mpdd-alignn 1.2.0__tar.gz → 1.3.0__tar.gz

{mpdd_alignn-1.2.0 → mpdd_alignn-1.3.0}/LICENSE.md

@@ -1,6 +1,6 @@
 The MIT License (MIT)
 
-Copyright (c) 2024 Adam M. Krajewski
+Copyright (c) 2024-2026 Adam M. Krajewski
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

{mpdd_alignn-1.2.0 → mpdd_alignn-1.3.0}/PKG-INFO

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: mpdd-alignn
-Version: 1.2.0
+Version: 1.3.0
 Summary: Our version of the NIST-JARVIS ALIGNN optimized in terms of model performance and to some extent reliability, for large-scale deployments over the MPDD infrastructure by Phases Research Lab.
-Author-email: "Adam M. Krajewski" <adam@phaseslab.org>
+Author-email: "Adam M. Krajewski" <adam@phaseslab.org>, Ricardo Amaral <rna5137@psu.edu>
 Project-URL: Homepage, https://github.com/PhasesResearchLab/mpdd-alignn
 Project-URL: Bug Tracker, https://github.com/PhasesResearchLab/mpdd-alignn/issues
 Project-URL: Research Page, https://phaseslab.com/mpdd
@@ -34,6 +34,7 @@ Requires-Dist: dgl>=1.1.3
 Requires-Dist: scikit-learn>=1.3.1
 Requires-Dist: matplotlib>=3.4.1
 Requires-Dist: tqdm>=4.60.0
+Requires-Dist: natsort>=8.3.0
 Requires-Dist: pandas>=1.2.3
 Requires-Dist: pydantic_settings
 Requires-Dist: ruamel.yaml

{mpdd_alignn-1.2.0 → mpdd_alignn-1.3.0}/alignn/pretrained.py

@@ -5,7 +5,7 @@ import os
 import sys
 import json
 import time
-from typing import List, Dict, Union
+from typing import List, Dict, Union, Optional
 from importlib import resources
 
 # Extra utility imports
@@ -18,6 +18,7 @@ import pandas as pd
 tqdm.pandas()
 from pysmartdl2 import SmartDL
 from dgl.data.utils import save_graphs, load_graphs
+from natsort import natsorted
 
 # ML imports
 import torch
@@ -83,23 +84,46 @@ def unzip_default_models() -> None:
         with zipfile.ZipFile(model['model'], 'r') as zip_ref:
             zip_ref.extractall(model['model'].replace('.zip', ''))
 
-def run_models_from_directory(
-        directory: str, 
-        mode: str = "serial",
-        saveGraphs: bool = False):
-    """Run all default models on all structures in a directory that are either in POSCAR or CIF format."""
-    # Parse all structures into Atoms objects
-    atoms_array = []
-    outputs: List[Dict[str, Union[float, str]]] = []
-    for file in os.listdir(directory):
-        if file.endswith(("poscar", "POSCAR", "vasp", "VASP")):
-            outputs.append({"name": file})
-            atoms_array.append(Atoms.from_poscar(os.path.join(directory, file)))
-        elif file.endswith(("cif", "CIF")):
-            outputs.append({"name": file})
-            atoms_array.append(Atoms.from_cif(os.path.join(directory, file)))
-        else:
-            print(f"Skipping file {file} as it is not a POSCAR or CIF file!", flush=True)
+def _run_pretrained_models(
+    atoms_array: List[Atoms],
+    outputs: List[Dict[str, Union[float, str]]],
+    mode: str = "serial",
+    saveGraphs: bool = False,
+    models: Optional[List[str]] = None,
+):
+    """Internal helper to run default models on already parsed structures.
+    
+    Parameters
+    ----------
+    atoms_array : List[Atoms]
+        Array of atomic structures to run predictions on.
+    outputs : List[Dict[str, Union[float, str]]]
+        List of output dictionaries (typically containing file names).
+    mode : str, optional
+        Either "serial" or "parallel" for execution mode. Default is "serial".
+    saveGraphs : bool, optional
+        Whether to save computed graphs to disk. Default is False.
+    models : List[str], optional
+        List of model names to run. If None, all default models are run.
+        
+    Returns
+    -------
+    List[Dict[str, Union[float, str]]]
+        Output dictionaries with model predictions added as new keys.
+    """
+    # Validate and filter models
+    if models is None:
+        active_models = default_models
+    else:
+        model_names_set = {m["name"] for m in default_models}
+        invalid_models = [m for m in models if m not in model_names_set]
+        if invalid_models:
+            raise ValueError(
+                f"The following model names were not found in default models: {invalid_models}. "
+                f"Available models: {sorted(model_names_set)}"
+            )
+        active_models = [m for m in default_models if m["name"] in set(models)]
+    
     # Convert all Atoms to Graphs
     print(f"Converting {len(atoms_array)} structures to graphs...", flush=True)
     if mode == "serial":
