Merge branch 'devel' into D0812_devel_head_alias

Author: iProzd

deepmd/dpmodel/fitting/property_fitting.py

@@ -12,6 +12,10 @@
 from deepmd.dpmodel.fitting.invar_fitting import (
     InvarFitting,
 )
+from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
+    OutputVariableDef,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -108,6 +112,20 @@ def __init__(
             type_map=type_map,
         )
 
+    def output_def(self) -> FittingOutputDef:
+        return FittingOutputDef(
+            [
+                OutputVariableDef(
+                    self.var_name,
+                    [self.dim_out],
+                    reducible=True,
+                    r_differentiable=False,
+                    c_differentiable=False,
+                    intensive=self.intensive,
+                ),
+            ]
+        )
+
     @classmethod
     def deserialize(cls, data: dict) -> "PropertyFittingNet":
         data = data.copy()

deepmd/dpmodel/model/dp_model.py

@@ -45,3 +45,7 @@ def update_sel(
             train_data, type_map, local_jdata["descriptor"]
         )
         return local_jdata_cpy, min_nbor_dist
+
+    def get_fitting_net(self):
+        """Get the fitting network."""
+        return self.atomic_model.fitting

deepmd/dpmodel/model/make_model.py

@@ -355,6 +355,7 @@ def forward_common_atomic(
                 self.atomic_output_def(),
                 extended_coord,
                 do_atomic_virial=do_atomic_virial,
+                mask=atomic_ret["mask"] if "mask" in atomic_ret else None,
             )
 
         forward_lower = call_lower

deepmd/dpmodel/model/property_model.py

@@ -25,3 +25,7 @@ def __init__(
     ) -> None:
         DPModelCommon.__init__(self)
         DPPropertyModel_.__init__(self, *args, **kwargs)
+
+    def get_var_name(self) -> str:
+        """Get the name of the property."""
+        return self.get_fitting_net().var_name

deepmd/dpmodel/model/transform_output.py

@@ -1,5 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
+from typing import (
+    Optional,
+)
+
 import array_api_compat
 import numpy as np
 
@@ -24,6 +28,7 @@ def fit_output_to_model_output(
     fit_output_def: FittingOutputDef,
     coord_ext: np.ndarray,
     do_atomic_virial: bool = False,
+    mask: Optional[np.ndarray] = None,
 ) -> dict[str, np.ndarray]:
     """Transform the output of the fitting network to
     the model output.
@@ -38,9 +43,19 @@ def fit_output_to_model_output(
         if vdef.reducible:
             kk_redu = get_reduce_name(kk)
             # cast to energy prec before reduction
-            model_ret[kk_redu] = xp.sum(
-                vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis
-            )
+            if vdef.intensive:
+                if mask is not None:
+                    model_ret[kk_redu] = xp.sum(
+                        vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis
+                    ) / np.sum(mask, axis=-1, keepdims=True)
+                else:
+                    model_ret[kk_redu] = xp.mean(
+                        vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis
+                    )
+            else:
+                model_ret[kk_redu] = xp.sum(
+                    vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis
+                )
             if vdef.r_differentiable:
                 kk_derv_r, kk_derv_c = get_deriv_name(kk)
                 # name-holders

deepmd/jax/model/base_model.py

@@ -46,7 +46,16 @@ def forward_common_atomic(
         atom_axis = -(len(shap) + 1)
         if vdef.reducible:
             kk_redu = get_reduce_name(kk)
-            model_predict[kk_redu] = jnp.sum(vv, axis=atom_axis)
+            if vdef.intensive:
+                mask = atomic_ret["mask"] if "mask" in atomic_ret else None
+                if mask is not None:
+                    model_predict[kk_redu] = jnp.sum(vv, axis=atom_axis) / jnp.sum(
+                        mask, axis=-1, keepdims=True
+                    )
+                else:
+                    model_predict[kk_redu] = jnp.mean(vv, axis=atom_axis)
+            else:
+                model_predict[kk_redu] = jnp.sum(vv, axis=atom_axis)
             kk_derv_r, kk_derv_c = get_deriv_name(kk)
             if vdef.r_differentiable:
 

deepmd/pd/train/training.py

@@ -918,7 +918,9 @@ def log_loss_valid(_task_key="Default"):
                 self.t0 = current_time
                 if self.rank == 0 and self.timing_in_training:
                     eta = int(
-                        (self.num_steps - display_step_id) / self.disp_freq * train_time
+                        (self.num_steps - display_step_id)
+                        / min(self.disp_freq, display_step_id - self.start_step)
+                        * train_time
                     )
                     log.info(
                         format_training_message(

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -363,21 +363,25 @@ def compute_or_load_stat(
         self,
         merged: Union[Callable[[], list[dict]], list[dict]],
         stat_file_path: Optional[DPPath] = None,
+        compute_or_load_out_stat: bool = True,
     ) -> NoReturn:
         """
