feat: support serialize and deserialize for atomic model's compute output stat

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
 from typing import (
     Dict,
     List,
@@ -30,11 +31,44 @@
         type_map: List[str],
         atom_exclude_types: List[int] = [],
         pair_exclude_types: List[Tuple[int, int]] = [],
+        rcond: Optional[float] = None,
+        preset_out_bias: Optional[Dict[str, np.ndarray]] = None,
     ):
         super().__init__()
         self.type_map = type_map
         self.reinit_atom_exclude(atom_exclude_types)
         self.reinit_pair_exclude(pair_exclude_types)
+        self.rcond = rcond
+        self.preset_out_bias = preset_out_bias
+
+    def init_out_stat(self):
+        """Initialize the output bias."""
+        ntypes = self.get_ntypes()
+        self.bias_keys: List[str] = list(self.fitting_output_def().keys())
+        self.max_out_size = max(
+            [self.atomic_output_def()[kk].size for kk in self.bias_keys]
+        )
+        self.n_out = len(self.bias_keys)
+        out_bias_data = np.zeros([self.n_out, ntypes, self.max_out_size])
+        out_std_data = np.ones([self.n_out, ntypes, self.max_out_size])
+        self.out_bias = out_bias_data
+        self.out_std = out_std_data
+
+    def __setitem__(self, key, value):
+        if key in ["out_bias"]:
+            self.out_bias = value
+        elif key in ["out_std"]:
+            self.out_std = value
+        else:
+            raise KeyError(key)
+
+    def __getitem__(self, key):
+        if key in ["out_bias"]:
+            return self.out_bias
+        elif key in ["out_std"]:
+            return self.out_std
+        else:
+            raise KeyError(key)
 
     def get_type_map(self) -> List[str]:
         """Get the type map."""
@@ -132,6 +166,7 @@
             fparam=fparam,
             aparam=aparam,
         )
+        ret_dict = self.apply_out_stat(ret_dict, atype)
 
         # nf x nloc
         atom_mask = ext_atom_mask[:, :nloc].astype(np.int32)
@@ -150,6 +185,84 @@
 
     def serialize(self) -> dict:
         return {
+            "type_map": self.type_map,
             "atom_exclude_types": self.atom_exclude_types,
             "pair_exclude_types": self.pair_exclude_types,
+            "rcond": self.rcond,
+            "preset_out_bias": self.preset_out_bias,
+            "@variables": {
+                "out_bias": self.out_bias,
+                "out_std": self.out_std,
+            },
         }
+
+    @classmethod
+    def deserialize(cls, data: dict) -> "BaseAtomicModel":
+        data = copy.deepcopy(data)
+        variables = data.pop("@variables")
+        obj = cls(**data)
+        for kk in variables.keys():
+            obj[kk] = variables[kk]
+        return obj
+
+    def apply_out_stat(
+        self,
+        ret: Dict[str, np.ndarray],
+        atype: np.ndarray,
+    ):
+        """Apply the stat to each atomic output.
+        The developer may override the method to define how the bias is applied
+        to the atomic output of the model.
+
+        Parameters
+        ----------
+        ret
+            The returned dict by the forward_atomic method
+        atype
+            The atom types. nf x nloc
+
+        """
+        out_bias, out_std = self._fetch_out_stat(self.bias_keys)
+        for kk in self.bias_keys:
+            # nf x nloc x odims, out_bias: ntypes x odims
+            ret[kk] = ret[kk] + out_bias[kk][atype]
+        return ret
+
+    def _varsize(
+        self,
+        shape: List[int],
+    ) -> int:
+        output_size = 1
+        len_shape = len(shape)
+        for i in range(len_shape):
+            output_size *= shape[i]
+        return output_size
+
+    def _get_bias_index(
+        self,
+        kk: str,
+    ) -> int:
+        res: List[int] = []
+        for i, e in enumerate(self.bias_keys):
+            if e == kk:
+                res.append(i)
+        assert len(res) == 1
+        return res[0]
+
+    def _fetch_out_stat(
+        self,
+        keys: List[str],
+    ) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray]]:
+        ret_bias = {}
+        ret_std = {}
+        ntypes = self.get_ntypes()
+        for kk in keys:
+            idx = self._get_bias_index(kk)
+            isize = self._varsize(self.atomic_output_def()[kk].shape)
+            ret_bias[kk] = self.out_bias[idx, :, :isize].reshape(
+                [ntypes] + list(self.atomic_output_def()[kk].shape)  # noqa: RUF005
+            )
+            ret_std[kk] = self.out_std[idx, :, :isize].reshape(
+                [ntypes] + list(self.atomic_output_def()[kk].shape)  # noqa: RUF005
+            )
+        return ret_bias, ret_std

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -49,11 +49,12 @@
         type_map: List[str],
         **kwargs,
     ):
