fix: add numb_fparam & numb_aparam to dipole & polar fitting (#4405)

Fix #4396. Fix #4397. Fix #4398. Throw errors for TF.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

## Release Notes

- **New Features**
- Introduced new parameters `numb_fparam` and `numb_aparam` for improved
fitting configurations in both dipole and polar fitting classes.
- Added methods to retrieve the values of the new parameters and
enhanced input requirement management.

- **Documentation**
- Updated training documentation to clarify the handling of new
parameters and their limitations in the TensorFlow backend.

- **Bug Fixes**
- Updated test configurations to reflect the new parameter structure,
ensuring consistency across tests for dipole and polar models.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: njzjz

deepmd/tf/fit/dipole.py

@@ -29,6 +29,9 @@
     one_layer,
     one_layer_rand_seed_shift,
 )
+from deepmd.utils.data import (
+    DataRequirementItem,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -51,6 +54,10 @@ class DipoleFittingSeA(Fitting):
     resnet_dt : bool
             Time-step `dt` in the resnet construction:
             y = x + dt * \phi (Wx + b)
+    numb_fparam
+            Number of frame parameters
+    numb_aparam
+            Number of atomic parameters
     sel_type : list[int]
             The atom types selected to have an atomic dipole prediction. If is None, all atoms are selected.
     seed : int
@@ -75,6 +82,8 @@ def __init__(
         embedding_width: int,
         neuron: list[int] = [120, 120, 120],
         resnet_dt: bool = True,
+        numb_fparam: int = 0,
+        numb_aparam: int = 0,
         sel_type: Optional[list[int]] = None,
         seed: Optional[int] = None,
         activation_function: str = "tanh",
@@ -108,6 +117,18 @@ def __init__(
         self.mixed_prec = None
         self.mixed_types = mixed_types
         self.type_map = type_map
+        self.numb_fparam = numb_fparam
+        self.numb_aparam = numb_aparam
+        if numb_fparam > 0:
+            raise ValueError("numb_fparam is not supported in the dipole fitting")
+        if numb_aparam > 0:
+            raise ValueError("numb_aparam is not supported in the dipole fitting")
+        self.fparam_avg = None
+        self.fparam_std = None
+        self.fparam_inv_std = None
+        self.aparam_avg = None
+        self.aparam_std = None
+        self.aparam_inv_std = None
 
     def get_sel_type(self) -> int:
         """Get selected type."""
@@ -372,6 +393,8 @@ def serialize(self, suffix: str) -> dict:
             "dim_out": 3,
             "neuron": self.n_neuron,
             "resnet_dt": self.resnet_dt,
+            "numb_fparam": self.numb_fparam,
+            "numb_aparam": self.numb_aparam,
             "activation_function": self.activation_function_name,
             "precision": self.fitting_precision.name,
             "exclude_types": [],
@@ -412,3 +435,29 @@ def deserialize(cls, data: dict, suffix: str):
             suffix=suffix,
         )
         return fitting
+
+    @property
+    def input_requirement(self) -> list[DataRequirementItem]:
+        """Return data requirements needed for the model input."""
+        data_requirement = []
+        if self.numb_fparam > 0:
+            data_requirement.append(
+                DataRequirementItem(
+                    "fparam", self.numb_fparam, atomic=False, must=True, high_prec=False
+                )
+            )
+        if self.numb_aparam > 0:
+            data_requirement.append(
+                DataRequirementItem(
+                    "aparam", self.numb_aparam, atomic=True, must=True, high_prec=False
+                )
+            )
+        return data_requirement
+
+    def get_numb_fparam(self) -> int:
+        """Get the number of frame parameters."""
+        return self.numb_fparam
+
+    def get_numb_aparam(self) -> int:
+        """Get the number of atomic parameters."""
+        return self.numb_aparam

deepmd/tf/fit/polar.py

@@ -34,6 +34,9 @@
     one_layer,
     one_layer_rand_seed_shift,
 )
+from deepmd.utils.data import (
+    DataRequirementItem,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -56,6 +59,10 @@ class PolarFittingSeA(Fitting):
     resnet_dt : bool
             Time-step `dt` in the resnet construction:
             y = x + dt * \phi (Wx + b)
+    numb_fparam
+            Number of frame parameters
+    numb_aparam
+            Number of atomic parameters
     sel_type : list[int]
             The atom types selected to have an atomic polarizability prediction. If is None, all atoms are selected.
     fit_diag : bool
@@ -86,6 +93,8 @@ def __init__(
         embedding_width: int,
         neuron: list[int] = [120, 120, 120],
         resnet_dt: bool = True,
+        numb_fparam: int = 0,
+        numb_aparam: int = 0,
         sel_type: Optional[list[int]] = None,
         fit_diag: bool = True,
         scale: Optional[list[float]] = None,
@@ -151,6 +160,18 @@ def __init__(
         self.mixed_prec = None
         self.mixed_types = mixed_types
         self.type_map = type_map
+        self.numb_fparam = numb_fparam
+        self.numb_aparam = numb_aparam
+        if numb_fparam > 0:
+            raise ValueError("numb_fparam is not supported in the dipole fitting")
+        if numb_aparam > 0:
+            raise ValueError("numb_aparam is not supported in the dipole fitting")
+        self.fparam_avg = None
+        self.fparam_std = None
+        self.fparam_inv_std = None
+        self.aparam_avg = None
+        self.aparam_std = None
+        self.aparam_inv_std = None
 
     def get_sel_type(self) -> list[int]:
         """Get selected atom types."""
@@ -565,6 +586,8 @@ def serialize(self, suffix: str) -> dict:
             "dim_out": 3,
             "neuron": self.n_neuron,
             "resnet_dt": self.resnet_dt,
+            "numb_fparam": self.numb_fparam,
+            "numb_aparam": self.numb_aparam,
             "activation_function": self.activation_function_name,
             "precision": self.fitting_precision.name,
             "exclude_types": [],
@@ -777,3 +800,29 @@ def get_loss(self, loss: dict, lr) -> Loss:
             atomic=False,
             label_name="polarizability",
         )
+
+    @property
+    def input_requirement(self) -> list[DataRequirementItem]:
+        """Return data requirements needed for the model input."""
+        data_requirement = []
+        if self.numb_fparam > 0:
+            data_requirement.append(
+                DataRequirementItem(
+                    "fparam", self.numb_fparam, atomic=False, must=True, high_prec=False
+                )
+            )
+        if self.numb_aparam > 0:
+            data_requirement.append(
+                DataRequirementItem(
+                    "aparam", self.numb_aparam, atomic=True, must=True, high_prec=False
+                )
+            )
+        return data_requirement
+
+    def get_numb_fparam(self) -> int:
+        """Get the number of frame parameters."""
+        return self.numb_fparam
+
+    def get_numb_aparam(self) -> int:
+        """Get the number of atomic parameters."""
+        return self.numb_aparam

deepmd/utils/argcheck.py

@@ -1595,6 +1595,8 @@ def fitting_property():
 
 @fitting_args_plugin.register("polar", doc=doc_polar)
 def fitting_polar():
+    doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams."
+    doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams."
     doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built."
     doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.'
     doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection'
@@ -1609,6 +1611,20 @@ def fitting_polar():
     doc_shift_diag = "Whether to shift the diagonal of polar, which is beneficial to training. Default is true."
 
     return [
+        Argument(
+            "numb_fparam",
+            int,
+            optional=True,
+            default=0,
+            doc=doc_only_pt_supported + doc_numb_fparam,
+        ),
+        Argument(
+            "numb_aparam",
+            int,
+            optional=True,
+            default=0,
+            doc=doc_only_pt_supported + doc_numb_aparam,
+        ),
         Argument(
             "neuron",
             list[int],
@@ -1649,13 +1665,29 @@ def fitting_polar():
 
 @fitting_args_plugin.register("dipole", doc=doc_dipole)
 def fitting_dipole():
+    doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams."
+    doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams."
     doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built."
     doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.'
     doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection'
     doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision."
     doc_sel_type = "The atom types for which the atomic dipole will be provided. If not set, all types will be selected."
     doc_seed = "Random seed for parameter initialization of the fitting net"
     return [
+        Argument(
+            "numb_fparam",
+            int,
+            optional=True,
+            default=0,
+            doc=doc_only_pt_supported + doc_numb_fparam,
+        ),
+        Argument(
+            "numb_aparam",
+            int,
+            optional=True,
+            default=0,
+            doc=doc_only_pt_supported + doc_numb_aparam,
+        ),
         Argument(
             "neuron",
             list[int],

doc/model/train-fitting-tensor.md

@@ -247,3 +247,4 @@ During training, at each step when the `lcurve.out` is printed, the system used
 
 To only fit against a subset of atomic types, in the TensorFlow backend, {ref}`fitting_net/sel_type <model[standard]/fitting_net[dipole]/sel_type>` should be set to selected types;
 in other backends, {ref}`atom_exclude_types <model/atom_exclude_types>` should be set to excluded types.
+The TensorFlow backend does not support {ref}`numb_fparam <model[standard]/fitting_net[dipole]/numb_fparam>` and {ref}`numb_aparam <model[standard]/fitting_net[dipole]/numb_aparam>`.

source/tests/consistent/model/test_dipole.py

@@ -62,8 +62,7 @@ def data(self) -> dict:
                 "type": "dipole",
                 "neuron": [4, 4, 4],
                 "resnet_dt": True,
-                # TODO: add numb_fparam argument to dipole fitting
-                "_numb_fparam": 0,
+                "numb_fparam": 0,
                 "precision": "float64",
                 "seed": 1,
             },

source/tests/consistent/model/test_polar.py

@@ -62,8 +62,7 @@ def data(self) -> dict:
                 "type": "polar",
                 "neuron": [4, 4, 4],
                 "resnet_dt": True,
-                # TODO: add numb_fparam argument to polar fitting
-                "_numb_fparam": 0,
+                "numb_fparam": 0,
                 "precision": "float64",
                 "seed": 1,
             },

source/tests/consistent/model/test_property.py

@@ -57,8 +57,7 @@ def data(self) -> dict:
                 "type": "property",
                 "neuron": [4, 4, 4],
                 "resnet_dt": True,
-                # TODO: add numb_fparam argument to property fitting
-                "_numb_fparam": 0,
+                "numb_fparam": 0,
                 "precision": "float64",
                 "seed": 1,
             },