support numerical precision and env_mat (#3114)

- change of numerical precision is supported
- environment matrix is supported.

---------

Co-authored-by: Han Wang <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

deepmd_utils/model_format/__init__.py

@@ -1,4 +1,10 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from .common import (
+    PRECISION_DICT,
+)
+from .env_mat import (
+    EnvMat,
+)
 from .network import (
     EmbeddingNet,
     NativeLayer,
@@ -9,10 +15,12 @@
 )
 
 __all__ = [
+    "EnvMat",
     "EmbeddingNet",
     "NativeLayer",
     "NativeNet",
     "load_dp_model",
     "save_dp_model",
     "traverse_model_dict",
+    "PRECISION_DICT",
 ]

deepmd_utils/model_format/common.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from abc import (
+    ABC,
+)
+
+import numpy as np
+
+PRECISION_DICT = {
+    "float16": np.float16,
+    "float32": np.float32,
+    "float64": np.float64,
+    "half": np.float16,
+    "single": np.float32,
+    "double": np.float64,
+}
+DEFAULT_PRECISION = "float64"
+
+
+class NativeOP(ABC):
+    """The unit operation of a native model."""
+
+    def call(self, *args, **kwargs):
+        """Forward pass in NumPy implementation."""
+        raise NotImplementedError

deepmd_utils/model_format/env_mat.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Optional,
+    Union,
+)
+
+import numpy as np
+
+from .common import (
+    NativeOP,
+)
+
+
+def compute_smooth_weight(
+    distance: np.ndarray,
+    rmin: float,
+    rmax: float,
+):
+    """Compute smooth weight for descriptor elements."""
+    min_mask = distance <= rmin
+    max_mask = distance >= rmax
+    mid_mask = np.logical_not(np.logical_or(min_mask, max_mask))
+    uu = (distance - rmin) / (rmax - rmin)
+    vv = uu * uu * uu * (-6.0 * uu * uu + 15.0 * uu - 10.0) + 1.0
+    return vv * mid_mask + min_mask
+
+
+def _make_env_mat(
+    nlist,
+    coord,
+    rcut: float,
+    ruct_smth: float,
+):
+    """Make smooth environment matrix."""
+    nf, nloc, nnei = nlist.shape
+    # nf x nall x 3
+    coord = coord.reshape(nf, -1, 3)
+    mask = nlist >= 0
+    nlist = nlist * mask
+    # nf x (nloc x nnei) x 3
+    index = np.tile(nlist.reshape(nf, -1, 1), (1, 1, 3))
+    coord_r = np.take_along_axis(coord, index, 1)
+    # nf x nloc x nnei x 3
+    coord_r = coord_r.reshape(nf, nloc, nnei, 3)
+    # nf x nloc x 1 x 3
+    coord_l = coord[:, :nloc].reshape(nf, -1, 1, 3)
+    # nf x nloc x nnei x 3
+    diff = coord_r - coord_l
+    # nf x nloc x nnei
+    length = np.linalg.norm(diff, axis=-1, keepdims=True)
+    # for index 0 nloc atom
+    length = length + ~np.expand_dims(mask, -1)
+    t0 = 1 / length
+    t1 = diff / length**2
+    weight = compute_smooth_weight(length, ruct_smth, rcut)
+    env_mat_se_a = np.concatenate([t0, t1], axis=-1) * weight * np.expand_dims(mask, -1)
+    return env_mat_se_a, diff * np.expand_dims(mask, -1), weight
+
+
+class EnvMat(NativeOP):
+    def __init__(
+        self,
+        rcut,
+        rcut_smth,
+    ):
+        self.rcut = rcut
+        self.rcut_smth = rcut_smth
+
+    def call(
+        self,
+        nlist: np.ndarray,
+        coord_ext: np.ndarray,
+        atype_ext: np.ndarray,
+        davg: Optional[np.ndarray] = None,
+        dstd: Optional[np.ndarray] = None,
+    ) -> Union[np.ndarray, np.ndarray]:
+        """Compute the environment matrix.
+
+        Parameters
+        ----------
+        nlist
+            The neighbor list. shape: nf x nloc x nnei
+        coord_ext
+            The extended coordinates of atoms. shape: nf x (nallx3)
+        atype_ext
+            The extended aotm types. shape: nf x nall
+        davg
+            The data avg. shape: nt x nnei x 4
+        dstd
+            The inverse of data std. shape: nt x nnei x 4
+
+        Returns
+        -------
+        env_mat
+            The environment matrix. shape: nf x nloc x nnei x 4
+        switch
+            The value of switch function. shape: nf x nloc x nnei
+        """
+        em, sw = self._call(nlist, coord_ext)
+        nf, nloc, nnei = nlist.shape
+        atype = atype_ext[:, :nloc]
+        if davg is not None:
+            em -= davg[atype]
+        if dstd is not None:
+            em /= dstd[atype]
+        return em, sw
+
+    def _call(
+        self,
+        nlist,
+        coord_ext,
+    ):
+        em, diff, ww = _make_env_mat(nlist, coord_ext, self.rcut, self.rcut_smth)
+        return em, ww
+
+    def serialize(
+        self,
+    ) -> dict:
+        return {
+            "rcut": self.rcut,
+            "rcut_smth": self.rcut_smth,
+        }
+
+    @classmethod
+    def deserialize(
+        cls,
+        data: dict,
+    ) -> "EnvMat":
+        return cls(**data)

