Refactor property fitting interface

deepmd/dpmodel/output_def.py

@@ -197,6 +197,7 @@ def __init__(
         r_hessian: bool = False,
         magnetic: bool = False,
         intensive: bool = False,
+        sub_var_name: list[str] = [],
     ) -> None:
         self.name = name
         self.shape = list(shape)
@@ -220,6 +221,7 @@ def __init__(
         self.r_hessian = r_hessian
         self.magnetic = magnetic
         self.intensive = intensive
+        self.sub_var_name = sub_var_name
         if self.r_hessian:
             if not self.reducible:
                 raise ValueError("only reducible variable can calculate hessian")

deepmd/pt/loss/property.py

@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import logging
+from typing import (
+    Union,
+)
 
 import torch
 import torch.nn.functional as F
@@ -24,6 +27,8 @@ def __init__(
         loss_func: str = "smooth_mae",
         metric: list = ["mae"],
         beta: float = 1.00,
+        property_name: Union[str, list] = "property",
+        property_dim: Union[int, list] = 1,
         **kwargs,
     ) -> None:
         r"""Construct a layer to compute loss on property.
@@ -44,6 +49,13 @@ def __init__(
         self.loss_func = loss_func
         self.metric = metric
         self.beta = beta
+        if isinstance(property_name, str):
+            property_name = [property_name]
+        if isinstance(property_dim, int):
+            property_dim = [property_dim]
+        self.property_name = property_name
+        assert self.task_dim == sum(property_dim)
+        self.property_name_dim_mapping = dict(zip(property_name, property_dim))
 
     def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False):
         """Return loss on properties .
@@ -69,34 +81,52 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False
             Other losses for display.
         """
         model_pred = model(**input_dict)
-        assert label["property"].shape[-1] == self.task_dim
-        assert model_pred["property"].shape[-1] == self.task_dim
+        nbz = model_pred["property"].shape[0]
+        assert model_pred["property"].shape == (nbz, self.task_dim)
+
+        concat_property = []
+        for property_name in self.property_name:
+            assert label[property_name].shape == (
+                nbz,
+                self.property_name_dim_mapping[property_name],
+            )
+            concat_property.append(label[property_name])
+        label["property"] = torch.cat(concat_property, dim=1)
+        assert label["property"].shape == (nbz, self.task_dim)
+
+        out_std = model.atomic_model.out_std[0][0]
+        out_bias = model.atomic_model.out_bias[0][0]
+        assert len(out_std.shape) == 1
+        assert out_std.shape[0] == self.task_dim
+
         loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0]
         more_loss = {}
 
         # loss
         if self.loss_func == "smooth_mae":
             loss += F.smooth_l1_loss(
-                label["property"],
-                model_pred["property"],
+                (label["property"] - out_bias) / out_std,
+                (model_pred["property"] - out_bias) / out_std,
                 reduction="sum",
                 beta=self.beta,
             )
         elif self.loss_func == "mae":
             loss += F.l1_loss(
-                label["property"], model_pred["property"], reduction="sum"
+                (label["property"] - out_bias) / out_std,
+                (model_pred["property"] - out_bias) / out_std,
+                reduction="sum",
             )
         elif self.loss_func == "mse":
             loss += F.mse_loss(
-                label["property"],
-                model_pred["property"],
+                (label["property"] - out_bias) / out_std,
+                (model_pred["property"] - out_bias) / out_std,
                 reduction="sum",
             )
         elif self.loss_func == "rmse":
             loss += torch.sqrt(
                 F.mse_loss(
-                    label["property"],
-                    model_pred["property"],
+                    (label["property"] - out_bias) / out_std,
+                    (model_pred["property"] - out_bias) / out_std,
                     reduction="mean",
                 )
             )
@@ -138,13 +168,14 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False
     def label_requirement(self) -> list[DataRequirementItem]:
         """Return data label requirements needed for this loss calculation."""
         label_requirement = []
-        label_requirement.append(
-            DataRequirementItem(
-                "property",
-                ndof=self.task_dim,
-                atomic=False,
-                must=False,
-                high_prec=True,
+        for property_name in self.property_name:
+            label_requirement.append(
+                DataRequirementItem(
+                    property_name,
+                    ndof=self.property_name_dim_mapping[property_name],
+                    atomic=False,
+                    must=True,
+                    high_prec=True,
+                )
             )
-        )
         return label_requirement

deepmd/pt/model/atomic_model/property_atomic_model.py

@@ -35,15 +35,7 @@ def apply_out_stat(
             The atom types. nf x nloc
 
         """
-        if self.fitting_net.get_bias_method() == "normal":
-            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
-        elif self.fitting_net.get_bias_method() == "no_bias":
-            return ret
-        else:
-            raise NotImplementedError(
-                "Only 'normal' and 'no_bias' is supported for parameter 'bias_method'."
-            )
+        out_bias, out_std = self._fetch_out_stat(self.bias_keys)
+        for kk in self.bias_keys:
+            ret[kk] = ret[kk] * out_std[kk][0] + out_bias[kk][0]
+        return ret

deepmd/pt/model/task/property.py

@@ -2,6 +2,7 @@
 import logging
 from typing import (
     Optional,
+    Union,
 )
 
 import torch
@@ -82,6 +83,7 @@ def __init__(
         bias_atom_p: Optional[torch.Tensor] = None,
         intensive: bool = False,
         bias_method: str = "normal",
+        property_name: Union[str, list] = "property",
         resnet_dt: bool = True,
         numb_fparam: int = 0,
         numb_aparam: int = 0,
@@ -95,6 +97,7 @@ def __init__(
         self.task_dim = task_dim
         self.intensive = intensive
         self.bias_method = bias_method
+        self.property_name = property_name
         super().__init__(
             var_name="property",
             ntypes=ntypes,
@@ -126,6 +129,7 @@ def output_def(self) -> FittingOutputDef:
                     r_differentiable=False,
                     c_differentiable=False,
                     intensive=self.intensive,
+                    sub_var_name=self.property_name,
                 ),
             ]
         )

deepmd/pt/utils/stat.py

@@ -29,6 +29,7 @@
 from deepmd.utils.out_stat import (
     compute_stats_from_atomic,
     compute_stats_from_redu,
+    compute_stats_property,
 )
 from deepmd.utils.path import (
     DPPath,
@@ -290,6 +291,12 @@
         # remove the keys that are not in the sample
         keys = [keys] if isinstance(keys, str) else keys
         assert isinstance(keys, list)
+        if "property" in atomic_output.var_defs:
+            sub_keys = []
+            for key in keys:
+                if atomic_output.var_defs[key].sub_var_name is not None:
+                    sub_keys.extend(atomic_output.var_defs[key].sub_var_name)
+            keys.extend(sub_keys)
         new_keys = [
             ii
             for ii in keys
@@ -373,6 +380,16 @@
 
         # merge global/atomic bias
         bias_atom_e, std_atom_e = {}, {}
+        keys = (
+            ["property"]
+            if (
+                "property" in atomic_output.var_defs
+                and (
+                    ii in keys for ii in atomic_output.var_defs["property"].sub_var_name
+                )
+            )
+            else keys
+        )
         for kk in keys:
             # use atomic bias whenever available
             if kk in bias_atom_a:
@@ -476,26 +493,44 @@
     std_atom_e = {}
     for kk in keys:
         if kk in stats_input:
-            if atomic_output is not None and atomic_output.get_data()[kk].intensive:
-                task_dim = stats_input[kk].shape[1]
-                assert merged_natoms[kk].shape == (nf[kk], ntypes)
-                stats_input[kk] = (
-                    merged_natoms[kk].sum(axis=1).reshape(-1, 1) * stats_input[kk]
+            if "property" in atomic_output.var_defs:
+                bias_atom_e[kk], std_atom_e[kk] = compute_stats_property(
+                    stats_input[kk],
+                    merged_natoms[kk],
+                    assigned_bias=assigned_atom_ener[kk],
+                )
+            else:
+                bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu(
+                    stats_input[kk],
+                    merged_natoms[kk],
+                    assigned_bias=assigned_atom_ener[kk],
+                    rcond=rcond,
                 )
-                assert stats_input[kk].shape == (nf[kk], task_dim)
-            bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu(
-                stats_input[kk],
-                merged_natoms[kk],
-                assigned_bias=assigned_atom_ener[kk],
-                rcond=rcond,
-            )
         else:
             # this key does not have global labels, skip it.
             continue
+    if "property" in atomic_output.var_defs:
+        concat_bias = []
+        concat_std = []
+        for ii in atomic_output.var_defs["property"].sub_var_name:
+            assert ii in bias_atom_e.keys()
+            assert ii in std_atom_e.keys()
+            concat_bias.append(bias_atom_e[ii])
+            concat_std.append(std_atom_e[ii])
+        del bias_atom_e, std_atom_e
+        bias_atom_e = {}
+        std_atom_e = {}
+        bias_atom_e["property"] = np.concatenate(concat_bias, axis=-1)
+        std_atom_e["property"] = np.concatenate(concat_std, axis=-1)
+        std_atom_e["property"] = np.tile(
+            std_atom_e["property"], (bias_atom_e["property"].shape[0], 1)
+        )
+
+        return bias_atom_e, std_atom_e
+
     bias_atom_e, std_atom_e = _post_process_stat(bias_atom_e, std_atom_e)
 
     # unbias_e is only used for print rmse
-
     if model_pred is None:
         unbias_e = {
             kk: merged_natoms[kk] @ bias_atom_e[kk].reshape(ntypes, -1)

deepmd/utils/argcheck.py

@@ -1581,6 +1581,7 @@ def fitting_property():
     doc_task_dim = "The dimension of outputs of fitting net"
     doc_intensive = "Whether the fitting property is intensive"
     doc_bias_method = "The method of applying the bias to each atomic output, user can select 'normal' or 'no_bias'. If 'no_bias' is used, no bias will be added to the atomic output."
+    doc_property_name = "TODO"
     return [
         Argument("numb_fparam", int, optional=True, default=0, doc=doc_numb_fparam),
         Argument("numb_aparam", int, optional=True, default=0, doc=doc_numb_aparam),
@@ -1614,6 +1615,13 @@ def fitting_property():
         Argument(
             "bias_method", str, optional=True, default="normal", doc=doc_bias_method
         ),
+        Argument(
+            "property_name",
+            [str, list],
+            optional=True,
+            default="property",
+            doc=doc_property_name,
+        ),
     ]
 
 
@@ -2481,6 +2489,8 @@ def loss_property():
     doc_loss_func = "The loss function to minimize, such as 'mae','smooth_mae'."
     doc_metric = "The metric for display. This list can include 'smooth_mae', 'mae', 'mse' and 'rmse'."
     doc_beta = "The 'beta' parameter in 'smooth_mae' loss."
+    doc_property_name = "The names of fitting properties, which should be consistent with the property names in the dataset."
+    doc_property_dim = "The dimensions of fitting properties, which should be consistent with the property dimensions in the dataset."
     return [
         Argument(
             "loss_func",
@@ -2503,6 +2513,20 @@ def loss_property():
             default=1.00,
             doc=doc_beta,
         ),
+        Argument(
+            "property_name",
+            [str, list],
+            optional=True,
+            default="property",
+            doc=doc_property_name,
+        ),
+        Argument(
+            "property_dim",
+            [int, list],
+            optional=True,
+            default=1,
+            doc=doc_property_dim,
+        ),
     ]
 
 

deepmd/utils/out_stat.py

@@ -33,8 +33,6 @@ def compute_stats_from_redu(
         The assigned output bias, shape is [ntypes, *(odim0, odim1, ...)].
         Set to a tensor of shape (odim0, odim1, ...) filled with nan if the bias
         of the type is not assigned.
-    rcond
-        Cut-off ratio for small singular values of a.
 
     Returns
     -------
@@ -130,3 +128,52 @@ def compute_stats_from_atomic(
             output[mask].std(axis=0) if output[mask].size > 0 else np.nan
         )
     return output_bias, output_std
+
+
+def compute_stats_property(
+    output_redu: np.ndarray,
+    natoms: np.ndarray,
+    assigned_bias: Optional[np.ndarray] = None,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Compute the output statistics.
+
+    Given the reduced output value and the number of atoms for each atom,
+    compute the least-squares solution as the atomic output bias and std.
+
+    Parameters
+    ----------
+    output_redu
+        The reduced output value, shape is [nframes, *(odim0, odim1, ...)].
+    natoms
+        The number of atoms for each atom, shape is [nframes, ntypes].
+    assigned_bias
+        The assigned output bias, shape is [ntypes, *(odim0, odim1, ...)].
+        Set to a tensor of shape (odim0, odim1, ...) filled with nan if the bias
+        of the type is not assigned.
+
+    Returns
+    -------
+    np.ndarray
+        The computed output bias, shape is [ntypes, *(odim0, odim1, ...)].
+    np.ndarray
+        The computed output std, shape is [*(odim0, odim1, ...)].
+    """
+    natoms = np.array(natoms)  # [nf, ntypes]
+    nf, ntypes = natoms.shape
+    output_redu = np.array(output_redu)
+    var_shape = list(output_redu.shape[1:])
+    output_redu = output_redu.reshape(nf, -1)
+    # check shape
+    assert output_redu.ndim == 2
+    assert natoms.ndim == 2
+    assert output_redu.shape[0] == natoms.shape[0]  # [nf,1]
+
+    computed_output_bias = np.repeat(
+        np.mean(output_redu, axis=0)[np.newaxis, :], ntypes, axis=0
+    )
+    output_std = np.std(output_redu, axis=0)
+
+    computed_output_bias = computed_output_bias.reshape([natoms.shape[1]] + var_shape)  # noqa: RUF005
+    output_std = output_std.reshape(var_shape)
+
+    return computed_output_bias, output_std