Add basic element filter

ml_peg/analysis/molecular_crystal/DMC_ICE13/analyse_DMC_ICE13.py

@@ -10,6 +10,7 @@
 from ml_peg.analysis.utils.decorators import build_table, plot_parity
 from ml_peg.analysis.utils.utils import (
     build_dispersion_name_map,
+    get_struct_info,
     load_metrics_config,
     mae,
 )
@@ -28,27 +29,12 @@
     METRICS_CONFIG_PATH
 )
 
-
-def get_polymorph_names() -> list[str]:
-    """
-    Get list of polymorph names.
-
-    Returns
-    -------
-    list[str]
-        List of polymorph names from structure files.
-    """
-    polymorph_names = []
-    for model_name in MODELS:
-        model_dir = CALC_PATH / model_name
-        if model_dir.exists():
-            xyz_files = sorted(model_dir.glob("*_polymorph.xyz"))
-            if xyz_files:
-                for xyz_file in xyz_files:
-                    atoms = read(xyz_file)
-                    polymorph_names.append(atoms.info["polymorph"])
-                break
-    return polymorph_names
+INFO = get_struct_info(
+    calc_path=CALC_PATH,
+    glob_pattern="*_polymorph.xyz",
+    info_keys=["polymorph"],
+    out_path=OUT_PATH,
+)
 
 
 @pytest.fixture
@@ -58,12 +44,12 @@ def get_polymorph_names() -> list[str]:
     x_label="Predicted lattice energy / meV",
     y_label="Reference lattice energy / meV",
     hoverdata={
-        "Polymorph": get_polymorph_names(),
+        "Polymorph": INFO["polymorph"],
     },
 )
 def lattice_energies() -> dict[str, list]:
     """
-    Get lattice energies for all DMC-ICE13 polymorphs.
+    Get lattice energies for all DMC-ICE13 polymorphs and plot as scatter.
 
     Returns
     -------
@@ -109,8 +95,7 @@ def lattice_energies() -> dict[str, list]:
     return results
 
 
-@pytest.fixture
-def dmc_ice13_errors(lattice_energies) -> dict[str, float]:
+def get_errors(lattice_energies: dict[str, list]) -> dict[str, float]:
     """
     Get mean absolute error for lattice energies.
 
@@ -126,7 +111,7 @@ def dmc_ice13_errors(lattice_energies) -> dict[str, float]:
     """
     results = {}
     for model_name in MODELS:
-        if lattice_energies[model_name]:
+        if lattice_energies.get(model_name):
             results[model_name] = mae(
                 lattice_energies["ref"], lattice_energies[model_name]
             )
@@ -135,30 +120,47 @@ def dmc_ice13_errors(lattice_energies) -> dict[str, float]:
     return results
 
 
+def get_metrics(lattice_energies: dict[str, list]) -> dict[str, dict]:
+    """
+    Get all DMC-ICE13 metrics.
+
+    Parameters
+    ----------
+    lattice_energies
+        Dictionary of reference and predicted lattice energies.
+
+    Returns
+    -------
+    dict[str, dict]
+        Metric names and values for all models.
+    """
+    return {
+        "MAE": get_errors(lattice_energies),
+    }
+
+
 @pytest.fixture
 @build_table(
     filename=OUT_PATH / "dmc_ice13_metrics_table.json",
     metric_tooltips=DEFAULT_TOOLTIPS,
     thresholds=DEFAULT_THRESHOLDS,
     mlip_name_map=DISPERSION_NAME_MAP,
 )
-def metrics(dmc_ice13_errors: dict[str, float]) -> dict[str, dict]:
+def metrics(lattice_energies: dict[str, list]) -> dict[str, dict]:
     """
     Get all DMC-ICE13 metrics.
 
     Parameters
     ----------
-    dmc_ice13_errors
-        Mean absolute errors for all polymorphs.
+    lattice_energies
+        Dictionary of reference and predicted lattice energies.
 
     Returns
     -------
     dict[str, dict]
         Metric names and values for all models.
     """
-    return {
-        "MAE": dmc_ice13_errors,
-    }
+    return get_metrics(lattice_energies)
 
 
 def test_dmc_ice13(metrics: dict[str, dict]) -> None:

ml_peg/analysis/molecular_crystal/X23/analyse_X23.py

@@ -11,6 +11,7 @@
 from ml_peg.analysis.utils.decorators import build_table, plot_parity
 from ml_peg.analysis.utils.utils import (
     build_dispersion_name_map,
+    get_struct_info,
     load_metrics_config,
     mae,
 )
@@ -32,27 +33,13 @@
 # Unit conversion
 EV_TO_KJ_PER_MOL = units.mol / units.kJ
 
-
-def get_system_names() -> list[str]:
-    """
-    Get list of X23 system names.
-
-    Returns
-    -------
-    list[str]
-        List of system names from structure files.
-    """
-    system_names = []
-    for model_name in MODELS:
-        model_dir = CALC_PATH / model_name
-        if model_dir.exists():
-            xyz_files = sorted(model_dir.glob("*.xyz"))
-            if xyz_files:
-                for xyz_file in xyz_files:
-                    atoms = read(xyz_file)
-                    system_names.append(atoms.info["system"])
-                break
-    return system_names
+INFO = get_struct_info(
+    calc_path=CALC_PATH,
+    glob_pattern="*.xyz",
+    index="0",
+    info_keys=["system"],
+    out_path=OUT_PATH,
+)
 
 
 @pytest.fixture
@@ -62,7 +49,7 @@ def get_system_names() -> list[str]:
     x_label="Predicted lattice energy / kJ/mol",
     y_label="Reference lattice energy / kJ/mol",
     hoverdata={
-        "System": get_system_names(),
+        "System": INFO["system"],
     },
 )
 def lattice_energies() -> dict[str, list]:
@@ -112,8 +99,7 @@ def lattice_energies() -> dict[str, list]:
     return results
 
 
-@pytest.fixture
-def x23_errors(lattice_energies) -> dict[str, float]:
+def get_errors(lattice_energies: dict[str, list]) -> dict[str, float]:
     """
     Get mean absolute error for lattice energies.
 
