From af517212d19b93941a76991d0ffbc62bc9f1a7da Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Mon, 26 Feb 2024 12:00:36 +0800
Subject: [PATCH 01/20] feat: add se_r descriptor

---
 deepmd/pt/model/descriptor/__init__.py        |   6 +
 deepmd/pt/model/descriptor/env_mat.py         |  50 +++
 deepmd/pt/model/descriptor/se_r.py            | 399 ++++++++++++++++++
 deepmd/pt/utils/env_mat_stat.py               | 136 ++++++
 source/tests/pt/model/test_descriptor_se_r.py | 174 ++++++++
 5 files changed, 765 insertions(+)
 create mode 100644 deepmd/pt/model/descriptor/se_r.py
 create mode 100644 source/tests/pt/model/test_descriptor_se_r.py

diff --git a/deepmd/pt/model/descriptor/__init__.py b/deepmd/pt/model/descriptor/__init__.py
index 1c2e943369..445537e2c5 100644
--- a/deepmd/pt/model/descriptor/__init__.py
+++ b/deepmd/pt/model/descriptor/__init__.py
@@ -13,6 +13,7 @@
 )
 from .env_mat import (
     prod_env_mat_se_a,
+    prod_env_mat_se_r,
 )
 from .gaussian_lcc import (
     DescrptGaussianLcc,
@@ -27,6 +28,9 @@
     DescrptBlockSeA,
     DescrptSeA,
 )
+from .se_r import(
+    DescrptSeR
+)
 
 __all__ = [
     "Descriptor",
@@ -35,9 +39,11 @@
     "DescrptBlockSeA",
     "DescrptBlockSeAtten",
     "DescrptSeA",
+    "DescrptSeR",
     "DescrptDPA1",
     "DescrptDPA2",
     "prod_env_mat_se_a",
+    "prod_env_mat_se_r",
     "DescrptGaussianLcc",
     "DescrptBlockHybrid",
     "DescrptBlockRepformers",
diff --git a/deepmd/pt/model/descriptor/env_mat.py b/deepmd/pt/model/descriptor/env_mat.py
index b3235de175..be6c2437d8 100644
--- a/deepmd/pt/model/descriptor/env_mat.py
+++ b/deepmd/pt/model/descriptor/env_mat.py
@@ -29,6 +29,27 @@ def _make_env_mat_se_a(nlist, coord, rcut: float, ruct_smth: float):
     env_mat_se_a = torch.cat([t0, t1], dim=-1) * weight
     return env_mat_se_a, diff * mask.unsqueeze(-1), weight
 
+def _make_env_mat_se_r(nlist, coord, rcut: float, ruct_smth: float):
+    """Make smooth environment matrix."""
+    bsz, natoms, nnei = nlist.shape
+    coord = coord.view(bsz, -1, 3)
+    nall = coord.shape[1]
+    mask = nlist >= 0
+    # nlist = nlist * mask  ## this impl will contribute nans in Hessian calculation.
+    nlist = torch.where(mask, nlist, nall - 1)
+    coord_l = coord[:, :natoms].view(bsz, -1, 1, 3)
+    index = nlist.view(bsz, -1).unsqueeze(-1).expand(-1, -1, 3)
+    coord_r = torch.gather(coord, 1, index)
+    coord_r = coord_r.view(bsz, natoms, nnei, 3)
+    diff = coord_r - coord_l
+    length = torch.linalg.norm(diff, dim=-1, keepdim=True)
+    # for index 0 nloc atom
+    length = length + ~mask.unsqueeze(-1)
+    t0 = 1 / length
+    weight = compute_smooth_weight(length, ruct_smth, rcut)
+    weight = weight * mask.unsqueeze(-1)
+    env_mat_se_r = t0 * weight
+    return env_mat_se_r, diff * mask.unsqueeze(-1), weight
 
 def prod_env_mat_se_a(
     extended_coord, nlist, atype, mean, stddev, rcut: float, rcut_smth: float
@@ -58,3 +79,32 @@ def prod_env_mat_se_a(
     t_std = stddev[atype]  # [n_atom, dim, 4]
     env_mat_se_a = (_env_mat_se_a - t_avg) / t_std
     return env_mat_se_a, diff, switch
+
+def prod_env_mat_se_r(
+    extended_coord, nlist, atype, mean, stddev, rcut: float, rcut_smth: float
+):
+    """Generate smooth environment matrix from atom coordinates and other context.
+
+    Args:
+    - extended_coord: Copied atom coordinates with shape [nframes, nall*3].
+    - atype: Atom types with shape [nframes, nloc].
+    - natoms: Batched atom statisics with shape [len(sec)+2].
+    - box: Batched simulation box with shape [nframes, 9].
+    - mean: Average value of descriptor per element type with shape [len(sec), nnei, 1].
+    - stddev: Standard deviation of descriptor per element type with shape [len(sec), nnei, 1].
+    - deriv_stddev:  StdDev of descriptor derivative per element type with shape [len(sec), nnei, 1, 3].
+    - rcut: Cut-off radius.
+    - rcut_smth: Smooth hyper-parameter for pair force & energy.
+
+    Returns
+    -------
+    - env_mat_se_r: Shape is [nframes, natoms[1]*nnei*1].
+    """
+    nframes = extended_coord.shape[0]
+    _env_mat_se_r, diff, switch = _make_env_mat_se_r(
+        nlist, extended_coord, rcut, rcut_smth
+    )  # shape [n_atom, dim, 1]
+    t_avg = mean[atype]  # [n_atom, dim, 1]
+    t_std = stddev[atype]  # [n_atom, dim, 1]
+    env_mat_se_r = (_env_mat_se_r - t_avg) / t_std
+    return env_mat_se_r, diff, switch
diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
new file mode 100644
index 0000000000..461667364b
--- /dev/null
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -0,0 +1,399 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Dict,
+    List,
+    Optional,
+    Tuple,
+)
+
+import numpy as np
+import torch
+
+from deepmd.pt.model.descriptor import (
+    Descriptor,
+    prod_env_mat_se_r,
+)
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.pt.utils.env import (
+    PRECISION_DICT,
+    RESERVED_PRECISON_DICT,
+)
+from deepmd.pt.utils.env_mat_stat import (
+    EnvMatStatSeR,
+)
+from deepmd.utils.env_mat_stat import (
+    StatItem,
+)
+from deepmd.utils.path import (
+    DPPath,
+)
+
+try:
+    from typing import (
+        Final,
+    )
+except ImportError:
+    from torch.jit import Final
+
+from deepmd.dpmodel.utils import EnvMat as DPEnvMat
+from deepmd.pt.model.network.mlp import (
+    EmbeddingNet,
+    NetworkCollection,
+)
+from deepmd.pt.model.network.network import (
+    TypeFilter,
+)
+from deepmd.pt.utils.exclude_mask import (
+    PairExcludeMask,
+)
+
+
+@Descriptor.register("se_e2_r")
+class DescrptSeR(Descriptor):
+    def __init__(
+        self,
+        rcut,
+        rcut_smth,
+        sel,
+        neuron=[25, 50, 100],
+        set_davg_zero: bool = False,
+        activation_function: str = "tanh",
+        precision: str = "float64",
+        resnet_dt: bool = False,
+        exclude_types: List[Tuple[int, int]] = [],
+        old_impl: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.rcut = rcut
+        self.rcut_smth = rcut_smth
+        self.neuron = neuron
+        self.filter_neuron = self.neuron
+        self.set_davg_zero = set_davg_zero
+        self.activation_function = activation_function
+        self.precision = precision
+        self.prec = PRECISION_DICT[self.precision]
+        self.resnet_dt = resnet_dt
+        self.old_impl = old_impl
+        self.exclude_types = exclude_types
+        self.ntypes = len(sel)
+        self.emask = PairExcludeMask(len(sel), exclude_types=exclude_types)
+
+        self.sel = sel
+        self.sec = torch.tensor(
+            np.append([0], np.cumsum(self.sel)), dtype=int, device=env.DEVICE
+        )
+        self.split_sel = self.sel
+        self.nnei = sum(sel)
+        self.ndescrpt = self.nnei * 1
+
+        wanted_shape = (self.ntypes, self.nnei, 1)
+        mean = torch.zeros(wanted_shape, dtype=self.prec, device=env.DEVICE)
+        stddev = torch.ones(wanted_shape, dtype=self.prec, device=env.DEVICE)
+        self.register_buffer("mean", mean)
+        self.register_buffer("stddev", stddev)
+        self.filter_layers_old = None
+        self.filter_layers = None
+
+        if self.old_impl:
+            filter_layers = []
+            # TODO: remove
+            start_index = 0
+            for type_i in range(self.ntypes):
+                one = TypeFilter(start_index, sel[type_i], self.filter_neuron)
+                filter_layers.append(one)
+                start_index += sel[type_i]
+            self.filter_layers_old = torch.nn.ModuleList(filter_layers)
+        else:
+            filter_layers = NetworkCollection(
+                ndim=1, ntypes=len(sel), network_type="embedding_network"
+            )
+            # TODO: ndim=2 if type_one_side=False
+            for ii in range(self.ntypes):
+                filter_layers[(ii,)] = EmbeddingNet(
+                    1,
+                    self.filter_neuron,
+                    activation_function=self.activation_function,
+                    precision=self.precision,
+                    resnet_dt=self.resnet_dt,
+                )
+            self.filter_layers = filter_layers
+        self.stats = None
+
+    def get_rcut(self) -> float:
+        """Returns the cut-off radius."""
+        return self.rcut
+
+    def get_nsel(self) -> int:
+        """Returns the number of selected atoms in the cut-off radius."""
+        return sum(self.sel)
+
+    def get_sel(self) -> List[int]:
+        """Returns the number of selected atoms for each type."""
+        return self.sel
+
+    def get_ntypes(self) -> int:
+        """Returns the number of element types."""
+        return self.ntypes
+
+    def get_dim_out(self) -> int:
+        """Returns the output dimension."""
+        return self.neuron[-1]
+
+    def get_dim_emb(self) -> int:
+        """Returns the output dimension."""
+        raise NotImplementedError
+
+    def get_dim_in(self) -> int:
+        """Returns the input dimension."""
+        return 0
+
+    def mixed_types(self) -> bool:
+        """If true, the discriptor
+        1. assumes total number of atoms aligned across frames;
+        2. requires a neighbor list that does not distinguish different atomic types.
+
+        If false, the discriptor
+        1. assumes total number of atoms of each atom type aligned across frames;
+        2. requires a neighbor list that distinguishes different atomic types.
+
+        """
+        return False
+
+    def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None):
+        """Update mean and stddev for descriptor elements."""
+        env_mat_stat = EnvMatStatSeR(self)
+        if path is not None:
+            path = path / env_mat_stat.get_hash()
+        env_mat_stat.load_or_compute_stats(merged, path)
+        self.stats = env_mat_stat.stats
+        mean, stddev = env_mat_stat()
+        if not self.set_davg_zero:
+            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
+        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
+        if not self.set_davg_zero:
+            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
+        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
+
+    def get_stats(self) -> Dict[str, StatItem]:
+        """Get the statistics of the descriptor."""
+        if self.stats is None:
+            raise RuntimeError(
+                "The statistics of the descriptor has not been computed."
+            )
+        return self.stats
+    
+    def __setitem__(self, key, value):
+        if key in ("avg", "data_avg", "davg"):
+            self.mean = value
+        elif key in ("std", "data_std", "dstd"):
+            self.stddev = value
+        else:
+            raise KeyError(key)
+
+    def __getitem__(self, key):
+        if key in ("avg", "data_avg", "davg"):
+            return self.mean
+        elif key in ("std", "data_std", "dstd"):
+            return self.stddev
+        else:
+            raise KeyError(key)
+
+    @classmethod
+    def get_data_process_key(cls, config):
+        """
+        Get the keys for the data preprocess.
+        Usually need the information of rcut and sel.
+        TODO Need to be deprecated when the dataloader has been cleaned up.
+        """
+        descrpt_type = config["type"]
+        assert descrpt_type in ["se_e2_r"]
+        return {"sel": config["sel"], "rcut": config["rcut"]}
+
+    @property
+    def data_stat_key(self):
+        """
+        Get the keys for the data statistic of the descriptor.
+        Return a list of statistic names needed, such as "sumr", "sumr2" or "sumn".
+        """
+        return ["sumr", "sumn", "sumr2"]
+
+    def forward(
+        self,
+        coord_ext: torch.Tensor,
+        atype_ext: torch.Tensor,
+        nlist: torch.Tensor,
+        mapping: Optional[torch.Tensor] = None,
+    ):
+        """Compute the descriptor.
+
+        Parameters
+        ----------
+        coord_ext
+            The extended coordinates of atoms. shape: nf x (nallx3)
+        atype_ext
+            The extended aotm types. shape: nf x nall
+        nlist
+            The neighbor list. shape: nf x nloc x nnei
+        mapping
+            The index mapping, not required by this descriptor.
+
+        Returns
+        -------
+        descriptor
+            The descriptor. shape: nf x nloc x (ng x axis_neuron)
+        gr
+            The rotationally equivariant and permutationally invariant single particle
+            representation. shape: nf x nloc x ng x 3
+        g2
+            The rotationally invariant pair-partical representation.
+            this descriptor returns None
+        h2
+            The rotationally equivariant pair-partical representation.
+            this descriptor returns None
+        sw
+            The smooth switch function.
+
+        """
+        
+        del mapping
+        nloc = nlist.shape[1]
+        atype = atype_ext[:, :nloc]
+        dmatrix, diff, sw = prod_env_mat_se_r(
+            coord_ext,
+            nlist,
+            atype,
+            self.mean,
+            self.stddev,
+            self.rcut,
+            self.rcut_smth,
+        )
+        assert dmatrix.shape == (2,3,7,1)
+
+        if self.old_impl:
+            assert self.filter_layers_old is not None
+            dmatrix = dmatrix.view(
+                -1, self.ndescrpt
+            )  # shape is [nframes*nall, self.ndescrpt]
+            xyz_scatter = torch.empty(
+                1,
+                device=env.DEVICE,
+            )
+            ret = self.filter_layers_old[0](dmatrix)
+            xyz_scatter = ret
+            for ii, transform in enumerate(self.filter_layers_old[1:]):
+                # shape is [nframes*nall, 1, self.filter_neuron[-1]]
+                ret = transform.forward(dmatrix)
+                xyz_scatter = xyz_scatter + ret
+        else:
+            assert self.filter_layers is not None
+            dmatrix = dmatrix.view(-1, self.nnei, 1)
+            dmatrix = dmatrix.to(dtype=self.prec)
+            nfnl = dmatrix.shape[0]
+            # pre-allocate a shape to pass jit
+            xyz_scatter = torch.zeros(
+                [nfnl, 1, self.filter_neuron[-1]], dtype=self.prec, device=env.DEVICE
+            )
+
+            # nfnl x nnei
+            exclude_mask = self.emask(nlist, atype_ext).view(nfnl, -1)
+            for ii, ll in enumerate(self.filter_layers.networks):
+                # nfnl x nt
+                mm = exclude_mask[:, self.sec[ii] : self.sec[ii + 1]]
+                # nfnl x nt x 1
+                rr = dmatrix[:, self.sec[ii] : self.sec[ii + 1], :]
+                rr = rr * mm[:, :, None]
+                ss = rr[:, :, :1]
+                # nfnl x nt x ng
+                gg = ll.forward(ss)
+                # nfnl x 1 x ng
+                gr = torch.matmul(rr.permute(0, 2, 1), gg)
+                xyz_scatter += gr
+
+        xyz_scatter /= self.nnei
+        xyz_scatter_1 = xyz_scatter.permute(0, 2, 1)
+
+        result = torch.matmul(
+            xyz_scatter_1, xyz_scatter
+        )  # shape is [nframes*nall, self.filter_neuron[-1], 1]
+        result = result.view(-1, nloc, self.filter_neuron[-1] * 1)
+        return (
+            result.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION),
+            None,
+            None,
+            None,
+            sw,
+        )
+
+    def set_stat_mean_and_stddev(
+        self,
+        mean: torch.Tensor,
+        stddev: torch.Tensor,
+    ) -> None:
+        self.mean = mean
+        self.stddev = stddev
+
+    def serialize(self) -> dict:
+        return {
+            "rcut": self.rcut,
+            "rcut_smth": self.rcut_smth,
+            "sel": self.sel,
+            "neuron": self.neuron,
+            "resnet_dt": self.resnet_dt,
+            "set_davg_zero": self.set_davg_zero,
+            "activation_function": self.activation_function,
+            # make deterministic
+            "precision": RESERVED_PRECISON_DICT[self.prec],
+            "embeddings": self.filter_layers.serialize(),
+            "env_mat": DPEnvMat(self.rcut, self.rcut_smth).serialize(),
+            "exclude_types": self.exclude_types,
+            "@variables": {
+                "davg": self["davg"].detach().cpu().numpy(),
+                "dstd": self["dstd"].detach().cpu().numpy(),
+            },
+            ## to be updated when the options are supported.
+            "trainable": True,
+            "type_one_side": True,
+            "spin": None,
+        }
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "DescrptSeR":
+        data = data.copy()
+        variables = data.pop("@variables")
+        embeddings = data.pop("embeddings")
+        env_mat = data.pop("env_mat")
+        obj = cls(**data)
+
+        def t_cvt(xx):
+            return torch.tensor(xx, dtype=obj.prec, device=env.DEVICE)
+
+        obj["davg"] = t_cvt(variables["davg"])
+        obj["dstd"] = t_cvt(variables["dstd"])
+        obj.filter_layers = NetworkCollection.deserialize(embeddings)
+        return obj
+
+def analyze_descrpt(matrix, ndescrpt, natoms):
+    """Collect avg, square avg and count of descriptors in a batch."""
+    ntypes = natoms.shape[1] - 2
+    start_index = 0
+    sysr = []
+    sysn = []
+    sysr2 = []
+    for type_i in range(ntypes):
+        end_index = start_index + natoms[0, 2 + type_i]
+        dd = matrix[:, start_index:end_index]  # all descriptors for this element
+        start_index = end_index
+        dd = np.reshape(
+            dd, [-1, 1]
+        )  # Shape is [nframes*natoms[2+type_id]*self.nnei, 1]
+        ddr = dd[:, :1]
+        sumr = np.sum(ddr)
+        sumn = dd.shape[0]  # Value is nframes*natoms[2+type_id]*self.nnei
+        sumr2 = np.sum(np.multiply(ddr, ddr))
+        sysr.append(sumr)
+        sysn.append(sumn)
+        sysr2.append(sumr2)
+    return sysr, sysr2, sysn
diff --git a/deepmd/pt/utils/env_mat_stat.py b/deepmd/pt/utils/env_mat_stat.py
index 2f3c728c99..61f1dccb38 100644
--- a/deepmd/pt/utils/env_mat_stat.py
+++ b/deepmd/pt/utils/env_mat_stat.py
@@ -14,6 +14,7 @@
 )
 from deepmd.pt.model.descriptor.env_mat import (
     prod_env_mat_se_a,
+    prod_env_mat_se_r,
 )
 from deepmd.pt.utils import (
     env,
@@ -193,3 +194,138 @@ def __call__(self):
         mean = np.stack(all_davg)
         stddev = np.stack(all_dstd)
         return mean, stddev
+
+
+class EnvMatStatSeR(EnvMatStat):
+    """Environmental matrix statistics for the se_r environemntal matrix.
+
+    Parameters
+    ----------
+    descriptor : DescriptorBlock
+        The descriptor of the model.
+    """
+
+    def __init__(self, descriptor: "DescriptorBlock"):
+        super().__init__()
+        self.descriptor = descriptor
+
+    def iter(
+        self, data: List[Dict[str, torch.Tensor]]
+    ) -> Iterator[Dict[str, StatItem]]:
+        """Get the iterator of the environment matrix.
+
+        Parameters
+        ----------
+        data : List[Dict[str, torch.Tensor]]
+            The environment matrix.
+
+        Yields
+        ------
+        Dict[str, StatItem]
+            The statistics of the environment matrix.
+        """
+        zero_mean = torch.zeros(
+            self.descriptor.get_ntypes(),
+            self.descriptor.get_nsel(),
+            1,
+            dtype=env.GLOBAL_PT_FLOAT_PRECISION,
+            device=env.DEVICE,
+        )
+        one_stddev = torch.ones(
+            self.descriptor.get_ntypes(),
+            self.descriptor.get_nsel(),
+            1,
+            dtype=env.GLOBAL_PT_FLOAT_PRECISION,
+            device=env.DEVICE,
+        )
+        for system in data:
+            coord, atype, box, natoms = (
+                system["coord"],
+                system["atype"],
+                system["box"],
+                system["natoms"],
+            )
+            (
+                extended_coord,
+                extended_atype,
+                mapping,
+                nlist,
+            ) = extend_input_and_build_neighbor_list(
+                coord,
+                atype,
+                self.descriptor.get_rcut(),
+                self.descriptor.get_sel(),
+                mixed_types=self.descriptor.mixed_types(),
+                box=box,
+            )
+            env_mat, _, _ = prod_env_mat_se_r(
+                extended_coord,
+                nlist,
+                atype,
+                zero_mean,
+                one_stddev,
+                self.descriptor.get_rcut(),
+                # TODO: export rcut_smth from DescriptorBlock
+                self.descriptor.rcut_smth,
+            )
+            # reshape to nframes * nloc at the atom level,
+            # so nframes/mixed_type do not matter
+            env_mat = env_mat.view(
+                coord.shape[0] * coord.shape[1], self.descriptor.get_nsel(), 1
+            )
+            atype = atype.view(coord.shape[0] * coord.shape[1])
+            # (1, nloc) eq (ntypes, 1), so broadcast is possible
+            # shape: (ntypes, nloc)
+            type_idx = torch.eq(
+                atype.view(1, -1),
+                torch.arange(
+                    self.descriptor.get_ntypes(), device=env.DEVICE, dtype=torch.int32
+                ).view(-1, 1),
+            )
+            for type_i in range(self.descriptor.get_ntypes()):
+                dd = env_mat[type_idx[type_i]]
+                dd = dd.reshape([-1, 1])  # typen_atoms * nnei, 4
+                env_mats = {}
+                env_mats[f"r_{type_i}"] = dd[:, :1]
+                yield self.compute_stat(env_mats)
+
+    def get_hash(self) -> str:
+        """Get the hash of the environment matrix.
+
+        Returns
+        -------
+        str
+            The hash of the environment matrix.
+        """
+        return get_hash(
+            {
+                "type": "se_r",
+                "ntypes": self.descriptor.get_ntypes(),
+                "rcut": round(self.descriptor.get_rcut(), 2),
+                "rcut_smth": round(self.descriptor.rcut_smth, 2),
+                "nsel": self.descriptor.get_nsel(),
+                "sel": self.descriptor.get_sel(),
+                "mixed_types": self.descriptor.mixed_types(),
+            }
+        )
+
+    def __call__(self):
+        avgs = self.get_avg()
+        stds = self.get_std()
+
+        all_davg = []
+        all_dstd = []
+        for type_i in range(self.descriptor.get_ntypes()):
+            davgunit = [[avgs[f"r_{type_i}"]]]
+            dstdunit = [
+                [
+                    stds[f"r_{type_i}"]
+                ]
+            ]
+            davg = np.tile(davgunit, [self.descriptor.get_nsel(), 1])
+            dstd = np.tile(dstdunit, [self.descriptor.get_nsel(), 1])
+            all_davg.append(davg)
+            all_dstd.append(dstd)
+        mean = np.stack(all_davg)
+        stddev = np.stack(all_dstd)
+        return mean, stddev
\ No newline at end of file
diff --git a/source/tests/pt/model/test_descriptor_se_r.py b/source/tests/pt/model/test_descriptor_se_r.py
new file mode 100644
index 0000000000..abf3a050d3
--- /dev/null
+++ b/source/tests/pt/model/test_descriptor_se_r.py
@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import itertools
+import unittest
+
+import numpy as np
+import torch
+
+# from deepmd.dpmodel.descriptor import DescrptSeR as DPDescrptSeR
+from deepmd.pt.model.descriptor.se_r import (
+    DescrptSeR,
+)
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.pt.utils.env import (
+    PRECISION_DICT,
+)
+
+from .test_env_mat import (
+    TestCaseSingleFrameWithNlist,
+)
+from .test_mlp import (
+    get_tols,
+)
+
+dtype = env.GLOBAL_PT_FLOAT_PRECISION
+
+
+# to be merged with the tf test case
+class TestDescrptSeR(unittest.TestCase, TestCaseSingleFrameWithNlist):
+    def setUp(self):
+        TestCaseSingleFrameWithNlist.setUp(self)
+
+    def test_consistency(
+        self,
+    ):
+        rng = np.random.default_rng()
+        nf, nloc, nnei = self.nlist.shape
+        davg = rng.normal(size=(self.nt, nnei, 1))
+        dstd = rng.normal(size=(self.nt, nnei, 1))
+        dstd = 0.1 + np.abs(dstd)
+
+        for idt, prec, em in itertools.product(
+            [False, True],
+            ["float64", "float32"],
+            [[], [[0, 1]], [[1, 1]]],
+        ):
+            dtype = PRECISION_DICT[prec]
+            rtol, atol = get_tols(prec)
+            err_msg = f"idt={idt} prec={prec}"
+            # sea new impl
+            dd0 = DescrptSeR(
+                self.rcut,
+                self.rcut_smth,
+                self.sel,
+                precision=prec,
+                resnet_dt=idt,
+                old_impl=False,
+                exclude_mask=em,
+            ).to(env.DEVICE)
+            dd0.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE)
+            dd0.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE)
+
+            rd0, _, _, _, _ = dd0(
+                torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE),
+                torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE),
+                torch.tensor(self.nlist, dtype=int, device=env.DEVICE),
+            )
+            # serialization
+            dd1 = DescrptSeR.deserialize(dd0.serialize())
+            rd1, gr1, _, _, sw1 = dd1(
+                torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE),
+                torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE),
+                torch.tensor(self.nlist, dtype=int, device=env.DEVICE),
+            )
+            np.testing.assert_allclose(
+                rd0.detach().cpu().numpy(),
+                rd1.detach().cpu().numpy(),
+                rtol=rtol,
+                atol=atol,
+                err_msg=err_msg,
+            )
+            np.testing.assert_allclose(
+                rd0.detach().cpu().numpy()[0][self.perm[: self.nloc]],
+                rd0.detach().cpu().numpy()[1],
+                rtol=rtol,
+                atol=atol,
+                err_msg=err_msg,
+            )
+            # dp impl
+            # dd2 = DPDescrptSeR.deserialize(dd0.serialize())
+            # rd2, gr2, _, _, sw2 = dd2.call(
+            #     self.coord_ext,
+            #     self.atype_ext,
+            #     self.nlist,
+            # )
+            # for aa, bb in zip([rd1, gr1, sw1], [rd2, gr2, sw2]):
+            #     np.testing.assert_allclose(
+            #         aa.detach().cpu().numpy(),
+            #         bb,
+            #         rtol=rtol,
+            #         atol=atol,
+            #         err_msg=err_msg,
+            #     )
+            # old impl
+            if idt is False and prec == "float64":
+                dd3 = DescrptSeR(
+                    self.rcut,
+                    self.rcut_smth,
+                    self.sel,
+                    precision=prec,
+                    resnet_dt=idt,
+                    old_impl=True,
+                ).to(env.DEVICE)
+                dd0_state_dict = dd0.state_dict()
+                dd3_state_dict = dd3.state_dict()
+                for i in dd3_state_dict:
+                    dd3_state_dict[i] = (
+                        dd0_state_dict[
+                            i.replace(".deep_layers.", ".layers.").replace(
+                                "filter_layers_old.", "filter_layers.networks."
+                            )
+                        ]
+                        .detach()
+                        .clone()
+                    )
+                    if ".bias" in i:
+                        dd3_state_dict[i] = dd3_state_dict[i].unsqueeze(0)
+                dd3.load_state_dict(dd3_state_dict)
+
+                rd3, gr3, _, _, sw3 = dd3(
+                    torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE),
+                    torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE),
+                    torch.tensor(self.nlist, dtype=int, device=env.DEVICE),
+                )
+                for aa, bb in zip([rd1, gr1, sw1], [rd3, gr3, sw3]):
+                    np.testing.assert_allclose(
+                        aa.detach().cpu().numpy(),
+                        bb.detach().cpu().numpy(),
+                        rtol=rtol,
+                        atol=atol,
+                        err_msg=err_msg,
+                    )
+
+    def test_jit(
+        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)
+
+        for idt, prec in itertools.product(
+            [False, True],
+            ["float64", "float32"],
+        ):
+            dtype = PRECISION_DICT[prec]
+            rtol, atol = get_tols(prec)
+            err_msg = f"idt={idt} prec={prec}"
+            # sea new impl
+            dd0 = DescrptSeR(
+                self.rcut,
+                self.rcut_smth,
+                self.sel,
+                precision=prec,
+                resnet_dt=idt,
+                old_impl=False,
+            )
+            dd0.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE)
+            dd0.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE)
+            dd1 = DescrptSeR.deserialize(dd0.serialize())
+            model = torch.jit.script(dd0)
+            model = torch.jit.script(dd1)

