PLUMED

Official website	How is metatomic supported?
https://www.plumed.org/	In the official (development) version

The metatomic interface for PLUMED allows using an exported AtomisticModel to compute arbitrary functions of the atomic coordinates, using them as collective variables to perform advanced sampling such as metadynamics. Below we also provide a minimal example of the implementation of a minimalist, model that is compatible with this interface. See also this recipe for more realistic, complex demonstrations.

How to install the code

See the official installation instructions in the documentation of PLUMED.

Supported model outputs

The model must provide a features output, and it is important that this output has a fixed size, and that the size can be determined by executing the model with an empty system (as this is how PLUMED determines internally the size of a CV). A minimal example of a model that computes the distance between two atoms is given below. Note how the capabilities and outputs of the model are defined to create an AtomisticModel from a bare torch.nn.Module, before exporting it as a torchscript file.

from typing import Dict, List, Optional

import metatensor.torch as mts
import torch

import metatomic.torch as mta


class Distance(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def requested_neighbor_lists(self) -> List[mta.NeighborListOptions]:
        # request a neighbor list to be computed and stored in the
        # system passed to `forward`. In this case we return an empty
        # list because we don't use neighbor lists
        return []

    def forward(
        self,
        systems: List[mta.System],
        outputs: Dict[str, mta.ModelOutput],
        selected_atoms: Optional[mts.Labels],
    ) -> Dict[str, mts.TensorMap]:
        if "features" not in outputs:
            return {}

        if outputs["features"].per_atom:
            raise ValueError("per-atoms features are not supported in this model")

        # PLUMED will first call the model with 0 atoms to get the size of the
        # output, so we need to handle this case first
        if len(systems[0]) == 0:
            # prepares an empty TensorMap with the correct metadata
            keys = mts.Labels("_", torch.tensor([[0]]))
            block = mts.TensorBlock(
                # this shape will be used by PLUMED to know that the CV only has 1 entry
                torch.zeros((0, 1), dtype=torch.float64),
                samples=mts.Labels("structure", torch.zeros((0, 1), dtype=torch.int32)),
                components=[],
                properties=mts.Labels(
                    "distance", torch.zeros((1, 1), dtype=torch.int32)
                ),
            )

            return {"features": mts.TensorMap(keys, [block])}

        if selected_atoms is None:
            raise ValueError("this model requires selected_atoms to be set")

        if len(selected_atoms) != 2:
            raise ValueError(
                "the model should be given two atoms to compute a distance"
            )

        atom_ids = selected_atoms.column("atom").squeeze()
        values = []
        for system in systems:
            if len(system.positions) < selected_atoms.values.max():
                raise ValueError(
                    "System size is too small for the selected atoms indices"
                )
            distance = torch.sqrt(
                (
                    (system.positions[atom_ids[0]] - system.positions[atom_ids[1]]) ** 2
                ).sum()
            )
            values.append(distance)

        # creates a tensor map to contain the values and return
        keys = mts.Labels("_", torch.tensor([[0]]))
        block = mts.TensorBlock(
            torch.stack(values, dim=0).reshape(-1, 1),
            samples=mts.Labels(
                "system",
                torch.arange(len(systems), dtype=torch.int32).reshape((-1, 1)),
            ),
            components=[],
            properties=mts.Labels("distance", torch.zeros((1, 1), dtype=torch.int32)),
        )

        return {"features": mts.TensorMap(keys, [block])}


# instantiates the model, describes its metadata, and export
distance = Distance()

# metatdata about the model itself
metadata = mta.ModelMetadata(
    name="Distance",
    description="Computes the distance between two selected atoms",
)

# metatdata about what the model can do
capabilities = mta.ModelCapabilities(
    length_unit="Angstrom",
    outputs={"features": mta.ModelOutput(per_atom=False)},
    atomic_types=[0],
    interaction_range=torch.inf,
    supported_devices=["cpu", "cuda"],
    dtype="float64",
)

model = mta.AtomisticModel(
    module=distance.eval(),
    metadata=metadata,
    capabilities=capabilities,
)

model.save("mta-distance.pt")

How to use the model in PLUMED

See the official syntax reference in the PLUMED documentation. An example of a PLUMED input to load the model above could read

dist: METATOMIC ...
    MODEL=mta-distance.pt
    SPECIES1=1-416  # no need for species
    SPECIES_TO_TYPES=0  # map everything to zero
    SELECTED_ATOMS=401,402  # indices of atoms (1-based)
...