diff --git a/ml_peg/calcs/conftest.py b/ml_peg/calcs/conftest.py new file mode 100644 index 000000000..7b2855061 --- /dev/null +++ b/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") diff --git a/ml_peg/cli/cli.py b/ml_peg/cli/cli.py index 352b68b10..e727a37a8 100644 --- a/ml_peg/cli/cli.py +++ b/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]) diff --git a/ml_peg/models/__init__.py b/ml_peg/models/__init__.py index 9e29c36ff..d630ac690 100644 --- a/ml_peg/models/__init__.py +++ b/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 diff --git a/ml_peg/models/get_models.py b/ml_peg/models/get_models.py index 5700b902b..223fb57ff 100644 --- a/ml_peg/models/get_models.py +++ b/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) diff --git a/ml_peg/models/mock.py b/ml_peg/models/mock.py new file mode 100644 index 000000000..ddc94c37b --- /dev/null +++ b/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.") diff --git a/ml_peg/models/models.py b/ml_peg/models/models.py index e8980b13d..ecc592deb 100644 --- a/ml_peg/models/models.py +++ b/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()