Unified torch-based short-range MD interface

[agoscinski]

MLP packages

allegro (e3nn based)

Python

def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
    # AtomicDataDict.Type is Dict[str, torch.Tensor] 

https://github.com/mir-group/allegro/blob/71ea80992f70a84fc65ecfb5a5ca45e8a82c6031/allegro/nn/_allegro.py#L418

LAMMPS interface

c10::Dict<std::string, torch::Tensor> input;
input.insert("pos", pos_tensor.to(device));
input.insert("edge_index", edges_tensor.to(device));// (2, n_edges) pair list, basically reformatted NL
input.insert("atom_types", ij2type_tensor.to(device));
std::vector<torch::IValue> input_vector(1, input);
auto output = model.forward(input_vector).toGenericDict();

https://github.com/mir-group/pair_allegro/blob/a7899a9c4c6be0620e11bef0bca8d06d9f7f32a8/pair_allegro.cpp#L426-L432

For backwards compatability (torch <=1.10), they reimplement some buggy torch::jit::freeze code from torch, that freezes the model if it is not already so.
They need to do so, because their network is based on e3nn objects that pytorch model can only handle frozen.

TorchANI

def forward(self, species_coordinates: Tuple[Tensor, Tensor], # tuple of species and positions
            cell: Optional[Tensor] = None,
            pbc: Optional[Tensor] = None) # they compute NL internally, so they need pbc
            -> SpeciesEnergies:

https://github.com/aiqm/torchani/blob/40cf334dbdf71903f30be8560784ca793b830221/torchani/models.py#L83-L85

LAMMPS
TODO haven't found anything, but haven't searched much. There are ANI interfaces with OpenKIM so it is definetly available with LAMMPS

OpenMM-ML
https://github.com/openmm/openmm-ml/blob/main/openmmml/models/anipotential.py
OpenMM-Torch
https://github.com/openmm/openmm-torch/blob/master/tutorials/openmm-torch-nnpops.ipynb

nequip (message-passing, e3nn based)

Message-passing
TODO don't know exactly, the use some edge_cell_shifts

def with_edge_vectors(data: Type, with_lengths: bool = True) -> Type:
    """Compute the edge displacement vectors for a graph.

    If ``data.pos.requires_grad`` and/or ``data.cell.requires_grad``, this
    method will return edge vectors correctly connected in the autograd graph.

    Returns:
        Tensor [n_edges, 3] edge displacement vectors
    """
    ...

https://github.com/mir-group/nequip/blob/c56f48fcc9b4018a84e1ed28f762fadd5bc763f1/nequip/data/AtomicDataDict.py#L47-L55

Python
the same as in allegro

c10::Dict<std::string, torch::Tensor> input;
input.insert("pos", pos_tensor.to(device));
input.insert("edge_index", edges_tensor.to(device)); // (2, n_edges) pair list, basically reformatted NL
input.insert("edge_cell_shift", edge_cell_shifts_tensor.to(device)); // (n_edges, 3) TODO think, needed for message passing
//  edge_cell_shift (Tensor [n_edges, 3], optional): which periodic image
//    of the target point each edge goes to, relative to the source point.
input.insert("cell", cell_tensor.to(device));
input.insert("atom_types", tag2type_tensor.to(device));
std::vector<torch::IValue> input_vector(1, input);
auto output = model.forward(input_vector).toGenericDict();

https://github.com/mir-group/pair_nequip/blob/b3680ef1288747fd8694693c1a747be4dd6a76b9/pair_nequip.cpp#L445-L451

OpenMM interface
open issue mir-group/nequip#288

Schnet (message-passing)

Message-passing
I don't know how they handle it

def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    ...

https://github.com/atomistic-machine-learning/schnetpack/blob/3f6eecee4b35ea5bb6e5f883ae70b789f2bd3e75/src/schnetpack/model/base.py#L172
https://github.com/atomistic-machine-learning/schnetpack/blob/3f6eecee4b35ea5bb6e5f883ae70b789f2bd3e75/src/schnetpack/representation/schnet.py#L138-L142

