dp: pt: rearrange folder structure of atomic_model and model.

deepmd/dpmodel/__init__.py

@@ -5,7 +5,6 @@
     NativeOP,
 )
 from .model import (
-    DPAtomicModel,
     DPModel,
 )
 from .output_def import (
@@ -20,7 +19,6 @@
 
 __all__ = [
     "DPModel",
-    "DPAtomicModel",
     "PRECISION_DICT",
     "DEFAULT_PRECISION",
     "NativeOP",

deepmd/dpmodel/atomic_model/__init__.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+from .linear_atomic_model import (
+    DPZBLLinearAtomicModel,
+    LinearAtomicModel,
+)
+from .make_base_atomic_model import (
+    make_base_atomic_model,
+)
+from .pairtab_atomic_model import (
+    PairTabAtomicModel,
+)
+
+__all__ = [
+    "make_base_atomic_model",
+    "BaseAtomicModel",
+    "DPAtomicModel",
+    "PairTabAtomicModel",
+    "LinearAtomicModel",
+    "DPZBLLinearAtomicModel",
+]

deepmd/dpmodel/model/linear_atomic_model.py → deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -13,24 +13,24 @@
 
 import numpy as np
 
-from deepmd.dpmodel import (
-    FittingOutputDef,
-    OutputVariableDef,
-)
 from deepmd.dpmodel.utils.nlist import (
     build_multiple_neighbor_list,
     get_multiple_nlist_key,
     nlist_distinguish_types,
 )
 
+from ..output_def import (
+    FittingOutputDef,
+    OutputVariableDef,
+)
 from .base_atomic_model import (
     BaseAtomicModel,
 )
 from .dp_atomic_model import (
     DPAtomicModel,
 )
 from .pairtab_atomic_model import (
-    PairTabModel,
+    PairTabAtomicModel,
 )
 
 
@@ -39,8 +39,8 @@
 
     Parameters
     ----------
-    models : list[DPAtomicModel or PairTabModel]
-        A list of models to be combined. PairTabModel must be used together with a DPAtomicModel.
+    models : list[DPAtomicModel or PairTabAtomicModel]
+        A list of models to be combined. PairTabAtomicModel must be used together with a DPAtomicModel.
     """
 
     def __init__(
@@ -212,7 +212,7 @@
     def __init__(
         self,
         dp_model: DPAtomicModel,
-        zbl_model: PairTabModel,
+        zbl_model: PairTabAtomicModel,
         sw_rmin: float,
         sw_rmax: float,
         smin_alpha: Optional[float] = 0.1,
@@ -277,7 +277,9 @@
         # use the larger rr based on nlist
         nlist_larger = zbl_nlist if zbl_nnei >= dp_nnei else dp_nlist
         masked_nlist = np.clip(nlist_larger, 0, None)
-        pairwise_rr = PairTabModel._get_pairwise_dist(extended_coord, masked_nlist)
+        pairwise_rr = PairTabAtomicModel._get_pairwise_dist(
+            extended_coord, masked_nlist
+        )
 
         numerator = np.sum(
             pairwise_rr * np.exp(-pairwise_rr / self.smin_alpha), axis=-1

deepmd/dpmodel/model/pairtab_atomic_model.py → deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -21,7 +21,7 @@
 )
 
 
-class PairTabModel(BaseAtomicModel):
+class PairTabAtomicModel(BaseAtomicModel):
     """Pairwise tabulation energy model.
 
     This model can be used to tabulate the pairwise energy between atoms for either
@@ -99,7 +99,7 @@ def serialize(self) -> dict:
         return {"tab": self.tab.serialize(), "rcut": self.rcut, "sel": self.sel}
 
     @classmethod
-    def deserialize(cls, data) -> "PairTabModel":
+    def deserialize(cls, data) -> "PairTabAtomicModel":
         rcut = data["rcut"]
         sel = data["sel"]
         tab = PairTab.deserialize(data["tab"])

deepmd/dpmodel/model/__init__.py

@@ -1,16 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from .dp_atomic_model import (
-    DPAtomicModel,
-)
 from .dp_model import (
     DPModel,
 )
-from .make_base_atomic_model import (
-    make_base_atomic_model,
-)
 
 __all__ = [
     "DPModel",
-    "DPAtomicModel",
-    "make_base_atomic_model",
 ]

deepmd/dpmodel/model/dp_model.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from .dp_atomic_model import (
+from deepmd.dpmodel.atomic_model import (
     DPAtomicModel,
 )
+
 from .make_model import (
     make_model,
 )

deepmd/pt/model/atomic_model/__init__.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.atomic_model import (
+    make_base_atomic_model,
+)
+
+from .base_atomic_model import (
+    BaseAtomicModel,
+)
+from .dp_atomic_model import (
+    DPAtomicModel,
+)
+from .linear_atomic_model import (
+    DPZBLLinearAtomicModel,
+    LinearAtomicModel,
+)
+from .pairtab_atomic_model import (
+    PairTabAtomicModel,
+)
+
+__all__ = [
+    "make_base_atomic_model",
+    "BaseAtomicModel",
+    "DPAtomicModel",
+    "PairTabAtomicModel",
+    "LinearAtomicModel",
+    "DPZBLLinearAtomicModel",
+]

deepmd/pt/model/model/base_atomic_model.py → deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from deepmd.dpmodel.model import (
+from deepmd.dpmodel.atomic_model import (
     make_base_atomic_model,
 )
 

deepmd/pt/model/model/dp_atomic_model.py → deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -28,15 +28,12 @@
 from .base_atomic_model import (
     BaseAtomicModel,
 )
-from .model import (
-    BaseModel,
-)
 
 log = logging.getLogger(__name__)
 
 
-class DPAtomicModel(BaseModel, BaseAtomicModel):
+class DPAtomicModel(torch.nn.Module, BaseAtomicModel):
     """Model give atomic prediction of some physical property.
 
     Parameters
     ----------

deepmd/pt/model/model/linear_atomic_model.py → deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -28,21 +28,18 @@
 from .dp_atomic_model import (
     DPAtomicModel,
 )
-from .model import (
-    BaseModel,
-)
 from .pairtab_atomic_model import (
-    PairTabModel,
+    PairTabAtomicModel,
 )
 
 
-class LinearAtomicModel(BaseModel, BaseAtomicModel):
+class LinearAtomicModel(torch.nn.Module, BaseAtomicModel):
     """Linear model make linear combinations of several existing models.
 
     Parameters
     ----------
-    models : list[DPAtomicModel or PairTabModel]
-        A list of models to be combined. PairTabModel must be used together with a DPAtomicModel.
+    models : list[DPAtomicModel or PairTabAtomicModel]
+        A list of models to be combined. PairTabAtomicModel must be used together with a DPAtomicModel.
     """
 
     def __init__(
@@ -226,7 +223,7 @@
     def __init__(
         self,
         dp_model: DPAtomicModel,
-        zbl_model: PairTabModel,
+        zbl_model: PairTabAtomicModel,
         sw_rmin: float,
         sw_rmax: float,
         smin_alpha: Optional[float] = 0.1,
@@ -294,7 +291,9 @@
         # use the larger rr based on nlist
         nlist_larger = zbl_nlist if zbl_nnei >= dp_nnei else dp_nlist
         masked_nlist = torch.clamp(nlist_larger, 0)
-        pairwise_rr = PairTabModel._get_pairwise_dist(extended_coord, masked_nlist)
+        pairwise_rr = PairTabAtomicModel._get_pairwise_dist(
+            extended_coord, masked_nlist
+        )
         numerator = torch.sum(
             pairwise_rr * torch.exp(-pairwise_rr / self.smin_alpha), dim=-1
         )  # masked nnei will be zero, no need to handle

deepmd/pt/model/model/pairtab_atomic_model.py → deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -19,12 +19,9 @@
 from .base_atomic_model import (
     BaseAtomicModel,
 )
-from .model import (
-    BaseModel,
-)
 
 
-class PairTabModel(BaseModel, BaseAtomicModel):
+class PairTabAtomicModel(torch.nn.Module, BaseAtomicModel):
     """Pairwise tabulation energy model.
 
     This model can be used to tabulate the pairwise energy between atoms for either
@@ -117,7 +114,7 @@
         return {"tab": self.tab.serialize(), "rcut": self.rcut, "sel": self.sel}
 
     @classmethod
-    def deserialize(cls, data) -> "PairTabModel":
+    def deserialize(cls, data) -> "PairTabAtomicModel":
         rcut = data["rcut"]
         sel = data["sel"]
         tab = PairTab.deserialize(data["tab"])

deepmd/pt/model/model/__init__.py

@@ -1,22 +1,25 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import copy
 
-from deepmd.pt.model.descriptor.descriptor import (
-    Descriptor,
-)
-from deepmd.pt.model.model.dp_atomic_model import (
+from deepmd.pt.model.atomic_model import (
     DPAtomicModel,
+    PairTabAtomicModel,
 )
-from deepmd.pt.model.model.pairtab_atomic_model import (
-    PairTabModel,
+from deepmd.pt.model.descriptor.descriptor import (
+    Descriptor,
 )
 from deepmd.pt.model.task import (
     Fitting,
 )
 
-from .ener import (
+from .dp_model import (
+    DPModel,
+)
+from .dp_zbl_model import (
+    DPZBLModel,
+)
+from .ener_model import (
     EnergyModel,
-    ZBLModel,
 )
 from .model import (
     BaseModel,
@@ -44,13 +47,13 @@
     dp_model = DPAtomicModel(descriptor, fitting, type_map=model_params["type_map"])
     # pairtab
     filepath = model_params["use_srtab"]
-    pt_model = PairTabModel(
+    pt_model = PairTabAtomicModel(
         filepath, model_params["descriptor"]["rcut"], model_params["descriptor"]["sel"]
     )
 
     rmin = model_params["sw_rmin"]
     rmax = model_params["sw_rmax"]
-    return ZBLModel(
+    return DPZBLModel(
         dp_model,
         pt_model,
         rmin,
@@ -84,4 +87,5 @@
     "BaseModel",
     "EnergyModel",
     "get_model",
+    "DPModel",
 ]

deepmd/pt/model/model/dp_model.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.pt.model.atomic_model import (
+    DPAtomicModel,
+)
+
+from .make_model import (
+    make_model,
+)
+
+DPModel = make_model(DPAtomicModel)