@@ -115,86 +139,182 @@ def run_models_from_directory(
         )
     else:
         raise ValueError(f"Mode {mode} not implemented!")
-    
+
     if saveGraphs:
         print("Saving graphs to disk...", flush=True)
         input_files = [i["name"] for i in outputs]
         for g, name in zip(graph_array, input_files):
-            save_graphs(f"graphs/{name}.graph.bin", list(g), formats='coo')
+            save_graphs(f"graphs/{name}.graph.bin", list(g), formats="coo")
         print("Graphs saved!", flush=True)
 
     modelArray = []
-    for model in default_models:
-        modelPath = str(resources.files('alignn').joinpath(model['model']))
-        zp = zipfile.ZipFile(modelPath, 'r')
-        # Get the full path of checkpoint_300.pt or best_model.pt in the zip. Pick the first one found.
-        # This is a workaround for the fact that some models (new ones) have a different naming convention
-        # for the ALIGNN checkpoint files.
+    for model in active_models:
+        modelPath = str(resources.files("alignn").joinpath(model["model"]))
+        zp = zipfile.ZipFile(modelPath, "r")
+
         modelCheckpoint = [
-            i for i in zp.namelist() 
+            i for i in zp.namelist()
             if ("checkpoint_" in i and "pt" in i) or "best_model.pt" in i
-            ]
+        ]
         if len(modelCheckpoint) == 0:
-            raise ValueError(f"No model identifier found in {modelPath} for model {model['name']}!", flush=True)
+            raise ValueError(
+                f"No model identifier found in {modelPath} for model {model['name']}!"
+            )
         if len(modelCheckpoint) > 1:
-            print(f"Multiple model identifiers ({len(modelCheckpoint)}) found in {modelPath} for model {model['name']}: {modelCheckpoint}. Using the first one in the list.", flush=True)
-            modelCheckpoint = modelCheckpoint[0]
-        else:
-            modelCheckpoint = modelCheckpoint[0]
-        config = json.loads(zp.read([i for i in zp.namelist() if "config.json" in i][0]))
+            print(
+                f"Multiple model identifiers ({len(modelCheckpoint)}) found in "
+                f"{modelPath} for model {model['name']}: {modelCheckpoint}. "
+                "Using the first one in the list.",
+                flush=True,
+            )
+        modelCheckpoint = modelCheckpoint[0]
+
+        config = json.loads(
+            zp.read([i for i in zp.namelist() if "config.json" in i][0])
+        )
         data = zipfile.ZipFile(modelPath).read(modelCheckpoint)
-        
-        # Create model based on the model name in config
+
         model_type = config["model"]["name"]
         if model_type == "alignn":
             loadedModel = ALIGNN(ALIGNNConfig(**config["model"]))
         elif model_type == "alignn_atomwise":
             loadedModel = ALIGNNAtomWise(ALIGNNAtomWiseConfig(**config["model"]))
         else:
-            raise ValueError(f"Unknown model type: {model_type} for model {model['name']}")
+            raise ValueError(
+                f"Unknown model type: {model_type} for model {model['name']}"
+            )
 
-        _, filename = tempfile.mkstemp()
+        fd, filename = tempfile.mkstemp()
+        os.close(fd)
         with open(filename, "wb") as f:
             f.write(data)
-        
-        # Load checkpoint and handle different checkpoint formats
-        checkpoint = torch.load(filename, map_location='cpu')
+
+        checkpoint = torch.load(filename, map_location="cpu")
         if "model" in checkpoint:
             loadedModel.load_state_dict(checkpoint["model"])
         else:
-            # Some checkpoints store the model state dict directly
             loadedModel.load_state_dict(checkpoint)
-            
-        loadedModel.to('cpu')
+
+        loadedModel.to("cpu")
         loadedModel.eval()
         modelArray.append(loadedModel)
-        
+
         print(f"Model {model['name']} loaded!", flush=True)
-    
-    print(f"Running {len(default_models)} models on {len(graph_array)} structures...", flush=True)
-    # Run all models on all graphs
-    for model, loaded_model in zip(default_models, modelArray):
+
+    print(
+        f"Running {len(active_models)} models on {len(graph_array)} structures...",
+        flush=True,
+    )
+    for model, loaded_model in zip(active_models, modelArray):
         for g, out in zip(graph_array, outputs):
             model_output = loaded_model([g[0], g[1]])
