Adding get_good_linestyles()

changelog.md

@@ -4,6 +4,7 @@
 
 ### [v0.1.7] (10.08.2022)
 
+- Adding new function `puma.utils.get_good_linestyles()` for easier linestyle management [#116](https://github.com/umami-hep/puma/pull/116)
 - Adding new option to place rejection label legend in ROC plots [#109](https://github.com/umami-hep/puma/pull/109)
 
 ### [v0.1.6] (26.07.2022)

puma/tests/utils/test_utils.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python
+"""
+Unit test script for the functions in utils/__init__.py
+"""
+
+import os
+import tempfile
+import unittest
+
+import numpy as np
+from matplotlib.testing.compare import compare_images
+
+from puma import Line2D, Line2DPlot
+from puma.utils import get_good_linestyles, logger, set_log_level
+
+set_log_level(logger, "DEBUG")
+
+
+class Linestyles_TestCase(unittest.TestCase):
+    """Test class for the linestyle management of puma."""
+
+    def setUp(self):
+        # Set up directories for comparison plots
+        self.tmp_dir = tempfile.TemporaryDirectory()
+        self.actual_plots_dir = f"{self.tmp_dir.name}/"
+        self.expected_plots_dir = os.path.join(
+            os.path.dirname(__file__), "expected_plots"
+        )
+
+    def test_raise_value_error(self):
+        """Test if ValueError is raised for wrong argument type."""
+        with self.assertRaises(ValueError):
+            get_good_linestyles(4)
+
+    def test_get_good_linestyles(self):
+        """Test if the default linestyles obtained are the correct ones."""
+
+        with self.subTest("Testing default linestyles"):
+            expected_linestyles = [
+                "solid",
+                (0, (1, 1)),
+                (0, (5, 1)),
+                (0, (3, 1, 1, 1)),
+                (0, (3, 1, 1, 1, 1, 1)),
+                (0, (1, 1)),
+                (0, (5, 5)),
+                "dashdot",
+                (0, (1, 10)),
+                (0, (5, 10)),
+                (0, (3, 10, 1, 10)),
+                (0, (3, 10, 1, 10, 1, 10)),
+                (0, (3, 5, 1, 5)),
+                (0, (3, 5, 1, 5, 1, 5)),
+            ]
+            actual_linestyles = get_good_linestyles()
+            self.assertListEqual(expected_linestyles, actual_linestyles)
+
+        with self.subTest("Test specifying list of names"):
+            linestyle_names = ["densely dashed", "dashdotted"]
+            expected_linestyles = [(0, (5, 1)), (0, (3, 5, 1, 5))]
+            actual_linestyles = get_good_linestyles(linestyle_names)
+            self.assertListEqual(expected_linestyles, actual_linestyles)
+
+        with self.subTest("Test case of only one name given (not list)"):
+            expected_linestyle = (0, (5, 1))
+            actual_linestyle = get_good_linestyles("densely dashed")
+            self.assertTupleEqual(expected_linestyle, actual_linestyle)
+
+    def test_linestyles_accepted_by_mpl(self):
+        """Test if all the linestyles from get_good_linestyles() are accepted by
+        matplotlib."""
+
+        test_plot = Line2DPlot()
+        for i, linestyle in enumerate(get_good_linestyles()):
+            test_plot.add(
+                Line2D(
+                    np.linspace(0, 10, 10),
+                    i * np.linspace(0, 10, 10),
+                    linestyle=linestyle,
+                )
+            )
+        test_plot.draw()
+        plotname = "test_linestyles.png"
+        test_plot.savefig(f"{self.actual_plots_dir}/{plotname}")
+        # Uncomment line below to update expected image
+        # test_plot.savefig(f"{self.expected_plots_dir}/{plotname}")
+        self.assertIsNone(
+            compare_images(
+                f"{self.actual_plots_dir}/{plotname}",
+                f"{self.expected_plots_dir}/{plotname}",
+                tol=1,
+            )
+        )

puma/utils/__init__.py

@@ -138,6 +138,70 @@ def get_good_colours(colour_scheme=None):
         return Dark2_8.mpl_colors
 
 
+def get_good_linestyles(names=None):
+    """Returns a list of good linestyles
+
+    Parameters
+    ----------
+    names : list or str, optional
+        List or string of the name(s) of the linestyle(s) you want to retrieve, e.g.
+        "densely dotted" or ["solid", "dashdot", "densely dashed"], by default None
+
+    Returns
+    -------
+    list
+        List of good linestyles. Either the specified selection or the whole list in
+        the predefined order.
+
+    Raises
+    ------
+    ValueError
+        If `names` is not a str or list.
+    """
+    linestyle_tuples = {
+        "solid": "solid",
+        "densely dashed": (0, (5, 1)),
+        "densely dotted": (0, (1, 1)),
+        "densely dashdotted": (0, (3, 1, 1, 1)),
+        "densely dashdotdotted": (0, (3, 1, 1, 1, 1, 1)),
+        "dotted": (0, (1, 1)),
+        "dashed": (0, (5, 5)),
+        "dashdot": "dashdot",
+        "loosely dashed": (0, (5, 10)),
+        "loosely dotted": (0, (1, 10)),
+        "loosely dashdotted": (0, (3, 10, 1, 10)),
+        "loosely dashdotdotted": (0, (3, 10, 1, 10, 1, 10)),
+        "dashdotted": (0, (3, 5, 1, 5)),
+        "dashdotdotted": (0, (3, 5, 1, 5, 1, 5)),
+    }
+
+    default_order = [
+        "solid",
+        "densely dotted",
+        "densely dashed",
+        "densely dashdotted",
+        "densely dashdotdotted",
+        "dotted",
+        "dashed",
+        "dashdot",
+        "loosely dotted",
+        "loosely dashed",
+        "loosely dashdotted",
+        "loosely dashdotdotted",
+        "dashdotted",
+        "dashdotdotted",
+    ]
+    if names is None:
+        names = default_order
+    elif isinstance(names, str):
+        return linestyle_tuples[names]
+    elif not isinstance(names, list):
+        raise ValueError(
+            "Invalid type of `names`, has to be a list of strings or a sting."
+        )
+    return [linestyle_tuples[name] for name in names]
+
+
 global_config = {
     "flavour_categories": {
         "bjets": {