From c0af6faf09b5b78b4d5e1621da53577d814ee2b2 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 26 Feb 2024 04:01:53 +0000
Subject: [PATCH 02/20] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt/model/descriptor/__init__.py | 4 ++--
 deepmd/pt/model/descriptor/env_mat.py  | 3 +++
 deepmd/pt/model/descriptor/se_r.py     | 8 ++++----
 deepmd/pt/utils/env_mat_stat.py        | 8 ++------
 4 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/deepmd/pt/model/descriptor/__init__.py b/deepmd/pt/model/descriptor/__init__.py
index 445537e2c5..b69e5a82b1 100644
--- a/deepmd/pt/model/descriptor/__init__.py
+++ b/deepmd/pt/model/descriptor/__init__.py
@@ -28,8 +28,8 @@
     DescrptBlockSeA,
     DescrptSeA,
 )
-from .se_r import(
-    DescrptSeR
+from .se_r import (
+    DescrptSeR,
 )
 
 __all__ = [
diff --git a/deepmd/pt/model/descriptor/env_mat.py b/deepmd/pt/model/descriptor/env_mat.py
index be6c2437d8..01972a5d81 100644
--- a/deepmd/pt/model/descriptor/env_mat.py
+++ b/deepmd/pt/model/descriptor/env_mat.py
@@ -29,6 +29,7 @@ def _make_env_mat_se_a(nlist, coord, rcut: float, ruct_smth: float):
     env_mat_se_a = torch.cat([t0, t1], dim=-1) * weight
     return env_mat_se_a, diff * mask.unsqueeze(-1), weight
 
+
 def _make_env_mat_se_r(nlist, coord, rcut: float, ruct_smth: float):
     """Make smooth environment matrix."""
     bsz, natoms, nnei = nlist.shape
@@ -51,6 +52,7 @@ def _make_env_mat_se_r(nlist, coord, rcut: float, ruct_smth: float):
     env_mat_se_r = t0 * weight
     return env_mat_se_r, diff * mask.unsqueeze(-1), weight
 
+
 def prod_env_mat_se_a(
     extended_coord, nlist, atype, mean, stddev, rcut: float, rcut_smth: float
 ):
@@ -80,6 +82,7 @@ def prod_env_mat_se_a(
     env_mat_se_a = (_env_mat_se_a - t_avg) / t_std
     return env_mat_se_a, diff, switch
 
+
 def prod_env_mat_se_r(
     extended_coord, nlist, atype, mean, stddev, rcut: float, rcut_smth: float
 ):
diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index 461667364b..d2cb5b4e8e 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -35,7 +35,7 @@
         Final,
     )
 except ImportError:
-    from torch.jit import Final
+    pass
 
 from deepmd.dpmodel.utils import EnvMat as DPEnvMat
 from deepmd.pt.model.network.mlp import (
@@ -184,7 +184,7 @@ def get_stats(self) -> Dict[str, StatItem]:
                 "The statistics of the descriptor has not been computed."
             )
         return self.stats
-    
+
     def __setitem__(self, key, value):
         if key in ("avg", "data_avg", "davg"):
             self.mean = value
@@ -257,7 +257,6 @@ def forward(
             The smooth switch function.
 
         """
-        
         del mapping
         nloc = nlist.shape[1]
         atype = atype_ext[:, :nloc]
@@ -270,7 +269,7 @@ def forward(
             self.rcut,
             self.rcut_smth,
         )
-        assert dmatrix.shape == (2,3,7,1)
+        assert dmatrix.shape == (2, 3, 7, 1)
 
         if self.old_impl:
             assert self.filter_layers_old is not None
@@ -375,6 +374,7 @@ def t_cvt(xx):
         obj.filter_layers = NetworkCollection.deserialize(embeddings)
         return obj
 
+
 def analyze_descrpt(matrix, ndescrpt, natoms):
     """Collect avg, square avg and count of descriptors in a batch."""
     ntypes = natoms.shape[1] - 2
diff --git a/deepmd/pt/utils/env_mat_stat.py b/deepmd/pt/utils/env_mat_stat.py
index 61f1dccb38..757e954018 100644
--- a/deepmd/pt/utils/env_mat_stat.py
+++ b/deepmd/pt/utils/env_mat_stat.py
@@ -317,15 +317,11 @@ def __call__(self):
         all_dstd = []
         for type_i in range(self.descriptor.get_ntypes()):
             davgunit = [[avgs[f"r_{type_i}"]]]
-            dstdunit = [
-                [
-                    stds[f"r_{type_i}"]
-                ]
-            ]
+            dstdunit = [[stds[f"r_{type_i}"]]]
             davg = np.tile(davgunit, [self.descriptor.get_nsel(), 1])
             dstd = np.tile(dstdunit, [self.descriptor.get_nsel(), 1])
             all_davg.append(davg)
             all_dstd.append(dstd)
         mean = np.stack(all_davg)
         stddev = np.stack(all_dstd)
-        return mean, stddev
\ No newline at end of file
+        return mean, stddev

From 8a8107c10cb518beb17597ea6104c096016c8db8 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Mon, 26 Feb 2024 12:29:36 +0800
Subject: [PATCH 03/20] fix: UTs, removed old impl

---
 deepmd/pt/model/descriptor/se_r.py            | 66 +++++++------------
 source/tests/pt/model/test_descriptor_se_r.py | 40 +----------
 2 files changed, 26 insertions(+), 80 deletions(-)

diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index d2cb5b4e8e..40e293f374 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -76,7 +76,7 @@ def __init__(
         self.precision = precision
         self.prec = PRECISION_DICT[self.precision]
         self.resnet_dt = resnet_dt
-        self.old_impl = old_impl
+        self.old_impl = False # this does not support old implementation.
         self.exclude_types = exclude_types
         self.ntypes = len(sel)
         self.emask = PairExcludeMask(len(sel), exclude_types=exclude_types)
@@ -269,47 +269,31 @@ def forward(
             self.rcut,
             self.rcut_smth,
         )
-        assert dmatrix.shape == (2, 3, 7, 1)
 
-        if self.old_impl:
-            assert self.filter_layers_old is not None
-            dmatrix = dmatrix.view(
-                -1, self.ndescrpt
-            )  # shape is [nframes*nall, self.ndescrpt]
-            xyz_scatter = torch.empty(
-                1,
-                device=env.DEVICE,
-            )
-            ret = self.filter_layers_old[0](dmatrix)
-            xyz_scatter = ret
-            for ii, transform in enumerate(self.filter_layers_old[1:]):
-                # shape is [nframes*nall, 1, self.filter_neuron[-1]]
-                ret = transform.forward(dmatrix)
-                xyz_scatter = xyz_scatter + ret
-        else:
-            assert self.filter_layers is not None
-            dmatrix = dmatrix.view(-1, self.nnei, 1)
-            dmatrix = dmatrix.to(dtype=self.prec)
-            nfnl = dmatrix.shape[0]
-            # pre-allocate a shape to pass jit
-            xyz_scatter = torch.zeros(
-                [nfnl, 1, self.filter_neuron[-1]], dtype=self.prec, device=env.DEVICE
-            )
+        
+        assert self.filter_layers is not None
+        dmatrix = dmatrix.view(-1, self.nnei, 1)
+        dmatrix = dmatrix.to(dtype=self.prec)
+        nfnl = dmatrix.shape[0]
+        # pre-allocate a shape to pass jit
+        xyz_scatter = torch.zeros(
+            [nfnl, 1, self.filter_neuron[-1]], dtype=self.prec, device=env.DEVICE
+        )
 
-            # nfnl x nnei
-            exclude_mask = self.emask(nlist, atype_ext).view(nfnl, -1)
-            for ii, ll in enumerate(self.filter_layers.networks):
-                # nfnl x nt
-                mm = exclude_mask[:, self.sec[ii] : self.sec[ii + 1]]
-                # nfnl x nt x 1
-                rr = dmatrix[:, self.sec[ii] : self.sec[ii + 1], :]
-                rr = rr * mm[:, :, None]
-                ss = rr[:, :, :1]
-                # nfnl x nt x ng
-                gg = ll.forward(ss)
-                # nfnl x 1 x ng
-                gr = torch.matmul(rr.permute(0, 2, 1), gg)
-                xyz_scatter += gr
+        # nfnl x nnei
+        exclude_mask = self.emask(nlist, atype_ext).view(nfnl, -1)
+        for ii, ll in enumerate(self.filter_layers.networks):
+            # nfnl x nt
+            mm = exclude_mask[:, self.sec[ii] : self.sec[ii + 1]]
+            # nfnl x nt x 1
+            rr = dmatrix[:, self.sec[ii] : self.sec[ii + 1], :]
+            rr = rr * mm[:, :, None]
+            ss = rr[:, :, :1]
+            # nfnl x nt x ng
+            gg = ll.forward(ss)
+            # nfnl x 1 x ng
+            gr = torch.matmul(rr.permute(0, 2, 1), gg)
+            xyz_scatter += gr
 
         xyz_scatter /= self.nnei
         xyz_scatter_1 = xyz_scatter.permute(0, 2, 1)
@@ -317,7 +301,7 @@ def forward(
         result = torch.matmul(
             xyz_scatter_1, xyz_scatter
         )  # shape is [nframes*nall, self.filter_neuron[-1], 1]
-        result = result.view(-1, nloc, self.filter_neuron[-1] * 1)
+        result = result.view(-1, nloc, self.filter_neuron[-1] * self.filter_neuron[-1])
         return (
             result.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION),
             None,
diff --git a/source/tests/pt/model/test_descriptor_se_r.py b/source/tests/pt/model/test_descriptor_se_r.py
index abf3a050d3..1b0a4d3dd1 100644
--- a/source/tests/pt/model/test_descriptor_se_r.py
+++ b/source/tests/pt/model/test_descriptor_se_r.py
@@ -102,45 +102,7 @@ def test_consistency(
             #         atol=atol,
             #         err_msg=err_msg,
             #     )
-            # old impl
-            if idt is False and prec == "float64":
-                dd3 = DescrptSeR(
-                    self.rcut,
-                    self.rcut_smth,
-                    self.sel,
-                    precision=prec,
-                    resnet_dt=idt,
-                    old_impl=True,
-                ).to(env.DEVICE)
-                dd0_state_dict = dd0.state_dict()
-                dd3_state_dict = dd3.state_dict()
-                for i in dd3_state_dict:
-                    dd3_state_dict[i] = (
-                        dd0_state_dict[
-                            i.replace(".deep_layers.", ".layers.").replace(
-                                "filter_layers_old.", "filter_layers.networks."
-                            )
-                        ]
-                        .detach()
-                        .clone()
-                    )
-                    if ".bias" in i:
-                        dd3_state_dict[i] = dd3_state_dict[i].unsqueeze(0)
-                dd3.load_state_dict(dd3_state_dict)
-
-                rd3, gr3, _, _, sw3 = dd3(
-                    torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE),
-                    torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE),
-                    torch.tensor(self.nlist, dtype=int, device=env.DEVICE),
-                )
-                for aa, bb in zip([rd1, gr1, sw1], [rd3, gr3, sw3]):
-                    np.testing.assert_allclose(
-                        aa.detach().cpu().numpy(),
-                        bb.detach().cpu().numpy(),
-                        rtol=rtol,
-                        atol=atol,
-                        err_msg=err_msg,
-                    )
+            
 
     def test_jit(
         self,

From fb6340b22dbfbf43fa658794f1a4919ea1f85284 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 26 Feb 2024 04:32:59 +0000
Subject: [PATCH 04/20] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt/model/descriptor/se_r.py            | 3 +--
 source/tests/pt/model/test_descriptor_se_r.py | 1 -
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index 40e293f374..85d0779e50 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -76,7 +76,7 @@ def __init__(
         self.precision = precision
         self.prec = PRECISION_DICT[self.precision]
         self.resnet_dt = resnet_dt
-        self.old_impl = False # this does not support old implementation.
+        self.old_impl = False  # this does not support old implementation.
         self.exclude_types = exclude_types
         self.ntypes = len(sel)
         self.emask = PairExcludeMask(len(sel), exclude_types=exclude_types)
@@ -270,7 +270,6 @@ def forward(
             self.rcut_smth,
         )
 
-        
         assert self.filter_layers is not None
         dmatrix = dmatrix.view(-1, self.nnei, 1)
         dmatrix = dmatrix.to(dtype=self.prec)
diff --git a/source/tests/pt/model/test_descriptor_se_r.py b/source/tests/pt/model/test_descriptor_se_r.py
index 1b0a4d3dd1..7e78ea95c9 100644
--- a/source/tests/pt/model/test_descriptor_se_r.py
+++ b/source/tests/pt/model/test_descriptor_se_r.py
@@ -102,7 +102,6 @@ def test_consistency(
             #         atol=atol,
             #         err_msg=err_msg,
             #     )
-            
 
     def test_jit(
         self,

From 2eb4041c8cf8b0e61801cf159514d804897b5818 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Mon, 26 Feb 2024 12:42:06 +0800
Subject: [PATCH 05/20] fix: pre-commit

---
 deepmd/pt/model/descriptor/env_mat.py         |  2 --
 deepmd/pt/model/descriptor/se_r.py            |  6 ------
 source/tests/pt/model/test_descriptor_se_r.py | 11 +++++------
 3 files changed, 5 insertions(+), 14 deletions(-)

diff --git a/deepmd/pt/model/descriptor/env_mat.py b/deepmd/pt/model/descriptor/env_mat.py
index 01972a5d81..81d0996ff2 100644
--- a/deepmd/pt/model/descriptor/env_mat.py
+++ b/deepmd/pt/model/descriptor/env_mat.py
@@ -73,7 +73,6 @@ def prod_env_mat_se_a(
     -------
     - env_mat_se_a: Shape is [nframes, natoms[1]*nnei*4].
     """
-    nframes = extended_coord.shape[0]
     _env_mat_se_a, diff, switch = _make_env_mat_se_a(
         nlist, extended_coord, rcut, rcut_smth
     )  # shape [n_atom, dim, 4]
@@ -103,7 +102,6 @@ def prod_env_mat_se_r(
     -------
     - env_mat_se_r: Shape is [nframes, natoms[1]*nnei*1].
     """
-    nframes = extended_coord.shape[0]
     _env_mat_se_r, diff, switch = _make_env_mat_se_r(
         nlist, extended_coord, rcut, rcut_smth
     )  # shape [n_atom, dim, 1]
diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index 85d0779e50..fb185dec81 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -30,12 +30,6 @@
     DPPath,
 )
 
-try:
-    from typing import (
-        Final,
-    )
-except ImportError:
-    pass
 
 from deepmd.dpmodel.utils import EnvMat as DPEnvMat
 from deepmd.pt.model.network.mlp import (
diff --git a/source/tests/pt/model/test_descriptor_se_r.py b/source/tests/pt/model/test_descriptor_se_r.py
index 7e78ea95c9..8220079638 100644
--- a/source/tests/pt/model/test_descriptor_se_r.py
+++ b/source/tests/pt/model/test_descriptor_se_r.py
@@ -35,7 +35,7 @@ def test_consistency(
         self,
     ):
         rng = np.random.default_rng()
-        nf, nloc, nnei = self.nlist.shape
+        _, _, nnei = self.nlist.shape
         davg = rng.normal(size=(self.nt, nnei, 1))
         dstd = rng.normal(size=(self.nt, nnei, 1))
         dstd = 0.1 + np.abs(dstd)
@@ -107,7 +107,7 @@ def test_jit(
         self,
     ):
         rng = np.random.default_rng()
-        nf, nloc, nnei = self.nlist.shape
+        _, _, 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)
@@ -117,8 +117,7 @@ def test_jit(
             ["float64", "float32"],
         ):
             dtype = PRECISION_DICT[prec]
-            rtol, atol = get_tols(prec)
-            err_msg = f"idt={idt} prec={prec}"
+
             # sea new impl
             dd0 = DescrptSeR(
                 self.rcut,
@@ -131,5 +130,5 @@ def test_jit(
             dd0.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE)
             dd0.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE)
             dd1 = DescrptSeR.deserialize(dd0.serialize())
-            model = torch.jit.script(dd0)
-            model = torch.jit.script(dd1)
+            torch.jit.script(dd0)
+            torch.jit.script(dd1)

From 6c5224ff05ded284af6c4f12e1d6b0e316164274 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 26 Feb 2024 04:43:02 +0000
Subject: [PATCH 06/20] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt/model/descriptor/se_r.py | 24 +++++++++++-------------
 1 file changed, 11 insertions(+), 13 deletions(-)

diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index fb185dec81..278210588b 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -9,10 +9,18 @@
 import numpy as np
 import torch
 
+from deepmd.dpmodel.utils import EnvMat as DPEnvMat
 from deepmd.pt.model.descriptor import (
     Descriptor,
     prod_env_mat_se_r,
 )
+from deepmd.pt.model.network.mlp import (
+    EmbeddingNet,
+    NetworkCollection,
+)
+from deepmd.pt.model.network.network import (
+    TypeFilter,
+)
 from deepmd.pt.utils import (
     env,
 )
@@ -23,6 +31,9 @@
 from deepmd.pt.utils.env_mat_stat import (
     EnvMatStatSeR,
 )
+from deepmd.pt.utils.exclude_mask import (
+    PairExcludeMask,
+)
 from deepmd.utils.env_mat_stat import (
     StatItem,
 )
@@ -31,19 +42,6 @@
 )
 
 
-from deepmd.dpmodel.utils import EnvMat as DPEnvMat
-from deepmd.pt.model.network.mlp import (
-    EmbeddingNet,
-    NetworkCollection,
-)
-from deepmd.pt.model.network.network import (
-    TypeFilter,
-)
-from deepmd.pt.utils.exclude_mask import (
-    PairExcludeMask,
-)
-
-
 @Descriptor.register("se_e2_r")
 class DescrptSeR(Descriptor):
     def __init__(

From b771db419cb61b163fa7511a49d53bee7ccab794 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Mon, 26 Feb 2024 13:28:54 +0800
Subject: [PATCH 07/20] fix: update se_r output

---
 deepmd/pt/model/descriptor/se_r.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index 278210588b..d7c16bbada 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -286,13 +286,9 @@ def forward(
             gr = torch.matmul(rr.permute(0, 2, 1), gg)
             xyz_scatter += gr
 
-        xyz_scatter /= self.nnei
-        xyz_scatter_1 = xyz_scatter.permute(0, 2, 1)
-
-        result = torch.matmul(
-            xyz_scatter_1, xyz_scatter
-        )  # shape is [nframes*nall, self.filter_neuron[-1], 1]
-        result = result.view(-1, nloc, self.filter_neuron[-1] * self.filter_neuron[-1])
+        res_rescale = 1.0 / 5.0
+        result = torch.mean(xyz_scatter, dim=1) * res_rescale
+        result = result.view(-1, nloc, self.filter_neuron[-1])
         return (
             result.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION),
             None,

From 8a1a86c650a750d651812c3ff82a06d83d35b2a4 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Mon, 26 Feb 2024 14:52:56 +0800
Subject: [PATCH 08/20] chore: refactor

---
 deepmd/pt/model/descriptor/__init__.py   |   7 +-
 deepmd/pt/model/descriptor/descriptor.py |   4 +-
 deepmd/pt/model/descriptor/env_mat.py    |  86 +++--------
 deepmd/pt/model/descriptor/repformers.py |   8 +-
 deepmd/pt/model/descriptor/se_a.py       |   8 +-
 deepmd/pt/model/descriptor/se_atten.py   |   8 +-
 deepmd/pt/model/descriptor/se_r.py       |  44 +++---
 deepmd/pt/utils/env_mat_stat.py          | 181 ++++-------------------
 source/tests/pt/model/test_descriptor.py |   4 +-
 source/tests/pt/model/test_env_mat.py    |   4 +-
 10 files changed, 88 insertions(+), 266 deletions(-)

diff --git a/deepmd/pt/model/descriptor/__init__.py b/deepmd/pt/model/descriptor/__init__.py
index b69e5a82b1..2c55346f94 100644
--- a/deepmd/pt/model/descriptor/__init__.py
+++ b/deepmd/pt/model/descriptor/__init__.py
@@ -12,8 +12,8 @@
     DescrptDPA2,
 )
 from .env_mat import (
-    prod_env_mat_se_a,
-    prod_env_mat_se_r,
+    prod_env_mat,
+
 )
 from .gaussian_lcc import (
     DescrptGaussianLcc,
@@ -42,8 +42,7 @@
     "DescrptSeR",
     "DescrptDPA1",
     "DescrptDPA2",
-    "prod_env_mat_se_a",
-    "prod_env_mat_se_r",
+    "prod_env_mat",
     "DescrptGaussianLcc",
     "DescrptBlockHybrid",
     "DescrptBlockRepformers",
diff --git a/deepmd/pt/model/descriptor/descriptor.py b/deepmd/pt/model/descriptor/descriptor.py
index 091f2b1e20..0611d8400b 100644
--- a/deepmd/pt/model/descriptor/descriptor.py
+++ b/deepmd/pt/model/descriptor/descriptor.py
@@ -24,7 +24,7 @@
     env,
 )
 from deepmd.pt.utils.env_mat_stat import (
-    EnvMatStatSeA,
+    EnvMatStatSe,
 )
 from deepmd.pt.utils.plugin import (
     Plugin,
@@ -222,7 +222,7 @@ def share_params(self, base_class, shared_level, resume=False):
             # link buffers
             if hasattr(self, "mean") and not resume:
                 # in case of change params during resume
-                base_env = EnvMatStatSeA(base_class)
+                base_env = EnvMatStatSe(base_class)
                 base_env.stats = base_class.stats
                 for kk in base_class.get_stats():
                     base_env.stats[kk] += self.get_stats()[kk]
diff --git a/deepmd/pt/model/descriptor/env_mat.py b/deepmd/pt/model/descriptor/env_mat.py
index 81d0996ff2..c3a1d38726 100644
--- a/deepmd/pt/model/descriptor/env_mat.py
+++ b/deepmd/pt/model/descriptor/env_mat.py
@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import Optional
 import torch
 
 from deepmd.pt.utils.preprocess import (
@@ -6,7 +7,7 @@
 )
 
 
-def _make_env_mat_se_a(nlist, coord, rcut: float, ruct_smth: float):
+def _make_env_mat(nlist, coord, rcut: float, ruct_smth: float, radial_only: bool=False):
     """Make smooth environment matrix."""
     bsz, natoms, nnei = nlist.shape
     coord = coord.view(bsz, -1, 3)
@@ -26,86 +27,35 @@ def _make_env_mat_se_a(nlist, coord, rcut: float, ruct_smth: float):
     t1 = diff / length**2
     weight = compute_smooth_weight(length, ruct_smth, rcut)
     weight = weight * mask.unsqueeze(-1)
-    env_mat_se_a = torch.cat([t0, t1], dim=-1) * weight
-    return env_mat_se_a, diff * mask.unsqueeze(-1), weight
+    if radial_only:
+        env_mat = t0 * weight
+    else:
+        env_mat = torch.cat([t0, t1], dim=-1) * weight
+    return env_mat, diff * mask.unsqueeze(-1), weight
 
 
-def _make_env_mat_se_r(nlist, coord, rcut: float, ruct_smth: float):
-    """Make smooth environment matrix."""
-    bsz, natoms, nnei = nlist.shape
-    coord = coord.view(bsz, -1, 3)
-    nall = coord.shape[1]
-    mask = nlist >= 0
-    # nlist = nlist * mask  ## this impl will contribute nans in Hessian calculation.
-    nlist = torch.where(mask, nlist, nall - 1)
-    coord_l = coord[:, :natoms].view(bsz, -1, 1, 3)
-    index = nlist.view(bsz, -1).unsqueeze(-1).expand(-1, -1, 3)
-    coord_r = torch.gather(coord, 1, index)
-    coord_r = coord_r.view(bsz, natoms, nnei, 3)
-    diff = coord_r - coord_l
-    length = torch.linalg.norm(diff, dim=-1, keepdim=True)
-    # for index 0 nloc atom
-    length = length + ~mask.unsqueeze(-1)
-    t0 = 1 / length
-    weight = compute_smooth_weight(length, ruct_smth, rcut)
-    weight = weight * mask.unsqueeze(-1)
-    env_mat_se_r = t0 * weight
-    return env_mat_se_r, diff * mask.unsqueeze(-1), weight
-
-
-def prod_env_mat_se_a(
-    extended_coord, nlist, atype, mean, stddev, rcut: float, rcut_smth: float
+def prod_env_mat(
+    extended_coord, nlist, atype, mean, stddev, rcut: float, rcut_smth: float, radial_only: bool=False
 ):
     """Generate smooth environment matrix from atom coordinates and other context.
 
     Args:
     - extended_coord: Copied atom coordinates with shape [nframes, nall*3].
     - atype: Atom types with shape [nframes, nloc].
-    - natoms: Batched atom statisics with shape [len(sec)+2].
-    - box: Batched simulation box with shape [nframes, 9].
-    - mean: Average value of descriptor per element type with shape [len(sec), nnei, 4].
-    - stddev: Standard deviation of descriptor per element type with shape [len(sec), nnei, 4].
-    - deriv_stddev:  StdDev of descriptor derivative per element type with shape [len(sec), nnei, 4, 3].
+    - mean: Average value of descriptor per element type with shape [len(sec), nnei, 4 or 1].
+    - stddev: Standard deviation of descriptor per element type with shape [len(sec), nnei, 4 or 1].
     - rcut: Cut-off radius.
     - rcut_smth: Smooth hyper-parameter for pair force & energy.
+    - radial_only: Whether to return a full description or a radial-only descriptor.
 
     Returns
     -------
-    - env_mat_se_a: Shape is [nframes, natoms[1]*nnei*4].
+    - env_mat: Shape is [nframes, natoms[1]*nnei*4].
     """
-    _env_mat_se_a, diff, switch = _make_env_mat_se_a(
-        nlist, extended_coord, rcut, rcut_smth
-    )  # shape [n_atom, dim, 4]
-    t_avg = mean[atype]  # [n_atom, dim, 4]
-    t_std = stddev[atype]  # [n_atom, dim, 4]
+    _env_mat_se_a, diff, switch = _make_env_mat(
+        nlist, extended_coord, rcut, rcut_smth, radial_only
+    )  # shape [n_atom, dim, 4 or 1]
+    t_avg = mean[atype]  # [n_atom, dim, 4 or 1]
+    t_std = stddev[atype]  # [n_atom, dim, 4 or 1]
     env_mat_se_a = (_env_mat_se_a - t_avg) / t_std
     return env_mat_se_a, diff, switch
-
-
-def prod_env_mat_se_r(
-    extended_coord, nlist, atype, mean, stddev, rcut: float, rcut_smth: float
-):
-    """Generate smooth environment matrix from atom coordinates and other context.
-
-    Args:
-    - extended_coord: Copied atom coordinates with shape [nframes, nall*3].
-    - atype: Atom types with shape [nframes, nloc].
-    - natoms: Batched atom statisics with shape [len(sec)+2].
-    - box: Batched simulation box with shape [nframes, 9].
-    - mean: Average value of descriptor per element type with shape [len(sec), nnei, 1].
-    - stddev: Standard deviation of descriptor per element type with shape [len(sec), nnei, 1].
-    - deriv_stddev:  StdDev of descriptor derivative per element type with shape [len(sec), nnei, 1, 3].
-    - rcut: Cut-off radius.
-    - rcut_smth: Smooth hyper-parameter for pair force & energy.
-
-    Returns
-    -------
-    - env_mat_se_r: Shape is [nframes, natoms[1]*nnei*1].
-    """
-    _env_mat_se_r, diff, switch = _make_env_mat_se_r(
-        nlist, extended_coord, rcut, rcut_smth
-    )  # shape [n_atom, dim, 1]
-    t_avg = mean[atype]  # [n_atom, dim, 1]
-    t_std = stddev[atype]  # [n_atom, dim, 1]
-    env_mat_se_r = (_env_mat_se_r - t_avg) / t_std
-    return env_mat_se_r, diff, switch
diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py
index 8aa1114fdc..dc4afcd07a 100644
--- a/deepmd/pt/model/descriptor/repformers.py
+++ b/deepmd/pt/model/descriptor/repformers.py
@@ -11,7 +11,7 @@
     DescriptorBlock,
 )
 from deepmd.pt.model.descriptor.env_mat import (
-    prod_env_mat_se_a,
+    prod_env_mat,
 )
 from deepmd.pt.model.network.network import (
     SimpleLinear,
@@ -20,7 +20,7 @@
     env,
 )
 from deepmd.pt.utils.env_mat_stat import (
-    EnvMatStatSeA,
+    EnvMatStatSe,
 )
 from deepmd.pt.utils.utils import (
     get_activation_fn,
@@ -222,7 +222,7 @@ def forward(
         nall = extended_coord.view(nframes, -1).shape[1] // 3
         atype = extended_atype[:, :nloc]
         # nb x nloc x nnei x 4, nb x nloc x nnei x 3, nb x nloc x nnei x 1
-        dmatrix, diff, sw = prod_env_mat_se_a(
+        dmatrix, diff, sw = prod_env_mat(
             extended_coord,
             nlist,
             atype,
@@ -279,7 +279,7 @@ def forward(
 
     def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None):
         """Update mean and stddev for descriptor elements."""
-        env_mat_stat = EnvMatStatSeA(self)
+        env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
         env_mat_stat.load_or_compute_stats(merged, path)
diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py
index 0550488ecf..9440881eae 100644
--- a/deepmd/pt/model/descriptor/se_a.py
+++ b/deepmd/pt/model/descriptor/se_a.py
@@ -13,7 +13,7 @@
 from deepmd.pt.model.descriptor import (
     Descriptor,
     DescriptorBlock,
-    prod_env_mat_se_a,
+    prod_env_mat,
 )
 from deepmd.pt.utils import (
     env,
@@ -23,7 +23,7 @@
     RESERVED_PRECISON_DICT,
 )
 from deepmd.pt.utils.env_mat_stat import (
-    EnvMatStatSeA,
+    EnvMatStatSe,
 )
 from deepmd.utils.env_mat_stat import (
     StatItem,
@@ -391,7 +391,7 @@ def __getitem__(self, key):
 
     def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None):
         """Update mean and stddev for descriptor elements."""
-        env_mat_stat = EnvMatStatSeA(self)
+        env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
         env_mat_stat.load_or_compute_stats(merged, path)
@@ -435,7 +435,7 @@ def forward(
         del extended_atype_embd, mapping
         nloc = nlist.shape[1]
         atype = extended_atype[:, :nloc]
-        dmatrix, diff, sw = prod_env_mat_se_a(
+        dmatrix, diff, sw = prod_env_mat(
             extended_coord,
             nlist,
             atype,
diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py
index 4a7469a804..0b32bd9341 100644
--- a/deepmd/pt/model/descriptor/se_atten.py
+++ b/deepmd/pt/model/descriptor/se_atten.py
@@ -12,7 +12,7 @@
     DescriptorBlock,
 )
 from deepmd.pt.model.descriptor.env_mat import (
-    prod_env_mat_se_a,
+    prod_env_mat,
 )
 from deepmd.pt.model.network.network import (
     NeighborWiseAttention,
@@ -22,7 +22,7 @@
     env,
 )
 from deepmd.pt.utils.env_mat_stat import (
-    EnvMatStatSeA,
+    EnvMatStatSe,
 )
 from deepmd.utils.env_mat_stat import (
     StatItem,
@@ -202,7 +202,7 @@ def dim_emb(self):
 
     def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None):
         """Update mean and stddev for descriptor elements."""
-        env_mat_stat = EnvMatStatSeA(self)
+        env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
         env_mat_stat.load_or_compute_stats(merged, path)
@@ -247,7 +247,7 @@ def forward(
         atype = extended_atype[:, :nloc]
         nb = nframes
         nall = extended_coord.view(nb, -1, 3).shape[1]
-        dmatrix, diff, sw = prod_env_mat_se_a(
+        dmatrix, diff, sw = prod_env_mat(
             extended_coord,
             nlist,
             atype,
diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index d7c16bbada..0ddbfd5c63 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -12,7 +12,7 @@
 from deepmd.dpmodel.utils import EnvMat as DPEnvMat
 from deepmd.pt.model.descriptor import (
     Descriptor,
-    prod_env_mat_se_r,
+    prod_env_mat,
 )
 from deepmd.pt.model.network.mlp import (
     EmbeddingNet,
@@ -29,7 +29,7 @@
     RESERVED_PRECISON_DICT,
 )
 from deepmd.pt.utils.env_mat_stat import (
-    EnvMatStatSeR,
+    EnvMatStatSe,
 )
 from deepmd.pt.utils.exclude_mask import (
     PairExcludeMask,
@@ -89,29 +89,20 @@ def __init__(
         self.filter_layers_old = None
         self.filter_layers = None
 
-        if self.old_impl:
-            filter_layers = []
-            # TODO: remove
-            start_index = 0
-            for type_i in range(self.ntypes):
-                one = TypeFilter(start_index, sel[type_i], self.filter_neuron)
-                filter_layers.append(one)
-                start_index += sel[type_i]
-            self.filter_layers_old = torch.nn.ModuleList(filter_layers)
-        else:
-            filter_layers = NetworkCollection(
-                ndim=1, ntypes=len(sel), network_type="embedding_network"
+        
+        filter_layers = NetworkCollection(
+            ndim=1, ntypes=len(sel), network_type="embedding_network"
+        )
+        # TODO: ndim=2 if type_one_side=False
+        for ii in range(self.ntypes):
+            filter_layers[(ii,)] = EmbeddingNet(
+                1,
+                self.filter_neuron,
+                activation_function=self.activation_function,
+                precision=self.precision,
+                resnet_dt=self.resnet_dt,
             )
-            # TODO: ndim=2 if type_one_side=False
-            for ii in range(self.ntypes):
-                filter_layers[(ii,)] = EmbeddingNet(
-                    1,
-                    self.filter_neuron,
-                    activation_function=self.activation_function,
-                    precision=self.precision,
-                    resnet_dt=self.resnet_dt,
-                )
-            self.filter_layers = filter_layers
+        self.filter_layers = filter_layers
         self.stats = None
 
     def get_rcut(self) -> float:
@@ -156,7 +147,7 @@ def mixed_types(self) -> bool:
 
     def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None):
         """Update mean and stddev for descriptor elements."""
-        env_mat_stat = EnvMatStatSeR(self)
+        env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
         env_mat_stat.load_or_compute_stats(merged, path)
@@ -252,7 +243,7 @@ def forward(
         del mapping
         nloc = nlist.shape[1]
         atype = atype_ext[:, :nloc]
-        dmatrix, diff, sw = prod_env_mat_se_r(
+        dmatrix, diff, sw = prod_env_mat(
             coord_ext,
             nlist,
             atype,
@@ -260,6 +251,7 @@ def forward(
             self.stddev,
             self.rcut,
             self.rcut_smth,
+            True,
         )
 
         assert self.filter_layers is not None
diff --git a/deepmd/pt/utils/env_mat_stat.py b/deepmd/pt/utils/env_mat_stat.py
index 757e954018..7f63f0a99e 100644
--- a/deepmd/pt/utils/env_mat_stat.py
+++ b/deepmd/pt/utils/env_mat_stat.py
@@ -13,8 +13,7 @@
     get_hash,
 )
 from deepmd.pt.model.descriptor.env_mat import (
-    prod_env_mat_se_a,
-    prod_env_mat_se_r,
+    prod_env_mat,
 )
 from deepmd.pt.utils import (
     env,
@@ -57,8 +56,8 @@ def compute_stat(self, env_mat: Dict[str, torch.Tensor]) -> Dict[str, StatItem]:
         return stats
 
 
-class EnvMatStatSeA(EnvMatStat):
-    """Environmental matrix statistics for the se_a environemntal matrix.
+class EnvMatStatSe(EnvMatStat):
+    """Environmental matrix statistics for the se_a/se_r environemntal matrix.
 
     Parameters
     ----------
@@ -69,6 +68,7 @@ class EnvMatStatSeA(EnvMatStat):
     def __init__(self, descriptor: "DescriptorBlock"):
         super().__init__()
         self.descriptor = descriptor
+        self.last_dim = self.descriptor.ndescrpt// self.descriptor.nnei # se_r=1, se_a=4
 
     def iter(
         self, data: List[Dict[str, torch.Tensor]]
@@ -88,14 +88,14 @@ def iter(
         zero_mean = torch.zeros(
             self.descriptor.get_ntypes(),
             self.descriptor.get_nsel(),
-            4,
+            self.last_dim,
             dtype=env.GLOBAL_PT_FLOAT_PRECISION,
             device=env.DEVICE,
         )
         one_stddev = torch.ones(
             self.descriptor.get_ntypes(),
             self.descriptor.get_nsel(),
-            4,
+            self.last_dim,
             dtype=env.GLOBAL_PT_FLOAT_PRECISION,
             device=env.DEVICE,
         )
@@ -119,7 +119,7 @@ def iter(
                 mixed_types=self.descriptor.mixed_types(),
                 box=box,
             )
-            env_mat, _, _ = prod_env_mat_se_a(
+            env_mat, _, _ = prod_env_mat(
                 extended_coord,
                 nlist,
                 atype,
@@ -132,7 +132,7 @@ def iter(
             # reshape to nframes * nloc at the atom level,
             # so nframes/mixed_type do not matter
             env_mat = env_mat.view(
-                coord.shape[0] * coord.shape[1], self.descriptor.get_nsel(), 4
+                coord.shape[0] * coord.shape[1], self.descriptor.get_nsel(), self.last_dim
             )
             atype = atype.view(coord.shape[0] * coord.shape[1])
             # (1, nloc) eq (ntypes, 1), so broadcast is possible
@@ -145,10 +145,11 @@ def iter(
             )
             for type_i in range(self.descriptor.get_ntypes()):
                 dd = env_mat[type_idx[type_i]]
-                dd = dd.reshape([-1, 4])  # typen_atoms * nnei, 4
+                dd = dd.reshape([-1, self.last_dim])  # typen_atoms * nnei, 4
                 env_mats = {}
                 env_mats[f"r_{type_i}"] = dd[:, :1]
-                env_mats[f"a_{type_i}"] = dd[:, 1:]
+                if self.last_dim == 4:
+                    env_mats[f"a_{type_i}"] = dd[:, 1:]
                 yield self.compute_stat(env_mats)
 
     def get_hash(self) -> str:
@@ -159,9 +160,10 @@ def get_hash(self) -> str:
         str
             The hash of the environment matrix.
         """
+        dscpt_type = "se_a" if self.last_dim == 4 else "se_r"
         return get_hash(
             {
-                "type": "se_a",
+                "type": dscpt_type,
                 "ntypes": self.descriptor.get_ntypes(),
                 "rcut": round(self.descriptor.get_rcut(), 2),
                 "rcut_smth": round(self.descriptor.rcut_smth, 2),
@@ -177,151 +179,30 @@ def __call__(self):
 
         all_davg = []
         all_dstd = []
+        
         for type_i in range(self.descriptor.get_ntypes()):
-            davgunit = [[avgs[f"r_{type_i}"], 0, 0, 0]]
-            dstdunit = [
-                [
-                    stds[f"r_{type_i}"],
-                    stds[f"a_{type_i}"],
-                    stds[f"a_{type_i}"],
-                    stds[f"a_{type_i}"],
+            if self.last_dim == 4:
+                davgunit = [[avgs[f"r_{type_i}"], 0, 0, 0]]
+                dstdunit = [
+                    [
+                        stds[f"r_{type_i}"],
+                        stds[f"a_{type_i}"],
+                        stds[f"a_{type_i}"],
+                        stds[f"a_{type_i}"],
+                    ]
+                ]
+            elif self.last_dim == 1:
+                davgunit = [[avgs[f"r_{type_i}"]]]
+                dstdunit = [
+                    [
+                        stds[f"r_{type_i}"],
+                    ]
                 ]
-            ]
             davg = np.tile(davgunit, [self.descriptor.get_nsel(), 1])
             dstd = np.tile(dstdunit, [self.descriptor.get_nsel(), 1])
             all_davg.append(davg)
             all_dstd.append(dstd)
-        mean = np.stack(all_davg)
-        stddev = np.stack(all_dstd)
-        return mean, stddev
-
-
-class EnvMatStatSeR(EnvMatStat):
-    """Environmental matrix statistics for the se_r environemntal matrix.
-
-    Parameters
-    ----------
-    descriptor : DescriptorBlock
-        The descriptor of the model.
-    """
-
-    def __init__(self, descriptor: "DescriptorBlock"):
-        super().__init__()
-        self.descriptor = descriptor
-
-    def iter(
-        self, data: List[Dict[str, torch.Tensor]]
-    ) -> Iterator[Dict[str, StatItem]]:
-        """Get the iterator of the environment matrix.
-
-        Parameters
-        ----------
-        data : List[Dict[str, torch.Tensor]]
-            The environment matrix.
-
-        Yields
-        ------
-        Dict[str, StatItem]
-            The statistics of the environment matrix.
-        """
-        zero_mean = torch.zeros(
-            self.descriptor.get_ntypes(),
-            self.descriptor.get_nsel(),
-            1,
-            dtype=env.GLOBAL_PT_FLOAT_PRECISION,
-            device=env.DEVICE,
-        )
-        one_stddev = torch.ones(
-            self.descriptor.get_ntypes(),
-            self.descriptor.get_nsel(),
-            1,
-            dtype=env.GLOBAL_PT_FLOAT_PRECISION,
-            device=env.DEVICE,
-        )
-        for system in data:
-            coord, atype, box, natoms = (
-                system["coord"],
-                system["atype"],
-                system["box"],
-                system["natoms"],
-            )
-            (
-                extended_coord,
-                extended_atype,
-                mapping,
-                nlist,
-            ) = extend_input_and_build_neighbor_list(
-                coord,
-                atype,
-                self.descriptor.get_rcut(),
-                self.descriptor.get_sel(),
-                mixed_types=self.descriptor.mixed_types(),
-                box=box,
-            )
-            env_mat, _, _ = prod_env_mat_se_r(
-                extended_coord,
-                nlist,
-                atype,
-                zero_mean,
-                one_stddev,
-                self.descriptor.get_rcut(),
-                # TODO: export rcut_smth from DescriptorBlock
-                self.descriptor.rcut_smth,
-            )
-            # reshape to nframes * nloc at the atom level,
-            # so nframes/mixed_type do not matter
-            env_mat = env_mat.view(
-                coord.shape[0] * coord.shape[1], self.descriptor.get_nsel(), 1
-            )
-            atype = atype.view(coord.shape[0] * coord.shape[1])
-            # (1, nloc) eq (ntypes, 1), so broadcast is possible
-            # shape: (ntypes, nloc)
-            type_idx = torch.eq(
-                atype.view(1, -1),
-                torch.arange(
-                    self.descriptor.get_ntypes(), device=env.DEVICE, dtype=torch.int32
-                ).view(-1, 1),
-            )
-            for type_i in range(self.descriptor.get_ntypes()):
-                dd = env_mat[type_idx[type_i]]
-                dd = dd.reshape([-1, 1])  # typen_atoms * nnei, 4
-                env_mats = {}
-                env_mats[f"r_{type_i}"] = dd[:, :1]
-                yield self.compute_stat(env_mats)
 
-    def get_hash(self) -> str:
-        """Get the hash of the environment matrix.
-
-        Returns
-        -------
-        str
-            The hash of the environment matrix.
-        """
-        return get_hash(
-            {
-                "type": "se_r",
-                "ntypes": self.descriptor.get_ntypes(),
-                "rcut": round(self.descriptor.get_rcut(), 2),
-                "rcut_smth": round(self.descriptor.rcut_smth, 2),
-                "nsel": self.descriptor.get_nsel(),
-                "sel": self.descriptor.get_sel(),
-                "mixed_types": self.descriptor.mixed_types(),
-            }
-        )
-
-    def __call__(self):
-        avgs = self.get_avg()
-        stds = self.get_std()
-
-        all_davg = []
-        all_dstd = []
-        for type_i in range(self.descriptor.get_ntypes()):
-            davgunit = [[avgs[f"r_{type_i}"]]]
-            dstdunit = [[stds[f"r_{type_i}"]]]
-            davg = np.tile(davgunit, [self.descriptor.get_nsel(), 1])
-            dstd = np.tile(dstdunit, [self.descriptor.get_nsel(), 1])
-            all_davg.append(davg)
-            all_dstd.append(dstd)
         mean = np.stack(all_davg)
         stddev = np.stack(all_dstd)
-        return mean, stddev
+        return mean, stddev
\ No newline at end of file
diff --git a/source/tests/pt/model/test_descriptor.py b/source/tests/pt/model/test_descriptor.py
index 529a83ac6d..ffad27201a 100644
--- a/source/tests/pt/model/test_descriptor.py
+++ b/source/tests/pt/model/test_descriptor.py
@@ -14,7 +14,7 @@
 )
 
 from deepmd.pt.model.descriptor import (
-    prod_env_mat_se_a,
+    prod_env_mat,
 )
 from deepmd.pt.utils import (
     dp_random,
@@ -155,7 +155,7 @@ def test_consistency(self):
             mixed_types=False,
             box=self.pt_batch["box"].to(env.DEVICE),
         )
-        my_d, _, _ = prod_env_mat_se_a(
+        my_d, _, _ = prod_env_mat(
             extended_coord,
             nlist,
             atype,
diff --git a/source/tests/pt/model/test_env_mat.py b/source/tests/pt/model/test_env_mat.py
index ee262e7ee5..fee3fd6fea 100644
--- a/source/tests/pt/model/test_env_mat.py
+++ b/source/tests/pt/model/test_env_mat.py
@@ -8,7 +8,7 @@
     EnvMat,
 )
 from deepmd.pt.model.descriptor.env_mat import (
-    prod_env_mat_se_a,
+    prod_env_mat,
 )
 from deepmd.pt.utils import (
     env,
@@ -99,7 +99,7 @@ def test_consistency(
         dstd = 0.1 + np.abs(dstd)
         em0 = EnvMat(self.rcut, self.rcut_smth)
         mm0, ww0 = em0.call(self.coord_ext, self.atype_ext, self.nlist, davg, dstd)
-        mm1, _, ww1 = prod_env_mat_se_a(
+        mm1, _, ww1 = prod_env_mat(
             torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE),
             torch.tensor(self.nlist, dtype=int, device=env.DEVICE),
             torch.tensor(self.atype_ext[:, :nloc], dtype=int, device=env.DEVICE),

From 50cdfe002034025a260344b2c3e7cc1fdf3c03db Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 26 Feb 2024 06:53:37 +0000
Subject: [PATCH 09/20] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt/model/descriptor/__init__.py |  1 -
 deepmd/pt/model/descriptor/env_mat.py  | 15 ++++++++++++---
 deepmd/pt/model/descriptor/se_r.py     |  4 ----
 deepmd/pt/utils/env_mat_stat.py        | 12 ++++++++----
 4 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/deepmd/pt/model/descriptor/__init__.py b/deepmd/pt/model/descriptor/__init__.py
index 2c55346f94..784db59be0 100644
--- a/deepmd/pt/model/descriptor/__init__.py
+++ b/deepmd/pt/model/descriptor/__init__.py
@@ -13,7 +13,6 @@
 )
 from .env_mat import (
     prod_env_mat,
-
 )
 from .gaussian_lcc import (
     DescrptGaussianLcc,
diff --git a/deepmd/pt/model/descriptor/env_mat.py b/deepmd/pt/model/descriptor/env_mat.py
index c3a1d38726..4e6ffb7785 100644
--- a/deepmd/pt/model/descriptor/env_mat.py
+++ b/deepmd/pt/model/descriptor/env_mat.py
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from typing import Optional
+
 import torch
 
 from deepmd.pt.utils.preprocess import (
@@ -7,7 +7,9 @@
 )
 
 
-def _make_env_mat(nlist, coord, rcut: float, ruct_smth: float, radial_only: bool=False):
+def _make_env_mat(
+    nlist, coord, rcut: float, ruct_smth: float, radial_only: bool = False
+):
     """Make smooth environment matrix."""
     bsz, natoms, nnei = nlist.shape
     coord = coord.view(bsz, -1, 3)
@@ -35,7 +37,14 @@ def _make_env_mat(nlist, coord, rcut: float, ruct_smth: float, radial_only: bool
 
 
 def prod_env_mat(
-    extended_coord, nlist, atype, mean, stddev, rcut: float, rcut_smth: float, radial_only: bool=False
+    extended_coord,
+    nlist,
+    atype,
+    mean,
+    stddev,
+    rcut: float,
+    rcut_smth: float,
+    radial_only: bool = False,
 ):
     """Generate smooth environment matrix from atom coordinates and other context.
 
diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index 0ddbfd5c63..4ccb6f8842 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -18,9 +18,6 @@
     EmbeddingNet,
     NetworkCollection,
 )
-from deepmd.pt.model.network.network import (
-    TypeFilter,
-)
 from deepmd.pt.utils import (
     env,
 )
@@ -89,7 +86,6 @@ def __init__(
         self.filter_layers_old = None
         self.filter_layers = None
 
-        
         filter_layers = NetworkCollection(
             ndim=1, ntypes=len(sel), network_type="embedding_network"
         )
diff --git a/deepmd/pt/utils/env_mat_stat.py b/deepmd/pt/utils/env_mat_stat.py
index 7f63f0a99e..3af03bda97 100644
--- a/deepmd/pt/utils/env_mat_stat.py
+++ b/deepmd/pt/utils/env_mat_stat.py
@@ -68,7 +68,9 @@ class EnvMatStatSe(EnvMatStat):
     def __init__(self, descriptor: "DescriptorBlock"):
         super().__init__()
         self.descriptor = descriptor
-        self.last_dim = self.descriptor.ndescrpt// self.descriptor.nnei # se_r=1, se_a=4
+        self.last_dim = (
+            self.descriptor.ndescrpt // self.descriptor.nnei
+        )  # se_r=1, se_a=4
 
     def iter(
         self, data: List[Dict[str, torch.Tensor]]
@@ -132,7 +134,9 @@ def iter(
             # reshape to nframes * nloc at the atom level,
             # so nframes/mixed_type do not matter
             env_mat = env_mat.view(
-                coord.shape[0] * coord.shape[1], self.descriptor.get_nsel(), self.last_dim
+                coord.shape[0] * coord.shape[1],
+                self.descriptor.get_nsel(),
+                self.last_dim,
             )
             atype = atype.view(coord.shape[0] * coord.shape[1])
             # (1, nloc) eq (ntypes, 1), so broadcast is possible
@@ -179,7 +183,7 @@ def __call__(self):
 
         all_davg = []
         all_dstd = []
-        
+
         for type_i in range(self.descriptor.get_ntypes()):
             if self.last_dim == 4:
                 davgunit = [[avgs[f"r_{type_i}"], 0, 0, 0]]
@@ -205,4 +209,4 @@ def __call__(self):
 
         mean = np.stack(all_davg)
         stddev = np.stack(all_dstd)
-        return mean, stddev
\ No newline at end of file
+        return mean, stddev

From 84d61da944267f7d11baccf6da43464607115fe6 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Mon, 26 Feb 2024 15:39:34 +0800
Subject: [PATCH 10/20] feat: add numpy impl

---
 deepmd/dpmodel/descriptor/__init__.py         |   4 +
 deepmd/dpmodel/descriptor/se_r.py             | 351 ++++++++++++++++++
 deepmd/dpmodel/utils/env_mat.py               |  25 +-
 deepmd/pt/model/descriptor/repformers.py      |   1 +
 deepmd/pt/model/descriptor/se_r.py            |   1 +
 source/tests/pt/model/test_descriptor_se_r.py |  32 +-
 6 files changed, 391 insertions(+), 23 deletions(-)
 create mode 100644 deepmd/dpmodel/descriptor/se_r.py

diff --git a/deepmd/dpmodel/descriptor/__init__.py b/deepmd/dpmodel/descriptor/__init__.py
index 5eca26acc5..9c085568ba 100644
--- a/deepmd/dpmodel/descriptor/__init__.py
+++ b/deepmd/dpmodel/descriptor/__init__.py
@@ -5,8 +5,12 @@
 from .se_e2_a import (
     DescrptSeA,
 )
+from .se_r import (
+    DescrptSeR
+)
 
 __all__ = [
     "DescrptSeA",
+    "DescrptSeR",
     "make_base_descriptor",
 ]
diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py
new file mode 100644
index 0000000000..fffdd06194
--- /dev/null
+++ b/deepmd/dpmodel/descriptor/se_r.py
@@ -0,0 +1,351 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import numpy as np
+
+from deepmd.utils.path import (
+    DPPath,
+)
+
+try:
+    from deepmd._version import version as __version__
+except ImportError:
+    __version__ = "unknown"
+
+import copy
+from typing import (
+    Any,
+    List,
+    Optional,
+)
+
+from deepmd.dpmodel import (
+    DEFAULT_PRECISION,
+    PRECISION_DICT,
+    NativeOP,
+)
+from deepmd.dpmodel.utils import (
+    EmbeddingNet,
+    EnvMat,
+    NetworkCollection,
+    PairExcludeMask,
+)
+
+from .base_descriptor import (
+    BaseDescriptor,
+)
+
+
+@BaseDescriptor.register("se_e2_r")
+@BaseDescriptor.register("se_r")
+class DescrptSeR(NativeOP, BaseDescriptor):
+    r"""DeepPot-SE constructed from all information (both angular and radial) of
+    atomic configurations. The embedding takes the distance between atoms as input.
+
+    The descriptor :math:`\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}` is given by [1]_
+
+    .. math::
+        \mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_<
+
+    where :math:`\mathcal{R}^i \in \mathbb{R}^{N \times 4}` is the coordinate
+    matrix, and each row of :math:`\mathcal{R}^i` can be constructed as follows
+
+    .. math::
+        (\mathcal{R}^i)_j = [
+        \begin{array}{c}
+            s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}}
+        \end{array}
+        ]
+
+    where :math:`\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})` is
+    the relative coordinate and :math:`r_{ji}=\lVert \mathbf{R}_{ji} \lVert` is its norm.
+    The switching function :math:`s(r)` is defined as:
+
+    .. math::
+        s(r)=
+        \begin{cases}
+        \frac{1}{r}, & r<r_s \\
+        \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\
+        0, & r \geq r_c
+        \end{cases}
+
+    Each row of the embedding matrix  :math:`\mathcal{G}^i \in \mathbb{R}^{N \times M_1}` consists of outputs
+    of a embedding network :math:`\mathcal{N}` of :math:`s(r_{ji})`:
+
+    .. math::
+        (\mathcal{G}^i)_j = \mathcal{N}(s(r_{ji}))
+
+    :math:`\mathcal{G}^i_< \in \mathbb{R}^{N \times M_2}` takes first :math:`M_2` columns of
+    :math:`\mathcal{G}^i`. The equation of embedding network :math:`\mathcal{N}` can be found at
+    :meth:`deepmd.tf.utils.network.embedding_net`.
+
+    Parameters
+    ----------
+    rcut
+            The cut-off radius :math:`r_c`
+    rcut_smth
+            From where the environment matrix should be smoothed :math:`r_s`
+    sel : list[str]
+            sel[i] specifies the maxmum number of type i atoms in the cut-off radius
+    neuron : list[int]
+            Number of neurons in each hidden layers of the embedding net :math:`\mathcal{N}`
+    axis_neuron
+            Number of the axis neuron :math:`M_2` (number of columns of the sub-matrix of the embedding matrix)
+    resnet_dt
+            Time-step `dt` in the resnet construction:
+            y = x + dt * \phi (Wx + b)
+    trainable
+            If the weights of embedding net are trainable.
+    type_one_side
+            Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
+    exclude_types : List[List[int]]
+            The excluded pairs of types which have no interaction with each other.
+            For example, `[[0, 1]]` means no interaction between type 0 and type 1.
+    set_davg_zero
+            Set the shift of embedding net input to zero.
+    activation_function
+            The activation function in the embedding net. Supported options are |ACTIVATION_FN|
+    precision
+            The precision of the embedding net parameters. Supported options are |PRECISION|
+    multi_task
+            If the model has multi fitting nets to train.
+    spin
+            The deepspin object.
+
+    Limitations
+    -----------
+    The currently implementation does not support the following features
+
+    1. type_one_side == False
+    2. exclude_types != []
+    3. spin is not None
+
+    References
+    ----------
+    .. [1] Linfeng Zhang, Jiequn Han, Han Wang, Wissam A. Saidi, Roberto Car, and E. Weinan. 2018.
+       End-to-end symmetry preserving inter-atomic potential energy model for finite and extended
+       systems. In Proceedings of the 32nd International Conference on Neural Information Processing
+       Systems (NIPS'18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.
+    """
+
+    def __init__(
+        self,
+        rcut: float,
+        rcut_smth: float,
+        sel: List[str],
+        neuron: List[int] = [24, 48, 96],
+        resnet_dt: bool = False,
+        trainable: bool = True,
+        type_one_side: bool = True,
+        exclude_types: List[List[int]] = [],
+        set_davg_zero: bool = False,
+        activation_function: str = "tanh",
+        precision: str = DEFAULT_PRECISION,
+        spin: Optional[Any] = None,
+        # consistent with argcheck, not used though
+        seed: Optional[int] = None,
+    ) -> None:
+        ## seed, uniform_seed, multi_task, not included.
+        if not type_one_side:
+            raise NotImplementedError("type_one_side == False not implemented")
+        if spin is not None:
+            raise NotImplementedError("spin is not implemented")
+
+        self.rcut = rcut
+        self.rcut_smth = rcut_smth
+        self.sel = sel
+        self.ntypes = len(self.sel)
+        self.neuron = neuron
+        self.resnet_dt = resnet_dt
+        self.trainable = trainable
+        self.type_one_side = type_one_side
+        self.exclude_types = exclude_types
+        self.set_davg_zero = set_davg_zero
+        self.activation_function = activation_function
+        self.precision = precision
+        self.spin = spin
+        self.emask = PairExcludeMask(self.ntypes, self.exclude_types)
+
+        in_dim = 1  # not considiering type embedding
+        self.embeddings = NetworkCollection(
+            ntypes=self.ntypes,
+            ndim=(1 if self.type_one_side else 2),
+            network_type="embedding_network",
+        )
+        if not self.type_one_side:
+            raise NotImplementedError("type_one_side == False not implemented")
+        for ii in range(self.ntypes):
+            self.embeddings[(ii,)] = EmbeddingNet(
+                in_dim,
+                self.neuron,
+                self.activation_function,
+                self.resnet_dt,
+                self.precision,
+            )
+        self.env_mat = EnvMat(self.rcut, self.rcut_smth)
+        self.nnei = np.sum(self.sel)
+        self.davg = np.zeros([self.ntypes, self.nnei, 1])
+        self.dstd = np.ones([self.ntypes, self.nnei, 1])
+        self.orig_sel = self.sel
+
+    def __setitem__(self, key, value):
+        if key in ("avg", "data_avg", "davg"):
+            self.davg = value
+        elif key in ("std", "data_std", "dstd"):
+            self.dstd = value
+        else:
+            raise KeyError(key)
+
+    def __getitem__(self, key):
+        if key in ("avg", "data_avg", "davg"):
+            return self.davg
+        elif key in ("std", "data_std", "dstd"):
+            return self.dstd
+        else:
+            raise KeyError(key)
+
+    @property
+    def dim_out(self):
+        """Returns the output dimension of this descriptor."""
+        return self.get_dim_out()
+
+    def get_dim_out(self):
+        """Returns the output dimension of this descriptor."""
+        return self.neuron[-1]
+
+    def get_dim_emb(self):
+        """Returns the embedding (g2) dimension of this descriptor."""
+        raise NotImplementedError
+
+    def get_rcut(self):
+        """Returns cutoff radius."""
+        return self.rcut
+
+    def get_sel(self):
+        """Returns cutoff radius."""
+        return self.sel
+
+    def mixed_types(self):
+        """Returns if the descriptor requires a neighbor list that distinguish different
+        atomic types or not.
+        """
+        return False
+
+    def get_ntypes(self) -> int:
+        """Returns the number of element types."""
+        return self.ntypes
+
+    def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None):
+        """Update mean and stddev for descriptor elements."""
+        raise NotImplementedError
+
+    def cal_g(
+        self,
+        ss,
+        ll,
+    ):
+        nf, nloc, nnei = ss.shape[0:3]
+        ss = ss.reshape(nf, nloc, nnei, 1)
+        # nf x nloc x nnei x ng
+        gg = self.embeddings[(ll,)].call(ss)
+        return gg
+
+    def call(
+        self,
+        coord_ext,
+        atype_ext,
+        nlist,
+        mapping: Optional[np.ndarray] = None,
+    ):
+        """Compute the descriptor.
+
+        Parameters
+        ----------
+        coord_ext
+            The extended coordinates of atoms. shape: nf x (nallx3)
+        atype_ext
+            The extended aotm types. shape: nf x nall
+        nlist
+            The neighbor list. shape: nf x nloc x nnei
+        mapping
+            The index mapping from extended to lcoal region. not used by this descriptor.
+
+        Returns
+        -------
+        descriptor
+            The descriptor. shape: nf x nloc x (ng x axis_neuron)
+        gr
+            The rotationally equivariant and permutationally invariant single particle
+            representation. shape: nf x nloc x ng x 3
+        g2
+            The rotationally invariant pair-partical representation.
+            this descriptor returns None
+        h2
+            The rotationally equivariant pair-partical representation.
+            this descriptor returns None
+        sw
+            The smooth switch function.
+        """
+        del mapping
+        # nf x nloc x nnei x 1
+        rr, ww = self.env_mat.call(coord_ext, atype_ext, nlist, self.davg, self.dstd,True)
+        nf, nloc, nnei, _ = rr.shape
+        sec = np.append([0], np.cumsum(self.sel))
+
+        ng = self.neuron[-1]
+        gr = np.zeros([nf, nloc, ng, 1])
+        exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)
+        for tt in range(self.ntypes):
+            mm = exclude_mask[:, :, sec[tt] : sec[tt + 1]]
+            tr = rr[:, :, sec[tt] : sec[tt + 1], :]
+            tr = tr * mm[:, :, :, None]
+            ss = tr[..., 0:1]
+            gg = self.cal_g(ss, tt)
+            # nf x nloc x ng x 1
+            gr += np.einsum("flni,flnj->flij", gg, tr)
+        
+        res_rescale = 1.0 / 5.0
+        res = gr * res_rescale
+        res = res.reshape(nf, nloc, -1)
+        return res, None, None, None, ww
+
+    def serialize(self) -> dict:
+        """Serialize the descriptor to dict."""
+        return {
+            "@class": "Descriptor",
+            "type": "se_r",
+            "rcut": self.rcut,
+            "rcut_smth": self.rcut_smth,
+            "sel": self.sel,
+            "neuron": self.neuron,
+            "resnet_dt": self.resnet_dt,
+            "trainable": self.trainable,
+            "type_one_side": self.type_one_side,
+            "exclude_types": self.exclude_types,
+            "set_davg_zero": self.set_davg_zero,
+            "activation_function": self.activation_function,
+            # make deterministic
+            "precision": np.dtype(PRECISION_DICT[self.precision]).name,
+            "spin": self.spin,
+            "env_mat": self.env_mat.serialize(),
+            "embeddings": self.embeddings.serialize(),
+            "@variables": {
+                "davg": self.davg,
+                "dstd": self.dstd,
+            },
+        }
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "DescrptSeR":
+        """Deserialize from dict."""
+        data = copy.deepcopy(data)
+        data.pop("@class", None)
+        data.pop("type", None)
+        variables = data.pop("@variables")
+        embeddings = data.pop("embeddings")
+        env_mat = data.pop("env_mat")
+        obj = cls(**data)
+
+        obj["davg"] = variables["davg"]
+        obj["dstd"] = variables["dstd"]
+        obj.embeddings = NetworkCollection.deserialize(embeddings)
+        obj.env_mat = EnvMat.deserialize(env_mat)
+        return obj
diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py
index 070b0e1549..129ae9c971 100644
--- a/deepmd/dpmodel/utils/env_mat.py
+++ b/deepmd/dpmodel/utils/env_mat.py
@@ -30,6 +30,7 @@ def _make_env_mat(
     coord,
     rcut: float,
     ruct_smth: float,
+    radial_only: bool=False,
 ):
     """Make smooth environment matrix."""
     nf, nloc, nnei = nlist.shape
@@ -54,8 +55,11 @@ def _make_env_mat(
     t1 = diff / length**2
     weight = compute_smooth_weight(length, ruct_smth, rcut)
     weight = weight * np.expand_dims(mask, -1)
-    env_mat_se_a = np.concatenate([t0, t1], axis=-1) * weight
-    return env_mat_se_a, diff * np.expand_dims(mask, -1), weight
+    if radial_only:
+        env_mat = t0 * weight
+    else:
+        env_mat = np.concatenate([t0, t1], axis=-1) * weight
+    return env_mat, diff * np.expand_dims(mask, -1), weight
 
 
 class EnvMat(NativeOP):
@@ -74,6 +78,7 @@ def call(
         nlist: np.ndarray,
         davg: Optional[np.ndarray] = None,
         dstd: Optional[np.ndarray] = None,
+        radial_only: bool = False,
     ) -> Union[np.ndarray, np.ndarray]:
         """Compute the environment matrix.
 
@@ -86,18 +91,23 @@ def call(
         atype_ext
             The extended aotm types. shape: nf x nall
         davg
-            The data avg. shape: nt x nnei x 4
+            The data avg. shape: nt x nnei x (4 or 1)
         dstd
-            The inverse of data std. shape: nt x nnei x 4
+            The inverse of data std. shape: nt x nnei x (4 or 1)
+        radial_only
+            Whether to only compute radial part of the environment matrix.
+            If True, the output will be of shape nf x nloc x nnei x 1.
+            Otherwise, the output will be of shape nf x nloc x nnei x 4.
+            Default: False.
 
         Returns
         -------
         env_mat
-            The environment matrix. shape: nf x nloc x nnei x 4
+            The environment matrix. shape: nf x nloc x nnei x (4 or 1)
         switch
             The value of switch function. shape: nf x nloc x nnei
         """
-        em, sw = self._call(nlist, coord_ext)
+        em, sw = self._call(nlist, coord_ext, radial_only)
         nf, nloc, nnei = nlist.shape
         atype = atype_ext[:, :nloc]
         if davg is not None:
@@ -110,8 +120,9 @@ def _call(
         self,
         nlist,
         coord_ext,
+        radial_only
     ):
-        em, diff, ww = _make_env_mat(nlist, coord_ext, self.rcut, self.rcut_smth)
+        em, diff, ww = _make_env_mat(nlist, coord_ext, self.rcut, self.rcut_smth, radial_only)
         return em, ww
 
     def serialize(
diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py
index dc4afcd07a..76df81d347 100644
--- a/deepmd/pt/model/descriptor/repformers.py
+++ b/deepmd/pt/model/descriptor/repformers.py
@@ -100,6 +100,7 @@ def __init__(
         self.nlayers = nlayers
         sel = [sel] if isinstance(sel, int) else sel
         self.nnei = sum(sel)
+        self.ndescrpt = self.nnei * 4 # use full descriptor.
         assert len(sel) == 1
         self.sel = sel
         self.sec = self.sel
diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index 4ccb6f8842..383f6ea25a 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -40,6 +40,7 @@
 
 
 @Descriptor.register("se_e2_r")
+@Descriptor.register("se_r")
 class DescrptSeR(Descriptor):
     def __init__(
         self,
diff --git a/source/tests/pt/model/test_descriptor_se_r.py b/source/tests/pt/model/test_descriptor_se_r.py
index 8220079638..c999f06863 100644
--- a/source/tests/pt/model/test_descriptor_se_r.py
+++ b/source/tests/pt/model/test_descriptor_se_r.py
@@ -5,7 +5,7 @@
 import numpy as np
 import torch
 
-# from deepmd.dpmodel.descriptor import DescrptSeR as DPDescrptSeR
+from deepmd.dpmodel.descriptor import DescrptSeR as DPDescrptSeR
 from deepmd.pt.model.descriptor.se_r import (
     DescrptSeR,
 )
@@ -68,7 +68,7 @@ def test_consistency(
             )
             # serialization
             dd1 = DescrptSeR.deserialize(dd0.serialize())
-            rd1, gr1, _, _, sw1 = dd1(
+            rd1, _, _, _, sw1 = dd1(
                 torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE),
                 torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE),
                 torch.tensor(self.nlist, dtype=int, device=env.DEVICE),
@@ -88,20 +88,20 @@ def test_consistency(
                 err_msg=err_msg,
             )
             # dp impl
-            # dd2 = DPDescrptSeR.deserialize(dd0.serialize())
-            # rd2, gr2, _, _, sw2 = dd2.call(
-            #     self.coord_ext,
-            #     self.atype_ext,
-            #     self.nlist,
-            # )
-            # for aa, bb in zip([rd1, gr1, sw1], [rd2, gr2, sw2]):
-            #     np.testing.assert_allclose(
-            #         aa.detach().cpu().numpy(),
-            #         bb,
-            #         rtol=rtol,
-            #         atol=atol,
-            #         err_msg=err_msg,
-            #     )
+            dd2 = DPDescrptSeR.deserialize(dd0.serialize())
+            rd2, _, _, _, sw2 = dd2.call(
+                self.coord_ext,
+                self.atype_ext,
+                self.nlist,
+            )
+            for aa, bb in zip([rd1, sw1], [rd2, sw2]):
+                np.testing.assert_allclose(
+                    aa.detach().cpu().numpy(),
+                    bb,
+                    rtol=rtol,
+                    atol=atol,
+                    err_msg=err_msg,
+                )
 
     def test_jit(
         self,

From 08a6988fbca0725ea147c148df804049f9db8bee Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 26 Feb 2024 07:40:35 +0000
Subject: [PATCH 11/20] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/dpmodel/descriptor/__init__.py    |  2 +-
 deepmd/dpmodel/descriptor/se_r.py        |  6 ++++--
 deepmd/dpmodel/utils/env_mat.py          | 13 +++++--------
 deepmd/pt/model/descriptor/repformers.py |  2 +-
 4 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/__init__.py b/deepmd/dpmodel/descriptor/__init__.py
index 9c085568ba..08f8eb5052 100644
--- a/deepmd/dpmodel/descriptor/__init__.py
+++ b/deepmd/dpmodel/descriptor/__init__.py
@@ -6,7 +6,7 @@
     DescrptSeA,
 )
 from .se_r import (
-    DescrptSeR
+    DescrptSeR,
 )
 
 __all__ = [
diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py
index fffdd06194..70b8655a24 100644
--- a/deepmd/dpmodel/descriptor/se_r.py
+++ b/deepmd/dpmodel/descriptor/se_r.py
@@ -286,7 +286,9 @@ def call(
         """
         del mapping
         # nf x nloc x nnei x 1
-        rr, ww = self.env_mat.call(coord_ext, atype_ext, nlist, self.davg, self.dstd,True)
+        rr, ww = self.env_mat.call(
+            coord_ext, atype_ext, nlist, self.davg, self.dstd, True
+        )
         nf, nloc, nnei, _ = rr.shape
         sec = np.append([0], np.cumsum(self.sel))
 
@@ -301,7 +303,7 @@ def call(
             gg = self.cal_g(ss, tt)
             # nf x nloc x ng x 1
             gr += np.einsum("flni,flnj->flij", gg, tr)
-        
+
         res_rescale = 1.0 / 5.0
         res = gr * res_rescale
         res = res.reshape(nf, nloc, -1)
diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py
index 129ae9c971..0e861d9f38 100644
--- a/deepmd/dpmodel/utils/env_mat.py
+++ b/deepmd/dpmodel/utils/env_mat.py
@@ -30,7 +30,7 @@ def _make_env_mat(
     coord,
     rcut: float,
     ruct_smth: float,
-    radial_only: bool=False,
+    radial_only: bool = False,
 ):
     """Make smooth environment matrix."""
     nf, nloc, nnei = nlist.shape
@@ -116,13 +116,10 @@ def call(
             em /= dstd[atype]
         return em, sw
 
-    def _call(
-        self,
-        nlist,
-        coord_ext,
-        radial_only
-    ):
-        em, diff, ww = _make_env_mat(nlist, coord_ext, self.rcut, self.rcut_smth, radial_only)
+    def _call(self, nlist, coord_ext, radial_only):
+        em, diff, ww = _make_env_mat(
+            nlist, coord_ext, self.rcut, self.rcut_smth, radial_only
+        )
         return em, ww
 
     def serialize(
diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py
index 76df81d347..ad523bcc2d 100644
--- a/deepmd/pt/model/descriptor/repformers.py
+++ b/deepmd/pt/model/descriptor/repformers.py
@@ -100,7 +100,7 @@ def __init__(
         self.nlayers = nlayers
         sel = [sel] if isinstance(sel, int) else sel
         self.nnei = sum(sel)
-        self.ndescrpt = self.nnei * 4 # use full descriptor.
+        self.ndescrpt = self.nnei * 4  # use full descriptor.
         assert len(sel) == 1
         self.sel = sel
         self.sec = self.sel

From 1f0fd9997233afdf45d6ebd1f68fb3b4f2bf35e7 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Tue, 27 Feb 2024 09:45:58 +0800
Subject: [PATCH 12/20] fix: UTs

---
 deepmd/tf/descriptor/se_r.py                  |  95 ++++++++-
 .../tests/consistent/descriptor/test_se_r.py  | 185 ++++++++++++++++++
 2 files changed, 279 insertions(+), 1 deletion(-)
 create mode 100644 source/tests/consistent/descriptor/test_se_r.py

diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py
index f790d0a8fb..8f9d429ba0 100644
--- a/deepmd/tf/descriptor/se_r.py
+++ b/deepmd/tf/descriptor/se_r.py
@@ -7,6 +7,10 @@
 
 import numpy as np
 
+from deepmd.dpmodel.utils.env_mat import (
+    EnvMat,
+)
+
 from deepmd.tf.common import (
     cast_precision,
     get_activation_func,
@@ -115,8 +119,9 @@ def __init__(
         self.seed_shift = embedding_net_rand_seed_shift(self.filter_neuron)
         self.trainable = trainable
         self.filter_activation_fn = get_activation_func(activation_function)
+        self.activation_function_name = activation_function
         self.filter_precision = get_precision(precision)
-        exclude_types = exclude_types
+        self.orig_exclude_types = exclude_types
         self.exclude_types = set()
         for tt in exclude_types:
             assert len(tt) == 2
@@ -698,3 +703,91 @@ def _filter_r(
             result = tf.reduce_mean(xyz_scatter, axis=1) * res_rescale
 
         return result
+
+    @classmethod
+    def deserialize(cls, data: dict, suffix: str = ""):
+        """Deserialize the model.
+
+        Parameters
+        ----------
+        data : dict
+            The serialized data
+
+        Returns
+        -------
+        Model
+            The deserialized model
+        """
+        if cls is not DescrptSeR:
+            raise NotImplementedError("Not implemented in class %s" % cls.__name__)
+        data = data.copy()
+        embedding_net_variables = cls.deserialize_network(
+            data.pop("embeddings"), suffix=suffix
+        )
+        data.pop("env_mat")
+        variables = data.pop("@variables")
+        descriptor = cls(**data)
+        descriptor.embedding_net_variables = embedding_net_variables
+        descriptor.davg = variables["davg"].reshape(
+            descriptor.ntypes, descriptor.ndescrpt
+        )
+        descriptor.dstd = variables["dstd"].reshape(
+            descriptor.ntypes, descriptor.ndescrpt
+        )
+        return descriptor
+
+    def serialize(self, suffix: str = "") -> dict:
+        """Serialize the model.
+
+        Parameters
+        ----------
+        suffix : str, optional
+            The suffix of the scope
+
+        Returns
+        -------
+        dict
+            The serialized data
+        """
+        if type(self) is not DescrptSeR:
+            raise NotImplementedError(
+                "Not implemented in class %s" % self.__class__.__name__
+            )
+        if self.embedding_net_variables is None:
+            raise RuntimeError("init_variables must be called before serialize")
+        if self.spin is not None:
+            raise NotImplementedError("spin is unsupported")
+        assert self.davg is not None
+        assert self.dstd is not None
+        # TODO: not sure how to handle type embedding - type embedding is not a model parameter,
+        # but instead a part of the input data. Maybe the interface should be refactored...
+        return {
+            "rcut": self.rcut,
+            "rcut_smth": self.rcut_smth,
+            "sel": self.sel_r,
+            "neuron": self.filter_neuron,
+            "resnet_dt": self.filter_resnet_dt,
+            "trainable": self.trainable,
+            "type_one_side": self.type_one_side,
+            "exclude_types": list(self.orig_exclude_types),
+            "set_davg_zero": self.set_davg_zero,
+            "activation_function": self.activation_function_name,
+            "precision": self.filter_precision.name,
+            "embeddings": self.serialize_network(
+                ntypes=self.ntypes,
+                ndim=(1 if self.type_one_side else 2),
+                in_dim=1,
+                neuron=self.filter_neuron,
+                activation_function=self.activation_function_name,
+                resnet_dt=self.filter_resnet_dt,
+                variables=self.embedding_net_variables,
+                excluded_types=self.exclude_types,
+                suffix=suffix,
+            ),
+            "env_mat": EnvMat(self.rcut, self.rcut_smth).serialize(),
+            "@variables": {
+                "davg": self.davg.reshape(self.ntypes, self.nnei_r, 1),
+                "dstd": self.dstd.reshape(self.ntypes, self.nnei_r, 1),
+            },
+            "spin": self.spin,
+        }
\ No newline at end of file
diff --git a/source/tests/consistent/descriptor/test_se_r.py b/source/tests/consistent/descriptor/test_se_r.py
new file mode 100644
index 0000000000..354ae1cc99
--- /dev/null
+++ b/source/tests/consistent/descriptor/test_se_r.py
@@ -0,0 +1,185 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+from typing import (
+    Any,
+    Tuple,
+)
+
+import numpy as np
+
+from deepmd.dpmodel.descriptor.se_r import DescrptSeR as DescrptSeRDP
+from deepmd.env import (
+    GLOBAL_NP_FLOAT_PRECISION,
+)
+
+from ..common import (
+    INSTALLED_PT,
+    INSTALLED_TF,
+    CommonTest,
+    parameterized,
+)
+from .common import (
+    DescriptorTest,
+)
+
+if INSTALLED_PT:
+    from deepmd.pt.model.descriptor.se_r import DescrptSeR as DescrptSeRPT
+else:
+    DescrptSeAPT = None
+if INSTALLED_TF:
+    from deepmd.tf.descriptor.se_r import DescrptSeR as DescrptSeRTF
+else:
+    DescrptSeATF = None
+from deepmd.utils.argcheck import (
+    descrpt_se_r_args,
+)
+
+
+@parameterized(
+    (True, False),  # resnet_dt
+    (True, False),  # type_one_side
+    ([], [[0, 1]]),  # excluded_types
+    ("float32", "float64"),  # precision
+)
+class TestSeA(CommonTest, DescriptorTest, unittest.TestCase):
+    @property
+    def data(self) -> dict:
+        (
+            resnet_dt,
+            type_one_side,
+            excluded_types,
+            precision,
+        ) = self.param
+        return {
+            "sel": [10, 10],
+            "rcut_smth": 5.80,
+            "rcut": 6.00,
+            "neuron": [6, 12, 24],
+            "resnet_dt": resnet_dt,
+            "type_one_side": type_one_side,
+            "exclude_types": excluded_types,
+            "precision": precision,
+            "seed": 1145141919810,
+        }
+
+    @property
+    def skip_pt(self) -> bool:
+        (
+            resnet_dt,
+            type_one_side,
+            excluded_types,
+            precision,
+        ) = self.param
+        return not type_one_side or CommonTest.skip_pt
+
+    @property
+    def skip_dp(self) -> bool:
+        (
+            resnet_dt,
+            type_one_side,
+            excluded_types,
+            precision,
+        ) = self.param
+        return not type_one_side or CommonTest.skip_dp
+
+    tf_class = DescrptSeRTF
+    dp_class = DescrptSeRDP
+    pt_class = DescrptSeRPT
+    args = descrpt_se_r_args()
+
+    def setUp(self):
+        CommonTest.setUp(self)
+
+        self.ntypes = 2
+        self.coords = np.array(
+            [
+                12.83,
+                2.56,
+                2.18,
+                12.09,
+                2.87,
+                2.74,
+                00.25,
+                3.32,
+                1.68,
+                3.36,
+                3.00,
+                1.81,
+                3.51,
+                2.51,
+                2.60,
+                4.27,
+                3.22,
+                1.56,
+            ],
+            dtype=GLOBAL_NP_FLOAT_PRECISION,
+        )
+        self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32)
+        self.box = np.array(
+            [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0],
+            dtype=GLOBAL_NP_FLOAT_PRECISION,
+        )
+        self.natoms = np.array([6, 6, 2, 4], dtype=np.int32)
+
+    def build_tf(self, obj: Any, suffix: str) -> Tuple[list, dict]:
+        return self.build_tf_descriptor(
+            obj,
+            self.natoms,
+            self.coords,
+            self.atype,
+            self.box,
+            suffix,
+        )
+
+    def eval_dp(self, dp_obj: Any) -> Any:
+        return self.eval_dp_descriptor(
+            dp_obj,
+            self.natoms,
+            self.coords,
+            self.atype,
+            self.box,
+        )
+
+    def eval_pt(self, pt_obj: Any) -> Any:
+        return self.eval_pt_descriptor(
+            pt_obj,
+            self.natoms,
+            self.coords,
+            self.atype,
+            self.box,
+        )
+
+    def extract_ret(self, ret: Any, backend) -> Tuple[np.ndarray, ...]:
+        return (ret[0],)
+
+    @property
+    def rtol(self) -> float:
+        """Relative tolerance for comparing the return value."""
+        (
+            resnet_dt,
+            type_one_side,
+            excluded_types,
+            precision,
+        ) = self.param
+        if precision == "float64":
+            return 1e-10
+        elif precision == "float32":
+            return 1e-4
+        else:
+            raise ValueError(f"Unknown precision: {precision}")
+
+    @property
+    def atol(self) -> float:
+        """Absolute tolerance for comparing the return value."""
+        (
+            resnet_dt,
+            type_one_side,
+            excluded_types,
+            precision,
+        ) = self.param
+        if precision == "float64":
+            return 1e-10
+        elif precision == "float32":
+            return 1e-4
+        else:
+            raise ValueError(f"Unknown precision: {precision}")

From c07e02c8beefffe1bb86f68e998ec4f295464e1c Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 27 Feb 2024 01:47:16 +0000
Subject: [PATCH 13/20] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/tf/descriptor/se_r.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py
index 8f9d429ba0..1436faea7f 100644
--- a/deepmd/tf/descriptor/se_r.py
+++ b/deepmd/tf/descriptor/se_r.py
@@ -10,7 +10,6 @@
 from deepmd.dpmodel.utils.env_mat import (
     EnvMat,
 )
-
 from deepmd.tf.common import (
     cast_precision,
     get_activation_func,
@@ -790,4 +789,4 @@ def serialize(self, suffix: str = "") -> dict:
                 "dstd": self.dstd.reshape(self.ntypes, self.nnei_r, 1),
             },
             "spin": self.spin,
-        }
\ No newline at end of file
+        }

From e5b074eec6e4f0e46798b698a4610259e12c98f3 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Tue, 27 Feb 2024 10:06:02 +0800
Subject: [PATCH 14/20] gix: match serialization

---
 deepmd/pt/model/descriptor/se_r.py | 2 ++
 deepmd/tf/descriptor/se_r.py       | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index 383f6ea25a..e7fcd2abfe 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -296,6 +296,8 @@ def set_stat_mean_and_stddev(
 
     def serialize(self) -> dict:
         return {
+            "@class": "Descriptor",
+            "type": "se_r",
             "rcut": self.rcut,
             "rcut_smth": self.rcut_smth,
             "sel": self.sel,
diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py
index 1436faea7f..1a12befdf0 100644
--- a/deepmd/tf/descriptor/se_r.py
+++ b/deepmd/tf/descriptor/se_r.py
@@ -761,6 +761,8 @@ def serialize(self, suffix: str = "") -> dict:
         # TODO: not sure how to handle type embedding - type embedding is not a model parameter,
         # but instead a part of the input data. Maybe the interface should be refactored...
         return {
+            "@class": "Descriptor",
+            "type": "se_r",
             "rcut": self.rcut,
             "rcut_smth": self.rcut_smth,
             "sel": self.sel_r,

From 3fe3ed4c9317eb5bdf54ae4081ed9ed107915c3e Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Tue, 27 Feb 2024 10:23:55 +0800
Subject: [PATCH 15/20] chore: refactor device

---
 deepmd/pt/model/descriptor/se_r.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index e7fcd2abfe..e965c5b6cb 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -257,7 +257,7 @@ def forward(
         nfnl = dmatrix.shape[0]
         # pre-allocate a shape to pass jit
         xyz_scatter = torch.zeros(
-            [nfnl, 1, self.filter_neuron[-1]], dtype=self.prec, device=env.DEVICE
+            [nfnl, 1, self.filter_neuron[-1]], dtype=self.prec, device=coord_ext.DEVICE
         )
 
         # nfnl x nnei

From 9d23e969e279b02241cbd0bf026003b7cd239867 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Tue, 27 Feb 2024 11:01:45 +0800
Subject: [PATCH 16/20] chore: refactor device

---
 deepmd/pt/model/descriptor/se_r.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index e965c5b6cb..b18619a7c2 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -257,7 +257,7 @@ def forward(
         nfnl = dmatrix.shape[0]
         # pre-allocate a shape to pass jit
         xyz_scatter = torch.zeros(
-            [nfnl, 1, self.filter_neuron[-1]], dtype=self.prec, device=coord_ext.DEVICE
+            [nfnl, 1, self.filter_neuron[-1]], dtype=self.prec, device=coord_ext.device
         )
 
         # nfnl x nnei

From c073241793be83d6a4352bae39f0f0fa874ee123 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Tue, 27 Feb 2024 12:21:44 +0800
Subject: [PATCH 17/20] fix: UTs

---
 deepmd/dpmodel/descriptor/se_r.py  | 12 ++++-----
 deepmd/pt/model/descriptor/se_r.py | 43 +++++++++---------------------
 2 files changed, 18 insertions(+), 37 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py
index 70b8655a24..bfddfecbb0 100644
--- a/deepmd/dpmodel/descriptor/se_r.py
+++ b/deepmd/dpmodel/descriptor/se_r.py
@@ -293,19 +293,19 @@ def call(
         sec = np.append([0], np.cumsum(self.sel))
 
         ng = self.neuron[-1]
-        gr = np.zeros([nf, nloc, ng, 1])
+        xyz_scatter = np.zeros([nf, nloc, ng])
         exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)
         for tt in range(self.ntypes):
             mm = exclude_mask[:, :, sec[tt] : sec[tt + 1]]
             tr = rr[:, :, sec[tt] : sec[tt + 1], :]
             tr = tr * mm[:, :, :, None]
-            ss = tr[..., 0:1]
-            gg = self.cal_g(ss, tt)
+            gg = self.cal_g(tr, tt)
+            gg = np.mean(gg, axis=2)
             # nf x nloc x ng x 1
-            gr += np.einsum("flni,flnj->flij", gg, tr)
+            xyz_scatter += gg
 
-        res_rescale = 1.0 / 5.0
-        res = gr * res_rescale
+        res_rescale = 1.0 / 10.0
+        res = xyz_scatter * res_rescale
         res = res.reshape(nf, nloc, -1)
         return res, None, None, None, ww
 
diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index b18619a7c2..c2179ea3dc 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -10,8 +10,10 @@
 import torch
 
 from deepmd.dpmodel.utils import EnvMat as DPEnvMat
+from .base_descriptor import (
+    BaseDescriptor
+)
 from deepmd.pt.model.descriptor import (
-    Descriptor,
     prod_env_mat,
 )
 from deepmd.pt.model.network.mlp import (
@@ -39,9 +41,9 @@
 )
 
 
-@Descriptor.register("se_e2_r")
-@Descriptor.register("se_r")
-class DescrptSeR(Descriptor):
+@BaseDescriptor.register("se_e2_r")
+@BaseDescriptor.register("se_r")
+class DescrptSeR(BaseDescriptor, torch.nn.Module):
     def __init__(
         self,
         rcut,
@@ -181,25 +183,6 @@ def __getitem__(self, key):
         else:
             raise KeyError(key)
 
-    @classmethod
-    def get_data_process_key(cls, config):
-        """
-        Get the keys for the data preprocess.
-        Usually need the information of rcut and sel.
-        TODO Need to be deprecated when the dataloader has been cleaned up.
-        """
-        descrpt_type = config["type"]
-        assert descrpt_type in ["se_e2_r"]
-        return {"sel": config["sel"], "rcut": config["rcut"]}
-
-    @property
-    def data_stat_key(self):
-        """
-        Get the keys for the data statistic of the descriptor.
-        Return a list of statistic names needed, such as "sumr", "sumr2" or "sumn".
-        """
-        return ["sumr", "sumn", "sumr2"]
-
     def forward(
         self,
         coord_ext: torch.Tensor,
@@ -266,17 +249,15 @@ def forward(
             # nfnl x nt
             mm = exclude_mask[:, self.sec[ii] : self.sec[ii + 1]]
             # nfnl x nt x 1
-            rr = dmatrix[:, self.sec[ii] : self.sec[ii + 1], :]
-            rr = rr * mm[:, :, None]
-            ss = rr[:, :, :1]
+            ss = dmatrix[:, self.sec[ii] : self.sec[ii + 1], :]
+            ss = ss * mm[:, :, None]
             # nfnl x nt x ng
             gg = ll.forward(ss)
-            # nfnl x 1 x ng
-            gr = torch.matmul(rr.permute(0, 2, 1), gg)
-            xyz_scatter += gr
+            gg = torch.mean(gg, dim=1).unsqueeze(1)
+            xyz_scatter += gg
 
-        res_rescale = 1.0 / 5.0
-        result = torch.mean(xyz_scatter, dim=1) * res_rescale
+        res_rescale = 1.0 / 10.0
+        result = xyz_scatter * res_rescale
         result = result.view(-1, nloc, self.filter_neuron[-1])
         return (
             result.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION),

From ce85c0f070954319496ceda34ed5b44000caae3e Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 27 Feb 2024 04:22:47 +0000
Subject: [PATCH 18/20] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt/model/descriptor/se_r.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index c2179ea3dc..655ac393a3 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -10,9 +10,6 @@
 import torch
 
 from deepmd.dpmodel.utils import EnvMat as DPEnvMat
-from .base_descriptor import (
-    BaseDescriptor
-)
 from deepmd.pt.model.descriptor import (
     prod_env_mat,
 )
@@ -40,6 +37,10 @@
     DPPath,
 )
 
+from .base_descriptor import (
+    BaseDescriptor,
+)
+
 
 @BaseDescriptor.register("se_e2_r")
 @BaseDescriptor.register("se_r")

From 72d7a3752690a81101ced7e9eafe675365f2e8f5 Mon Sep 17 00:00:00 2001
From: Anyang Peng <137014849+anyangml@users.noreply.github.com>
Date: Tue, 27 Feb 2024 16:43:57 +0800
Subject: [PATCH 19/20] fix: dtype

---
 deepmd/dpmodel/descriptor/se_e2_a.py | 12 ++++---
 deepmd/dpmodel/descriptor/se_r.py    | 50 ++++------------------------
 deepmd/pt/model/descriptor/se_a.py   |  4 +--
 deepmd/pt/model/descriptor/se_r.py   | 29 +---------------
 4 files changed, 15 insertions(+), 80 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py
index 97ab719c62..009d762f82 100644
--- a/deepmd/dpmodel/descriptor/se_e2_a.py
+++ b/deepmd/dpmodel/descriptor/se_e2_a.py
@@ -6,7 +6,9 @@
 from deepmd.utils.path import (
     DPPath,
 )
-
+from deepmd.env import (
+    GLOBAL_NP_FLOAT_PRECISION
+)
 try:
     from deepmd._version import version as __version__
 except ImportError:
@@ -183,8 +185,8 @@ def __init__(
             )
         self.env_mat = EnvMat(self.rcut, self.rcut_smth)
         self.nnei = np.sum(self.sel)
-        self.davg = np.zeros([self.ntypes, self.nnei, 4])
-        self.dstd = np.ones([self.ntypes, self.nnei, 4])
+        self.davg = np.zeros([self.ntypes, self.nnei, 4], dtype=PRECISION_DICT[self.precision])
+        self.dstd = np.ones([self.ntypes, self.nnei, 4], dtype=PRECISION_DICT[self.precision])
         self.orig_sel = self.sel
 
     def __setitem__(self, key, value):
@@ -292,7 +294,7 @@ def call(
         sec = np.append([0], np.cumsum(self.sel))
 
         ng = self.neuron[-1]
-        gr = np.zeros([nf * nloc, ng, 4])
+        gr = np.zeros([nf * nloc, ng, 4], dtype=PRECISION_DICT[self.precision])
         exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)
         # merge nf and nloc axis, so for type_one_side == False,
         # we don't require atype is the same in all frames
@@ -322,7 +324,7 @@ def call(
         # nf x nloc x ng x ng1
         grrg = np.einsum("flid,fljd->flij", gr, gr1)
         # nf x nloc x (ng x ng1)
-        grrg = grrg.reshape(nf, nloc, ng * self.axis_neuron)
+        grrg = grrg.reshape(nf, nloc, ng * self.axis_neuron).astype(GLOBAL_NP_FLOAT_PRECISION)
         return grrg, gr[..., 1:], None, None, ww
 
     def serialize(self) -> dict:
diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py
index bfddfecbb0..99728d9108 100644
--- a/deepmd/dpmodel/descriptor/se_r.py
+++ b/deepmd/dpmodel/descriptor/se_r.py
@@ -16,6 +16,7 @@
     List,
     Optional,
 )
+from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
 
 from deepmd.dpmodel import (
     DEFAULT_PRECISION,
@@ -37,45 +38,8 @@
 @BaseDescriptor.register("se_e2_r")
 @BaseDescriptor.register("se_r")
 class DescrptSeR(NativeOP, BaseDescriptor):
-    r"""DeepPot-SE constructed from all information (both angular and radial) of
-    atomic configurations. The embedding takes the distance between atoms as input.
+    r"""DeepPot-SE_R constructed from only the radial imformation of atomic configurations.
 
-    The descriptor :math:`\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}` is given by [1]_
-
-    .. math::
-        \mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_<
-
-    where :math:`\mathcal{R}^i \in \mathbb{R}^{N \times 4}` is the coordinate
-    matrix, and each row of :math:`\mathcal{R}^i` can be constructed as follows
-
-    .. math::
-        (\mathcal{R}^i)_j = [
-        \begin{array}{c}
-            s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}}
-        \end{array}
-        ]
-
-    where :math:`\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})` is
-    the relative coordinate and :math:`r_{ji}=\lVert \mathbf{R}_{ji} \lVert` is its norm.
-    The switching function :math:`s(r)` is defined as:
-
-    .. math::
-        s(r)=
-        \begin{cases}
-        \frac{1}{r}, & r<r_s \\
-        \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\
-        0, & r \geq r_c
-        \end{cases}
-
-    Each row of the embedding matrix  :math:`\mathcal{G}^i \in \mathbb{R}^{N \times M_1}` consists of outputs
-    of a embedding network :math:`\mathcal{N}` of :math:`s(r_{ji})`:
-
-    .. math::
-        (\mathcal{G}^i)_j = \mathcal{N}(s(r_{ji}))
-
-    :math:`\mathcal{G}^i_< \in \mathbb{R}^{N \times M_2}` takes first :math:`M_2` columns of
-    :math:`\mathcal{G}^i`. The equation of embedding network :math:`\mathcal{N}` can be found at
-    :meth:`deepmd.tf.utils.network.embedding_net`.
 
     Parameters
     ----------
@@ -87,8 +51,6 @@ class DescrptSeR(NativeOP, BaseDescriptor):
             sel[i] specifies the maxmum number of type i atoms in the cut-off radius
     neuron : list[int]
             Number of neurons in each hidden layers of the embedding net :math:`\mathcal{N}`
-    axis_neuron
-            Number of the axis neuron :math:`M_2` (number of columns of the sub-matrix of the embedding matrix)
     resnet_dt
             Time-step `dt` in the resnet construction:
             y = x + dt * \phi (Wx + b)
@@ -182,8 +144,8 @@ def __init__(
             )
         self.env_mat = EnvMat(self.rcut, self.rcut_smth)
         self.nnei = np.sum(self.sel)
-        self.davg = np.zeros([self.ntypes, self.nnei, 1])
-        self.dstd = np.ones([self.ntypes, self.nnei, 1])
+        self.davg = np.zeros([self.ntypes, self.nnei, 1], dtype=PRECISION_DICT[self.precision])
+        self.dstd = np.ones([self.ntypes, self.nnei, 1], dtype=PRECISION_DICT[self.precision])
         self.orig_sel = self.sel
 
     def __setitem__(self, key, value):
@@ -293,7 +255,7 @@ def call(
         sec = np.append([0], np.cumsum(self.sel))
 
         ng = self.neuron[-1]
-        xyz_scatter = np.zeros([nf, nloc, ng])
+        xyz_scatter = np.zeros([nf, nloc, ng], dtype=PRECISION_DICT[self.precision])
         exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)
         for tt in range(self.ntypes):
             mm = exclude_mask[:, :, sec[tt] : sec[tt + 1]]
@@ -306,7 +268,7 @@ def call(
 
         res_rescale = 1.0 / 10.0
         res = xyz_scatter * res_rescale
-        res = res.reshape(nf, nloc, -1)
+        res = res.reshape(nf, nloc, -1).astype(GLOBAL_NP_FLOAT_PRECISION)
         return res, None, None, None, ww
 
     def serialize(self) -> dict:
diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py
index ce8522ae5c..eec6e54962 100644
--- a/deepmd/pt/model/descriptor/se_a.py
+++ b/deepmd/pt/model/descriptor/se_a.py
@@ -393,9 +393,7 @@ def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None)
         if not self.set_davg_zero:
             self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
         self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
-        if not self.set_davg_zero:
-            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
-        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
+
 
     def get_stats(self) -> Dict[str, StatItem]:
         """Get the statistics of the descriptor."""
diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index 655ac393a3..22c0f55e55 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -156,9 +156,6 @@ def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None)
         if not self.set_davg_zero:
             self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
         self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
-        if not self.set_davg_zero:
-            self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
-        self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
 
     def get_stats(self) -> Dict[str, StatItem]:
         """Get the statistics of the descriptor."""
@@ -316,28 +313,4 @@ def t_cvt(xx):
         obj["davg"] = t_cvt(variables["davg"])
         obj["dstd"] = t_cvt(variables["dstd"])
         obj.filter_layers = NetworkCollection.deserialize(embeddings)
-        return obj
-
-
-def analyze_descrpt(matrix, ndescrpt, natoms):
-    """Collect avg, square avg and count of descriptors in a batch."""
-    ntypes = natoms.shape[1] - 2
-    start_index = 0
-    sysr = []
-    sysn = []
-    sysr2 = []
-    for type_i in range(ntypes):
-        end_index = start_index + natoms[0, 2 + type_i]
-        dd = matrix[:, start_index:end_index]  # all descriptors for this element
-        start_index = end_index
-        dd = np.reshape(
-            dd, [-1, 1]
-        )  # Shape is [nframes*natoms[2+type_id]*self.nnei, 1]
-        ddr = dd[:, :1]
-        sumr = np.sum(ddr)
-        sumn = dd.shape[0]  # Value is nframes*natoms[2+type_id]*self.nnei
-        sumr2 = np.sum(np.multiply(ddr, ddr))
-        sysr.append(sumr)
-        sysn.append(sumn)
-        sysr2.append(sumr2)
-    return sysr, sysr2, sysn
+        return obj
\ No newline at end of file

From 54375bbeb70476b009e8c118def1b8e4dbbb727d Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 27 Feb 2024 08:44:26 +0000
Subject: [PATCH 20/20] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/dpmodel/descriptor/se_e2_a.py | 19 +++++++++++++------
 deepmd/dpmodel/descriptor/se_r.py    | 12 +++++++++---
 deepmd/pt/model/descriptor/se_a.py   |  1 -
 deepmd/pt/model/descriptor/se_r.py   |  2 +-
 4 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py
index 009d762f82..a28215c35a 100644
--- a/deepmd/dpmodel/descriptor/se_e2_a.py
+++ b/deepmd/dpmodel/descriptor/se_e2_a.py
@@ -3,12 +3,13 @@
 
 import numpy as np
 
+from deepmd.env import (
+    GLOBAL_NP_FLOAT_PRECISION,
+)
 from deepmd.utils.path import (
     DPPath,
 )
-from deepmd.env import (
-    GLOBAL_NP_FLOAT_PRECISION
-)
+
 try:
     from deepmd._version import version as __version__
 except ImportError:
@@ -185,8 +186,12 @@ def __init__(
             )
         self.env_mat = EnvMat(self.rcut, self.rcut_smth)
         self.nnei = np.sum(self.sel)
-        self.davg = np.zeros([self.ntypes, self.nnei, 4], dtype=PRECISION_DICT[self.precision])
-        self.dstd = np.ones([self.ntypes, self.nnei, 4], dtype=PRECISION_DICT[self.precision])
+        self.davg = np.zeros(
+            [self.ntypes, self.nnei, 4], dtype=PRECISION_DICT[self.precision]
+        )
+        self.dstd = np.ones(
+            [self.ntypes, self.nnei, 4], dtype=PRECISION_DICT[self.precision]
+        )
         self.orig_sel = self.sel
 
     def __setitem__(self, key, value):
@@ -324,7 +329,9 @@ def call(
         # nf x nloc x ng x ng1
         grrg = np.einsum("flid,fljd->flij", gr, gr1)
         # nf x nloc x (ng x ng1)
-        grrg = grrg.reshape(nf, nloc, ng * self.axis_neuron).astype(GLOBAL_NP_FLOAT_PRECISION)
+        grrg = grrg.reshape(nf, nloc, ng * self.axis_neuron).astype(
+            GLOBAL_NP_FLOAT_PRECISION
+        )
         return grrg, gr[..., 1:], None, None, ww
 
     def serialize(self) -> dict:
diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py
index 99728d9108..77e43f7d85 100644
--- a/deepmd/dpmodel/descriptor/se_r.py
+++ b/deepmd/dpmodel/descriptor/se_r.py
@@ -16,7 +16,6 @@
     List,
     Optional,
 )
-from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
 
 from deepmd.dpmodel import (
     DEFAULT_PRECISION,
@@ -29,6 +28,9 @@
     NetworkCollection,
     PairExcludeMask,
 )
+from deepmd.env import (
+    GLOBAL_NP_FLOAT_PRECISION,
+)
 
 from .base_descriptor import (
     BaseDescriptor,
@@ -144,8 +146,12 @@ def __init__(
             )
         self.env_mat = EnvMat(self.rcut, self.rcut_smth)
         self.nnei = np.sum(self.sel)
-        self.davg = np.zeros([self.ntypes, self.nnei, 1], dtype=PRECISION_DICT[self.precision])
-        self.dstd = np.ones([self.ntypes, self.nnei, 1], dtype=PRECISION_DICT[self.precision])
+        self.davg = np.zeros(
+            [self.ntypes, self.nnei, 1], dtype=PRECISION_DICT[self.precision]
+        )
+        self.dstd = np.ones(
+            [self.ntypes, self.nnei, 1], dtype=PRECISION_DICT[self.precision]
+        )
         self.orig_sel = self.sel
 
     def __setitem__(self, key, value):
diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py
index eec6e54962..033d640ad8 100644
--- a/deepmd/pt/model/descriptor/se_a.py
+++ b/deepmd/pt/model/descriptor/se_a.py
@@ -394,7 +394,6 @@ def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None)
             self.mean.copy_(torch.tensor(mean, device=env.DEVICE))
         self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
 
-
     def get_stats(self) -> Dict[str, StatItem]:
         """Get the statistics of the descriptor."""
         if self.stats is None:
diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index 22c0f55e55..c685640426 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -313,4 +313,4 @@ def t_cvt(xx):
         obj["davg"] = t_cvt(variables["davg"])
         obj["dstd"] = t_cvt(variables["dstd"])
         obj.filter_layers = NetworkCollection.deserialize(embeddings)
-        return obj
\ No newline at end of file
+        return obj