Init branch

deepmd/pt/loss/property.py

@@ -85,35 +85,40 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False
         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([label["dipole_moment"],label["homo"]],dim=1)
+        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",
                 )
             )

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/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,15 +291,18 @@ def compute_output_stats(
         # remove the keys that are not in the sample
         keys = [keys] if isinstance(keys, str) else keys
         assert isinstance(keys, list)
+        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
             if (ii in sampled[0].keys()) or ("atom_" + ii in sampled[0].keys())
         ]
         del keys
         keys = new_keys
-        from IPython import embed
-        embed()
         # split system based on label
         atomic_sampled_idx = defaultdict(list)
         global_sampled_idx = defaultdict(list)
@@ -375,6 +379,7 @@ def compute_output_stats(
 
         # 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:
@@ -478,26 +483,42 @@ def compute_output_stats_global(
     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/out_stat.py

@@ -130,3 +130,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.
+    rcond
+        Cut-off ratio for small singular values of a.
+
+    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