+        super().__init__(type_map, **kwargs)
         self.type_map = type_map
         self.descriptor = descriptor
         self.fitting = fitting
         self.type_map = type_map
-        super().__init__(type_map, **kwargs)
+        super().init_out_stat()
 
     def fitting_output_def(self) -> FittingOutputDef:
         """Get the output def of the fitting net."""
@@ -157,7 +158,7 @@
             {
                 "@class": "Model",
                 "type": "standard",
-                "@version": 1,
+                "@version": 2,
                 "type_map": self.type_map,
                 "descriptor": self.descriptor.serialize(),
                 "fitting": self.fitting.serialize(),
@@ -168,13 +169,14 @@
     @classmethod
     def deserialize(cls, data) -> "DPAtomicModel":
         data = copy.deepcopy(data)
-        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        check_version_compatibility(data.pop("@version", 2), 2, 2)
         data.pop("@class")
         data.pop("type")
         descriptor_obj = BaseDescriptor.deserialize(data.pop("descriptor"))
         fitting_obj = BaseFitting.deserialize(data.pop("fitting"))
-        type_map = data.pop("type_map")
-        obj = cls(descriptor_obj, fitting_obj, type_map=type_map, **data)
+        data["descriptor"] = descriptor_obj
+        data["fitting"] = fitting_obj
+        obj = super().deserialize(data)
         return obj
 
     def get_dim_fparam(self) -> int:

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -52,6 +52,8 @@
         type_map: List[str],
         **kwargs,
     ):
+        super().__init__(type_map, **kwargs)
+        super().init_out_stat()
         self.models = models
         sub_model_type_maps = [md.get_type_map() for md in models]
         err_msg = []
@@ -66,7 +68,6 @@
             self.mapping_list.append(self.remap_atype(tpmp, self.type_map))
         assert len(err_msg) == 0, "\n".join(err_msg)
         self.mixed_types_list = [model.mixed_types() for model in self.models]
