feat(tf): support training stat_file

deepmd/tf/descriptor/se_a.py

@@ -77,6 +77,9 @@
 from .se import (
     DescrptSe,
 )
+from .stat import (
+    load_or_compute_se_input_stats,
+)
 
 
 @Descriptor.register("se_e2_a")
@@ -374,7 +377,8 @@ def compute_input_stats(
         **kwargs
             Additional keyword arguments.
         """
-        if True:
+
+        def compute_stats() -> dict[str, Any]:
             sumr = []
             suma = []
             sumn = []
@@ -398,7 +402,16 @@ def compute_input_stats(
                 "sumr2": sumr2,
                 "suma2": suma2,
             }
-            self.merge_input_stats(stat_dict)
+            return stat_dict
+
+        stat_dict = load_or_compute_se_input_stats(
+            self,
+            kwargs.get("stat_file_path"),
+            last_dim=4,
+            compute=compute_stats,
+            mixed_types=False,
+        )
+        self.merge_input_stats(stat_dict)
 
     def merge_input_stats(self, stat_dict: dict[str, Any]) -> None:
         """Merge the statistics computed from compute_input_stats to obtain the self.davg and self.dstd.

deepmd/tf/descriptor/se_atten.py

@@ -90,6 +90,9 @@
 from .se_a import (
     DescrptSeA,
 )
+from .stat import (
+    load_or_compute_se_input_stats,
+)
 
 log = logging.getLogger(__name__)
 
@@ -373,7 +376,8 @@ def compute_input_stats(
         **kwargs
             Additional keyword arguments.
         """
-        if True:
+
+        def compute_stats() -> dict[str, Any]:
             sumr = []
             suma = []
             sumn = []
@@ -418,7 +422,16 @@ def compute_input_stats(
                 "sumr2": sumr2,
                 "suma2": suma2,
             }
-            self.merge_input_stats(stat_dict)
+            return stat_dict
+
+        stat_dict = load_or_compute_se_input_stats(
+            self,
+            kwargs.get("stat_file_path"),
+            last_dim=4,
+            compute=compute_stats,
+            mixed_types=True,
+        )
+        self.merge_input_stats(stat_dict)
 
     def enable_compression(
         self,

deepmd/tf/descriptor/se_r.py

@@ -50,6 +50,9 @@
 from .se import (
     DescrptSe,
 )
+from .stat import (
+    load_or_compute_se_input_stats,
+)
 
 
 @Descriptor.register("se_e2_r")
@@ -274,17 +277,27 @@ def compute_input_stats(
         **kwargs
             Additional keyword arguments.
         """
-        sumr = []
-        sumn = []
-        sumr2 = []
-        for cc, bb, tt, nn, mm in zip(
-            data_coord, data_box, data_atype, natoms_vec, mesh, strict=True
-        ):
-            sysr, sysr2, sysn = self._compute_dstats_sys_se_r(cc, bb, tt, nn, mm)
-            sumr.append(sysr)
-            sumn.append(sysn)
-            sumr2.append(sysr2)
-        stat_dict = {"sumr": sumr, "sumn": sumn, "sumr2": sumr2}
+
+        def compute_stats() -> dict[str, Any]:
+            sumr = []
+            sumn = []
+            sumr2 = []
+            for cc, bb, tt, nn, mm in zip(
+                data_coord, data_box, data_atype, natoms_vec, mesh, strict=True
+            ):
+                sysr, sysr2, sysn = self._compute_dstats_sys_se_r(cc, bb, tt, nn, mm)
+                sumr.append(sysr)
+                sumn.append(sysn)
+                sumr2.append(sysr2)
+            return {"sumr": sumr, "sumn": sumn, "sumr2": sumr2}
+
+        stat_dict = load_or_compute_se_input_stats(
+            self,
+            kwargs.get("stat_file_path"),
+            last_dim=1,
+            compute=compute_stats,
+            mixed_types=False,
+        )
         self.merge_input_stats(stat_dict)
 
     def merge_input_stats(self, stat_dict: dict[str, Any]) -> None:

deepmd/tf/descriptor/se_t.py

@@ -52,6 +52,9 @@
 from .se import (
     DescrptSe,
 )
+from .stat import (
+    load_or_compute_se_input_stats,
+)
 
 
 @Descriptor.register("se_e3")
@@ -257,7 +260,8 @@ def compute_input_stats(
         **kwargs
             Additional keyword arguments.
         """
-        if True:
+
+        def compute_stats() -> dict[str, Any]:
             sumr = []
             suma = []
             sumn = []
@@ -281,7 +285,16 @@ def compute_input_stats(
                 "sumr2": sumr2,
                 "suma2": suma2,
             }
-            self.merge_input_stats(stat_dict)
+            return stat_dict
+
+        stat_dict = load_or_compute_se_input_stats(
+            self,
+            kwargs.get("stat_file_path"),
+            last_dim=4,
+            compute=compute_stats,
+            mixed_types=False,
+        )
+        self.merge_input_stats(stat_dict)
 
     def merge_input_stats(self, stat_dict: dict[str, Any]) -> None:
         """Merge the statistics computed from compute_input_stats to obtain the self.davg and self.dstd.

deepmd/tf/descriptor/stat.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from collections.abc import (
+    Callable,
+)
+from typing import (
+    Any,
+)
+
+import numpy as np
+
+from deepmd.common import (
+    get_hash,
+)
+from deepmd.utils.path import (
+    DPPath,
+)
+
+
+def _descriptor_rcut_smth(descrpt: Any) -> float:
+    if hasattr(descrpt, "rcut_smth"):
+        return descrpt.rcut_smth
+    return descrpt.rcut_r_smth
+
+
+def _descriptor_sel(descrpt: Any, last_dim: int) -> list[int]:
+    if hasattr(descrpt, "get_sel"):
+        sel = descrpt.get_sel()
+    elif last_dim == 1:
+        sel = descrpt.sel_r
+    else:
+        sel = descrpt.sel_a
+    if isinstance(sel, np.ndarray):
+        sel = sel.tolist()
+    elif isinstance(sel, int):
+        sel = [sel]
+    return [int(ii) for ii in sel]
+
+
+def _descriptor_stat_path(
+    descrpt: Any,
+    stat_file_path: DPPath | None,
+    last_dim: int,
+    mixed_types: bool,
+) -> DPPath | None:
+    if stat_file_path is None:
+        return None
+    sel = _descriptor_sel(descrpt, last_dim)
+    stat_hash = get_hash(
+        {
+            "type": "se_a" if last_dim == 4 else "se_r",
+            "ntypes": descrpt.get_ntypes(),
+            "rcut": round(descrpt.get_rcut(), 2),
+            "rcut_smth": round(_descriptor_rcut_smth(descrpt), 2),
+            "nsel": sum(sel),
+            "sel": sel,
+            "mixed_types": mixed_types,
+        }
+    )
+    return stat_file_path / stat_hash
+
+
+def _stat_keys(ntypes: int, angular: bool) -> list[str]:
+    keys = [f"r_{ii}" for ii in range(ntypes)]
+    if angular:
+        keys.extend(f"a_{ii}" for ii in range(ntypes))
+    return keys
+
+
+def _load_se_input_stats(
+    path: DPPath | None,
+    ntypes: int,
+    angular: bool,
+) -> dict[str, list[list[float]]] | None:
+    if path is None or not path.is_dir():
+        return None
+    if any(not (path / kk).is_file() for kk in _stat_keys(ntypes, angular)):
+        return None
+
+    sumr = []
+    sumn = []
+    sumr2 = []
+    suma = []
+    suma2 = []
+    for type_i in range(ntypes):
+        r_stat = (path / f"r_{type_i}").load_numpy()
+        sumn.append(float(r_stat[0]))
+        sumr.append(float(r_stat[1]))
+        sumr2.append(float(r_stat[2]))
+        if angular:
+            a_stat = (path / f"a_{type_i}").load_numpy()
+            suma.append(float(a_stat[1]) / 3.0)
+            suma2.append(float(a_stat[2]) / 3.0)
+
+    ret = {
+        "sumr": [sumr],
+        "sumn": [sumn],
+        "sumr2": [sumr2],
+    }
+    if angular:
+        ret["suma"] = [suma]
+        ret["suma2"] = [suma2]
+    return ret
+
+
+def _save_se_input_stats(
+    path: DPPath | None,
+    stat_dict: dict[str, Any],
+    ntypes: int,
+    angular: bool,
+) -> None:
+    if path is None:
+        return
+    path.mkdir(parents=True, exist_ok=True)
+
+    sumr = np.sum(stat_dict["sumr"], axis=0)
+    sumn = np.sum(stat_dict["sumn"], axis=0)
+    sumr2 = np.sum(stat_dict["sumr2"], axis=0)
+    if angular:
+        suma = np.sum(stat_dict["suma"], axis=0)
+        suma2 = np.sum(stat_dict["suma2"], axis=0)
+
+    for type_i in range(ntypes):
+        (path / f"r_{type_i}").save_numpy(
+            np.array([sumn[type_i], sumr[type_i], sumr2[type_i]])
+        )
+        if angular:
+            (path / f"a_{type_i}").save_numpy(
+                np.array([3.0 * sumn[type_i], 3.0 * suma[type_i], 3.0 * suma2[type_i]])
+            )
+
+
+def load_or_compute_se_input_stats(
+    descrpt: Any,
+    stat_file_path: DPPath | None,
+    last_dim: int,
+    compute: Callable[[], dict[str, Any]],
+    mixed_types: bool = False,
+) -> dict[str, Any]:
+    """Load or compute SE descriptor input statistics using EnvMatStatSe format."""
+    angular = last_dim == 4
+    stat_path = _descriptor_stat_path(descrpt, stat_file_path, last_dim, mixed_types)
+    stat_dict = _load_se_input_stats(stat_path, descrpt.get_ntypes(), angular)
+    if stat_dict is not None:
+        return stat_dict
+
+    stat_dict = compute()
+    _save_se_input_stats(stat_path, stat_dict, descrpt.get_ntypes(), angular)
+    return stat_dict

deepmd/tf/entrypoints/train.py

@@ -8,10 +8,14 @@
 import json
 import logging
 import time
+from pathlib import (
+    Path,
+)
 from typing import (
     Any,
 )
 
+import h5py
 import numpy as np
 
 from deepmd.common import (
@@ -50,6 +54,9 @@
 from deepmd.utils.data_system import (
     get_data,
 )
+from deepmd.utils.path import (
+    DPPath,
+)
 
 __all__ = ["train"]
 
@@ -232,6 +239,21 @@ def _do_work(
     # setup data modifier
     modifier = get_modifier(jdata["model"].get("modifier", None))
 
+    # extract stat_file from training parameters
+    stat_file_path = None
+    if not is_compress:
+        stat_file_raw = jdata["training"].get("stat_file", None)
+        if stat_file_raw is not None and run_opt.is_chief:
+            stat_file_target = Path(stat_file_raw)
+            stat_file_target.parent.mkdir(parents=True, exist_ok=True)
+            if not stat_file_target.exists():
+                if stat_file_raw.endswith((".h5", ".hdf5")):
+                    with h5py.File(stat_file_raw, "w") as f:
+                        pass
+                else:
+                    stat_file_target.mkdir(parents=True, exist_ok=True)
+            stat_file_path = DPPath(stat_file_raw, "a")
+
     # decouple the training data from the model compress process
     train_data = None
     valid_data = None
@@ -289,7 +311,12 @@ def _do_work(
         origin_type_map = get_data(
             jdata["training"]["training_data"], rcut, None, modifier
         ).get_type_map()
-    model.build(train_data, stop_batch, origin_type_map=origin_type_map)
+    model.build(
+        train_data,
+        stop_batch,
+        origin_type_map=origin_type_map,
+        stat_file_path=stat_file_path,
+    )
 
     if not is_compress:
         # train the model with the provided systems in a cyclic way