LAMMPS interface (the pair potential is a based on the one from nequip)

  c10::Dict<std::string, torch::Tensor> input;
  input.insert("_positions", pos_tensor.to(device));
  input.insert("_idx_i", edges_tensor[0].to(device)); // (n_edges) i part of pair list, basically reformatted NL
  input.insert("_idx_j", edges_tensor[1].to(device)); // (n_edges) j part pair list, basically reformatted NL
  input.insert("_idx_m", idx_m.to(device)); // TODO haven't seen any usage of this in code
  input.insert("_offsets", edge_cell_shifts_tensor.to(device)); // TODO haven't seen any usage of this in code
  input.insert("_cell", cell_tensor.to(device));
  input.insert("_n_atoms", n_atoms_tensor.to(device));
  input.insert("_atomic_numbers", tag2type_tensor.to(device));
  std::vector<torch::IValue> input_vector(1, input);

  auto output = model.forward(input_vector).toGenericDict();

https://github.com/atomistic-machine-learning/schnetpack/blob/3f6eecee4b35ea5bb6e5f883ae70b789f2bd3e75/interfaces/lammps/pair_schnetpack.cpp#L323-L332
https://schnetpack.readthedocs.io/en/latest/howtos/lammps.html

MACE (message-passing)

Message-passing
https://github.com/ACEsuit/lammps/blob/4ee087496324a8ce320459dc0bdf58b8e5e74ab0/src/ML-MACE/pair_mace.cpp#L362-L363

Python

def forward(
    self,
    data: Dict[str, torch.Tensor],
    training: bool = False,
    compute_force: bool = True,
    compute_virials: bool = False,
    compute_stress: bool = False,
    compute_displacement: bool = False,
) -> Dict[str, Optional[torch.Tensor]]:

https://github.com/ACEsuit/mace/blob/44b6da4a5edaa4c3ef867a11728555403d1d475d/mace/modules/models.py#L150-L158

LAMMPS interface

  c10::Dict<std::string, torch::Tensor> input;
  input.insert("batch", batch);
  input.insert("cell", cell);
  input.insert("edge_index", edge_index); // (2, n_edges) pair list, basically reformatted NL
  input.insert("energy", energy);
  input.insert("forces", forces);
  input.insert("node_attrs", node_attrs); // (n_nodes, n_feats) feature matrix
  input.insert("positions", positions);
  input.insert("ptr", ptr); // TODO contains numbers 0 and number of total atoms, haven't seen any usage of this in code
  input.insert("shifts", shifts); // TODO related to domain decomposition shifts, think it is needed for message-passing
  input.insert("unit_shifts", unit_shifts); // TODO related to domain decomposition shifts, think it is needed for message-passing
  input.insert("weight", weight); // TODO related to domain decomposition shifts, think it is needed for message-passing
  auto output = model.forward({input, mask, true, true, false}).toGenericDict(); // mask for ghost atoms

https://github.com/ACEsuit/lammps/blob/4ee087496324a8ce320459dc0bdf58b8e5e74ab0/src/ML-MACE/pair_mace.cpp#L204-L216
https://mace-docs.readthedocs.io/en/latest/guide/lammps.html

OpenMM Python interface
https://github.com/jharrymoore/openmm-ml/blob/main/openmmml/models/mace_potential.py
Neighborlist https://github.com/felixmusil/torch_nl

MD packages

GPUMD

no interface

GROMACS

no interface

LAMMPS

You get a neighbour list by list of atoms and a flat neighbour list with offsets to determine which section corresponds to which atom
CPU (dynamics) <--> GPU (potential)

• does not recompute neighbourlist in each step

OpenMM

It has several relevant modules for us

• core module https://github.com/openmm
• ML: python API for MLP https://github.com/openmm/openmm-ml
  • does not pass any NL
  • API https://github.com/openmm/openmm-ml/blob/main/openmmml/mlpotential.py
  • Implementation for ANI https://github.com/openmm/openmm-ml/blob/main/openmmml/models/anipotential.py
