fix: address comments

deepmd/pt/model/atomic_model/__init__.py

@@ -33,10 +33,19 @@
 from .polar_atomic_model import (
     DPPolarAtomicModel,
 )
+from .dos_atomic_model import (
+    DPDOSAtomicModel
+)
+from .energy_atomic_model import (
+    DPEnergyAtomicModel
+)
+
 
 __all__ = [
     "BaseAtomicModel",
     "DPAtomicModel",
+    "DPDOSAtomicModel",
+    "DPEnergyAtomicModel",
     "PairTabAtomicModel",
     "LinearEnergyAtomicModel",
     "DPPolarAtomicModel",

deepmd/pt/model/atomic_model/dos_atomic_model.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+class DPDOSAtomicModel(DPAtomicModel):
+    pass

deepmd/pt/model/atomic_model/energy_atomic_model.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+
+class DPEnergyAtomicModel(DPAtomicModel):
+    pass

deepmd/pt/model/model/dipole_model.py

@@ -14,15 +14,15 @@
 )
 
 from .dp_model import (
-    DPModel,
+    DPModelCommon,
 )
 from .make_model import (
     make_model,
 )
 
 
 @BaseModel.register("dipole")
-class DipoleModel(DPModel, make_model(DPDipoleAtomicModel)):
+class DipoleModel(DPModelCommon, make_model(DPDipoleAtomicModel)):
     model_type = "dipole"
 
     def __init__(
@@ -68,14 +68,6 @@ def forward(
             model_predict["updated_coord"] += coord
         return model_predict
 
-    def get_fitting_net(self):
-        """Get the fitting network."""
-        return self.atomic_model.fitting_net
-
-    def get_descriptor(self):
-        """Get the descriptor."""
-        return self.atomic_model.descriptor
-
     @torch.jit.export
     def forward_lower(
         self,

deepmd/pt/model/model/dos_model.py

@@ -7,22 +7,22 @@
 import torch
 
 from deepmd.pt.model.atomic_model import (
-    DPAtomicModel,
+    DPDOSAtomicModel,
 )
 from deepmd.pt.model.model.model import (
     BaseModel,
 )
 
 from .dp_model import (
-    DPModel,
+    DPModelCommon,
 )
 from .make_model import (
     make_model,
 )
 
 
 @BaseModel.register("dos")
-class DOSModel(DPModel, make_model(DPAtomicModel)):
+class DOSModel(DPModelCommon, make_model(DPDOSAtomicModel)):
     model_type = "dos"
 
     def __init__(
@@ -61,14 +61,6 @@ def forward(
             model_predict["updated_coord"] += coord
         return model_predict
 
-    def get_fitting_net(self):
-        """Get the fitting network."""
-        return self.atomic_model.fitting_net
-
-    def get_descriptor(self):
-        """Get the descriptor."""
-        return self.atomic_model.descriptor
-
     @torch.jit.export
     def get_numb_dos(self) -> int:
         """Get the number of  DOS for DOSFittingNet."""

deepmd/pt/model/model/dp_model.py

@@ -4,7 +4,7 @@
 )
 
 
-class DPModel:
+class DPModelCommon:
     """A base class to implement common methods for all the Models."""
 
     @classmethod
@@ -23,3 +23,12 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict):
             global_jdata, local_jdata["descriptor"]
         )
         return local_jdata_cpy
+
+
+    def get_fitting_net(self):
+        """Get the fitting network."""
+        return self.atomic_model.fitting_net
+
+    def get_descriptor(self):
+        """Get the descriptor."""
+        return self.atomic_model.descriptor

deepmd/pt/model/model/dp_zbl_model.py

@@ -14,7 +14,7 @@
 )
 
 from .dp_model import (
-    DPModel,
+    DPModelCommon,
 )
 from .make_model import (
     make_model,
@@ -24,7 +24,7 @@
 
 
 @BaseModel.register("zbl")
-class DPZBLModel(DPModel, DPZBLModel_):
+class DPZBLModel(DPModelCommon, DPZBLModel_):
     model_type = "ener"
 
     def __init__(

deepmd/pt/model/model/ener_model.py

@@ -7,22 +7,22 @@
 import torch
 
 from deepmd.pt.model.atomic_model import (
-    DPAtomicModel,
+    DPEnergyAtomicModel,
 )
 from deepmd.pt.model.model.model import (
     BaseModel,
 )
 
 from .dp_model import (
-    DPModel,
+    DPModelCommon,
 )
 from .make_model import (
     make_model,
 )
 
 
-@BaseModel.register("standard")
-class EnergyModel(DPModel, make_model(DPAtomicModel)):
+@BaseModel.register("energy")
+class EnergyModel(DPModelCommon, make_model(DPEnergyAtomicModel)):
     model_type = "ener"
 
     def __init__(
@@ -70,14 +70,6 @@ def forward(
             model_predict["updated_coord"] += coord
         return model_predict
 
-    def get_fitting_net(self):
-        """Get the fitting network."""
-        return self.atomic_model.fitting_net
-
-    def get_descriptor(self):
-        """Get the descriptor."""
-        return self.atomic_model.descriptor
-
     @torch.jit.export
     def forward_lower(
         self,

deepmd/pt/model/model/polar_model.py

@@ -14,15 +14,15 @@
 )
 
 from .dp_model import (
-    DPModel,
+    DPModelCommon,
 )
 from .make_model import (
     make_model,
 )
 
 
 @BaseModel.register("polar")
-class PolarModel(DPModel, make_model(DPPolarAtomicModel)):
+class PolarModel(DPModelCommon, make_model(DPPolarAtomicModel)):
     model_type = "polar"
 
     def __init__(
@@ -60,14 +60,6 @@ def forward(
             model_predict["updated_coord"] += coord
         return model_predict
 
-    def get_fitting_net(self):
-        """Get the fitting network."""
-        return self.atomic_model.fitting_net
-
-    def get_descriptor(self):
-        """Get the descriptor."""
-        return self.atomic_model.descriptor
-
     @torch.jit.export
     def forward_lower(
         self,