rearrange the def of atomic model. provide base class for fitting. re…

…moved the task base.

deepmd/dpmodel/fitting/__init__.py

@@ -1,8 +1,12 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from .fitting import (
+from .invar_fitting import (
     InvarFitting,
 )
+from .make_base_fitting import (
+    make_base_fitting,
+)
 
 __all__ = [
     "InvarFitting",
+    "make_base_fitting",
 ]

deepmd/dpmodel/fitting/base_fitting.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import numpy as np
+
+from .make_base_fitting import (
+    make_base_fitting,
+)
+
+BaseFitting = make_base_fitting(np.ndarray, "call")

deepmd/dpmodel/fitting/fitting.py → deepmd/dpmodel/fitting/invar_fitting.py

@@ -2,6 +2,7 @@
 import copy
 from typing import (
     Any,
+    Dict,
     List,
     Optional,
 )
@@ -22,10 +23,14 @@
     NetworkCollection,
 )
 
+from .base_fitting import (
+    BaseFitting,
+)
+
 
 @fitting_check_output
-class InvarFitting(NativeOP):
-    r"""Fitting the energy (or a porperty of `dim_out`) of the system. The force and the virial can also be trained.
+class InvarFitting(NativeOP, BaseFitting):
+    r"""Fitting the energy (or a rotationally invariant porperty of `dim_out`) of the system. The force and the virial can also be trained.
 
     Lets take the energy fitting task as an example.
     The potential energy :math:`E` is a fitting network function of the descriptor :math:`\mathcal{D}`:
@@ -279,7 +284,7 @@ def call(
         h2: Optional[np.array] = None,
         fparam: Optional[np.array] = None,
         aparam: Optional[np.array] = None,
-    ):
+    ) -> Dict[str, np.array]:
         """Calculate the fitting.
 
         Parameters

deepmd/dpmodel/fitting/make_base_fitting.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from abc import (
+    ABC,
+    abstractclassmethod,
+    abstractmethod,
+)
+from typing import (
+    Dict,
+    Optional,
+)
+
+from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
+)
+
+
+def make_base_fitting(
+    T_Tensor,
+    FWD_Method: str = "call",
+):
+    """Make the base class for the fitting."""
+
+    class BF(ABC):
+        """Base fitting provides the interfaces of fitting net."""
+
+        @abstractmethod
+        def output_def(self) -> FittingOutputDef:
+            pass
+
+        @abstractmethod
+        def fwd(
+            self,
+            descriptor: T_Tensor,
+            atype: T_Tensor,
+            gr: Optional[T_Tensor] = None,
+            g2: Optional[T_Tensor] = None,
+            h2: Optional[T_Tensor] = None,
+            fparam: Optional[T_Tensor] = None,
+            aparam: Optional[T_Tensor] = None,
+        ) -> Dict[str, T_Tensor]:
+            pass
+
+        @abstractmethod
+        def serialize(self) -> dict:
+            pass
+
+        @abstractclassmethod
+        def deserialize(cls):
+            pass
+
+    setattr(BF, FWD_Method, BF.fwd)
+    delattr(BF, "fwd")
+
+    return BF

deepmd/dpmodel/model/__init__.py

@@ -5,7 +5,7 @@
 from .dp_model import (
     DPModel,
 )
-from .make_atomic_model import (
+from .make_base_atomic_model import (
     make_base_atomic_model,
 )
 

deepmd/dpmodel/model/base_atomic_model.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import numpy as np
 
-from .make_atomic_model import (
+from .make_base_atomic_model import (
     make_base_atomic_model,
 )
 

deepmd/dpmodel/model/dp_atomic_model.py

@@ -50,7 +50,7 @@ def __init__(
         self.descriptor = descriptor
         self.fitting = fitting
 
-    def get_fitting_output_def(self) -> FittingOutputDef:
+    def fitting_output_def(self) -> FittingOutputDef:
         """Get the output def of the fitting net."""
         return self.fitting.output_def()
 

deepmd/dpmodel/model/make_atomic_model.py → deepmd/dpmodel/model/make_base_atomic_model.py

@@ -20,7 +20,7 @@ class BAM(ABC):
         """Base Atomic Model provides the interfaces of an atomic model."""
 
         @abstractmethod
-        def get_fitting_output_def(self) -> FittingOutputDef:
+        def fitting_output_def(self) -> FittingOutputDef:
             pass
 
         @abstractmethod
@@ -55,4 +55,29 @@ def serialize(self) -> dict:
         def deserialize(cls):
             pass
 
+        def do_grad(
+            self,
+            var_name: Optional[str] = None,
+        ) -> bool:
+            """Tell if the output variable `var_name` is differentiable.
+            if var_name is None, returns if any of the variable is differentiable.
+
+            """
+            odef = self.fitting_output_def()
+            if var_name is None:
+                require: List[bool] = []
+                for vv in odef.keys():
+                    require.append(self.do_grad_(vv))
+                return any(require)
+            else:
+                return self.do_grad_(var_name)
+
+        def do_grad_(
+            self,
+            var_name: str,
+        ) -> bool:
+            """Tell if the output variable `var_name` is differentiable."""
+            assert var_name is not None
+            return self.fitting_output_def()[var_name].differentiable
+
     return BAM

deepmd/dpmodel/model/make_model.py

@@ -56,9 +56,9 @@ def __init__(
                 **kwargs,
             )
 
-        def get_model_output_def(self):
+        def model_output_def(self):
             """Get the output def for the model."""
-            return ModelOutputDef(self.get_fitting_output_def())
+            return ModelOutputDef(self.fitting_output_def())
 
         def call(
             self,
@@ -125,7 +125,7 @@ def call(
             )
             model_predict = communicate_extended_output(
                 model_predict_lower,
-                self.get_model_output_def(),
+                self.model_output_def(),
                 mapping,
                 do_atomic_virial=do_atomic_virial,
             )
@@ -182,7 +182,7 @@ def call_lower(
             )
             model_predict = fit_output_to_model_output(
                 atomic_ret,
-                self.get_fitting_output_def(),
+                self.fitting_output_def(),
                 extended_coord,
                 do_atomic_virial=do_atomic_virial,
             )

deepmd/pt/model/model/base_atomic_model.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+
+import torch
+
+from deepmd.dpmodel.model import (
+    make_base_atomic_model,
+)
+
+BaseAtomicModel = make_base_atomic_model(torch.Tensor)

deepmd/pt/model/model/dp_atomic_model.py

@@ -19,15 +19,15 @@
     InvarFitting,
 )
 
-from .atomic_model import (
-    AtomicModel,
+from .base_atomic_model import (
+    BaseAtomicModel,
 )
 from .model import (
     BaseModel,
 )
 
 
-class DPAtomicModel(BaseModel, AtomicModel):
+class DPAtomicModel(BaseModel, BaseAtomicModel):
     """Model give atomic prediction of some physical property.
 
     Parameters
