Add example for constructing flexible water model

Author: PicoCentauri

examples/water-model/.gitignore

@@ -0,0 +1 @@
+*.pt

examples/water-model/README.rst

@@ -0,0 +1,5 @@
+Empircal water models
+=====================
+
+This example shows implemenetations of three and four point flexible water models and
+uses them to run molecular dynamics simulations of water.

examples/water-model/_model.py

@@ -0,0 +1,387 @@
+from typing import Dict, List, Optional
+
+import torch
+from metatensor.torch import Labels, TensorBlock, TensorMap, sum_over_samples
+from metatensor.torch.atomistic import ModelOutput, NeighborListOptions, System
+from torchpme import Calculator, CoulombPotential, P3MCalculator
+
+
+def lennard_jones_pair(
+    distances: torch.Tensor,
+    sigma: torch.Tensor,
+    epsilon: torch.Tensor,
+    cutoff: torch.Tensor,
+):
+    """Lennard-Jones potential for pair terms."""
+    c6 = (sigma**2 / distances**2) ** 3
+    c12 = c6**2
+    lj = 4 * epsilon * (c12 - c6)
+
+    cutoff = 1 / cutoff
+    offset = 4 * epsilon * (sigma**12 * cutoff**12 - sigma**6 * cutoff**6)
+
+    return lj - offset
+
+
+def harmonic_distance_pair(
+    distances: torch.Tensor,
+    coefficient: torch.Tensor,
+    equilibrium_distance: torch.Tensor,
+):
+    """Harmonic potential for bond terms."""
+    r2 = (distances - equilibrium_distance) ** 2
+
+    return 0.5 * coefficient * r2
+
+
+def harmonic_angular(
+    angles: torch.Tensor, coefficient: torch.Tensor, equilibrium_angle: torch.Tensor
+):
+    """Harmonic potential for angular terms."""
+    return 0.5 * coefficient * (angles - equilibrium_angle) ** 2
+
+
+def compute_angles(positions: torch.Tensor, neighbor_indices: torch.Tensor):
+    """Compute the angles formed by triplet of atoms based on their positions."""
+    atom_is = neighbor_indices[:, 0]
+    atom_js = neighbor_indices[:, 1]
+    atom_ks = neighbor_indices[:, 2]
+
+    pos_is = positions[atom_is]
+    pos_js = positions[atom_js]
+    pos_ks = positions[atom_ks]
+
+    R_ij = pos_js - pos_is
+    R_ik = pos_ks - pos_is
+
+    return angle(R_ij, R_ik)
+
+
+def angle(a: torch.Tensor, b: torch.Tensor, dim: int = -1):
+    """Compute the angle between two vectors a and b.
+
+    Code is taken from https://github.com/pytorch/pytorch/issues/59194
+    """
+    a_norm = a.norm(p=2, dim=dim, keepdim=True)
+    b_norm = a.norm(p=2, dim=dim, keepdim=True)
+    angles = 2 * torch.atan2(
+        (a * b_norm - a_norm * b).norm(p=2, dim=dim),
+        (a * b_norm + a_norm * b).norm(p=2, dim=dim),
+    )
+
+    return angles
+
+
+class WaterModel(torch.nn.Module):
+    def __init__(
+        self,
+        cutoff: float,
+        O_sigma: float,
+        O_epsilon: float,
+        O_charge: float,
+        OH_bond_coefficient: float,
+        OH_equilibrium_distance: float,
+        HOH_angle_coefficient: float,
+        HOH_equilibrium_angle: float,
+        pme_smearing: float,
+        pme_mesh_spacing: float,
+        pme_interpolation_nodes: int = 4,
+        pme_prefactor: float = 1,
+        four_point_model: bool = False,
+        dtype: Optional[float] = None,
+    ):
+        """
+        Flexible water model for three and four point models.
+
+        The model contains Lennard Jones interactions between the oxygens as well as
+        intra molecular bond and angle terms. The electrostatics are computed using the
+        P3M method. For a four point model the fourth side for the charge interaction is
+        computed implicitly based on the position of the other atoms.
+
+        :param cutoff: Cutoff for the Lennard-Jones interactions.
+        :param O_sigma: Sigma parameter for the oxygen Lennard-Jones interactions.
+        :param O_epsilon: Epsilon parameter for the oxygen Lennard-Jones interactions.
+        :param O_charge: Oxygen's atom charge; hydrogen is computed accordingly.
+        :param OH_bond_coefficient: Harmonic coefficient for the OH bond.
+        :param OH_equilibrium_distance: Equilibrium distance for the OH bond.
+        :param HOH_angle_coefficient: Harmonic coefficient for the HOH angle.
+        :param HOH_equilibrium_angle: Equilibrium angle for the HOH angle in degrees.
+        :param pme_smearing: Smearing parameter for the PME.
+        :param pme_mesh_spacing: Mesh spacing for the PME.
+        :param pme_interpolation_nodes: Number of interpolation nodes for the PME.
+        :param pme_prefactor: Prefactor for the PME.
+        :param four_point_model: If :py:obj:`True`, use the four-point model for the
+            electrostatics. The fourth point M is implicitly derived from the other
+            atoms or each water molecule and used during the force computation. See
+            10.1063/1.3167790 for details on its derivation.
+        :param dtype: Floating point precision for the model. If :py:obj:`None`, the
+        :param dtype: default
+            dtype is used.
+        """
+        super().__init__()
+
+        self.dtype = dtype if dtype is not None else torch.get_default_dtype()
+        self.four_point_model = four_point_model
+
+        self.register_buffer("cutoff", torch.tensor(cutoff, dtype=self.dtype))
+        self.register_buffer("O_sigma", torch.tensor(O_sigma, dtype=self.dtype))
+        self.register_buffer("O_epsilon", torch.tensor(O_epsilon, dtype=self.dtype))
+        self.register_buffer(
+            "OH_bond_coefficient", torch.tensor(OH_bond_coefficient, dtype=self.dtype)
+        )
+        self.register_buffer(
+            "OH_equilibrium_distance",
+            torch.tensor(OH_equilibrium_distance, dtype=self.dtype),
+        )
+        self.register_buffer(
+            "HOH_angle_coefficient",
+            torch.tensor(HOH_angle_coefficient, dtype=self.dtype),
+        )
+
+        # Convert degree angle to radians
+        self.register_buffer(
+            "HOH_equilibrium_angle",
+            torch.tensor(HOH_equilibrium_angle * torch.pi / 180, dtype=self.dtype),
+        )
+
+        # Register charges for water model
+        H_charge = -O_charge / 2
+        self.register_buffer(
+            "OHH_charges",
+            torch.tensor([O_charge, H_charge, H_charge], dtype=self.dtype),
+        )
+
+        self.pme_calculator = P3MCalculator(
+            potential=CoulombPotential(pme_smearing),
+            mesh_spacing=pme_mesh_spacing,
+            interpolation_nodes=pme_interpolation_nodes,
+            prefactor=pme_prefactor,
+        )
+        self.coulomb_calculator = Calculator(
+            CoulombPotential(), prefactor=pme_prefactor
+        )
+
+        self.nl = NeighborListOptions(cutoff=cutoff, full_list=False, strict=False)
+
+    def requested_neighbor_lists(self):
+        return [self.nl]
+
+    def _setup_systems(
+        self,
+        systems: list[System],
+        selected_atoms: Optional[Labels] = None,
+    ) -> tuple[System, TensorBlock]:
+        """Remove possible ghost atoms and add charges to the system."""
+        if len(systems) > 1:
+            raise ValueError(f"only one system supported, got {len(systems)}")
+
+        system_i = 0
+        system = systems[system_i]
+
+        # select only real atoms and discard ghosts
+        if selected_atoms is not None:
+            current_system_mask = selected_atoms.column("system") == system_i
+            current_atoms = selected_atoms.column("atom")
+            current_atoms = current_atoms[current_system_mask].to(torch.long)
+
+            types = system.types[current_atoms]
+            positions = system.positions[current_atoms]
+        else:
+            types = system.types
+            positions = system.positions
+
+        system_final = System(types, positions, system.cell, system.pbc)
+
+        return system_final, system.get_neighbor_list(self.nl)
+
+    def forward(
+        self,
+        systems: List[System],  # noqa
+        outputs: Dict[str, ModelOutput],  # noqa
+        selected_atoms: Optional[Labels] = None,
+    ) -> Dict[str, TensorMap]:  # noqa
+        """
+        Compute the energy of the water model.
+
+        Water molecules have to be in order OHH and whole across the system.
+        """
+        if list(outputs.keys()) != ["energy"]:
+            raise ValueError(
+                f"`outputs` keys ({', '.join(outputs.keys())}) contain unsupported "
+                "keys. Only 'energy' is supported."
+            )
+
+        system, neighbors = self._setup_systems(systems, selected_atoms)
+        species = system.types
+
+        if system.positions.dtype != self.dtype:
+            raise ValueError(
+                f"system.positions.dtype ({system.positions.dtype}) must be "
+                f"equal to dtype at initilization ({self.dtype})"
+            )
+
+        device = system.positions.device
+        n_atoms = len(system)
+
+        neighbor_indices = neighbors.samples.view(["first_atom", "second_atom"]).values
+
+        if device == "cpu":
+            # move data to 64-bit integers, for some reason indexing with 64-bit
+            # is a lot faster than using 32-bit integers on CPU. CUDA seems fine
+            # with either types
+            neighbor_indices = neighbor_indices.to(
+                torch.int64, memory_format=torch.contiguous_format
+            )
+
+        neighbor_distances = torch.linalg.norm(neighbors.values, dim=1).squeeze(1)
+
+        # Verify that system only contains water molecules in the correct order
+        if n_atoms % 3 != 0:
+            raise ValueError(
+                "system must be water containing a multiple of 3 atoms. "
+                f"Found {n_atoms} atoms!"
+            )
+
+        reference_types = torch.tensor(
+            [8, 1, 1], dtype=self.dtype, device=device
+        ).repeat(n_atoms // 3)
+        if not torch.all(system.types == reference_types):
+            raise ValueError(
+                "system must contain only water molecules in the order OHH"
+            )
+
+        energies = torch.zeros(n_atoms, dtype=self.dtype, device=device)
+        ###################
+        # O-O Lennard-Jones
+        ###################
+        i = neighbor_indices[:, 0]
+        j = neighbor_indices[:, 1]
+        lj_mask = (species[i] == 8) & (species[j] == 8)
+        lj_neighbor_indices = neighbor_indices[lj_mask]
+        lj_neighbor_distances = neighbor_distances[lj_mask]
+
+        lj = lennard_jones_pair(
+            distances=lj_neighbor_distances,
+            sigma=self.O_sigma,
+            epsilon=self.O_epsilon,
+            cutoff=self.cutoff,
+        )
+
+        energies.index_add_(0, lj_neighbor_indices[:, 0], lj)
+        energies.index_add_(0, lj_neighbor_indices[:, 1], lj)
+
+        ##########
+        # O-H bond
+        ##########
+        # select pairs within the same molecule
+        mol_mask = (i // 3) == (j // 3)
+
+        bond_mask = mol_mask & (species[i] == 8)
+        bond_neighbor_indices = neighbor_indices[bond_mask]
+        bond_neighbor_distances = neighbor_distances[bond_mask]
+
+        cell_dimensions = torch.linalg.norm(system.cell, dim=1)
+        min_dimension = float(torch.min(cell_dimensions))
+        half_cell = min_dimension / 2.0
+
+        if torch.any(bond_neighbor_distances > half_cell):
+            raise ValueError(
+                "Bond distances are larger than half of the cell size. "
+                "Most likely molecules are not whole."
+                "This is not supported by the model."
+            )
+
+        bond = harmonic_distance_pair(
+            distances=bond_neighbor_distances,
+            coefficient=self.OH_bond_coefficient,
+            equilibrium_distance=self.OH_equilibrium_distance,
+        )
+
+        energies.index_add_(0, bond_neighbor_indices[:, 0], bond)
+        energies.index_add_(0, bond_neighbor_indices[:, 1], bond)
+
+        #############
+        # H-O-H angle
+        #############
+        all_idx = torch.arange(n_atoms, device=device)
+        angle_indices = torch.vstack([all_idx[0::3], all_idx[1::3], all_idx[2::3]]).T
+        angle_values = compute_angles(system.positions, angle_indices)
+
+        angles = harmonic_angular(
+            angles=angle_values,
+            coefficient=self.HOH_angle_coefficient,
+            equilibrium_angle=self.HOH_equilibrium_angle,
+        )
+
+        energies.index_add_(0, angle_indices[:, 0], angles)
+        energies.index_add_(0, angle_indices[:, 1], angles)
+        energies.index_add_(0, angle_indices[:, 2], angles)
+
+        ################
+        # Electrostatics
+        ################
+
+        # fourth point is computed according to eq. 2 in 10.1063/1.3167790
+        if self.four_point_model:
+            positions_coul = torch.vstack(
+                [
+                    (
+                        (system.positions[1::3] + system.positions[2::3]) * 0.5
+                        + system.positions[0::3] * 3
+                    )
+                    / 4,
+                    system.positions[1::3],
+                    system.positions[2::3],
+                ]
+            )
+        else:
+            positions_coul = system.positions
+
+        charges = self.OHH_charges.tile((n_atoms // 3,)).unsqueeze(-1)
+
+        # all to all interactions
+        potential = self.pme_calculator(
+            positions=positions_coul,
+            cell=system.cell,
+            charges=charges,
+            neighbor_indices=neighbor_indices,
+            neighbor_distances=neighbor_distances,
+        )
+
+        potential_exclusion = self.coulomb_calculator(
+            positions=positions_coul,
+            cell=system.cell,
+            charges=charges,
+            neighbor_indices=bond_neighbor_indices,
+            neighbor_distances=bond_neighbor_distances,
+        )
+
+        potential -= potential_exclusion
+        energies += (potential * charges).flatten()
+
+        #####################
+        # Wrap into TensorMap
+        #####################
+        samples = torch.zeros((n_atoms, 2), device=device, dtype=torch.int32)
+        samples[:, 0] = 0
+        samples[:, 1] = torch.arange(n_atoms, device=device, dtype=torch.int32)
+
+        properties = torch.tensor([[0]], device=device, dtype=torch.int32)
+
+        block = TensorBlock(
+            values=energies.unsqueeze(-1),
+            samples=Labels(["system", "atom"], samples),
+            components=[],
+            properties=Labels("energy", properties),
+        )
+
+        keys = Labels("_", torch.zeros(1, 1, dtype=torch.int32, device=device))
+
+        energy_tensor = TensorMap(keys=keys, blocks=[block])
+
+        if outputs["energy"].per_atom:
+            energy = energy_tensor
+        else:
+            energy = sum_over_samples(energy_tensor, sample_names="atom")
+
+        return {"energy": energy}