@@ -129,7 +115,7 @@ def x23_errors(lattice_energies) -> dict[str, float]:
     """
     results = {}
     for model_name in MODELS:
-        if lattice_energies[model_name]:
+        if lattice_energies.get(model_name):
             results[model_name] = mae(
                 lattice_energies["ref"], lattice_energies[model_name]
             )
@@ -138,30 +124,47 @@ def x23_errors(lattice_energies) -> dict[str, float]:
     return results
 
 
+def get_metrics(lattice_energies: dict[str, list]) -> dict[str, dict]:
+    """
+    Get all X23 metrics.
+
+    Parameters
+    ----------
+    lattice_energies
+        Dictionary of reference and predicted lattice energies.
+
+    Returns
+    -------
+    dict[str, dict]
+        Metric names and values for all models.
+    """
+    return {
+        "MAE": get_errors(lattice_energies),
+    }
+
+
 @pytest.fixture
 @build_table(
     filename=OUT_PATH / "x23_metrics_table.json",
     metric_tooltips=DEFAULT_TOOLTIPS,
     thresholds=DEFAULT_THRESHOLDS,
     mlip_name_map=DISPERSION_NAME_MAP,
 )
-def metrics(x23_errors: dict[str, float]) -> dict[str, dict]:
+def metrics(lattice_energies: dict[str, list]) -> dict[str, dict]:
     """
     Get all X23 metrics.
 
     Parameters
     ----------
-    x23_errors
-        Mean absolute errors for all systems.
+    lattice_energies
+        Dictionary of reference and predicted lattice energies.
 
     Returns
     -------
     dict[str, dict]
         Metric names and values for all models.
     """
-    return {
-        "MAE": x23_errors,
-    }
+    return get_metrics(lattice_energies)
 
 
 def test_x23(metrics: dict[str, dict]) -> None:

ml_peg/analysis/nebs/li_diffusion/analyse_li_diffusion.py

@@ -9,7 +9,7 @@
 import pytest
 
 from ml_peg.analysis.utils.decorators import build_table, plot_scatter
-from ml_peg.analysis.utils.utils import load_metrics_config
+from ml_peg.analysis.utils.utils import load_metrics_config, write_struct_info
 from ml_peg.app import APP_ROOT
 from ml_peg.calcs import CALCS_ROOT
 from ml_peg.models import current_models
@@ -27,20 +27,20 @@
 REF_VALUES = {"path_b": 0.27, "path_c": 2.5}
 
 
-def plot_nebs(model: str, path: Literal["b", "c"]) -> None:
+def plot_nebs(model_name: str, path: Literal["b", "c"]) -> None:
     """
     Plot NEB paths and save all structure files.
 
     Parameters
     ----------
-    model
+    model_name
         Name of MLIP.
     path
         Path "b" or "c" for NEB.
     """
 
     @plot_scatter(
-        filename=OUT_PATH / f"figure_{model}_neb_{path.lower()}.json",
+        filename=OUT_PATH / model_name / f"figure_neb_{path.lower()}.json",
         title=f"NEB path {path.upper()}",
         x_label="Image",
         y_label="Energy / eV",
@@ -57,16 +57,16 @@ def plot_neb() -> dict[str, tuple[list[float], list[float]]]:
         """
         results = {}
         structs = read(
-            CALC_PATH / f"li_diffusion_{path.lower()}-{model}-neb-band.extxyz",
+            CALC_PATH / model_name / f"li_diffusion_{path.lower()}-neb-band.extxyz",
             index=":",
         )
-        results[model] = [
+        results[model_name] = [
             list(range(len(structs))),
             [struct.get_potential_energy() for struct in structs],
         ]
-        structs_dir = OUT_PATH / model
+        structs_dir = OUT_PATH / model_name
         structs_dir.mkdir(parents=True, exist_ok=True)
-        write(structs_dir / f"{model}-{path.lower()}-neb-band.extxyz", structs)
+        write(structs_dir / f"{path.lower()}-neb-band.extxyz", structs)
 
         return results
 
@@ -88,7 +88,7 @@ def path_b_error() -> dict[str, float]:
     for model_name in MODELS:
         plot_nebs(model_name, "b")
         with open(
-            CALC_PATH / f"li_diffusion_b-{model_name}-neb-results.dat", encoding="utf8"
+            CALC_PATH / model_name / "li_diffusion_b-neb-results.dat", encoding="utf8"
         ) as f:
             data = f.readlines()
             pred_barrier, _, _ = tuple(float(x) for x in data[1].split())
@@ -111,7 +111,7 @@ def path_c_error() -> dict[str, float]:
     for model_name in MODELS:
         plot_nebs(model_name, "c")
         with open(
-            CALC_PATH / f"li_diffusion_c-{model_name}-neb-results.dat", encoding="utf8"
+            CALC_PATH / model_name / "li_diffusion_c-neb-results.dat", encoding="utf8"
         ) as f:
             data = f.readlines()
             pred_barrier, _, _ = tuple(float(x) for x in data[1].split())
@@ -158,4 +158,8 @@ def test_li_diffusion(metrics: dict[str, dict]) -> None:
     metrics
         All Li diffusion metrics.
     """
-    return
+    write_struct_info(
+        data_path=CALC_PATH / "mock" / "li_diffusion_b-neb-band.extxyz",
+        out_path=OUT_PATH,
+        index=0,
+    )