• Torch: torch API for MLP https://github.com/openmm/openmm-torch
  • does not pass any NL
  • API https://github.com/openmm/openmm-torch/tree/master/openmmapi/src
  • Implementation for ANI https://github.com/openmm/openmm-torch/blob/master/tutorials/openmm-torch-nnpops.ipynb
• NNPOps: Optimized MD operations https://github.com/openmm/NNPOps
  • their NL seems to be super fast 🌟 [FEATURE] OpenMM mir-group/nequip#288 (comment)

TorchMD

• core https://github.com/torchmd/torchmd/tree/master
• Torch https://github.com/torchmd/torchmd-net/tree/main
  I dunno

ACEMD

I dunno, seems closed source

Prototyping unified MD interface

Requirements

Core idea:

• You implement your torch model with this interface and you can use "all" MD software (LAMMPS and OpenMM as current goals)
• Depending on the MD interface we need have different presets for neighbourlist and units

Why not OpenKIM?

• they offer no interface with a MD software that supports GPU
• no GPU neighborlist
• by focusing only on torch-based models we possibly simplify our workload, OpenKIM supports everything

Neighbourlist

• support for neighbourlist
• support no neighbourlist, only positions (people want to compute their own neighbourlists)

Most of the time is spent computing the angular features, so it's important to have a neighbor list with the shorter cutoff just for them
openmm/NNPOps#2 (comment)

---> We need to offer CPU and GPU neighbourlist for potential

Atomic numbers

• MD software packages do not always provide them (OpenMM), so we need an option to store them

Unit conversion

• support for unit conversions: model input (e.g. OpenMM in nanometers), output <---> MD input, output

Message-passing

• support for additional parameters for message-passing like cell shifts

Pseudo Code

Mainly inspired from Allegro and ANI, beause the others are entangled with message-passing options that I don't fully understand

class OpenMD(metaclass=MetaClass):

    def __init__(self,
            model: torch.nn.Moule,
            model_units_type : str
            atom_types: torch.Tensor = None,
            ):
        pass

    def forward(self, **kwargs):
        if self._md_nl_type == 'nl':
            energies = self._forward_nl(**kwargs)
        elif self._md_nl_type == 'no_nl':
            energies = self._forward_no_nl(**kwargs)

        # needs some units module supported with torch
        energies = units.convert(energies
                self.model_units_type,
                self._md_units_type)

    @abstractmethod
    def _forward_no_nl(self, positions: torch.Tensor, cell: torch.Tensor):
        # ... some individual implementation

    @abstractmethod
    def _forward_nl(self, positions: torch.Tensor, edges: torch.Tensor, cell: torch.Tensor):
        # ... some individual implementation
		
    @abstractmethod
    def _forward_ase(self, frames: ase.Atoms):
        # ... some individual implementation
        
    @abstractmethod
    @property
    def cutoff(self):
        # ... requires some individual implementation
        pass
		
    def freeze_md_interface(self,
            md_interface : str,
            **input_kwargs):
        if self.md_interface == 'lammps':
            self._md_units_type = 'angstroem'
            self._forward_nl(**input_kwargs)
        elif self.md_interface == 'openmm':
            self._md_units_type = 'nanometer'
            self._forward_no_nl(**input_kwargs)

        self.model.eval()
        torch.jit.freeze(self)

### Implementation Allegro
from mlpops import neighbourlist
class AllegroOpenMD(OpenMD):
    def _forward_nl(self, positions: torch.Tensor, edges: torch.Tensor, cell: torch.Tensor):
        self.model.forward(positions, edges, cell)

    def _forward_no_nl(self, positions: torch.Tensor, cell: torch.Tensor):
        # does not work like this, but the idea is that we provide something
        edges = neighbourlist(positions, cell, cutoff)
        self._forward_nl(positions, edges, cell)
        
    def _forward_ase(self, frames: ase.Atoms):
		#...
		
    def cutoff(self):
        self.model.cutoff