minor code improvement based on @coderabbitai

Author: ChiahsinChu

deepmd/pt/modifier/dipole_charge.py

@@ -129,13 +129,14 @@ def forward(
             Atom parameters with shape (nframes, natoms, nap), by default None
         do_atomic_virial : bool, optional
             Whether to compute atomic virial, by default False
+            Note: This parameter is currently not implemented and is ignored
 
         Returns
         -------
         dict[str, torch.Tensor]
             Dictionary containing the correction terms:
             - energy: Energy correction tensor with shape (nframes, 1)
-            - force: Force correction tensor with shape (nframes, natoms+nsel, 3)
+            - force: Force correction tensor with shape (nframes, natoms, 3)
             - virial: Virial correction tensor with shape (nframes, 3, 3)
         """
         if box is None:
@@ -305,7 +306,10 @@ def extend_system_coord(
 
         # nframe x natoms x 3
         dipole = torch.cat(all_dipole, dim=0)
-        assert dipole.shape[0] == nframes
+        if dipole.shape[0] != nframes:
+            raise RuntimeError(
+                f"Dipole shape mismatch: expected {nframes} frames, got {dipole.shape[0]}"
+            )
 
         dipole_reshaped = dipole.reshape(nframes, natoms, 3)
         coord_reshaped = coord.reshape(nframes, natoms, 3)

source/tests/pt/modifier/test_dipole_charge.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import json
 import os
+import tempfile
 import unittest
 from pathlib import (
     Path,
@@ -50,6 +51,9 @@ def ref_data():
 
 class TestDipoleChargeModifier(unittest.TestCase):
     def setUp(self) -> None:
+        self.test_dir = tempfile.TemporaryDirectory()
+        self.orig_dir = os.getcwd()
+        os.chdir(self.test_dir.name)
         # setup parameter
         # numerical consistency can only be achieved with high prec
         self.ewald_h = 0.1
@@ -170,7 +174,7 @@ def test_serialize(self):
             to_numpy_array(ret0["virial"]), to_numpy_array(ret1["virial"])
         )
 
-    def test_box_none_warning(self):
+    def test_box_none_error(self):
         """Test that a RuntimeError is raised when box is None."""
         coord, _b, atype = ref_data()
         # consistent with the input shape from BaseModifier.modify_data
@@ -209,12 +213,5 @@ def test_train(self):
         trainer.run()
 
     def tearDown(self) -> None:
-        for f in os.listdir("."):
-            if f.startswith("frozen_model") and f.endswith(".pth"):
-                os.remove(f)
-            if f.startswith("dw_model") and (f.endswith(".pth") or f.endswith(".pb")):
-                os.remove(f)
-            if f.startswith("model.ckpt") and f.endswith(".pt"):
-                os.remove(f)
-            if f in ["lcurve.out", "checkpoint"]:
-                os.remove(f)
+        os.chdir(self.orig_dir)
+        self.test_dir.cleanup()