Repository: RemDelaporteMathurin/h-transport-materials

Issue:
PropertiesGroup plot add a colour_by argument
Users should be able to colour by properties (default), materials, isotopes, authors with a `colour_by` argument.

Additionally, say users pass a colour_by argument, they should be able to pass a `colours` dict.

For example:

```python
import h_transport_materials as htm

group = htm.diffusitivites
htm.plotting.plot(group, colour_by="material", colours={"tungsten": "tab:grey", "copper": "tab:blue"})

```



Discussion:

Patch:
diff --git a/docs/user/plotting.rst b/docs/user/plotting.rst
index c94eb73..dbc4fa6 100644
--- a/docs/user/plotting.rst
+++ b/docs/user/plotting.rst
@@ -171,6 +171,24 @@ Calculate the mean value and plot it too:
    plt.yscale("log")
    plt.show()
 
+The properties can be coloured according to different attributes like ``materials``, ``author``...
+
+.. plot::
+   :include-source: true
+
+   import h_transport_materials as htm
+   from h_transport_materials.plotting import plot
+   import matplotlib.pyplot as plt
+
+   # filter only tungsten and H
+   diffusivities = htm.diffusivities.filter(material="tungsten")
+
+   plot(diffusivities, colour_by="author")
+
+   plt.title("Tungsten diffusivity")
+   plt.yscale("log")
+   plt.show()
+
 Interactive visualisation
 -------------------------
 
diff --git a/h_transport_materials/material.py b/h_transport_materials/material.py
index fd40742..4965a07 100644
--- a/h_transport_materials/material.py
+++ b/h_transport_materials/material.py
@@ -30,6 +30,9 @@ class Material:
     def __str__(self) -> str:
         return self.name
 
+    def __hash__(self):
+        return hash(self.name)
+
 
 class PlasmaFacing(Material):
     family = "plasma facing"
@@ -68,6 +71,9 @@ class PureMetal(Metal):
 
         return super().__eq__(mat)
 
+    def __hash__(self):
+        return hash(self.name)
+
 
 class Hastelloy(Alloy):
     family = "hastelloy"
diff --git a/h_transport_materials/plotting.py b/h_transport_materials/plotting.py
index fdbef74..15ed71b 100644
--- a/h_transport_materials/plotting.py
+++ b/h_transport_materials/plotting.py
@@ -15,6 +15,7 @@ def plot(
     auto_label=True,
     show_datapoints=True,
     scatter_kwargs={},
+    colour_by="property",
     **kwargs
 ):
     """Plots a Property object on a temperature plot
@@ -34,6 +35,8 @@ def plot(
             scattered too. Defaults to True.
         scatter_kwargs (dict, optional): other matplotlib.pyplot.scatter arguments.
             Defaults to {}.
+        colour_by (str, optional): a property attribute to colour by (eg. "author", "isotope",
+            "material"). Defaults to "property".
         kwargs: other matplotlib.pyplot.plot arguments
     Returns:
         matplotlib.lines.Line2D: the Line2D artist
@@ -52,8 +55,18 @@ def plot(
         group = prop
         if prop.units == "mixed units":
             raise ValueError("Cannot plot group with mixed units")
+
+        # compute the prop to colour mapping
+        if colour_by != "property":
+            prop_to_color = get_prop_to_color(group, colour_by)
+
         lines = []
         for single_prop in group:
+            # change colour from kwargs if need be
+            current_kwargs = kwargs.copy()
+            if colour_by != "property" and "color" not in kwargs:
+                current_kwargs["color"] = prop_to_color[single_prop]
+
             l = plot_property(
                 single_prop,
                 T_bounds=T_bounds,
@@ -61,7 +74,7 @@ def plot(
                 auto_label=auto_label,
                 show_datapoints=show_datapoints,
                 scatter_kwargs=scatter_kwargs,
-                **kwargs
+                **current_kwargs
             )
             lines.append(l)
         return lines
@@ -109,6 +122,28 @@ def plot_property(
     return l
 
 
+def get_prop_to_color(group: PropertiesGroup, colour_by: str):
+    """Returns a dictionary mapping Property objects to a colour based on
+    a property attribute
+
+    Args:
+        group (PropertiesGroup): a group of properties
+        colour_by (str): a property attribute to colour by (eg. "author", "isotope", "material")
+
+    Returns:
+        dict: a dictionary mapping properties to colours
+    """
+    colour_cycle = plt.rcParams["axes.prop_cycle"].by_key()["color"]
+
+    all_keys = list(set([getattr(prop, colour_by) for prop in group]))
+    key_to_colour = {
+        key: colour_cycle[i % len(colour_cycle)] for i, key in enumerate(all_keys)
+    }
+    prop_to_colour = {prop: key_to_colour[getattr(prop, colour_by)] for prop in group}
+
+    return prop_to_colour
+
+
 def line_labels(
     ax=None,
     min_label_distance: float or str = "auto",


Test patch:
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
index b74b762..2f582d5 100644
--- a/tests/test_plotting.py
+++ b/tests/test_plotting.py
@@ -47,3 +47,10 @@ def test_plot_group_with_mixed_units():
     """Tests that plotting a group with mixed units raises an error"""
     with pytest.raises(ValueError):
         htm.plotting.plot(htm.solubilities)
+
+
+def test_plot_group_with_colour_by():
+    """Tests that a group can be plotted with a non-default colour_by argument"""
+
+    htm.plotting.plot(htm.diffusivities, colour_by="material")
+    plt.clf()