-            
-            # Handle different output formats
+
             if isinstance(model_output, dict):
-                # For atomwise models that return dict with 'out' key
-                if 'out' in model_output:
-                    out_data = model_output['out'].item()
+                if "out" in model_output:
+                    out_data = model_output["out"].item()
                 else:
-                    # Take the first value if dict has other keys
                     out_data = list(model_output.values())[0].item()
             else:
-                # For regular models that return tensor directly
                 out_data = model_output.item()
-                
-            out[model['name']] = round(out_data, 6)
-    
-    print("All models runs complete!", flush=True)
 
+            out[model["name"]] = round(out_data, 6)
+
+    print("All models runs complete!", flush=True)
     return outputs
 
+
+def run_models_from_directory(
+    directory: str,
+    mode: str = "serial",
+    saveGraphs: bool = False,
+    models: Optional[List[str]] = None,
+):
+    """Run default models on all structures in a directory that are either in POSCAR or CIF format.
+    
+    Parameters
+    ----------
+    directory : str
+        Path to directory containing structure files.
+    mode : str, optional
+        Either "serial" or "parallel" for execution mode. Default is "serial".
+    saveGraphs : bool, optional
+        Whether to save computed graphs to disk. Default is False.
+    models : List[str], optional
+        List of model names to run. If None, all default models are run.
+        
+    Returns
+    -------
+    List[Dict[str, Union[float, str]]]
+        Output dictionaries with model predictions added as new keys.
+    """
+    atoms_array = []
+    outputs: List[Dict[str, Union[float, str]]] = []
+
+    for file in natsorted(os.listdir(directory)):
+        if file.endswith(("poscar", "POSCAR", "vasp", "VASP")):
+            outputs.append({"name": file})
+            atoms_array.append(Atoms.from_poscar(os.path.join(directory, file)))
+        elif file.endswith(("cif", "CIF")):
+            outputs.append({"name": file})
+            atoms_array.append(Atoms.from_cif(os.path.join(directory, file)))
+        else:
+            print(f"Skipping file {file} as it is not a POSCAR or CIF file!", flush=True)
+
+    return _run_pretrained_models(
+        atoms_array=atoms_array,
+        outputs=outputs,
+        mode=mode,
+        saveGraphs=saveGraphs,
+        models=models,
+    )
+
+
+def run_models_from_structure(
+    structure: str,
+    mode: str = "serial",
+    saveGraphs: bool = False,
+    models: Optional[List[str]] = None,
+):
+    """Run default models on a single structure file in POSCAR or CIF format.
+    
+    Parameters
+    ----------
+    structure : str
+        Path to structure file.
+    mode : str, optional
+        Either "serial" or "parallel" for execution mode. Default is "serial".
+    saveGraphs : bool, optional
+        Whether to save computed graphs to disk. Default is False.
+    models : List[str], optional
+        List of model names to run. If None, all default models are run.
+        
+    Returns
+    -------
+    List[Dict[str, Union[float, str]]]
+        List containing one output dictionary for the input structure, with
+        model predictions added as new keys.
+    """
+    atoms_array = []
+    outputs: List[Dict[str, Union[float, str]]] = []
+
+    file = os.path.basename(structure)
+    if file.endswith(("poscar", "POSCAR", "vasp", "VASP")):
+        outputs.append({"name": file})
+        atoms_array.append(Atoms.from_poscar(structure))
+    elif file.endswith(("cif", "CIF")):
+        outputs.append({"name": file})
+        atoms_array.append(Atoms.from_cif(structure))
+    else:
+        raise ValueError(f"File {structure} is not a POSCAR/VASP or CIF file!")
+
+    return _run_pretrained_models(
+        atoms_array=atoms_array,
+        outputs=outputs,
+        mode=mode,
+        saveGraphs=saveGraphs,
+        models=models,
+    )
+
 # ******* Old method to download models *******
 """
 Name of the model, figshare link, number of outputs,

mpdd_alignn-1.3.0/alignn/tests/test_mpddalignn.py

@@ -0,0 +1,184 @@
+"""Standalone regression tests for pretrained MPDD ALIGNN wrappers."""
+
+from pathlib import Path
+from unittest import TestCase, mock
+
+import pytest
+
+from alignn import pretrained
+from alignn.pretrained import _run_pretrained_models
+
+
+def test_invalid_model_name_raises_error():
+    """Regression test: passing invalid model names should raise a ValueError
+    from `_run_pretrained_models` with the expected message.
+    """
+    invalid_models_list = ["fake_model_name_xyz", "another_bad_model"]
+
+    with pytest.raises(ValueError, match="The following model names were not found in default models"):
+        _run_pretrained_models(
+            atoms_array=[],
+            outputs=[],
+            models=invalid_models_list,
+        )
+
+
+def _fake_run_pretrained_models(
+    atoms_array, outputs, mode="serial", saveGraphs=False, models=None
+):
+    """Return deterministic payloads so the wrapper behavior can be compared.
+
+    Includes models information in response for validation.
+    """
+    return [
+        {
+            **out,
+            "mock_pred": f"{out['name']}|{mode}|{saveGraphs}",
+            "models_used": models,  # Include for verification
+        }
+        for out in outputs
+    ]
+
+
+# Get model names from default models
+_default_models = pretrained.get_default_models()
+_model_names = [m["name"] for m in _default_models]
+_first_model = _model_names[0] if _model_names else "model_1"
+_second_model = _model_names[1] if len(_model_names) > 1 else "model_2"
+_third_model = _model_names[2] if len(_model_names) > 2 else "model_3"
+
+
+class TestMPDDAlignNPretrainedWrappers(TestCase):
+    def test_run_models_from_directory_and_structure_agree_for_sigma_phase(self):
+        sigma_phase_dir = Path(__file__).resolve().parents[2] / "example.SigmaPhase"
+        poscar_files = sorted(
+            path
+            for path in sigma_phase_dir.iterdir()
+            if path.is_file() and path.suffix.lower() == ".poscar"
+        )
+
+        self.assertGreater(
+            len(poscar_files),
+            0,
+            "Expected at least one POSCAR file under example.SigmaPhase.",
+        )
+
+        with mock.patch.object(
+            pretrained,
+            "_run_pretrained_models",
+            side_effect=_fake_run_pretrained_models,
+        ):
+            directory_results = pretrained.run_models_from_directory(
+                str(sigma_phase_dir),
+                mode="serial",
+                saveGraphs=False,
+            )
+
+            directory_by_name = {result["name"]: result for result in directory_results}
+            expected_names = {path.name for path in poscar_files}
+
+            self.assertEqual(set(directory_by_name), expected_names)
+
+            for path in poscar_files:
+                with self.subTest(structure=path.name):
+                    structure_result = pretrained.run_models_from_structure(
+                        str(path),
+                        mode="serial",
+                        saveGraphs=False,
+                    )
+
+                    self.assertEqual(len(structure_result), 1)
+                    self.assertEqual(structure_result[0], directory_by_name[path.name])
+
+
+@pytest.mark.parametrize(
+    "models_selection,mode,saveGraphs",
+    [
+        (None, "serial", False),                                    # Default: all models
+        ([_first_model, _second_model], "serial", False),           # Two models
+        ([_first_model], "serial", True),                           # Single model with saveGraphs
+        ([], "serial", False),                                      # Empty list (no models)
+        ([_third_model, _first_model], "parallel", False),          # Multiple models, parallel mode
+    ],
+)
+class TestModelFiltering:
+    """Parametrized tests for model filtering functionality."""
+
+    def test_run_models_from_structure_with_model_filter(
+        self, models_selection, mode, saveGraphs
+    ):
+        """Test that run_models_from_structure correctly passes and validates model filtering."""
+        sigma_phase_dir = Path(__file__).resolve().parents[2] / "example.SigmaPhase"
+        poscar_files = sorted(
+            path
+            for path in sigma_phase_dir.iterdir()
+            if path.is_file() and path.suffix.lower() == ".poscar"
+        )
+
+        assert len(poscar_files) > 0, "Expected at least one POSCAR file in example.SigmaPhase"
+
+        test_path = str(poscar_files[0])
+
+        with mock.patch.object(
+            pretrained,
+            "_run_pretrained_models",
+            side_effect=_fake_run_pretrained_models,
+        ) as mock_run:
+            results = pretrained.run_models_from_structure(
+                test_path,
+                mode=mode,
+                saveGraphs=saveGraphs,
+                models=models_selection,
+            )
+
+            # Verify the mock was called with correct parameters
+            mock_run.assert_called_once()
+            call_kwargs = mock_run.call_args.kwargs
+
+            # Verify models parameter was passed correctly
+            assert call_kwargs["models"] == models_selection, (
+                f"Expected models={models_selection}, got {call_kwargs['models']}"
+            )
+
+            # Verify other parameters
+            assert call_kwargs["mode"] == mode
+            assert call_kwargs["saveGraphs"] == saveGraphs
+
+            # Verify the models information is in the returned results
+            assert len(results) == 1
+            assert results[0]["models_used"] == models_selection
+
+    def test_run_models_from_directory_with_model_filter(
+        self, models_selection, mode, saveGraphs
+    ):
+        """Test that run_models_from_directory correctly passes and validates model filtering."""
+        sigma_phase_dir = Path(__file__).resolve().parents[2] / "example.SigmaPhase"
+
+        with mock.patch.object(
+            pretrained,
+            "_run_pretrained_models",
+            side_effect=_fake_run_pretrained_models,
+        ) as mock_run:
+            results = pretrained.run_models_from_directory(
+                str(sigma_phase_dir),
+                mode=mode,
+                saveGraphs=saveGraphs,
+                models=models_selection,
+            )
+
+            # Verify the mock was called with correct parameters
+            mock_run.assert_called_once()
+            call_kwargs = mock_run.call_args.kwargs
+
+            # Verify models parameter was passed correctly
+            assert call_kwargs["models"] == models_selection, (
+                f"Expected models={models_selection}, got {call_kwargs['models']}"
+            )
+
+            # Verify other parameters
+            assert call_kwargs["mode"] == mode
+            assert call_kwargs["saveGraphs"] == saveGraphs
+
+            # Verify the models information is in all returned results
+            for result in results:
+                assert result["models_used"] == models_selection

{mpdd_alignn-1.2.0 → mpdd_alignn-1.3.0}/mpdd_alignn.egg-info/PKG-INFO

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: mpdd-alignn
-Version: 1.2.0
+Version: 1.3.0
 Summary: Our version of the NIST-JARVIS ALIGNN optimized in terms of model performance and to some extent reliability, for large-scale deployments over the MPDD infrastructure by Phases Research Lab.
-Author-email: "Adam M. Krajewski" <adam@phaseslab.org>
+Author-email: "Adam M. Krajewski" <adam@phaseslab.org>, Ricardo Amaral <rna5137@psu.edu>
 Project-URL: Homepage, https://github.com/PhasesResearchLab/mpdd-alignn
 Project-URL: Bug Tracker, https://github.com/PhasesResearchLab/mpdd-alignn/issues
 Project-URL: Research Page, https://phaseslab.com/mpdd
@@ -34,6 +34,7 @@ Requires-Dist: dgl>=1.1.3
 Requires-Dist: scikit-learn>=1.3.1
 Requires-Dist: matplotlib>=3.4.1
 Requires-Dist: tqdm>=4.60.0
+Requires-Dist: natsort>=8.3.0
 Requires-Dist: pandas>=1.2.3
 Requires-Dist: pydantic_settings
 Requires-Dist: ruamel.yaml

{mpdd_alignn-1.2.0 → mpdd_alignn-1.3.0}/mpdd_alignn.egg-info/SOURCES.txt

@@ -68,6 +68,7 @@ alignn/scripts/train_edos_pdos.py
 alignn/scripts/train_megnet.py
 alignn/tests/test_alignn_ff.py
 alignn/tests/test_force_reduction.py
+alignn/tests/test_mpddalignn.py
 alignn/tests/test_prop.py
 example.SigmaPhase/0-Cr8Fe18Ni4.POSCAR
 example.SigmaPhase/1-Cr16Fe8Ni6.POSCAR

{mpdd_alignn-1.2.0 → mpdd_alignn-1.3.0}/mpdd_alignn.egg-info/requires.txt

@@ -8,6 +8,7 @@ dgl>=1.1.3
 scikit-learn>=1.3.1
 matplotlib>=3.4.1
 tqdm>=4.60.0
+natsort>=8.3.0
 pandas>=1.2.3
 pydantic_settings
 ruamel.yaml

{mpdd_alignn-1.2.0 → mpdd_alignn-1.3.0}/pyproject.toml

@@ -4,9 +4,10 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mpdd-alignn"
-version = "1.2.0"
+version = "1.3.0"
 authors = [
-  { name="Adam M. Krajewski", email="adam@phaseslab.org" }
+  { name="Adam M. Krajewski", email="adam@phaseslab.org" },
+  { name="Ricardo Amaral", email="rna5137@psu.edu" },
 ]
 classifiers = [
     "Programming Language :: Python :: 3",
@@ -37,6 +38,7 @@ dependencies = [
     "scikit-learn>=1.3.1",
     "matplotlib>=3.4.1",
     "tqdm>=4.60.0",
+    "natsort>=8.3.0",
     "pandas>=1.2.3",
     "pydantic_settings",
     "ruamel.yaml",