From cf415a6db0009e8aeebf34adfc8dec7bebf017b9 Mon Sep 17 00:00:00 2001
From: Matthew Evans <git@ml-evs.science>
Date: Tue, 13 Sep 2022 03:02:32 +0100
Subject: [PATCH] Add function to turn formula into colour

---
 matador/utils/viz_utils.py | 31 +++++++++++++++++++++++++++----
 tests/test_viz_utils.py    | 18 ++++++++++++++++++
 2 files changed, 45 insertions(+), 4 deletions(-)
 create mode 100755 tests/test_viz_utils.py

diff --git a/matador/utils/viz_utils.py b/matador/utils/viz_utils.py
index 75f65220..f263c02f 100644
--- a/matador/utils/viz_utils.py
+++ b/matador/utils/viz_utils.py
@@ -155,7 +155,7 @@ def fresnel_view(
     images: Union[bool, float] = True,
     pad_cell: bool = True,
     lights: Optional[Callable] = None,
-    **camera_kwargs
+    **camera_kwargs,
 ) -> "fresnel.Scene":
     """Return a fresnel scene visualising the input crystal.
 
@@ -649,7 +649,7 @@ def fresnel_plot(
     labels: Union[bool, List[str]] = True,
     renderer: Optional[Callable] = None,
     camera_patches: Optional[List[Optional[Dict]]] = None,
-    **fresnel_view_kwargs
+    **fresnel_view_kwargs,
 ):
     """Visualize a series of structures as a grid of matplotlib plots.
 
@@ -739,6 +739,27 @@ def fresnel_plot(
     return fig, axes, scenes
 
 
+def formula_to_colour(formula: str) -> List[float]:
+    """Return an RGBA colour for the given chemical formula, provided
+    the formula has 3 or fewer species.
+
+    """
+    from matador.utils.chem_utils import get_stoich_from_formula, get_concentration
+
+    stoichiometry = get_stoich_from_formula(formula)
+    elements = [d[0] for d in stoichiometry]
+    if len(elements) > 3:
+        raise RuntimeError(
+            f"Cannot mix a colour for more than 3 elements, received: {stoichiometry}"
+        )
+
+    concentration = get_concentration(
+        stoichiometry, elements=elements, include_end=True
+    )
+
+    return colour_from_ternary_concentration(concentration, species=elements)
+
+
 def colour_from_ternary_concentration(
     conc: Union[Tuple[float, float], Tuple[float, float, float]],
     species: List[str],
@@ -751,10 +772,12 @@ def colour_from_ternary_concentration(
         RGBA array.
 
     """
-    if len(conc) == 2:
+    if len(conc) == 1:
+        return get_element_colours()[species[0]] + [alpha]
+    elif len(conc) == 2:
         x, y = conc
         z = 1 - x - y
     else:
         x, y, z = conc
-    colours = [np.asarray(get_element_colours()[s]) for s in species]
+    colours = [np.asarray(get_element_colours()[s]) for s in species + ["H"]]
     return ((x * colours[0] + y * colours[1] + z * colours[2])).tolist() + [alpha]
diff --git a/tests/test_viz_utils.py b/tests/test_viz_utils.py
new file mode 100755
index 00000000..395cce2c
--- /dev/null
+++ b/tests/test_viz_utils.py
@@ -0,0 +1,18 @@
+import unittest
+import numpy as np
+
+
+class VizUtilTest(unittest.TestCase):
+    """Tests viz util functions."""
+
+    def test_formula_to_colour(self):
+        from matador.utils.viz_utils import formula_to_colour
+
+        self.assertListEqual(formula_to_colour("K"), [0.4, 0.14, 0.43, 1])
+        self.assertListEqual(formula_to_colour("KP"), [0.62, 0.275, 0.255, 1])
+        self.assertListEqual(formula_to_colour("P"), [0.84, 0.41, 0.08, 1])
+        np.testing.assert_array_almost_equal(
+            formula_to_colour("KSnP"), [0.60666, 0.366666, 0.3999999, 1], decimal=3
+        )
+        with self.assertRaises(RuntimeError):
+            formula_to_colour("KSnPSb")