pairtab_atomic_model.py

# SPDX-License-Identifier: LGPL-3.0-or-later
from collections.abc import (
    Callable,
)
from typing import (
    Any,
    NoReturn,
)

import array_api_compat
import numpy as np

from deepmd.dpmodel.array_api import (
    Array,
    xp_take_along_axis,
)
from deepmd.dpmodel.output_def import (
    FittingOutputDef,
    OutputVariableDef,
)
from deepmd.dpmodel.utils.safe_gradient import (
    safe_for_sqrt,
)
from deepmd.utils.pair_tab import (
    PairTab,
)
from deepmd.utils.path import (
    DPPath,
)
from deepmd.utils.version import (
    check_version_compatibility,
)

from .base_atomic_model import (
    BaseAtomicModel,
)


@BaseAtomicModel.register("pairtab")
class PairTabAtomicModel(BaseAtomicModel):
    r"""Pairwise tabulation energy model.

    This model can be used to tabulate the pairwise energy between atoms for either
    short-range or long-range interactions, such as D3, LJ, ZBL, etc. It should not
    be used alone, but rather as one submodel of a linear (sum) model, such as
    DP+D3.

    Do not put the model on the first model of a linear model, since the linear
    model fetches the type map from the first model.

    At this moment, the model does not smooth the energy at the cutoff radius, so
    one needs to make sure the energy has been smoothed to zero.

    The pairwise energy is computed by table lookup and interpolation:

    .. math::
        E^i = \frac{1}{2} \sum_{j \in \mathcal{N}(i)} E_{t_i, t_j}(r_{ij}),

    where :math:`E_{t_i, t_j}(r)` is the tabulated pairwise energy between atom types
    :math:`t_i` and :math:`t_j` at distance :math:`r`, obtained via cubic spline
    interpolation from the table data. The factor of :math:`\frac{1}{2}` avoids
    double-counting of pairwise interactions.

    Parameters
    ----------
    tab_file : str
        The path to the tabulation file.
    rcut : float
        The cutoff radius.
    sel : int or list[int]
        The maxmum number of atoms in the cut-off radius.
    type_map : list[str]
        Mapping atom type to the name (str) of the type.
        For example `type_map[1]` gives the name of the type 1.
    """

    def __init__(
        self,
        tab_file: str,
        rcut: float,
        sel: int | list[int],
        type_map: list[str],
        rcond: float | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(type_map, **kwargs)
        super().init_out_stat()
        self.tab_file = tab_file
        self.rcut = rcut

        self.tab = PairTab(self.tab_file, rcut=rcut)
        self.ntypes = len(type_map)
        self.rcond = rcond

        if self.tab_file is not None:
            tab_info, tab_data = self.tab.get()
            nspline, ntypes_tab = tab_info[-2:].astype(int)
            self.tab_info = tab_info
            self.tab_data = tab_data.reshape(ntypes_tab, ntypes_tab, nspline, 4)
            if self.ntypes != ntypes_tab:
                raise ValueError(
                    "The `type_map` provided does not match the number of columns in the table."
                )
        else:
            self.tab_info, self.tab_data = None, None

        if isinstance(sel, int):
            self.sel = sel
        elif isinstance(sel, list):
            self.sel = sum(sel)
        else:
            raise TypeError("sel must be int or list[int]")

    def fitting_output_def(self) -> FittingOutputDef:
        return FittingOutputDef(
            [
                OutputVariableDef(
                    name="energy",
                    shape=[1],
                    reducible=True,
                    r_differentiable=True,
                    c_differentiable=True,
                )
            ]
        )

    def get_rcut(self) -> float:
        return self.rcut

    def get_type_map(self) -> list[str]:
        return self.type_map

    def get_sel(self) -> list[int]:
        return [self.sel]

    def set_case_embd(self, case_idx: int) -> NoReturn:
        """
        Set the case embedding of this atomic model by the given case_idx,
        typically concatenated with the output of the descriptor and fed into the fitting net.
        """
        raise NotImplementedError(
            "Case identification not supported for PairTabAtomicModel!"
        )

    def get_nsel(self) -> int:
        return self.sel

    def mixed_types(self) -> bool:
        """If true, the model
        1. assumes total number of atoms aligned across frames;
        2. uses a neighbor list that does not distinguish different atomic types.

        If false, the model
        1. assumes total number of atoms of each atom type aligned across frames;
        2. uses a neighbor list that distinguishes different atomic types.

        """
        # to match DPA1 and DPA2.
        return True

    def has_message_passing(self) -> bool:
        """Returns whether the atomic model has message passing."""
        return False

    def need_sorted_nlist_for_lower(self) -> bool:
        """Returns whether the atomic model needs sorted nlist when using `forward_lower`."""
        return False

    def change_type_map(
        self, type_map: list[str], model_with_new_type_stat: Any = None
    ) -> None:
        """Change the type related params to new ones, according to `type_map` and the original one in the model.
        If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types.
        """
        assert type_map == self.type_map, (
            "PairTabAtomicModel does not support changing type map now. "
            "This feature is currently not implemented because it would require additional work to change the tab file. "
            "We may consider adding this support in the future if there is a clear demand for it."
        )

    def serialize(self) -> dict:
        dd = BaseAtomicModel.serialize(self)
        dd.update(
            {
                "@class": "Model",
                "type": "pairtab",
                "@version": 2,
                "tab": self.tab.serialize(),
                "rcut": self.rcut,
                "sel": self.sel,
                "type_map": self.type_map,
            }
        )
        return dd

    @classmethod
    def deserialize(cls, data: dict) -> "PairTabAtomicModel":
        data = data.copy()
        check_version_compatibility(data.pop("@version", 1), 2, 2)
        data.pop("@class")
        data.pop("type")
        tab = PairTab.deserialize(data.pop("tab"))
        data["tab_file"] = None
        tab_model = super().deserialize(data)

        tab_model.tab = tab
        # Extract nspline/ntypes from the numpy source before setting on the
        # model, because dpmodel_setattr may convert to torch tensor.
        nspline, ntypes = tab.tab_info[-2:].astype(int)
        tab_model.tab_info = tab.tab_info
        tab_model.tab_data = tab.tab_data.reshape(ntypes, ntypes, nspline, 4)
        return tab_model

    def compute_or_load_stat(
        self,
        sampled_func: Callable[[], list[dict]],
        stat_file_path: DPPath | None = None,
        compute_or_load_out_stat: bool = True,
        preset_observed_type: list[str] | None = None,
    ) -> None:
        """Compute or load the statistics parameters of the model.

        PairTabAtomicModel has no descriptor or fitting net input stats,
        so this only computes output stats (energy bias) when requested.

        Parameters
        ----------
        sampled_func
            The lazy sampled function to get data frames from different data systems.
        stat_file_path
            The path to the stat file.
        compute_or_load_out_stat : bool
            Whether to compute the output statistics.
        """
        if compute_or_load_out_stat:
            wrapped_sampler = self._make_wrapped_sampler(sampled_func)
            self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)

        if stat_file_path is not None and self.type_map is not None:
            stat_file_path /= " ".join(self.type_map)

        self._collect_and_set_observed_type(
            sampled_func if callable(sampled_func) else lambda: sampled_func,
            stat_file_path,
            preset_observed_type,
        )

    def forward_atomic(
        self,
        extended_coord: Array,
        extended_atype: Array,
        nlist: Array,
        mapping: Array | None = None,
        fparam: Array | None = None,
        aparam: Array | None = None,
        comm_dict: dict | None = None,
        charge_spin: Array | None = None,
    ) -> dict[str, Array]:
        del comm_dict  # pairtab is local; no MPI ghost exchange needed.
        xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlist)
        nframes, nloc, nnei = nlist.shape
        extended_coord = xp.reshape(extended_coord, (nframes, -1, 3))

        # this will mask all -1 in the nlist
        mask = nlist >= 0
        masked_nlist = nlist * mask

        atype = extended_atype[:, :nloc]  # (nframes, nloc)
        pairwise_rr = self._get_pairwise_dist(
            extended_coord, masked_nlist
        )  # (nframes, nloc, nnei)

        # (nframes, nloc, nnei), index type is int64.
        dev = array_api_compat.device(extended_atype)
        j_type = extended_atype[
            xp.arange(extended_atype.shape[0], dtype=xp.int64, device=dev)[
                :, None, None
            ],
            masked_nlist,
        ]

        raw_atomic_energy = self._pair_tabulated_inter(
            nlist, atype, j_type, pairwise_rr
        )
        atomic_energy = 0.5 * xp.sum(
            xp.where(nlist != -1, raw_atomic_energy, xp.zeros_like(raw_atomic_energy)),
            axis=-1,
        )
        atomic_energy = xp.reshape(atomic_energy, (nframes, nloc, 1))

        return {"energy": atomic_energy}

    def _pair_tabulated_inter(
        self,
        nlist: Array,
        i_type: Array,
        j_type: Array,
        rr: Array,
    ) -> Array:
        """Pairwise tabulated energy.

        Parameters
        ----------
        nlist : Array
            The unmasked neighbour list. (nframes, nloc)
        i_type : Array
            The integer representation of atom type for all local atoms for all frames. (nframes, nloc)
        j_type : Array
            The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei)
        rr : Array
            The salar distance vector between two atoms. (nframes, nloc, nnei)

        Returns
        -------
        np.ndarray
            The masked atomic energy for all local atoms for all frames. (nframes, nloc, nnei)

        Raises
        ------
        Exception
            If the distance is beyond the table.

        Notes
        -----
        This function is used to calculate the pairwise energy between two atoms.
        It uses a table containing cubic spline coefficients calculated in PairTab.
        """
        xp = array_api_compat.array_namespace(nlist, i_type, j_type, rr)
        nframes, nloc, nnei = nlist.shape
        rmin = self.tab_info[0]
        hh = self.tab_info[1]
        hi = 1.0 / hh

        # jax jit does not support convert to a Python int, so we need to convert to xp.int64.
        nspline = xp.astype(self.tab_info[2] + 0.1, xp.int64)

        uu = (rr - rmin) * hi  # this is broadcasted to (nframes,nloc,nnei)

        # if nnei of atom 0 has -1 in the nlist, uu would be 0.
        # this is to handle the nlist where the mask is set to 0, so that we don't raise exception for those atoms.
        uu = xp.where(nlist != -1, uu, nspline + 1)

        # unsupported by jax
        # if xp.any(uu < 0):
        #     raise Exception("coord go beyond table lower boundary")

        idx = xp.astype(uu, xp.int64)

        uu -= idx
        table_coef = self._extract_spline_coefficient(
            i_type, j_type, idx, self.tab_data[...], nspline
        )
        table_coef = xp.reshape(table_coef, (nframes, nloc, nnei, 4))
        ener = self._calculate_ener(table_coef, uu)
        # here we need to overwrite energy to zero at rcut and beyond.
        mask_beyond_rcut = rr >= self.rcut
        # also overwrite values beyond extrapolation to zero
        extrapolation_mask = rr >= self.tab.rmin + nspline * self.tab.hh
        ener = xp.where(
            xp.logical_or(mask_beyond_rcut, extrapolation_mask),
            xp.zeros_like(ener),
            ener,
        )

        return ener

    @staticmethod
    def _get_pairwise_dist(coords: Array, nlist: Array) -> Array:
        """Get pairwise distance `dr`.

        Parameters
        ----------
        coords : Array
            The coordinate of the atoms, shape of (nframes, nall, 3).
        nlist
            The masked nlist, shape of (nframes, nloc, nnei).

        Returns
        -------
        np.ndarray
            The pairwise distance between the atoms (nframes, nloc, nnei).
        """
        xp = array_api_compat.array_namespace(coords, nlist)
        dev = array_api_compat.device(nlist)
        # index type is int64
        batch_indices = xp.arange(nlist.shape[0], dtype=xp.int64, device=dev)[
            :, None, None
        ]
        neighbor_atoms = coords[batch_indices, nlist]
        loc_atoms = coords[:, : nlist.shape[1], :]
        pairwise_dr = loc_atoms[:, :, None, :] - neighbor_atoms
        pairwise_rr = safe_for_sqrt(xp.sum(pairwise_dr**2, axis=-1))

        return pairwise_rr

    @staticmethod
    def _extract_spline_coefficient(
        i_type: Array,
        j_type: Array,
        idx: Array,
        tab_data: Array,
        nspline: np.int64,
    ) -> Array:
        """Extract the spline coefficient from the table.

        Parameters
        ----------
        i_type : Array
            The integer representation of atom type for all local atoms for all frames. (nframes, nloc)
        j_type : Array
            The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei)
        idx : Array
            The index of the spline coefficient. (nframes, nloc, nnei)
        tab_data : Array
            The table storing all the spline coefficient. (ntype, ntype, nspline, 4)
        nspline : int
            The number of splines in the table.

        Returns
        -------
        np.ndarray
            The spline coefficient. (nframes, nloc, nnei, 4), shape may be squeezed.
        """
        xp = array_api_compat.array_namespace(i_type, j_type, idx, tab_data)
        # (nframes, nloc, nnei)
        expanded_i_type = xp.broadcast_to(
            i_type[:, :, xp.newaxis],
            (i_type.shape[0], i_type.shape[1], j_type.shape[-1]),
        )

        # (nframes, nloc, nnei, nspline, 4)
        expanded_tab_data = tab_data[expanded_i_type, j_type]

        # (nframes, nloc, nnei, 1, 4)
        expanded_idx = xp.broadcast_to(
            idx[..., xp.newaxis, xp.newaxis], (*idx.shape, 1, 4)
        )
        clipped_indices = xp.astype(xp.clip(expanded_idx, 0, nspline - 1), xp.int64)

        # (nframes, nloc, nnei, 4)
        final_coef = xp.squeeze(
            xp_take_along_axis(expanded_tab_data, clipped_indices, 3), axis=3
        )

        # when the spline idx is beyond the table, all spline coefficients are set to `0`, and the resulting ener corresponding to the idx is also `0`.
        final_coef = xp.where(
            xp.squeeze(expanded_idx, axis=3) > nspline,
            xp.zeros_like(final_coef),
            final_coef,
        )
        return final_coef

    @staticmethod
    def _calculate_ener(coef: Array, uu: Array) -> Array:
        """Calculate energy using spline coeeficients.

        Parameters
        ----------
        coef : Array
            The spline coefficients. (nframes, nloc, nnei, 4)
        uu : Array
            The atom displancemnt used in interpolation and extrapolation (nframes, nloc, nnei)

        Returns
        -------
        np.ndarray
            The atomic energy for all local atoms for all frames. (nframes, nloc, nnei)
        """
        a3, a2, a1, a0 = coef[..., 0], coef[..., 1], coef[..., 2], coef[..., 3]
        etmp = (a3 * uu + a2) * uu + a1  # this should be elementwise operations.
        ener = etmp * uu + a0  # this energy has the extrapolated value when rcut > rmax
        return ener

    def get_dim_fparam(self) -> int:
        """Get the number (dimension) of frame parameters of this atomic model."""
        return 0

    def get_dim_aparam(self) -> int:
        """Get the number (dimension) of atomic parameters of this atomic model."""
        return 0

    def get_sel_type(self) -> list[int]:
        """Get the selected atom types of this model.

        Only atoms with selected atom types have atomic contribution
        to the result of the model.
        If returning an empty list, all atom types are selected.
        """
        return []

    def is_aparam_nall(self) -> bool:
        """Check whether the shape of atomic parameters is (nframes, nall, ndim).

        If False, the shape is (nframes, nloc, ndim).
        """
        return False

    def enable_compression(
        self,
        min_nbor_dist: float,
        table_extrapolate: float = 5,
        table_stride_1: float = 0.01,
        table_stride_2: float = 0.1,
        check_frequency: int = -1,
    ) -> None:
        """Pairtab model does not support compression."""
        pass

    def compression_needs_min_nbor_dist(self) -> bool:
        """Return whether compression consumes the minimum neighbor distance.

        Returns
        -------
        bool
            Always ``False``. The tabulated pair potential carries its own
            domain, so compression is a no-op here.
        """
        return False