Feat: add dipole consistency test

deepmd/dpmodel/fitting/dipole_fitting.py

@@ -53,8 +53,6 @@ class DipoleFitting(GeneralFitting):
             If the weights of fitting net are trainable.
             Suppose that we have :math:`N_l` hidden layers in the fitting net,
             this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable.
-    atom_ener
-            Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set.
     activation_function
             The activation function :math:`\boldsymbol{\phi}` in the embedding net. Supported options are |ACTIVATION_FN|
     precision
@@ -91,7 +89,6 @@ def __init__(
         rcond: Optional[float] = None,
         tot_ener_zero: bool = False,
         trainable: Optional[List[bool]] = None,
-        atom_ener: Optional[List[Optional[float]]] = None,
         activation_function: str = "tanh",
         precision: str = DEFAULT_PRECISION,
         layer_name: Optional[List[Optional[str]]] = None,
@@ -102,6 +99,8 @@ def __init__(
         r_differentiable: bool = True,
         c_differentiable: bool = True,
         old_impl=False,
+        # not used
+        seed: Optional[int] = None,
     ):
         # seed, uniform_seed are not included
         if tot_ener_zero:
@@ -112,8 +111,6 @@ def __init__(
             raise NotImplementedError("use_aparam_as_mask is not implemented")
         if layer_name is not None:
             raise NotImplementedError("layer_name is not implemented")
-        if atom_ener is not None and atom_ener != []:
-            raise NotImplementedError("atom_ener is not implemented")
 
         self.embedding_width = embedding_width
         self.r_differentiable = r_differentiable
@@ -129,7 +126,6 @@ def __init__(
             rcond=rcond,
             tot_ener_zero=tot_ener_zero,
             trainable=trainable,
-            atom_ener=atom_ener,
             activation_function=activation_function,
             precision=precision,
             layer_name=layer_name,

deepmd/dpmodel/fitting/general_fitting.py

@@ -55,8 +55,6 @@ class GeneralFitting(NativeOP, BaseFitting):
             If the weights of fitting net are trainable.
             Suppose that we have :math:`N_l` hidden layers in the fitting net,
             this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable.
-    atom_ener
-            Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set.
     activation_function
             The activation function :math:`\boldsymbol{\phi}` in the embedding net. Supported options are |ACTIVATION_FN|
     precision
@@ -87,7 +85,6 @@ def __init__(
         rcond: Optional[float] = None,
         tot_ener_zero: bool = False,
         trainable: Optional[List[bool]] = None,
-        atom_ener: Optional[List[float]] = None,
         activation_function: str = "tanh",
         precision: str = DEFAULT_PRECISION,
         layer_name: Optional[List[Optional[str]]] = None,
@@ -110,7 +107,6 @@ def __init__(
             self.trainable = [True for ii in range(len(self.neuron) + 1)]
         if isinstance(self.trainable, bool):
             self.trainable = [self.trainable] * (len(self.neuron) + 1)
-        self.atom_ener = atom_ener
         self.activation_function = activation_function
         self.precision = precision
         self.layer_name = layer_name
@@ -236,7 +232,6 @@ def serialize(self) -> dict:
             # not supported
             "tot_ener_zero": self.tot_ener_zero,
             "trainable": self.trainable,
-            "atom_ener": self.atom_ener,
             "layer_name": self.layer_name,
             "use_aparam_as_mask": self.use_aparam_as_mask,
             "spin": self.spin,

deepmd/dpmodel/fitting/invar_fitting.py

@@ -115,7 +115,7 @@
         rcond: Optional[float] = None,
         tot_ener_zero: bool = False,
         trainable: Optional[List[bool]] = None,
-        atom_ener: Optional[List[float]] = None,
+        atom_ener: Optional[List[float]] = [],
         activation_function: str = "tanh",
         precision: str = DEFAULT_PRECISION,
         layer_name: Optional[List[Optional[str]]] = None,
@@ -139,6 +139,7 @@
             raise NotImplementedError("atom_ener is not implemented")
 
         self.dim_out = dim_out
+        self.atom_ener = atom_ener
         super().__init__(
             var_name=var_name,
             ntypes=ntypes,
@@ -150,7 +151,6 @@
             rcond=rcond,
             tot_ener_zero=tot_ener_zero,
             trainable=trainable,
-            atom_ener=atom_ener,
             activation_function=activation_function,
             precision=precision,
             layer_name=layer_name,
@@ -163,6 +163,7 @@
     def serialize(self) -> dict:
         data = super().serialize()
         data["dim_out"] = self.dim_out
+        data["atom_ener"] = self.atom_ener
         return data
 
     def _net_out_dim(self):

deepmd/dpmodel/fitting/polarizability_fitting.py

@@ -56,8 +56,6 @@ class PolarFitting(GeneralFitting):
             If the weights of fitting net are trainable.
             Suppose that we have :math:`N_l` hidden layers in the fitting net,
             this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable.
-    atom_ener
-            Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set.
     activation_function
             The activation function :math:`\boldsymbol{\phi}` in the embedding net. Supported options are |ACTIVATION_FN|
     precision
@@ -93,7 +91,6 @@ def __init__(
         rcond: Optional[float] = None,
         tot_ener_zero: bool = False,
         trainable: Optional[List[bool]] = None,
-        atom_ener: Optional[List[Optional[float]]] = None,
         activation_function: str = "tanh",
         precision: str = DEFAULT_PRECISION,
         layer_name: Optional[List[Optional[str]]] = None,
@@ -115,8 +112,6 @@ def __init__(
             raise NotImplementedError("use_aparam_as_mask is not implemented")
         if layer_name is not None:
             raise NotImplementedError("layer_name is not implemented")
-        if atom_ener is not None and atom_ener != []:
-            raise NotImplementedError("atom_ener is not implemented")
 
         self.embedding_width = embedding_width
         self.fit_diag = fit_diag
@@ -142,7 +137,6 @@ def __init__(
             rcond=rcond,
             tot_ener_zero=tot_ener_zero,
             trainable=trainable,
-            atom_ener=atom_ener,
             activation_function=activation_function,
             precision=precision,
             layer_name=layer_name,

deepmd/pt/model/task/ener.py

@@ -78,6 +78,8 @@
         Random seed.
     exclude_types: List[int]
         Atomic contributions of the excluded atom types are set zero.
+    atom_ener
+            Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set.
 
     """
 
@@ -98,9 +100,11 @@
         rcond: Optional[float] = None,
         seed: Optional[int] = None,
         exclude_types: List[int] = [],
+        atom_ener: Optional[List[float]] = None,
         **kwargs,
     ):
         self.dim_out = dim_out
+        self.atom_ener = atom_ener
         super().__init__(
             var_name=var_name,
             ntypes=ntypes,
@@ -126,6 +130,7 @@
     def serialize(self) -> dict:
         data = super().serialize()
         data["dim_out"] = self.dim_out
+        data["atom_ener"] = self.atom_ener
         return data
 
     @property

deepmd/pt/model/task/fitting.py

@@ -428,7 +428,6 @@ def serialize(self) -> dict:
             ## NOTICE:  not supported by far
             "tot_ener_zero": False,
             "trainable": [True] * (len(self.neuron) + 1),
-            "atom_ener": [],
             "layer_name": None,
             "use_aparam_as_mask": False,
             "spin": None,

deepmd/tf/descriptor/se_r.py

@@ -204,6 +204,10 @@
         """Returns the output dimension of this descriptor."""
         return self.filter_neuron[-1]
 
+    def get_dim_rot_mat_1(self) -> int:
+        """To accommodate dipole fitting, which needs embedding_width."""
+        return self.get_dim_out()
+
     def get_nlist(self):
         """Returns neighbor information.
 

deepmd/tf/descriptor/se_t.py

@@ -194,6 +194,10 @@
         """Returns the output dimension of this descriptor."""
         return self.filter_neuron[-1]
 
+    def get_dim_rot_mat_1(self) -> int:
+        """To accommodate dipole fitting, which needs embedding_width."""
+        return self.get_dim_out()
+
     def get_nlist(self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]:
         """Returns neighbor information.
 

deepmd/tf/fit/dipole.py

@@ -38,8 +38,12 @@
 
     Parameters
     ----------
-    descrpt : tf.Tensor
-            The descrptor
+    ntypes
+            The ntypes of the descrptor :math:`\mathcal{D}`
+    dim_descrpt
+            The dimension of the descrptor :math:`\mathcal{D}`
+    embedding_width
+            The rotation matrix dimension of the descrptor :math:`\mathcal{D}`
     neuron : List[int]
             Number of neurons in each hidden layer of the fitting net
     resnet_dt : bool
@@ -59,7 +63,9 @@
 
     def __init__(
         self,
-        descrpt: tf.Tensor,
+        ntypes: int,
+        dim_descrpt: int,
+        embedding_width: int,
         neuron: List[int] = [120, 120, 120],
         resnet_dt: bool = True,
         sel_type: Optional[List[int]] = None,
@@ -70,8 +76,8 @@
         **kwargs,
     ) -> None:
         """Constructor."""
-        self.ntypes = descrpt.get_ntypes()
-        self.dim_descrpt = descrpt.get_dim_out()
+        self.ntypes = ntypes
+        self.dim_descrpt = dim_descrpt
         self.n_neuron = neuron
         self.resnet_dt = resnet_dt
         self.sel_type = sel_type
@@ -83,9 +89,10 @@
         self.seed = seed
         self.uniform_seed = uniform_seed
         self.seed_shift = one_layer_rand_seed_shift()
+        self.activation_function_name = activation_function
         self.fitting_activation_fn = get_activation_func(activation_function)
         self.fitting_precision = get_precision(precision)
-        self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1()
+        self.dim_rot_mat_1 = embedding_width
         self.dim_rot_mat = self.dim_rot_mat_1 * 3
         self.useBN = False
         self.fitting_net_variables = None
@@ -327,3 +334,61 @@
             tensor_size=3,
             label_name="dipole",
         )
+
+    def serialize(self, suffix: str) -> dict:
+        """Serialize the model.
+
+        Returns
+        -------
+        dict
+            The serialized data
+        """
+        data = {
+            "var_name": "dipole",
+            "ntypes": self.ntypes,
+            "dim_descrpt": self.dim_descrpt,
+            "embedding_width": self.dim_rot_mat_1,
+            # very bad design: type embedding is not passed to the class
+            # TODO: refactor the class
+            "mixed_types": False,
+            "dim_out": 3,
+            "neuron": self.n_neuron,
+            "resnet_dt": self.resnet_dt,
+            "activation_function": self.activation_function_name,
+            "precision": self.fitting_precision.name,
+            "exclude_types": [],
+            "nets": self.serialize_network(
+                ntypes=self.ntypes,
+                # TODO: consider type embeddings
+                ndim=1,
+                in_dim=self.dim_descrpt,
+                out_dim=self.dim_rot_mat_1,
+                neuron=self.n_neuron,
+                activation_function=self.activation_function_name,
+                resnet_dt=self.resnet_dt,
+                variables=self.fitting_net_variables,
+                suffix=suffix,
+            ),
+        }
+        return data
+
+    @classmethod
+    def deserialize(cls, data: dict, suffix: str):
+        """Deserialize the model.
+
+        Parameters
+        ----------
+        data : dict
+            The serialized data
+
+        Returns
+        -------
+        Model
+            The deserialized model
+        """
+        fitting = cls(**data)
+        fitting.fitting_net_variables = cls.deserialize_network(
+            data["nets"],
+            suffix=suffix,
+        )
+        return fitting

deepmd/tf/fit/ener.py

@@ -95,8 +95,10 @@ class EnerFitting(Fitting):
 
     Parameters
     ----------
-    descrpt
-            The descrptor :math:`\mathcal{D}`
+    ntypes
+            The ntypes of the descrptor :math:`\mathcal{D}`
+    dim_descrpt
+            The dimension of the descrptor :math:`\mathcal{D}`
     neuron
             Number of neurons :math:`N` in each hidden layer of the fitting net
     resnet_dt

deepmd/tf/fit/fitting.py

@@ -6,6 +6,7 @@
 from typing import (
     Callable,
     List,
+    Optional,
     Type,
 )
 
@@ -175,6 +176,7 @@ def serialize_network(
         activation_function: str,
         resnet_dt: bool,
         variables: dict,
+        out_dim: Optional[int] = 1,
         suffix: str = "",
     ) -> dict:
         """Serialize network.
@@ -197,6 +199,8 @@ def serialize_network(
             The input variables
         suffix : str, optional
             The suffix of the scope
+        out_dim : int, optional
+            The output dimension
 
         Returns
         -------
@@ -231,7 +235,7 @@ def serialize_network(
                 # initialize the network if it is not initialized
                 fittings[network_idx] = FittingNet(
                     in_dim=in_dim,
-                    out_dim=1,
+                    out_dim=out_dim,
                     neuron=neuron,
                     activation_function=activation_function,
                     resnet_dt=resnet_dt,

deepmd/tf/model/model.py

@@ -693,6 +693,7 @@ def __init__(
                 spin=self.spin,
                 ntypes=self.descrpt.get_ntypes(),
                 dim_descrpt=self.descrpt.get_dim_out(),
+                embedding_width=self.descrpt.get_dim_rot_mat_1(),
             )
         self.rcut = self.descrpt.get_rcut()
         self.ntypes = self.descrpt.get_ntypes()

deepmd/tf/model/multi.py

@@ -139,6 +139,7 @@ def __init__(
                     spin=self.spin,
                     ntypes=self.descrpt.get_ntypes(),
                     dim_descrpt=self.descrpt.get_dim_out(),
+                    embedding_width=self.descrpt.get_dim_rot_mat_1(),
                 )
 
         # type embedding