fix: move imports to top-level and remove try/except from tests

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>

Author: Copilot

deepmd/tf/entrypoints/train.py

@@ -8,6 +8,9 @@
 import json
 import logging
 import time
+from pathlib import (
+    Path,
+)
 from typing import (
     Any,
     Optional,
@@ -48,6 +51,9 @@
 from deepmd.utils.data_system import (
     get_data,
 )
+from deepmd.utils.path import (
+    DPPath,
+)
 
 __all__ = ["train"]
 
@@ -236,14 +242,6 @@ def _do_work(
     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:
-            from pathlib import (
-                Path,
-            )
-
-            from deepmd.utils.path import (
-                DPPath,
-            )
-
             if not Path(stat_file_raw).exists():
                 if stat_file_raw.endswith((".h5", ".hdf5")):
                     with h5py.File(stat_file_raw, "w") as f:

deepmd/tf/model/ener.py

@@ -21,6 +21,9 @@
 from deepmd.tf.utils.spin import (
     Spin,
 )
+from deepmd.tf.utils.stat import (
+    compute_output_stats,
+)
 from deepmd.tf.utils.type_embed import (
     TypeEmbedNet,
 )
@@ -174,10 +177,6 @@ def _compute_output_stat(
     ) -> None:
         if stat_file_path is not None:
             # Use the new stat functionality with file save/load
-            from deepmd.tf.utils.stat import (
-                compute_output_stats,
-            )
-
             # Merge system stats for compatibility
             m_all_stat = merge_sys_stat(all_stat)
 

source/tests/tf/test_stat_file.py

@@ -51,40 +51,20 @@ def test_stat_file_tf(self) -> None:
                 mpi_log="master",
             )
 
-            try:
-                # Run training - this should create the stat file
-                _do_work(self.jdata, run_opt, is_compress=False)
+            # Run training - this should create the stat file
+            _do_work(self.jdata, run_opt, is_compress=False)
 
-                # Check if stat files were created
-                stat_path = Path(stat_file_path)
-                self.assertTrue(
-                    stat_path.exists(), "Stat file directory should be created"
-                )
+            # Check if stat files were created
+            stat_path = Path(stat_file_path)
+            self.assertTrue(stat_path.exists(), "Stat file directory should be created")
 
-                # Check for energy bias and std files
-                bias_file = stat_path / "bias_atom_energy"
-                std_file = stat_path / "std_atom_energy"
+            # Check for energy bias and std files
+            bias_file = stat_path / "bias_atom_energy"
+            std_file = stat_path / "std_atom_energy"
 
-                # At minimum, the directory structure should be created
-                # Even if files aren't created due to insufficient data, the directory should exist
-                self.assertTrue(
-                    stat_path.is_dir(), "Stat file path should be a directory"
-                )
-
-            except Exception as e:
-                # Print the exception for debugging but don't fail the test
-                # since we're mainly testing that the stat_file parameter is accepted
-                print(
-                    f"Training encountered an exception (expected for minimal test data): {e}"
-                )
-
-                # Still check that the stat file directory was created
-                stat_path = Path(stat_file_path)
-                if stat_path.exists():
-                    self.assertTrue(
-                        stat_path.is_dir(),
-                        "Stat file path should be a directory if created",
-                    )
+            # At minimum, the directory structure should be created
+            # Even if files aren't created due to insufficient data, the directory should exist
+            self.assertTrue(stat_path.is_dir(), "Stat file path should be a directory")
 
 
 if __name__ == "__main__":

source/tests/tf/test_stat_file_integration.py

@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""
-Integration test to validate stat_file functionality end-to-end
-"""
+"""Integration test to validate stat_file functionality end-to-end."""
 
 import json
 import os
@@ -76,27 +74,22 @@ def test_stat_file_save_and_load(self) -> None:
             with open(config_file, "w") as f:
                 json.dump(config, f, indent=2)
 
-            try:
-                # Attempt to run training
-                # This will fail due to missing data but should still process stat_file parameter
-                train(
-                    INPUT=config_file,
-                    init_model=None,
-                    restart=None,
-                    output=os.path.join(temp_dir, "output.json"),
-                    init_frz_model=None,
-                    mpi_log="master",
-                    log_level=20,
-                    log_path=None,
-                    is_compress=False,
-                    skip_neighbor_stat=True,
-                    finetune=None,
-                    use_pretrain_script=False,
-                )
-            except Exception as e:
-                # Expected to fail due to missing training data
-                # But the stat_file parameter should have been processed
-                print(f"Expected training failure: {e}")
+            # Attempt to run training
+            # This will fail due to missing data but should still process stat_file parameter
+            train(
+                INPUT=config_file,
+                init_model=None,
+                restart=None,
+                output=os.path.join(temp_dir, "output.json"),
+                init_frz_model=None,
+                mpi_log="master",
+                log_level=20,
+                log_path=None,
+                is_compress=False,
+                skip_neighbor_stat=True,
+                finetune=None,
+                use_pretrain_script=False,
+            )
 
             # The main validation is that the code didn't crash with an unrecognized parameter
             # and that if the stat file directory was attempted to be created, it exists
@@ -105,7 +98,6 @@ def test_stat_file_save_and_load(self) -> None:
                 self.assertTrue(
                     stat_path.is_dir(), "Stat file path should be a directory"
                 )
-                print(f"Stat file directory was created: {stat_file_path}")
 
             # This test primarily validates that the stat_file parameter is accepted
             # and processed without errors in the TF pipeline