dp: pt: rearrange folder structure of atomic_model and model. (#3268)

- create models `dpmodel.atomic_model` and `pt.model.atomic_model`. 
- place all the atomic model related codes into the new modules. 
- `PairTabModel` -> `PairTabAtomicModel`
- `ZBLModel` -> `DPZBLModel`

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

deepmd/dpmodel/__init__.py

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

deepmd/dpmodel/atomic_model/__init__.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""The atomic model provides the prediction of some property on each
+atom.  All the atomic models are not supposed to be directly accessed
+by users, but it provides a convenient interface for the
+implementation of models.
+
+Taking the energy models for example, the developeres only needs to
+implement the atomic energy prediction via an atomic model, and the
+model can be automatically made by the `deepmd.dpmodel.make_model`
+method. The `DPModel` is made by
+```
+DPModel = make_model(DPAtomicModel)
+```
+
+"""
+
+
+from .base_atomic_model import (
+    BaseAtomicModel,
+)
+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 @@ class LinearAtomicModel(BaseAtomicModel):
 
     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__(
@@ -240,7 +240,7 @@ class DPZBLLinearAtomicModel(LinearAtomicModel):
     def __init__(
         self,
         dp_model: DPAtomicModel,
-        zbl_model: PairTabModel,
+        zbl_model: PairTabAtomicModel,
         sw_rmin: float,
         sw_rmax: float,
         smin_alpha: Optional[float] = 0.1,
@@ -305,7 +305,9 @@ def _compute_weight(
         # 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,23 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from .dp_atomic_model import (
-    DPAtomicModel,
-)
+"""The model that takes the coordinates, cell and atom types as input
+and predicts some property. The models are automatically generated from
+atomic models by the `deepmd.dpmodel.make_model` method.
+
+The `make_model` method does the reduction, auto-differentiation
+(dummy for dpmodels) and communication of the atomic properties
+according to output variable definition
+`deepmd.dpmodel.OutputVariableDef`.
+
+"""
+
 from .dp_model import (
     DPModel,
 )
-from .make_base_atomic_model import (
-    make_base_atomic_model,
+from .make_model import (
+    make_model,
 )
 
 __all__ = [
     "DPModel",
-    "DPAtomicModel",
-    "make_base_atomic_model",
+    "make_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,37 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""The atomic model provides the prediction of some property on each
+atom.  All the atomic models are not supposed to be directly accessed
+by users, but it provides a convenient interface for the
+implementation of models.
+
+Taking the energy models for example, the developeres only needs to
+implement the atomic energy prediction via an atomic model, and the
+model can be automatically made by the `deepmd.dpmodel.make_model`
+method. The `DPModel` is made by
+```
+DPModel = make_model(DPAtomicModel)
+```
+
+"""
+
+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__ = [
+    "BaseAtomicModel",
+    "DPAtomicModel",
+    "PairTabAtomicModel",
+    "LinearAtomicModel",
+    "DPZBLLinearAtomicModel",
+]

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

@@ -3,7 +3,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,14 +28,11 @@
 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__(
@@ -266,7 +263,7 @@ class DPZBLLinearAtomicModel(LinearAtomicModel):
     def __init__(
         self,
         dp_model: DPAtomicModel,
-        zbl_model: PairTabModel,
+        zbl_model: PairTabAtomicModel,
         sw_rmin: float,
         sw_rmax: float,
         smin_alpha: Optional[float] = 0.1,
@@ -334,7 +331,9 @@ def _compute_weight(
         # 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 @@ 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/pt/model/model/__init__.py

@@ -1,26 +1,42 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+"""The model that takes the coordinates, cell and atom types as input
+and predicts some property. The models are automatically generated from
+atomic models by the `deepmd.dpmodel.make_model` method.
+
+The `make_model` method does the reduction, auto-differentiation and
+communication of the atomic properties according to output variable
+definition `deepmd.dpmodel.OutputVariableDef`.
+
+"""
+
 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 .make_hessian_model import (
     make_hessian_model,
 )
+from .make_model import (
+    make_model,
+)
 from .model import (
     BaseModel,
 )
@@ -47,13 +63,13 @@ def get_zbl_model(model_params):
     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,
@@ -85,7 +101,11 @@ def get_model(model_params):
 
 __all__ = [
     "BaseModel",
-    "EnergyModel",
     "get_model",
+    "get_zbl_model",
+    "DPModel",
+    "EnergyModel",
+    "DPZBLModel",
+    "make_model",
     "make_hessian_model",
 ]