-        super().__init__(type_map, **kwargs)
 
     def mixed_types(self) -> bool:
         """If true, the model
@@ -86,7 +87,7 @@
 
     def get_type_map(self) -> List[str]:
         """Get the type map."""
-        raise self.type_map
+        return self.type_map
 
     def get_model_rcuts(self) -> List[float]:
         """Get the cut-off radius for each individual models."""
@@ -222,27 +223,30 @@
         )
 
     def serialize(self) -> dict:
-        return {
-            "@class": "Model",
-            "type": "linear",
-            "@version": 1,
-            "models": [model.serialize() for model in self.models],
-            "type_map": self.type_map,
-        }
+        dd = super().serialize()
+        dd.update(
+            {
+                "@class": "Model",
+                "@version": 2,
+                "type": "linear",
+                "models": [model.serialize() for model in self.models],
+                "type_map": self.type_map,
+            }
+        )
+        return dd
 
     @classmethod
     def deserialize(cls, data: dict) -> "LinearEnergyAtomicModel":
         data = copy.deepcopy(data)
-        check_version_compatibility(data.pop("@version", 1), 1, 1)
-        data.pop("@class")
-        data.pop("type")
-        type_map = data.pop("type_map")
+        check_version_compatibility(data.pop("@version", 2), 2, 2)
+        data.pop("@class", None)
+        data.pop("type", None)
         models = [
             BaseAtomicModel.get_class_by_type(model["type"]).deserialize(model)
             for model in data["models"]
         ]
-        data.pop("models")
-        return cls(models, type_map, **data)
+        data["models"] = models
+        return super().deserialize(data)
 
     def _compute_weight(
         self,
@@ -336,24 +340,21 @@
         **kwargs,
     ):
         models = [dp_model, zbl_model]
-        super().__init__(models, type_map, **kwargs)
-        self.dp_model = dp_model
-        self.zbl_model = zbl_model
+        kwargs["models"] = models
+        kwargs["type_map"] = type_map
+        super().__init__(**kwargs)
 
         self.sw_rmin = sw_rmin
         self.sw_rmax = sw_rmax
         self.smin_alpha = smin_alpha
 
     def serialize(self) -> dict:
-        dd = BaseAtomicModel.serialize(self)
+        dd = super().serialize()
         dd.update(
             {
                 "@class": "Model",
-                "type": "zbl",
                 "@version": 2,
-                "models": LinearEnergyAtomicModel(
-                    models=[self.models[0], self.models[1]], type_map=self.type_map
-                ).serialize(),
+                "type": "zbl",
                 "sw_rmin": self.sw_rmin,
                 "sw_rmax": self.sw_rmax,
                 "smin_alpha": self.smin_alpha,
@@ -364,25 +365,15 @@
     @classmethod
     def deserialize(cls, data) -> "DPZBLLinearEnergyAtomicModel":
         data = copy.deepcopy(data)
-        check_version_compatibility(data.pop("@version", 1), 2, 1)
-        data.pop("@class")
-        data.pop("type")
-        sw_rmin = data.pop("sw_rmin")
-        sw_rmax = data.pop("sw_rmax")
-        smin_alpha = data.pop("smin_alpha")
-        linear_model = LinearEnergyAtomicModel.deserialize(data.pop("models"))
-        dp_model, zbl_model = linear_model.models
-        type_map = linear_model.type_map
-
-        return cls(
-            dp_model=dp_model,
-            zbl_model=zbl_model,
-            sw_rmin=sw_rmin,
-            sw_rmax=sw_rmax,
-            type_map=type_map,
-            smin_alpha=smin_alpha,
-            **data,
-        )
+        check_version_compatibility(data.pop("@version", 2), 2, 2)
+        models = [
+            BaseAtomicModel.get_class_by_type(model["type"]).deserialize(model)
+            for model in data["models"]
+        ]
+        data["dp_model"], data["zbl_model"] = models[0], models[1]
+        data.pop("@class", None)
+        data.pop("type", None)
+        return super().deserialize(data)
 
     def _compute_weight(
         self,

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -64,6 +64,7 @@
         **kwargs,
     ):
         super().__init__(type_map, **kwargs)
+        super().init_out_stat()
         self.tab_file = tab_file
         self.rcut = rcut
         self.type_map = type_map
@@ -71,7 +72,7 @@
         self.tab = PairTab(self.tab_file, rcut=rcut)
         self.type_map = type_map
         self.ntypes = len(type_map)
         self.rcond = rcond
         self.atom_ener = atom_ener
 
         if self.tab_file is not None:
@@ -155,7 +156,7 @@
             {
                 "@class": "Model",
                 "type": "pairtab",
-                "@version": 1,
+                "@version": 2,
                 "tab": self.tab.serialize(),
                 "rcut": self.rcut,
                 "sel": self.sel,
@@ -167,14 +168,13 @@
     @classmethod
     def deserialize(cls, data) -> "PairTabAtomicModel":
         data = copy.deepcopy(data)
-        check_version_compatibility(data.pop("@version", 1), 1, 1)
+        check_version_compatibility(data.pop("@version", 1), 2, 2)
         data.pop("@class")
         data.pop("type")
-        rcut = data.pop("rcut")
-        sel = data.pop("sel")
-        type_map = data.pop("type_map")
         tab = PairTab.deserialize(data.pop("tab"))
-        tab_model = cls(None, rcut, sel, type_map, **data)
+        data["tab_file"] = None
+        tab_model = super().deserialize(data)
+
         tab_model.tab = tab
         tab_model.tab_info = tab_model.tab.tab_info
         nspline, ntypes = tab_model.tab_info[-2:].astype(int)