Add mock calculators

ml_peg/calcs/conftest.py

@@ -0,0 +1,44 @@
+"""Configure pytest for calculations."""
+
+from __future__ import annotations
+
+from pytest import Config, Parser
+
+from ml_peg import models
+
+
+def pytest_addoption(parser: Parser) -> None:
+    """
+    Add custom CLI inputs to pytest.
+
+    Parameters
+    ----------
+    parser
+        Pytest parser object.
+    """
+    parser.addoption(
+        "--run-mock",
+        action="store_true",
+        default=False,
+        help="Include mock model in tests",
+    )
+    parser.addoption(
+        "--mock-only",
+        action="store_true",
+        default=False,
+        help="Only run mock model, ignoring other models",
+    )
+
+
+def pytest_configure(config: Config) -> None:
+    """
+    Configure pytest to custom CLI inputs.
+
+    Parameters
+    ----------
+    config
+        Pytest configuration object.
+    """
+    # Set current models from CLI input
+    models.run_mock = config.getoption("--run-mock")
+    models.mock_only = config.getoption("--mock-only")

ml_peg/cli/cli.py

@@ -176,6 +176,12 @@ def run_calcs(
     test: Annotated[
         str, Option(help="Test to run calculations for. Default is all tests.")
     ] = "*",
+    run_mock: Annotated[
+        bool, Option(help="Whether to run with mock calculator in addition to models.")
+    ] = True,
+    mock_only: Annotated[
+        bool, Option(help="Whether to only run mock calculator with no models.")
+    ] = False,
     run_slow: Annotated[
         bool, Option(help="Whether to run calculations labelled slow.")
     ] = True,
@@ -204,6 +210,10 @@ def run_calcs(
     test
         Test to run calculation for. Default is `*`, corresponding to all tests in the
         category.
+    run_mock
+        Whether to run mock calculations. Default is `True`.
+    mock_only
+        Whether to only run mock calculations, with no models. Default is `False`.
     run_slow
         Whether to run slow calculations. Default is `True`.
     run_very_slow
@@ -230,6 +240,12 @@ def run_calcs(
     if run_very_slow:
         options.extend(["--run-very-slow"])
 
+    if run_mock:
+        options.extend(["--run-mock"])
+
+    if mock_only:
+        options.extend(["--mock-only"])
+
     if models:
         options.extend(["--models", models])
 

ml_peg/models/__init__.py

@@ -7,3 +7,5 @@
 MODELS_ROOT = Path(__file__).parent
 current_models = None
 models_file = MODELS_ROOT / "models.yml"
+run_mock = False
+mock_only = False

ml_peg/models/get_models.py

@@ -111,6 +111,8 @@ def get_subset(
 def load_models(
     models: None | str | Iterable = None,
     filepath: Path | str | None = None,
+    run_mock: bool | None = None,
+    mock_only: bool | None = None,
 ) -> dict[str, Any]:
     """
     Load models for use in calculations.
@@ -123,15 +125,35 @@ def load_models(
         this will be treated as a comma-separated list.
     filepath
         Path to YAML file with models. Default is `models_file`.
+    run_mock
+            Whether to include mock model in the loaded models. Default is False.
+    mock_only
+            Whether to load only mock model, ignoring `models`. Default is False.
 
     Returns
     -------
     dict[str, Any]
-        Loaded models from models.yml.
+        Loaded models from models.yml and/or loaded mock model.
     """
-    from ml_peg.models.models import FairChemCalc, GenericASECalc, OrbCalc, PetMadCalc
-
-    loaded_models = {}
+    from ml_peg.models.models import (
+        FairChemCalc,
+        GenericASECalc,
+        MockCalc,
+        OrbCalc,
+        PetMadCalc,
+    )
+
+    if run_mock is None:
+        from ml_peg.models import run_mock
+    if mock_only is None:
+        from ml_peg.models import mock_only
+
+    if mock_only and not run_mock:
+        raise ValueError("Cannot set `mock_only` without `run_mock`")
+
+    loaded_models = {"mock": MockCalc()} if run_mock else {}
+    if mock_only:
+        return loaded_models
 
     filepath = filepath if filepath else models_file
     all_models = _load_models_yaml(filepath)

ml_peg/models/mock.py

@@ -0,0 +1,66 @@
+"""Mock calculator class."""
+
+from __future__ import annotations
+
+from ase import Atoms
+from ase.calculators.calculator import Calculator, all_changes
+import numpy as np
+
+
+class MockCalculator(Calculator):
+    """A mock calculator that returns zero values for all properties."""
+
+    implemented_properties = ["energy", "forces", "stress"]
+
+    def calculate(
+        self,
+        atoms: Atoms | None = None,
+        properties: list[str] | None = None,
+        system_changes: list[str] = all_changes,
+        **kwargs,
+    ) -> None:
+        """
+        Define calculation method for mock calculator.
+
+        Parameters
+        ----------
+        atoms
+            Atoms object to calculate properties for.
+        properties
+            List of properties to calculate.
+        system_changes
+            List of system changes to consider for calculation.
+        **kwargs
+            Any additional keyword arguments.
+        """
+        super().calculate(atoms, properties, system_changes)
+
+        if "energy" in properties:
+            self.results["energy"] = 0.0
+
+        if "forces" in properties:
+            self.results["forces"] = np.zeros((len(self.atoms), 3))
+
+        if "stress" in properties:
+            self.results["stress"] = np.zeros(6)
+
+
+class MockErrorCalculator(Calculator):
+    """A mock calculator that raises an error for all properties."""
+
+    implemented_properties = ["energy", "forces", "stress"]
+    results = {}
+    parameters = {}
+
+    def calculate(self, *args, **kwargs) -> None:
+        """
+        Define calculation method for mock calculator.
+
+        Parameters
+        ----------
+        *args
+            Any additional positional arguments.
+        **kwargs
+            Any additional keyword arguments.
+        """
+        raise ValueError("This is a mock error calculator. All calculations fail.")

ml_peg/models/models.py

@@ -231,7 +231,7 @@ def get_calculator(self, **kwargs) -> Calculator:
         Returns
         -------
         Calculator
-            Loaded ASE Orb Calculator.
+            Loaded ASE fairchem Calculator.
         """
         from fairchem.core import FAIRChemCalculator, pretrained_mlip
         # torch.serialization.add_safe_globals([slice])
@@ -257,3 +257,29 @@ def available(self) -> bool:
             return self.model_name in pretrained_mlip._MODEL_CKPTS.checkpoints
         except Exception:
             return False
+
+
+@dataclasses.dataclass(kw_only=True)
+class MockCalc(SumCalc):
+    """Dataclass for mock calculator."""
+
+    model_name: str = "mock"
+    trained_on_dispersion: bool = True
+
+    def get_calculator(self, **kwargs) -> Calculator:
+        """
+        Prepare and load the calculator.
+
+        Parameters
+        ----------
+        **kwargs
+            Any additional keyword arguments passed to `get_calculator`.
+
+        Returns
+        -------
+        Calculator
+            Loaded mock ASE Calculator.
+        """
+        from ml_peg.models.mock import MockCalculator
+
+        return MockCalculator()