tidy elasticity, move to interactive cells (but required a bit more code to get the legend to update too)

Author: joehart2001

ml_peg/analysis/bulk_crystal/elasticity/analyse_elasticity.py

@@ -12,7 +12,11 @@
     build_table,
     plot_density_scatter,
 )
-from ml_peg.analysis.utils.utils import load_metrics_config, mae
+from ml_peg.analysis.utils.utils import (
+    build_density_inputs,
+    load_metrics_config,
+    mae,
+)
 from ml_peg.app import APP_ROOT
 from ml_peg.calcs import CALCS_ROOT
 from ml_peg.models.get_models import get_model_names
@@ -31,26 +35,6 @@
 G_COLUMN = "G_vrh"
 
 
-def _load_results(model_name: str) -> pd.DataFrame | None:
-    """
-    Load benchmark results for a model, if available.
-
-    Parameters
-    ----------
-    model_name
-        Name of the model whose results should be read.
-
-    Returns
-    -------
-    pd.DataFrame | None
-        Results dataframe if present, otherwise ``None``.
-    """
-    results_path = CALC_PATH / model_name / "moduli_results.csv"
-    if not results_path.exists():
-        return None
-    return pd.read_csv(results_path)
-
-
 def _filter_results(df: pd.DataFrame, model_name: str) -> tuple[pd.DataFrame, int]:
     """
     Filter outlier predictions and return remaining data with exclusion count.
@@ -85,22 +69,11 @@ def _collect_model_data() -> dict[str, dict[str, Any]]:
     """
     stats: dict[str, dict[str, Any]] = {}
     for model_name in MODELS:
-        df = _load_results(model_name)
-        if df is None:
-            stats[model_name] = {
-                "bulk": None,
-                "shear": None,
-                "excluded": None,
-            }
-            continue
+        results_path = CALC_PATH / model_name / "moduli_results.csv"
+        df = pd.read_csv(results_path)
+
         filtered, excluded = _filter_results(df, model_name)
