Feat: add ZBL weighted DP model

deepmd/dpmodel/model/linear_atomic_model.py

@@ -0,0 +1,288 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from abc import (
+    abstractmethod,
+)
+from typing import (
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import numpy as np
+
+from deepmd.dpmodel import (
+    FittingOutputDef,
+    OutputVariableDef,
+)
+from deepmd.dpmodel.utils.nlist import (
+    build_multiple_neighbor_list,
+    get_multiple_nlist_key,
+    nlist_distinguish_types,
+)
+
+from .base_atomic_model import (
+    BaseAtomicModel,
+)
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+from .pairtab_atomic_model import (
+    PairTabModel,
+)
+
+
+class LinearAtomicModel(BaseAtomicModel):
+    """Linear model make linear combinations of several existing models.
+
+    Parameters
+    ----------
+    models : list[DPAtomicModel or PairTabModel]
+        A list of models to be combined. PairTabModel must be used together with a DPAtomicModel.
+    """
+
+    def __init__(
+        self,
+        models: List[BaseAtomicModel],
+        **kwargs,
+    ):
+        super().__init__()
+        self.models = models
+        self.distinguish_type_list = [
+            model.distinguish_types() for model in self.models
+        ]
+
+    def distinguish_types(self) -> bool:
+        """If distinguish different types by sorting."""
+        return False
+
+    def get_rcut(self) -> float:
+        """Get the cut-off radius."""
+        return max(self.get_rcuts())
+
+    def get_rcuts(self) -> List[float]:
+        """Get the cut-off radius for each individual models."""
+        return [model.get_rcut() for model in self.models]
+
+    def get_sel(self) -> List[int]:
+        return [model.get_nsel() for model in self.models]
+
+    def get_original_sels(self) -> List[Union[int, List[int]]]:
+        """Get the sels for each individual models."""
+        return [model.get_sel() for model in self.models]
+
+    def _sort_rcuts_sels(self) -> Tuple[List[int], List[float]]:
+        # sort the pair of rcut and sels in ascending order, first based on sel, then on rcut.
+        zipped = sorted(
+            zip(self.get_rcuts(), self.get_sel()), key=lambda x: (x[1], x[0])
+        )
+        return [p[0] for p in zipped], [p[1] for p in zipped]
+
+    def forward_atomic(
+        self,
+        extended_coord,
+        extended_atype,
+        nlist,
+        mapping: Optional[np.ndarray] = None,
+        fparam: Optional[np.ndarray] = None,
+        aparam: Optional[np.ndarray] = None,
+    ) -> Dict[str, np.ndarray]:
+        """Return atomic prediction.
+
+        Parameters
+        ----------
+        extended_coord
+            coodinates in extended region, (nframes, nall * 3)
+        extended_atype
+            atomic type in extended region, (nframes, nall)
+        nlist
+            neighbor list, (nframes, nloc, nsel).
+        mapping
+            mapps the extended indices to local indices.
+        fparam
+            frame parameter. (nframes, ndf)
+        aparam
+            atomic parameter. (nframes, nloc, nda)
+
+        Returns
+        -------
+        result_dict
+            the result dict, defined by the fitting net output def.
+        """
+        nframes, nloc, nnei = nlist.shape
+        self.extended_coord = extended_coord.reshape(nframes, -1, 3)
+        sorted_rcuts, sorted_sels = self._sort_rcuts_sels()
+        nlists = build_multiple_neighbor_list(
+            self.extended_coord,
+            nlist,
+            sorted_rcuts,
+            sorted_sels,
+        )
+        raw_nlists = [
+            nlists[get_multiple_nlist_key(rcut, sel)]
+            for rcut, sel in zip(self.get_rcuts(), self.get_sel())
+        ]
+        self.nlists_ = [
+            nl if not dt else nlist_distinguish_types(nl, extended_atype, sel)
+            for dt, nl, sel in zip(
+                self.distinguish_type_list, raw_nlists, self.get_original_sels()
+            )
+        ]
+        ener_list = [
+            model.forward_atomic(
+                self.extended_coord,
+                extended_atype,
+                nl,
+                mapping,
+            )["energy"]
+            for model, nl in zip(self.models, self.nlists_)
+        ]
+        self.weights = self._compute_weight()
+        self.atomic_bias = None
+        if self.atomic_bias is not None:
+            raise NotImplementedError("Need to add bias in a future PR.")
+        else:
+            fit_ret = {
+                "energy": np.sum(np.stack(ener_list) * np.stack(self.weights), axis=0),
+            }  # (nframes, nloc, 1)
+        return fit_ret
+
+    def fitting_output_def(self) -> FittingOutputDef:
+        return FittingOutputDef(
+            [
+                OutputVariableDef(
+                    name="energy", shape=[1], reduciable=True, differentiable=True
+                )
+            ]
+        )
+
+    def serialize(self) -> dict:
+        return {
+            "models": [model.serialize() for model in self.models],
+        }
+
+    @classmethod
+    def deserialize(cls, data) -> "LinearAtomicModel":
+        models = [DPAtomicModel.deserialize(model) for model in data["models"]]
+        return cls(models)
+
+    @abstractmethod
+    def _compute_weight(self) -> np.ndarray:
+        """This should be a list of user defined weights that matches the number of models to be combined."""
+        raise NotImplementedError
+
+
+class DPZBLLinearAtomicModel(LinearAtomicModel):
+    """Model linearly combine a list of AtomicModels.
+
+    Parameters
+    ----------
+    models
+            This linear model should take a DPAtomicModel and a PairTable model.
+    """
+
+    def __init__(
+        self,
+        dp_model: DPAtomicModel,
+        zbl_model: PairTabModel,
+        sw_rmin: float,
+        sw_rmax: float,
+        smin_alpha: Optional[float] = 0.1,
+        **kwargs,
+    ):
+        models = [dp_model, zbl_model]
+        super().__init__(models, **kwargs)
+        self.dp_model = dp_model
+        self.zbl_model = zbl_model
+
+        self.sw_rmin = sw_rmin
+        self.sw_rmax = sw_rmax
+        self.smin_alpha = smin_alpha
+
+    def serialize(self) -> dict:
+        return {
+            "dp_model": self.dp_model.serialize(),
+            "zbl_model": self.zbl_model.serialize(),
+            "sw_rmin": self.sw_rmin,
+            "sw_rmax": self.sw_rmax,
+            "smin_alpha": self.smin_alpha,
+        }
+
+    @classmethod
+    def deserialize(cls, data) -> "DPZBLLinearAtomicModel":
+        sw_rmin = data["sw_rmin"]
+        sw_rmax = data["sw_rmax"]
+        smin_alpha = data["smin_alpha"]
+
+        dp_model = DPAtomicModel.deserialize(data["dp_model"])
+        zbl_model = PairTabModel.deserialize(data["zbl_model"])
+
+        return cls(
+            dp_model=dp_model,
+            zbl_model=zbl_model,
+            sw_rmin=sw_rmin,
+            sw_rmax=sw_rmax,
+            smin_alpha=smin_alpha,
+        )
+
+    def _compute_weight(self) -> List[np.ndarray]:
+        """ZBL weight.
+
+        Returns
+        -------
+        List[np.ndarray]
+            the atomic ZBL weight for interpolation. (nframes, nloc, 1)
+        """
+        assert (
+            self.sw_rmax > self.sw_rmin
+        ), "The upper boundary `sw_rmax` must be greater than the lower boundary `sw_rmin`."
+
+        dp_nlist = self.nlists_[0]
+        zbl_nlist = self.nlists_[1]
+
+        zbl_nnei = zbl_nlist.shape[-1]
+        dp_nnei = dp_nlist.shape[-1]
+
+        # use the larger rr based on nlist
+        nlist_larger = zbl_nlist if zbl_nnei >= dp_nnei else dp_nlist
+        nloc = nlist_larger.shape[1]
+        masked_nlist = np.clip(nlist_larger, 0, None)
+        pairwise_rr = np.sqrt(
+            np.sum(
+                np.power(
+                    (
+                        np.expand_dims(self.extended_coord, 2)
+                        - np.expand_dims(self.extended_coord, 1)
+                    ),
+                    2,
+                ),
+                axis=-1,
+            )
+        )
+
+        rr = np.take_along_axis(pairwise_rr[:, :nloc, :], masked_nlist, 2)
+
+        numerator = np.sum(
+            rr * np.exp(-rr / self.smin_alpha), axis=-1
+        )  # masked nnei will be zero, no need to handle
+        denominator = np.sum(
+            np.where(
+                nlist_larger != -1,
+                np.exp(-rr / self.smin_alpha),
+                np.zeros_like(nlist_larger),
+            ),
+            axis=-1,
+        )  # handle masked nnei.
+        sigma = numerator / denominator
+        u = (sigma - self.sw_rmin) / (self.sw_rmax - self.sw_rmin)
+        coef = np.zeros_like(u)
+        left_mask = sigma < self.sw_rmin
+        mid_mask = (self.sw_rmin <= sigma) & (sigma < self.sw_rmax)
+        right_mask = sigma >= self.sw_rmax
+        coef[left_mask] = 1
+        smooth = -6 * u**5 + 15 * u**4 - 10 * u**3 + 1
+        coef[mid_mask] = smooth[mid_mask]
+        coef[right_mask] = 0
+        self.zbl_weight = coef
+        return [1 - np.expand_dims(coef, -1), np.expand_dims(coef, -1)]

deepmd/dpmodel/model/pair_tab_model.py → deepmd/dpmodel/model/pairtab_atomic_model.py

@@ -78,7 +78,10 @@ def fitting_output_def(self) -> FittingOutputDef:
     def get_rcut(self) -> float:
         return self.rcut
 
-    def get_sel(self) -> int:
+    def get_sel(self) -> List[int]:
+        return [self.sel]
+
+    def get_nsel(self) -> int:
         return self.sel
 
     def distinguish_types(self) -> bool:
@@ -196,7 +199,7 @@ def _pair_tabulated_inter(
             i_type, j_type, idx, self.tab_data, self.nspline
         )
         table_coef = table_coef.reshape(self.nframes, self.nloc, self.nnei, 4)
-        ener = self._calcualte_ener(table_coef, uu)
+        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
@@ -275,7 +278,7 @@ def _extract_spline_coefficient(
         return final_coef
 
     @staticmethod
-    def _calcualte_ener(coef: np.ndarray, uu: np.ndarray) -> np.ndarray:
+    def _calculate_ener(coef: np.ndarray, uu: np.ndarray) -> np.ndarray:
         """Calculate energy using spline coeeficients.
 
         Parameters

deepmd/dpmodel/utils/nlist.py

@@ -169,7 +169,7 @@ def build_multiple_neighbor_list(
     nall = coord1.shape[1]
     coord0 = coord1[:, :nloc, :]
     nlist_mask = nlist == -1
-    tnlist_0 = nlist
+    tnlist_0 = nlist.copy()
     tnlist_0[nlist_mask] = 0
     index = np.tile(tnlist_0.reshape(nb, nloc * nsel, 1), [1, 1, 3])
     coord2 = np.take_along_axis(coord1, index, axis=1).reshape(nb, nloc, nsel, 3)

deepmd/pt/model/model/__init__.py

@@ -4,18 +4,62 @@
 from deepmd.pt.model.descriptor.descriptor import (
     Descriptor,
 )
+from deepmd.pt.model.model.dp_atomic_model import (
+    DPAtomicModel,
+)
+from deepmd.pt.model.model.pairtab_atomic_model import (
+    PairTabModel,
+)
 from deepmd.pt.model.task import (
     Fitting,
 )
 
 from .ener import (
     EnergyModel,
+    ZBLModel,
 )
 from .model import (
     BaseModel,
 )
 
 
+def get_zbl_model(model_params, sampled=None):
+    model_params = copy.deepcopy(model_params)
+    ntypes = len(model_params["type_map"])
+    # descriptor
+    model_params["descriptor"]["ntypes"] = ntypes
+    descriptor = Descriptor(**model_params["descriptor"])
+    # fitting
+    fitting_net = model_params.get("fitting_net", None)
+    fitting_net["type"] = fitting_net.get("type", "ener")
+    fitting_net["ntypes"] = descriptor.get_ntypes()
+    fitting_net["distinguish_types"] = descriptor.distinguish_types()
+    fitting_net["embedding_width"] = descriptor.get_dim_out()
+    grad_force = "direct" not in fitting_net["type"]
+    if not grad_force:
+        fitting_net["out_dim"] = descriptor.get_dim_emb()
+        if "ener" in fitting_net["type"]:
+            fitting_net["return_energy"] = True
+    fitting = Fitting(**fitting_net)
+    dp_model = DPAtomicModel(
+        descriptor, fitting, type_map=model_params["type_map"], resuming=True
+    )
+    # pairtab
+    filepath = model_params["use_srtab"]
+    pt_model = PairTabModel(
+        filepath, model_params["descriptor"]["rcut"], model_params["descriptor"]["sel"]
+    )
+
+    rmin = model_params["sw_rmin"]
+    rmax = model_params["sw_rmax"]
+    return ZBLModel(
+        dp_model,
+        pt_model,
+        rmin,
+        rmax,
+    )
+
+
 def get_model(model_params, sampled=None):
     model_params = copy.deepcopy(model_params)
     ntypes = len(model_params["type_map"])