@@ -84,7 +84,7 @@ def __init__(
             sampled=sampled,
         )
 
-    def get_fitting_output_def(self) -> FittingOutputDef:
+    def fitting_output_def(self) -> FittingOutputDef:
         """Get the output def of the fitting net."""
         return (
             self.fitting_net.output_def()

deepmd/pt/model/model/make_model.py

@@ -57,9 +57,9 @@ def __init__(
                 **kwargs,
             )
 
-        def get_model_output_def(self):
+        def model_output_def(self):
             """Get the output def for the model."""
-            return ModelOutputDef(self.get_fitting_output_def())
+            return ModelOutputDef(self.fitting_output_def())
 
         # cannot use the name forward. torch script does not work
         def forward_common(
@@ -123,7 +123,7 @@ def forward_common(
             )
             model_predict = communicate_extended_output(
                 model_predict_lower,
-                self.get_model_output_def(),
+                self.model_output_def(),
                 mapping,
                 do_atomic_virial=do_atomic_virial,
             )
@@ -176,7 +176,7 @@ def forward_common_lower(
             )
             model_predict = fit_output_to_model_output(
                 atomic_ret,
-                self.get_fitting_output_def(),
+                self.fitting_output_def(),
                 extended_coord,
                 do_atomic_virial=do_atomic_virial,
             )

deepmd/pt/model/model/pair_tab.py

@@ -19,12 +19,12 @@
     PairTab,
 )
 
-from .atomic_model import (
-    AtomicModel,
+from .base_atomic_model import (
+    BaseAtomicModel,
 )
 
 
-class PairTabModel(nn.Module, AtomicModel):
+class PairTabModel(nn.Module, BaseAtomicModel):
     """Pairwise tabulation energy model.
 
     This model can be used to tabulate the pairwise energy between atoms for either
@@ -72,7 +72,7 @@ def __init__(
         else:
             raise TypeError("sel must be int or list[int]")
 
-    def get_fitting_output_def(self) -> FittingOutputDef:
+    def fitting_output_def(self) -> FittingOutputDef:
         return FittingOutputDef(
             [
                 OutputVariableDef(

deepmd/pt/model/task/__init__.py

@@ -2,6 +2,9 @@
 from .atten_lcc import (
     FittingNetAttenLcc,
 )
+from .base_fitting import (
+    BaseFitting,
+)
 from .denoise import (
     DenoiseNet,
 )
@@ -15,9 +18,6 @@
 from .fitting import (
     Fitting,
 )
-from .task import (
-    TaskBaseMethod,
-)
 from .type_predict import (
     TypePredictNet,
 )
@@ -29,6 +29,6 @@
     "EnergyFittingNet",
     "EnergyFittingNetDirect",
     "Fitting",
-    "TaskBaseMethod",
+    "BaseFitting",
     "TypePredictNet",
 ]

deepmd/pt/model/task/atten_lcc.py

@@ -6,15 +6,15 @@
     EnergyHead,
     NodeTaskHead,
 )
-from deepmd.pt.model.task.task import (
-    TaskBaseMethod,
+from deepmd.pt.model.task.fitting import (
+    Fitting,
 )
 from deepmd.pt.utils import (
     env,
 )
 
 
-class FittingNetAttenLcc(TaskBaseMethod):
+class FittingNetAttenLcc(Fitting):
     def __init__(
         self, embedding_width, bias_atom_e, pair_embed_dim, attention_heads, **kwargs
     ):

deepmd/pt/model/task/base_fitting.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import torch
+
+from deepmd.dpmodel.fitting import (
+    make_base_fitting,
+)
+
+BaseFitting = make_base_fitting(torch.Tensor, "forward")