From 9d5de1fb83d060e001a2f08662133801a6fae2d5 Mon Sep 17 00:00:00 2001
From: joehart2001 <jh2536@cam.ac.uk>
Date: Wed, 3 Jun 2026 22:14:09 +0100
Subject: [PATCH] add differnet shapes for GMTKN55 categories

---
 .../molecular/GMTKN55/analyse_GMTKN55.py      | 10 ++++
 ml_peg/analysis/utils/decorators.py           | 47 +++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/ml_peg/analysis/molecular/GMTKN55/analyse_GMTKN55.py b/ml_peg/analysis/molecular/GMTKN55/analyse_GMTKN55.py
index ae39c5279..283740a32 100644
--- a/ml_peg/analysis/molecular/GMTKN55/analyse_GMTKN55.py
+++ b/ml_peg/analysis/molecular/GMTKN55/analyse_GMTKN55.py
@@ -84,6 +84,14 @@ def structure_info() -> dict[str, dict[str, float] | list | NDArray]:
 
 INFO = structure_info()
 
+_CATEGORY_LABELS = {
+    "Basic properties and reaction energies for small systems": "Small systems",
+    "Intermolecular noncovalent interactions": "Intermolecular NCIs",
+    "Intramolecular noncovalent interactions": "Intramolecular NCIs",
+    "Reaction barrier heights": "Barrier heights",
+    "Reaction energies for large systems and isomerisation reactions": "Large systems",
+}
+
 
 @pytest.fixture
 @plot_parity(
@@ -97,6 +105,8 @@ def structure_info() -> dict[str, dict[str, float] | list | NDArray]:
         "System": INFO["systems"],
         "Excluded": INFO["excluded"],
     },
+    symbol_by=INFO["categories"].tolist(),
+    symbol_labels=_CATEGORY_LABELS,
 )
 def rel_energies() -> dict[str, list[float]]:
     """
diff --git a/ml_peg/analysis/utils/decorators.py b/ml_peg/analysis/utils/decorators.py
index 820333f6b..d32c82b24 100644
--- a/ml_peg/analysis/utils/decorators.py
+++ b/ml_peg/analysis/utils/decorators.py
@@ -164,6 +164,8 @@ def plot_parity(
     y_label: str | None = None,
     hoverdata: dict | None = None,
     filename: str = "parity.json",
+    symbol_by: list | None = None,
+    symbol_labels: dict[str, str] | None = None,
 ) -> Callable:
     """
     Plot parity plot of MLIP results against reference data.
@@ -180,6 +182,13 @@ def plot_parity(
         Hover data dictionary. Default is `{}`.
     filename
         Filename to save plot as JSON. Default is "parity.json".
+    symbol_by
+        Per-point list of group values. When provided, each point receives a
+        marker symbol based on its group, while trace colours still represent
+        models. Legend-only traces above the plot show one marker symbol per group.
+    symbol_labels
+        Optional mapping from ``symbol_by`` values to shorter display names
+        used in the legend. Values absent from this dict are shown as-is.
 
     Returns
     -------
@@ -231,6 +240,17 @@ def plot_parity_wrapper(*args, **kwargs) -> dict[str, Any]:
                 customdata = list(zip(*hoverdata.values(), strict=True))
 
             fig = go.Figure()
+            marker_kwargs = {}
+            if symbol_by:
+                symbols = ["circle", "square", "diamond", "cross", "x"]
+                groups = list(dict.fromkeys(symbol_by))
+                group_symbol = {
+                    g: symbols[i % len(symbols)] for i, g in enumerate(groups)
+                }
+                marker_kwargs = {
+                    "marker": {"symbol": [group_symbol[g] for g in symbol_by]}
+                }
+
             for mlip, value in results.items():
                 if mlip == "ref":
                     continue
@@ -242,9 +262,26 @@ def plot_parity_wrapper(*args, **kwargs) -> dict[str, Any]:
                         mode="markers",
                         customdata=customdata,
                         hovertemplate=hovertemplate,
+                        **marker_kwargs,
                     )
                 )
 
+            if symbol_by:
+                for group in groups:
+                    label = (symbol_labels or {}).get(group, group)
+                    fig.add_trace(
+                        go.Scatter(
+                            x=[None],
+                            y=[None],
+                            name=label,
+                            mode="markers",
+                            marker={"symbol": group_symbol[group], "color": "black"},
+                            hoverinfo="skip",
+                            legend="legend2",
+                            showlegend=True,
+                        )
+                    )
+
             full_fig = fig.full_figure_for_development()
             x_range = full_fig.layout.xaxis.range
             y_range = full_fig.layout.yaxis.range
@@ -268,6 +305,16 @@ def plot_parity_wrapper(*args, **kwargs) -> dict[str, Any]:
                 xaxis={"title": {"text": x_label}},
                 yaxis={"title": {"text": y_label}},
             )
+            if symbol_by:
+                fig.update_layout(
+                    legend2={
+                        "orientation": "h",
+                        "yanchor": "bottom",
+                        "y": 1.02,
+                        "xanchor": "left",
+                        "x": 0,
+                    }
+                )
 
             fig.update_traces()