-        Compute the output statistics (e.g. energy bias) for the fitting net from packed data.
+        Compute or load the statistics parameters of the model,
+        such as mean and standard deviation of descriptors or the energy bias of the fitting net.
+        When `sampled` is provided, all the statistics parameters will be calculated (or re-calculated for update),
+        and saved in the `stat_file_path`(s).
+        When `sampled` is not provided, it will check the existence of `stat_file_path`(s)
+        and load the calculated statistics parameters.
 
         Parameters
         ----------
-        merged : Union[Callable[[], list[dict]], list[dict]]
-            - list[dict]: A list of data samples from various data systems.
-                Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor`
-                originating from the `i`-th data system.
-            - Callable[[], list[dict]]: A lazy function that returns data samples in the above format
-                only when needed. Since the sampling process can be slow and memory-intensive,
-                the lazy function helps by only sampling once.
-        stat_file_path : Optional[DPPath]
-            The path to the stat file.
+        merged
+            The lazy sampled function to get data frames from different data systems.
+        stat_file_path
+            The dictionary of paths to the statistics files.
+        compute_or_load_out_stat : bool
+            Whether to compute the output statistics.
+            If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors).
 
         """
         raise NotImplementedError

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -285,6 +285,7 @@ def compute_or_load_stat(
         self,
         sampled_func,
         stat_file_path: Optional[DPPath] = None,
+        compute_or_load_out_stat: bool = True,
     ) -> None:
         """
         Compute or load the statistics parameters of the model,
@@ -300,6 +301,9 @@ def compute_or_load_stat(
             The lazy sampled function to get data frames from different data systems.
         stat_file_path
             The dictionary of paths to the statistics files.
+        compute_or_load_out_stat : bool
+            Whether to compute the output statistics.
+            If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors).
         """
         if stat_file_path is not None and self.type_map is not None:
             # descriptors and fitting net with different type_map
@@ -323,7 +327,8 @@ def wrapped_sampler():
         self.fitting_net.compute_input_stats(
             wrapped_sampler, protection=self.data_stat_protect
         )
-        self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)
+        if compute_or_load_out_stat:
+            self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)
 
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import functools
 from typing import (
-    Callable,
     Optional,
     Union,
 )
@@ -319,6 +319,10 @@ def apply_out_stat(
             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
 
     @staticmethod
@@ -464,34 +468,11 @@ def is_aparam_nall(self) -> bool:
         """
         return False
 
-    def compute_or_load_out_stat(
-        self,
-        merged: Union[Callable[[], list[dict]], list[dict]],
-        stat_file_path: Optional[DPPath] = None,
-    ) -> None:
-        """
-        Compute the output statistics (e.g. energy bias) for the fitting net from packed data.
-
-        Parameters
-        ----------
-        merged : Union[Callable[[], list[dict]], list[dict]]
-            - list[dict]: A list of data samples from various data systems.
-                Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor`
-                originating from the `i`-th data system.
-            - Callable[[], list[dict]]: A lazy function that returns data samples in the above format
-                only when needed. Since the sampling process can be slow and memory-intensive,
-                the lazy function helps by only sampling once.
-        stat_file_path : Optional[DPPath]
-            The path to the stat file.
-
-        """
-        for md in self.models:
-            md.compute_or_load_out_stat(merged, stat_file_path)
-
     def compute_or_load_stat(
         self,
         sampled_func,
         stat_file_path: Optional[DPPath] = None,
+        compute_or_load_out_stat: bool = True,
     ) -> None:
         """
         Compute or load the statistics parameters of the model,
@@ -507,9 +488,34 @@ def compute_or_load_stat(
             The lazy sampled function to get data frames from different data systems.
         stat_file_path
             The dictionary of paths to the statistics files.
+        compute_or_load_out_stat : bool
+            Whether to compute the output statistics.
+            If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors).
         """
         for md in self.models:
-            md.compute_or_load_stat(sampled_func, stat_file_path)
+            md.compute_or_load_stat(
+                sampled_func, stat_file_path, compute_or_load_out_stat=False
+            )
+
+        if stat_file_path is not None and self.type_map is not None:
+            # descriptors and fitting net with different type_map
+            # should not share the same parameters
+            stat_file_path /= " ".join(self.type_map)
+
+        @functools.lru_cache
+        def wrapped_sampler():
+            sampled = sampled_func()
+            if self.pair_excl is not None:
+                pair_exclude_types = self.pair_excl.get_exclude_types()
+                for sample in sampled:
+                    sample["pair_exclude_types"] = list(pair_exclude_types)
+            if self.atom_excl is not None:
+                atom_exclude_types = self.atom_excl.get_exclude_types()
+                for sample in sampled:
+                    sample["atom_exclude_types"] = list(atom_exclude_types)
+            return sampled
+
+        self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)
 
 
 class DPZBLLinearEnergyAtomicModel(LinearEnergyAtomicModel):