Add an interface to PET (#68)

---------

Co-authored-by: frostedoyster <[email protected]>
Co-authored-by: Filippo Bigi <[email protected]>
Co-authored-by: Arslan Mazitov <[email protected]>

.github/workflows/pet-tests.yml

@@ -0,0 +1,36 @@
+name: PET tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    # Check all PR
+
+jobs:
+  tests:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        include:
+          - os: ubuntu-22.04
+            python-version: "3.12"
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+    - run: pip install tox
+
+    - name: run PET tests
+      run: tox -e pet-tests
+      env:
+        # Use the CPU only version of torch when building/running the code
+        PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu
+
+    - name: Upload codecoverage
+      uses: codecov/codecov-action@v4
+      with:
+        files: ./tests/coverage.xml

docs/src/dev-docs/utils/data/index.rst

@@ -10,3 +10,4 @@ API for handling data in ``metatensor-models``.
    dataset
    readers/index
    writers
+   systems_to_ase

docs/src/dev-docs/utils/data/systems_to_ase.rst

@@ -0,0 +1,11 @@
+Converting Systems to ASE
+#########################
+
+Some machine learning models might train on ``ase.Atoms`` objects.
+This module provides a function to convert a ``metatensor.torch.atomistic.System``
+object to an ``ase.Atoms`` object.
+
+.. automodule:: metatensor.models.utils.data.system_to_ase
+    :members:
+    :undoc-members:
+    :show-inheritance:

pyproject.toml

@@ -60,6 +60,9 @@ soap-bpnn = []
 alchemical-model = [
   "torch_alchemical @ git+https://github.com/abmazitov/torch_alchemical.git@fafb0bd",
 ]
+pet = [
+  "pet @ git+https://github.com/serfg/pet.git@5668bda",
+]
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/metatensor/models/cli/conf/architecture/experimental.pet.yaml

@@ -0,0 +1,59 @@
+ARCHITECTURAL_HYPERS:
+  CUTOFF_DELTA: 0.2
+  AVERAGE_POOLING: False
+  TRANSFORMERS_CENTRAL_SPECIFIC: False
+  HEADS_CENTRAL_SPECIFIC: False
+  ADD_TOKEN_FIRST: True
+  ADD_TOKEN_SECOND: True
+  N_GNN_LAYERS: 3
+  TRANSFORMER_D_MODEL: 128
+  TRANSFORMER_N_HEAD: 4
+  TRANSFORMER_DIM_FEEDFORWARD: 512
+  HEAD_N_NEURONS: 128
+  N_TRANS_LAYERS: 3
+  ACTIVATION: silu
+  USE_LENGTH: True
+  USE_ONLY_LENGTH: False
+  R_CUT: 5.0
+  R_EMBEDDING_ACTIVATION: False
+  COMPRESS_MODE: mlp
+  BLEND_NEIGHBOR_SPECIES: False
+  AVERAGE_BOND_ENERGIES: False
+  USE_BOND_ENERGIES: True
+  USE_ADDITIONAL_SCALAR_ATTRIBUTES: False
+  SCALAR_ATTRIBUTES_SIZE: null
+  TRANSFORMER_TYPE: PostLN # PostLN or PreLN
+  USE_LONG_RANGE: False
+  K_CUT: null # should be float; only used when USE_LONG_RANGE is True
+
+
+FITTING_SCHEME:  
+  INITIAL_LR: 1e-4
+  EPOCH_NUM_ATOMIC: 1000000000000000000
+  SCHEDULER_STEP_SIZE_ATOMIC: 500000000
+  EPOCHS_WARMUP_ATOMIC: 250000000
+  GLOBAL_AUG: True
+  SLIDING_FACTOR: 0.7
+  ATOMIC_BATCH_SIZE: 850
+  MAX_TIME: 234000
+  ENERGY_WEIGHT: 0.1 # only used when fitting MLIP
+  MULTI_GPU: False
+  RANDOM_SEED: 0
+  CUDA_DETERMINISTIC: False
+  MODEL_TO_START_WITH: null
+  SUPPORT_MISSING_VALUES: False
+  USE_WEIGHT_DECAY: False
+  WEIGHT_DECAY: 0.0
+  DO_GRADIENT_CLIPPING: False
+  GRADIENT_CLIPPING_MAX_NORM: null # must be overwritten if DO_GRADIENT_CLIPPING is True
+  USE_SHIFT_AGNOSTIC_LOSS: False # only used when fitting general target. Primary use case: EDOS
+  ENERGIES_LOSS: per_structure # per_structure or per_atom
+
+MLIP_SETTINGS: # only used when fitting MLIP
+  ENERGY_KEY: energy
+  FORCES_KEY: forces
+  USE_ENERGIES: True
+  USE_FORCES: True
+
+UTILITY_FLAGS: #for internal usage; do not change/overwrite
+  CALCULATION_TYPE: null

src/metatensor/models/experimental/pet/__init__.py

@@ -0,0 +1,2 @@
+from .model import Model, DEFAULT_HYPERS  # noqa: F401
+from .train import train  # noqa: F401

src/metatensor/models/experimental/pet/model.py

@@ -0,0 +1,101 @@
+from typing import Dict, List, Optional
+
+import torch
+from metatensor.torch import Labels, TensorBlock, TensorMap
+from metatensor.torch.atomistic import (
+    ModelCapabilities,
+    ModelOutput,
+    NeighborsListOptions,
+    System,
+)
+from omegaconf import OmegaConf
+from pet.hypers import Hypers
+from pet.pet import PET
+
+from ... import ARCHITECTURE_CONFIG_PATH
+from .utils import systems_to_batch_dict
+
+
+DEFAULT_HYPERS = OmegaConf.to_container(
+    OmegaConf.load(ARCHITECTURE_CONFIG_PATH / "experimental.pet.yaml")
+)
+
+DEFAULT_MODEL_HYPERS = DEFAULT_HYPERS["ARCHITECTURAL_HYPERS"]
+
+# We hardcode some of the hypers to make PET work as a MLIP.
+DEFAULT_MODEL_HYPERS.update(
+    {"D_OUTPUT": 1, "TARGET_TYPE": "structural", "TARGET_AGGREGATION": "sum"}
+)
+
+ARCHITECTURE_NAME = "experimental.pet"
+
+
+class Model(torch.nn.Module):
+    def __init__(
+        self, capabilities: ModelCapabilities, hypers: Dict = DEFAULT_MODEL_HYPERS
+    ) -> None:
+        super().__init__()
+        self.name = ARCHITECTURE_NAME
+        self.hypers = Hypers(hypers) if isinstance(hypers, dict) else hypers
+        self.cutoff = (
+            self.hypers["R_CUT"] if isinstance(self.hypers, dict) else self.hypers.R_CUT
+        )
+        self.all_species: List[int] = capabilities.species
+        self.capabilities = capabilities
+        self.pet = PET(self.hypers, 0.0, len(self.all_species))
+
+    def set_trained_model(self, trained_model: torch.nn.Module) -> None:
+        self.pet = trained_model
+
+    def requested_neighbors_lists(
+        self,
+    ) -> List[NeighborsListOptions]:
+        return [
+            NeighborsListOptions(
+                model_cutoff=self.cutoff,
+                full_list=True,
+            )
+        ]
+
+    def forward(
+        self,
+        systems: List[System],
+        outputs: Dict[str, ModelOutput],
+        selected_atoms: Optional[Labels] = None,
+    ) -> Dict[str, TensorMap]:
+        if selected_atoms is not None:
+            raise NotImplementedError("PET does not support selected atoms.")
+        options = self.requested_neighbors_lists()[0]
+        batch = systems_to_batch_dict(systems, options, self.all_species)
+        predictions = self.pet(batch)
+        total_energies: Dict[str, TensorMap] = {}
+        for output_name in outputs:
+            total_energies[output_name] = TensorMap(
+                keys=Labels(
+                    names=["_"],
+                    values=torch.tensor(
+                        [[0]],
+                        device=predictions.device,
+                    ),
+                ),
+                blocks=[
+                    TensorBlock(
+                        samples=Labels(
+                            names=["structure"],
+                            values=torch.arange(
+                                len(predictions),
+                                device=predictions.device,
+                            ).view(-1, 1),
+                        ),
+                        components=[],
+                        properties=Labels(
+                            names=["_"],
+                            values=torch.zeros(
+                                (1, 1), dtype=torch.int32, device=predictions.device
+                            ),
+                        ),
+                        values=predictions,
+                    )
+                ],
+            )
+        return total_energies

src/metatensor/models/experimental/pet/tests/__init__.py

@@ -0,0 +1,6 @@
+from pathlib import Path
+
+DATASET_PATH = str(
+    Path(__file__).parent.resolve()
+    / "../../../../../../tests/resources/qm9_reduced_100.xyz"
+)

src/metatensor/models/experimental/pet/tests/test_functionality.py

@@ -0,0 +1,47 @@
+import ase
+import rascaline.torch
+import torch
+from metatensor.torch.atomistic import (
+    MetatensorAtomisticModel,
+    ModelCapabilities,
+    ModelEvaluationOptions,
+    ModelOutput,
+)
+
+from metatensor.models.experimental.pet import DEFAULT_HYPERS, Model
+from metatensor.models.utils.neighbors_lists import get_system_with_neighbors_lists
+
+
+def test_prediction_subset():
+    """Tests that the model can predict on a subset
+    of the elements it was trained on."""
+
+    capabilities = ModelCapabilities(
+        length_unit="Angstrom",
+        species=[1, 6, 7, 8],
+        outputs={
+            "energy": ModelOutput(
+                quantity="energy",
+                unit="eV",
+            )
+        },
+    )
+
+    model = Model(capabilities, DEFAULT_HYPERS["ARCHITECTURAL_HYPERS"]).to(
+        torch.float64
+    )
+    structure = ase.Atoms("O2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
+    system = rascaline.torch.systems_to_torch(structure)
+    system = get_system_with_neighbors_lists(system, model.requested_neighbors_lists())
+
+    evaluation_options = ModelEvaluationOptions(
+        length_unit=capabilities.length_unit,
+        outputs=capabilities.outputs,
+    )
+
+    model = MetatensorAtomisticModel(model.eval(), model.capabilities)
+    model(
+        [system],
+        evaluation_options,
+        check_consistency=True,
+    )

src/metatensor/models/experimental/pet/tests/test_torchscript.py

@@ -0,0 +1,41 @@
+import torch
+from metatensor.torch.atomistic import ModelCapabilities, ModelOutput
+
+from metatensor.models.experimental.pet import DEFAULT_HYPERS, Model
+
+
+def test_torchscript():
+    """Tests that the model can be jitted."""
+
+    capabilities = ModelCapabilities(
+        length_unit="Angstrom",
+        species=[1, 6, 7, 8],
+        outputs={
+            "energy": ModelOutput(
+                quantity="energy",
+                unit="eV",
+            )
+        },
+    )
+    pet = Model(capabilities, DEFAULT_HYPERS["ARCHITECTURAL_HYPERS"])
+    torch.jit.script(pet)
+
+
+def test_torchscript_save():
+    """Tests that the model can be jitted and saved."""
+
+    capabilities = ModelCapabilities(
+        length_unit="Angstrom",
+        species=[1, 6, 7, 8],
+        outputs={
+            "energy": ModelOutput(
+                quantity="energy",
+                unit="eV",
+            )
+        },
+    )
+    pet = Model(capabilities, DEFAULT_HYPERS["ARCHITECTURAL_HYPERS"])
+    torch.jit.save(
+        torch.jit.script(pet),
+        "pet.pt",
+    )