Add fparam and aparam

iProzd · Feb 29, 2024 · ab35653 · ab35653
1 parent 90be50e
commit ab35653
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
@@ -861,6 +861,8 @@ def get_data(self, is_train=True, task_key="Default"):
             "atype",
             "box",
             "spin",
+            "fparam",
+            "aparam",
         ]
         input_dict = {item_key: None for item_key in input_keys}
         label_dict = {}

diff --git a/deepmd/utils/env_mat_stat.py b/deepmd/utils/env_mat_stat.py
@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import logging
 from abc import (
     ABC,
     abstractmethod,
@@ -19,6 +20,8 @@
     DPPath,
 )
 
+log = logging.getLogger(__name__)
+
 
 class StatItem:
     """A class to store the statistics of the environment matrix.
@@ -170,10 +173,12 @@ def load_or_compute_stats(
         """
         if path is not None and path.is_dir():
             self.load_stats(path)
+            log.info(f"Load stats from {path}.")
         else:
             self.compute_stats(data)
             if path is not None:
                 self.save_stats(path)
+                log.info(f"Save stats to {path}.")
 
     def get_avg(self, default: float = 0) -> Dict[str, float]:
         """Get the average of the environment matrix.