deepmd_utils/model_format/network.py

@@ -4,9 +4,6 @@
 See issue #2982 for more information.
 """
 import json
-from abc import (
-    ABC,
-)
 from typing import (
     List,
     Optional,
@@ -20,6 +17,12 @@
 except ImportError:
     __version__ = "unknown"
 
+from .common import (
+    DEFAULT_PRECISION,
+    PRECISION_DICT,
+    NativeOP,
+)
+
 
 def traverse_model_dict(model_obj, callback: callable, is_variable: bool = False):
     """Traverse a model dict and call callback on each variable.
@@ -124,14 +127,6 @@ def load_dp_model(filename: str) -> dict:
     return model_dict
 
 
-class NativeOP(ABC):
-    """The unit operation of a native model."""
-
-    def call(self, *args, **kwargs):
-        """Forward pass in NumPy implementation."""
-        raise NotImplementedError
-
-
 class NativeLayer(NativeOP):
     """Native representation of a layer.
 
@@ -156,12 +151,16 @@ def __init__(
         idt: Optional[np.ndarray] = None,
         activation_function: Optional[str] = None,
         resnet: bool = False,
+        precision: str = DEFAULT_PRECISION,
     ) -> None:
-        self.w = w
-        self.b = b
-        self.idt = idt
+        prec = PRECISION_DICT[precision.lower()]
+        self.precision = precision
+        self.w = w.astype(prec) if w is not None else None
+        self.b = b.astype(prec) if b is not None else None
+        self.idt = idt.astype(prec) if idt is not None else None
         self.activation_function = activation_function
         self.resnet = resnet
+        self.check_type_consistency()
 
     def serialize(self) -> dict:
         """Serialize the layer to a dict.
@@ -180,6 +179,7 @@ def serialize(self) -> dict:
         return {
             "activation_function": self.activation_function,
             "resnet": self.resnet,
+            "precision": self.precision,
             "@variables": data,
         }
 