-        if filtered.empty:
-            stats[model_name] = {
-                "bulk": None,
-                "shear": None,
-                "excluded": excluded,
-            }
-            continue
+
         stats[model_name] = {
             "bulk": {
                 "ref": filtered[f"{K_COLUMN}_DFT"].tolist(),
@@ -115,44 +88,6 @@ def _collect_model_data() -> dict[str, dict[str, Any]]:
     return stats
 
 
-def _density_inputs(property_key: str, model_stats: dict[str, dict[str, Any]]) -> dict:
-    """
-    Prepare mapping for density scatter decorator.
-
-    Parameters
-    ----------
-    property_key
-        Property key to extract (``"bulk"`` or ``"shear"``).
-    model_stats
-        Aggregated statistics for each model.
-
-    Returns
-    -------
-    dict
-        Mapping of model name to density-scatter inputs.
-    """
-    inputs = {}
-    for model_name in MODELS:
-        prop = model_stats.get(model_name, {}).get(property_key)
-        excluded = model_stats.get(model_name, {}).get("excluded")
-        if not prop:
-            inputs[model_name] = {
-                "ref": [],
-                "pred": [],
-                "meta": {"excluded": excluded},
-            }
-            continue
-        ref_vals = prop["ref"]
-        pred_vals = prop["pred"]
-        inputs[model_name] = {
-            "ref": ref_vals,
-            "pred": pred_vals,
-            "mae": mae(ref_vals, pred_vals) if ref_vals else None,
-            "meta": {"excluded": excluded},
-        }
-    return inputs
-
-
 @pytest.fixture
 def elasticity_stats() -> dict[str, dict[str, Any]]:
     """
@@ -185,9 +120,6 @@ def bulk_mae(elasticity_stats: dict[str, dict[str, Any]]) -> dict[str, float | N
     results: dict[str, float | None] = {}
     for model_name in MODELS:
         prop = elasticity_stats.get(model_name, {}).get("bulk")
-        if not prop:
-            results[model_name] = None
-            continue
         results[model_name] = mae(prop["ref"], prop["pred"])
     return results
 
@@ -210,17 +142,14 @@ def shear_mae(elasticity_stats: dict[str, dict[str, Any]]) -> dict[str, float |
     results: dict[str, float | None] = {}
     for model_name in MODELS:
         prop = elasticity_stats.get(model_name, {}).get("shear")
-        if not prop:
-            results[model_name] = None
-            continue
         results[model_name] = mae(prop["ref"], prop["pred"])
     return results
 
 
 @pytest.fixture
 @plot_density_scatter(
     filename=OUT_PATH / "figure_bulk_density.json",
-    title="Bulk modulus density",
+    title="Bulk modulus density plot",
     x_label="Reference bulk modulus / GPa",
     y_label="Predicted bulk modulus / GPa",
 )
@@ -238,13 +167,13 @@ def bulk_density(elasticity_stats: dict[str, dict[str, Any]]) -> dict[str, dict]
     dict[str, dict]
         Mapping of model name to density-scatter data.
     """
-    return _density_inputs("bulk", elasticity_stats)
+    return build_density_inputs(MODELS, elasticity_stats, "bulk", mae_fn=mae)
 
 
 @pytest.fixture
 @plot_density_scatter(
     filename=OUT_PATH / "figure_shear_density.json",
-    title="Shear modulus density",
+    title="Shear modulus density plot",
     x_label="Reference shear modulus / GPa",
     y_label="Predicted shear modulus / GPa",
 )
@@ -262,7 +191,7 @@ def shear_density(elasticity_stats: dict[str, dict[str, Any]]) -> dict[str, dict
     dict[str, dict]
         Mapping of model name to density-scatter data.
     """
-    return _density_inputs("shear", elasticity_stats)
+    return build_density_inputs(MODELS, elasticity_stats, "shear", mae_fn=mae)
 
 
 @pytest.fixture

ml_peg/analysis/utils/decorators.py

@@ -261,12 +261,12 @@ def plot_density_scatter(
     seed: int = 0,
 ) -> Callable:
     """
-    Plot density-coloured parity scatter with model selector.
+    Plot density-coloured parity scatter with legend-based model toggling.
 
     The decorated function must return a mapping of model name to a dictionary with
     ``ref`` and ``pred`` arrays (and optional ``mae``). Each model is rendered as a
     scatter trace with marker colours indicating local data density.
-    Only one model is shown at a time using a dropdown selector.
+    Only one model is shown at a time; use the legend to toggle models.
 
     Parameters
     ----------
@@ -486,41 +486,21 @@ def _downsample(
                 )
             )
 
-            buttons = []
-            n_models = len(model_names)
-            for idx, model in enumerate(model_names):
-                visible = [False] * n_models + [True]
-                visible[idx] = True
-                buttons.append(
-                    {
-                        "label": model,
-                        "method": "update",
-                        "args": [
-                            {"visible": visible},
-                            {
-                                "title": f"{title} - {model}" if title else None,
-                                "annotations": [annotations[idx]],
-                            },
-                        ],
-                    }
-                )
+            # Store all annotations and model order in layout meta so consumers
+            # can swap annotation text when filtering per-model on the frontend.
+            layout_meta = {
+                "annotations": annotations,
+                "models": model_names,
+            }
 
             fig.update_layout(
-                title={"text": f"{title} - {model_names[0]}" if title else None},
+                title={"text": title} if title else None,
                 xaxis={"title": {"text": x_label}},
                 yaxis={"title": {"text": y_label}},
-                updatemenus=[
-                    {
-                        "buttons": buttons,
-                        "direction": "down",
-                        "showactive": True,
-                        "x": 0.0,
-                        "xanchor": "left",
-                        "y": 1.15,
-                        "yanchor": "top",
-                    }
-                ],
                 annotations=[annotations[0]],
+                meta=layout_meta,
+                showlegend=True,
+                legend_title_text="Model",
             )
 
             Path(filename).parent.mkdir(parents=True, exist_ok=True)

ml_peg/analysis/utils/utils.py

@@ -4,6 +4,7 @@
 
 from collections.abc import Callable
 from pathlib import Path
+from typing import Any
 
 from matplotlib import cm
 from matplotlib.colors import Colormap
@@ -114,6 +115,54 @@ def rmse(ref: list, prediction: list) -> float:
     return mean_squared_error(ref, prediction)
 
 
+def build_density_inputs(
+    models: list[str],
+    model_stats: dict[str, dict[str, Any]],
+    property_key: str,
+    *,
+    mae_fn: Callable[[list, list], float] | None = None,
+) -> dict[str, dict[str, Any]]:
+    """
+    Prepare a model->data mapping for density scatter plots.
+
+    Parameters
+    ----------
+    models
+        Ordered list of model names to include.
+    model_stats
+        Mapping of model -> {"<property_key>": {"ref": [...], "pred": [...]},
+        "excluded": int}.
+    property_key
+        Key to extract from ``model_stats`` for each model (e.g. ``"bulk"`` or
+        ``"shear"``).
+    mae_fn
+        Optional callable to compute MAE. Defaults to :func:`mae` when None.
+
+    Returns
+    -------
+    dict[str, dict[str, Any]]
+        Mapping ready for ``plot_density_scatter``.
+    """
+    mae_fn = mae if mae_fn is None else mae_fn
+    inputs: dict[str, dict[str, Any]] = {}
+
+    for model_name in models:
+        stats = model_stats.get(model_name, {})
+        prop = stats.get(property_key)
+        excluded = stats.get("excluded")
+
+        ref_vals = prop.get("ref", [])
+        pred_vals = prop.get("pred", [])
+        inputs[model_name] = {
+            "ref": ref_vals,
+            "pred": pred_vals,
+            "mae": mae_fn(ref_vals, pred_vals) if ref_vals else None,
+            "meta": {"excluded": excluded} if excluded is not None else {},
+        }
+
+    return inputs
+
+
 def calc_metric_scores(
     metrics_data: list[MetricRow],
     thresholds: Thresholds | None = None,

ml_peg/app/bulk_crystal/molecular.yml renamed to ml_peg/app/bulk_crystal/bulk_crystal.yml

@@ -1,2 +1,2 @@
-title: Bulk Crystal Systems
+title: Bulk Crystals
 description: Bulk crystal properties, including elastic moduli, phonons, and lattice constants.

ml_peg/app/bulk_crystal/elasticity/app_elasticity.py

@@ -7,9 +7,12 @@
 
 from ml_peg.app import APP_ROOT
 from ml_peg.app.base_app import BaseApp
-from ml_peg.app.utils.build_callbacks import plot_from_table_column
-from ml_peg.app.utils.load import read_plot
+from ml_peg.app.utils.build_callbacks import plot_from_table_cell
+from ml_peg.app.utils.load import read_density_plot_for_model
+from ml_peg.models.get_models import get_model_names
+from ml_peg.models.models import current_models
 
+MODELS = get_model_names(current_models)
 BENCHMARK_NAME = "Elasticity"
 DOCS_URL = "https://ddmms.github.io/ml-peg/user_guide/benchmarks/bulk.html#elasticity"
 DATA_PATH = APP_ROOT / "data" / "bulk_crystal" / "elasticity"
@@ -20,22 +23,26 @@ class ElasticityApp(BaseApp):
 
     def register_callbacks(self) -> None:
         """Register callbacks to app."""
-        bulk_plot = read_plot(
-            DATA_PATH / "figure_bulk_density.json",
-            id=f"{BENCHMARK_NAME}-bulk-figure",
-        )
-        shear_plot = read_plot(
-            DATA_PATH / "figure_shear_density.json",
-            id=f"{BENCHMARK_NAME}-shear-figure",
-        )
+        density_plots = {
+            model: {
+                "Bulk modulus MAE": read_density_plot_for_model(
+                    filename=DATA_PATH / "figure_bulk_density.json",
+                    model=model,
+                    id=f"{BENCHMARK_NAME}-{model}-bulk-figure",
+                ),
+                "Shear modulus MAE": read_density_plot_for_model(
+                    filename=DATA_PATH / "figure_shear_density.json",
+                    model=model,
+                    id=f"{BENCHMARK_NAME}-{model}-shear-figure",
+                ),
+            }
+            for model in MODELS
+        }
 
-        plot_from_table_column(
+        plot_from_table_cell(
             table_id=self.table_id,
             plot_id=f"{BENCHMARK_NAME}-figure-placeholder",
-            column_to_plot={
-                "Bulk modulus MAE": bulk_plot,
-                "Shear modulus MAE": shear_plot,
-            },
+            cell_to_plot=density_plots,
         )