Fix bug and clean up codes based on coderabbit

Author: ChiahsinChu

deepmd/tf/infer/deep_eval.py

@@ -144,7 +144,7 @@ def __init__(
         self.dm = None
         if self.modifier_type is not None:
             modifier = BaseModifier.get_class_by_type(self.modifier_type)
-            modifier_params = modifier.get_params(self)
+            modifier_params = modifier.get_params_from_frozen_model(self)
             self.dm = modifier.get_modifier(modifier_params)
 
     def _init_tensors(self) -> None:
@@ -662,8 +662,8 @@ def _get_natoms_and_nframes(
         coords: np.ndarray,
         atom_types: list[int] | np.ndarray,
     ) -> tuple[int, int]:
-        atom_types = np.reshape(atom_types, (-1))
-        natoms = len(atom_types)
+        # (natoms,) or (nframes, natoms,)
+        natoms = np.shape(atom_types)[-1]
         if natoms == 0:
             assert coords.size == 0
         else:

deepmd/tf/modifier/base_modifier.py

@@ -1,4 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from abc import (
+    abstractmethod,
+)
+
 from deepmd.dpmodel.modifier.base_modifier import (
     make_base_modifier,
 )
@@ -11,3 +15,23 @@ class BaseModifier(DeepPot, make_base_modifier()):
     def __init__(self, *args, **kwargs) -> None:
         """Construct a basic model for different tasks."""
         DeepPot.__init__(self, *args, **kwargs)
+
+    @staticmethod
+    @abstractmethod
+    def get_params_from_frozen_model(model) -> dict:
+        """Extract the modifier parameters from a model.
+
+        This method should extract the necessary parameters from a model
+        to create an instance of this modifier.
+
+        Parameters
+        ----------
+        model
+            The model from which to extract parameters
+
+        Returns
+        -------
+        dict
+            The modifier parameters
+        """
+        pass

deepmd/tf/modifier/dipole_charge.py

@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import os
+from typing import (
+    TYPE_CHECKING,
+)
 
 import numpy as np
 
@@ -13,9 +16,6 @@
     op_module,
     tf,
 )
-from deepmd.tf.infer import (
-    DeepEval,
-)
 from deepmd.tf.infer.deep_dipole import DeepDipoleOld as DeepDipole
 from deepmd.tf.infer.ewald_recp import (
     EwaldRecp,
@@ -30,6 +30,11 @@
     run_sess,
 )
 
+if TYPE_CHECKING:
+    from deepmd.tf.infer import (
+        DeepEval,
+    )
+
 
 @BaseModifier.register("dipole_charge")
 class DipoleChargeModifier(DeepDipole, BaseModifier):
@@ -492,7 +497,24 @@ def modify_data(self, data: dict, data_sys: DeepmdData) -> None:
             data["virial"] -= tot_v.reshape(data["virial"].shape)
 
     @staticmethod
-    def get_params(model: DeepEval):
+    def get_params_from_frozen_model(model: "DeepEval") -> dict:
+        """Extract modifier parameters from a DeepEval model.
+
+        Parameters
+        ----------
+        model : DeepEval
+            The DeepEval model instance containing the modifier tensors.
+
+        Returns
+        -------
+        dict
+            Dictionary containing modifier parameters:
+            - model_name : str
+            - model_charge_map : list[int]
+            - sys_charge_map : list[int]
+            - ewald_h : float
+            - ewald_beta : float
+        """
         t_mdl_name = model._get_tensor("modifier_attr/mdl_name:0")
         t_mdl_charge_map = model._get_tensor("modifier_attr/mdl_charge_map:0")
         t_sys_charge_map = model._get_tensor("modifier_attr/sys_charge_map:0")