@@ -192,14 +192,28 @@ def deserialize(cls, data: dict) -> "NativeLayer":
         data : dict
             The dict to deserialize from.
         """
+        precision = data.get("precision", DEFAULT_PRECISION)
         return cls(
             w=data["@variables"]["w"],
             b=data["@variables"].get("b", None),
             idt=data["@variables"].get("idt", None),
             activation_function=data["activation_function"],
             resnet=data.get("resnet", False),
+            precision=precision,
         )
 
+    def check_type_consistency(self):
+        precision = self.precision
+
+        def check_var(var):
+            if var is not None:
+                # assertion "float64" == "double" would fail
+                assert PRECISION_DICT[var.dtype.name] is PRECISION_DICT[precision]
+
+        check_var(self.w)
+        check_var(self.b)
+        check_var(self.idt)
+
     def __setitem__(self, key, value):
         if key in ("w", "matrix"):
             self.w = value
@@ -211,6 +225,8 @@ def __setitem__(self, key, value):
             self.activation_function = value
         elif key == "resnet":
             self.resnet = value
+        elif key == "precision":
+            self.precision = value
         else:
             raise KeyError(key)
 
@@ -225,6 +241,8 @@ def __getitem__(self, key):
             return self.activation_function
         elif key == "resnet":
             return self.resnet
+        elif key == "precision":
+            return self.precision
         else:
             raise KeyError(key)
 
@@ -338,6 +356,7 @@ def __init__(
         neuron: List[int] = [24, 48, 96],
         activation_function: str = "tanh",
         resnet_dt: bool = False,
+        precision: str = DEFAULT_PRECISION,
     ):
         layers = []
         i_in = in_dim
@@ -351,6 +370,7 @@ def __init__(
                     idt=rng.normal(size=(ii)) if resnet_dt else None,
                     activation_function=activation_function,
                     resnet=True,
+                    precision=precision,
                 ).serialize()
             )
             i_in = i_ot

source/tests/test_model_format_utils.py

@@ -10,6 +10,7 @@
 
 from deepmd_utils.model_format import (
     EmbeddingNet,
+    EnvMat,
     NativeLayer,
     NativeNet,
     load_dp_model,
@@ -19,18 +20,22 @@
 
 class TestNativeLayer(unittest.TestCase):
     def test_serialize_deserize(self):
-        for (ni, no), bias, ut, activation_function, resnet, ashp in itertools.product(
+        for (
+            ni,
+            no,
+        ), bias, ut, activation_function, resnet, ashp, prec in itertools.product(
             [(5, 5), (5, 10), (5, 9), (9, 5)],
             [True, False],
             [True, False],
             ["tanh", "none"],
             [True, False],
             [None, [4], [3, 2]],
+            ["float32", "float64", "single", "double"],
         ):
             ww = np.full((ni, no), 3.0)
             bb = np.full((no,), 4.0) if bias else None
             idt = np.full((no,), 5.0) if ut else None
-            nl0 = NativeLayer(ww, bb, idt, activation_function, resnet)
+            nl0 = NativeLayer(ww, bb, idt, activation_function, resnet, prec)
             nl1 = NativeLayer.deserialize(nl0.serialize())
             inp_shap = [ww.shape[0]]
             if ashp is not None:
@@ -134,3 +139,47 @@ def test_save_load_model(self):
     def tearDown(self) -> None:
         if os.path.exists(self.filename):
             os.remove(self.filename)
+
+
+class TestEnvMat(unittest.TestCase):
+    def setUp(self):
+        # nloc == 3, nall == 4
+        self.nloc = 3
+        self.nall = 4
+        self.nf, self.nt = 1, 2
+        self.coord_ext = np.array(
+            [
+                [0, 0, 0],
+                [0, 1, 0],
+                [0, 0, 1],
+                [0, -2, 0],
+            ],
+            dtype=np.float64,
+        ).reshape([1, self.nall * 3])
+        self.atype_ext = np.array([0, 0, 1, 0], dtype=int).reshape([1, self.nall])
+        # sel = [5, 2]
+        self.nlist = np.array(
+            [
+                [1, 3, -1, -1, -1, 2, -1],
+                [0, -1, -1, -1, -1, 2, -1],
+                [0, 1, -1, -1, -1, 0, -1],
+            ],
+            dtype=int,
+        ).reshape([1, self.nloc, 7])
+        self.rcut = 0.4
+        self.rcut_smth = 2.2
+
+    def test_self_consistency(
+        self,
+    ):
+        rng = np.random.default_rng()
+        nf, nloc, nnei = self.nlist.shape
+        davg = rng.normal(size=(self.nt, nnei, 4))
+        dstd = rng.normal(size=(self.nt, nnei, 4))
+        dstd = 0.1 + np.abs(dstd)
+        em0 = EnvMat(self.rcut, self.rcut_smth)
+        em1 = EnvMat.deserialize(em0.serialize())
+        mm0, ww0 = em0.call(self.nlist, self.coord_ext, self.atype_ext, davg, dstd)
+        mm1, ww1 = em1.call(self.nlist, self.coord_ext, self.atype_ext, davg, dstd)
+        np.testing.assert_allclose(mm0, mm1)
+        np.testing.assert_allclose(ww0, ww1)