From aae48508e002d194446d71260fdab7622c83b88e Mon Sep 17 00:00:00 2001 From: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> Date: Thu, 15 Feb 2024 16:16:38 +0800 Subject: [PATCH] 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 --- deepmd/dpmodel/__init__.py | 2 - deepmd/dpmodel/atomic_model/__init__.py | 42 +++++++++ .../base_atomic_model.py | 0 .../dp_atomic_model.py | 0 .../linear_atomic_model.py | 20 ++-- .../make_base_atomic_model.py | 0 .../pairtab_atomic_model.py | 4 +- deepmd/dpmodel/model/__init__.py | 21 +++-- deepmd/dpmodel/model/dp_model.py | 3 +- deepmd/pt/model/atomic_model/__init__.py | 37 ++++++++ .../base_atomic_model.py | 2 +- .../dp_atomic_model.py | 5 +- .../linear_atomic_model.py | 17 ++-- .../pairtab_atomic_model.py | 7 +- deepmd/pt/model/model/__init__.py | 42 ++++++--- deepmd/pt/model/model/dp_model.py | 10 ++ deepmd/pt/model/model/dp_zbl_model.py | 94 +++++++++++++++++++ .../pt/model/model/{ener.py => ener_model.py} | 90 +----------------- .../common/dpmodel/test_dp_atomic_model.py | 6 +- .../dpmodel/test_linear_atomic_model.py | 22 ++--- .../dpmodel/test_pairtab_atomic_model.py | 10 +- source/tests/pt/model/test_dp_atomic_model.py | 8 +- source/tests/pt/model/test_dp_model.py | 2 +- .../pt/model/test_linear_atomic_model.py | 26 +++-- .../tests/pt/model/test_make_hessian_model.py | 2 +- .../pt/model/test_pairtab_atomic_model.py | 14 +-- 26 files changed, 300 insertions(+), 186 deletions(-) create mode 100644 deepmd/dpmodel/atomic_model/__init__.py rename deepmd/dpmodel/{model => atomic_model}/base_atomic_model.py (100%) rename deepmd/dpmodel/{model => atomic_model}/dp_atomic_model.py (100%) rename deepmd/dpmodel/{model => atomic_model}/linear_atomic_model.py (96%) rename deepmd/dpmodel/{model => atomic_model}/make_base_atomic_model.py (100%) rename deepmd/dpmodel/{model => atomic_model}/pairtab_atomic_model.py (99%) create mode 100644 deepmd/pt/model/atomic_model/__init__.py rename deepmd/pt/model/{model => atomic_model}/base_atomic_model.py (89%) rename deepmd/pt/model/{model => atomic_model}/dp_atomic_model.py (98%) rename deepmd/pt/model/{model => atomic_model}/linear_atomic_model.py (96%) rename deepmd/pt/model/{model => atomic_model}/pairtab_atomic_model.py (98%) create mode 100644 deepmd/pt/model/model/dp_model.py create mode 100644 deepmd/pt/model/model/dp_zbl_model.py rename deepmd/pt/model/model/{ener.py => ener_model.py} (51%) diff --git a/deepmd/dpmodel/__init__.py b/deepmd/dpmodel/__init__.py index 906aac662a..6a7bdb3585 100644 --- a/deepmd/dpmodel/__init__.py +++ b/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", diff --git a/deepmd/dpmodel/atomic_model/__init__.py b/deepmd/dpmodel/atomic_model/__init__.py new file mode 100644 index 0000000000..2cd20f54c1 --- /dev/null +++ b/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", +] diff --git a/deepmd/dpmodel/model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py similarity index 100% rename from deepmd/dpmodel/model/base_atomic_model.py rename to deepmd/dpmodel/atomic_model/base_atomic_model.py diff --git a/deepmd/dpmodel/model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py similarity index 100% rename from deepmd/dpmodel/model/dp_atomic_model.py rename to deepmd/dpmodel/atomic_model/dp_atomic_model.py diff --git a/deepmd/dpmodel/model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py similarity index 96% rename from deepmd/dpmodel/model/linear_atomic_model.py rename to deepmd/dpmodel/atomic_model/linear_atomic_model.py index 8a53cb9229..48aacacdee 100644 --- a/deepmd/dpmodel/model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -13,16 +13,16 @@ 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, ) @@ -30,7 +30,7 @@ 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 diff --git a/deepmd/dpmodel/model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py similarity index 100% rename from deepmd/dpmodel/model/make_base_atomic_model.py rename to deepmd/dpmodel/atomic_model/make_base_atomic_model.py diff --git a/deepmd/dpmodel/model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py similarity index 99% rename from deepmd/dpmodel/model/pairtab_atomic_model.py rename to deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index 6d2de0e126..d5f2bd2f1a 100644 --- a/deepmd/dpmodel/model/pairtab_atomic_model.py +++ b/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"]) diff --git a/deepmd/dpmodel/model/__init__.py b/deepmd/dpmodel/model/__init__.py index 5c0a32673d..dda174fa4e 100644 --- a/deepmd/dpmodel/model/__init__.py +++ b/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", ] diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 819d46450e..6196bcfe87 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/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, ) diff --git a/deepmd/pt/model/atomic_model/__init__.py b/deepmd/pt/model/atomic_model/__init__.py new file mode 100644 index 0000000000..75c1ce3c2e --- /dev/null +++ b/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", +] diff --git a/deepmd/pt/model/model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py similarity index 89% rename from deepmd/pt/model/model/base_atomic_model.py rename to deepmd/pt/model/atomic_model/base_atomic_model.py index 3caa9ed5ae..73b2d76a6d 100644 --- a/deepmd/pt/model/model/base_atomic_model.py +++ b/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, ) diff --git a/deepmd/pt/model/model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py similarity index 98% rename from deepmd/pt/model/model/dp_atomic_model.py rename to deepmd/pt/model/atomic_model/dp_atomic_model.py index 2d04a2b8cf..17b70e4701 100644 --- a/deepmd/pt/model/model/dp_atomic_model.py +++ b/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 diff --git a/deepmd/pt/model/model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py similarity index 96% rename from deepmd/pt/model/model/linear_atomic_model.py rename to deepmd/pt/model/atomic_model/linear_atomic_model.py index 7035426402..7ccc5e225b 100644 --- a/deepmd/pt/model/model/linear_atomic_model.py +++ b/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 diff --git a/deepmd/pt/model/model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py similarity index 98% rename from deepmd/pt/model/model/pairtab_atomic_model.py rename to deepmd/pt/model/atomic_model/pairtab_atomic_model.py index 8aa780effb..cda9071cda 100644 --- a/deepmd/pt/model/model/pairtab_atomic_model.py +++ b/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"]) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 25db37a3d7..8199a8490b 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/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", ] diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py new file mode 100644 index 0000000000..75d3820e45 --- /dev/null +++ b/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) diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py new file mode 100644 index 0000000000..8d71157b60 --- /dev/null +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Dict, + Optional, +) + +import torch + +from deepmd.pt.model.atomic_model import ( + DPZBLLinearAtomicModel, +) + +from .make_model import ( + make_model, +) + +DPZBLModel_ = make_model(DPZBLLinearAtomicModel) + + +class DPZBLModel(DPZBLModel_): + model_type = "ener" + + def __init__( + self, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + + def forward( + self, + coord, + atype, + box: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + do_atomic_virial: bool = False, + ) -> Dict[str, torch.Tensor]: + model_ret = self.forward_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + + model_predict = {} + model_predict["atom_energy"] = model_ret["energy"] + model_predict["energy"] = model_ret["energy_redu"] + if self.do_grad("energy"): + model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) + if do_atomic_virial: + model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3) + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + else: + model_predict["force"] = model_ret["dforce"] + return model_predict + + def forward_lower( + self, + extended_coord, + extended_atype, + nlist, + mapping: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + do_atomic_virial: bool = False, + ): + model_ret = self.forward_common_lower( + extended_coord, + extended_atype, + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + + model_predict = {} + model_predict["atom_energy"] = model_ret["energy"] + model_predict["energy"] = model_ret["energy_redu"] + if self.do_grad("energy"): + model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial: + model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze( + -2 + ) + else: + assert model_ret["dforce"] is not None + model_predict["dforce"] = model_ret["dforce"] + model_predict = model_ret + return model_predict diff --git a/deepmd/pt/model/model/ener.py b/deepmd/pt/model/model/ener_model.py similarity index 51% rename from deepmd/pt/model/model/ener.py rename to deepmd/pt/model/model/ener_model.py index 3c0b66edcd..2afeb2762b 100644 --- a/deepmd/pt/model/model/ener.py +++ b/deepmd/pt/model/model/ener_model.py @@ -6,95 +6,9 @@ import torch -from .dp_atomic_model import ( - DPAtomicModel, +from .dp_model import ( + DPModel, ) -from .linear_atomic_model import ( - DPZBLLinearAtomicModel, -) -from .make_model import ( - make_model, -) - -DPModel = make_model(DPAtomicModel) -ZBLModel_ = make_model(DPZBLLinearAtomicModel) - - -class ZBLModel(ZBLModel_): - model_type = "ener" - - def __init__( - self, - *args, - **kwargs, - ): - super().__init__(*args, **kwargs) - - def forward( - self, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - do_atomic_virial: bool = False, - ) -> Dict[str, torch.Tensor]: - model_ret = self.forward_common( - coord, - atype, - box, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - - model_predict = {} - model_predict["atom_energy"] = model_ret["energy"] - model_predict["energy"] = model_ret["energy_redu"] - if self.do_grad("energy"): - model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) - if do_atomic_virial: - model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3) - model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) - else: - model_predict["force"] = model_ret["dforce"] - return model_predict - - def forward_lower( - self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - do_atomic_virial: bool = False, - ): - model_ret = self.forward_common_lower( - extended_coord, - extended_atype, - nlist, - mapping=mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - - model_predict = {} - model_predict["atom_energy"] = model_ret["energy"] - model_predict["energy"] = model_ret["energy_redu"] - if self.do_grad("energy"): - model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) - model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) - if do_atomic_virial: - model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze( - -2 - ) - else: - assert model_ret["dforce"] is not None - model_predict["dforce"] = model_ret["dforce"] - model_predict = model_ret - return model_predict class EnergyModel(DPModel): diff --git a/source/tests/common/dpmodel/test_dp_atomic_model.py b/source/tests/common/dpmodel/test_dp_atomic_model.py index 684bea45c6..ea5d811b77 100644 --- a/source/tests/common/dpmodel/test_dp_atomic_model.py +++ b/source/tests/common/dpmodel/test_dp_atomic_model.py @@ -3,15 +3,15 @@ import numpy as np +from deepmd.dpmodel.atomic_model import ( + DPAtomicModel, +) from deepmd.dpmodel.descriptor import ( DescrptSeA, ) from deepmd.dpmodel.fitting import ( InvarFitting, ) -from deepmd.dpmodel.model import ( - DPAtomicModel, -) from .case_single_frame_with_nlist import ( TestCaseSingleFrameWithNlist, diff --git a/source/tests/common/dpmodel/test_linear_atomic_model.py b/source/tests/common/dpmodel/test_linear_atomic_model.py index 6dcff97d74..74f2daf7be 100644 --- a/source/tests/common/dpmodel/test_linear_atomic_model.py +++ b/source/tests/common/dpmodel/test_linear_atomic_model.py @@ -6,21 +6,21 @@ import numpy as np +from deepmd.dpmodel.atomic_model import ( + DPAtomicModel, +) +from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + DPZBLLinearAtomicModel, +) +from deepmd.dpmodel.atomic_model.pairtab_atomic_model import ( + PairTabAtomicModel, +) from deepmd.dpmodel.descriptor.se_e2_a import ( DescrptSeA, ) from deepmd.dpmodel.fitting.invar_fitting import ( InvarFitting, ) -from deepmd.dpmodel.model.dp_atomic_model import ( - DPAtomicModel, -) -from deepmd.dpmodel.model.linear_atomic_model import ( - DPZBLLinearAtomicModel, -) -from deepmd.dpmodel.model.pairtab_atomic_model import ( - PairTabModel, -) class TestWeightCalculation(unittest.TestCase): @@ -53,7 +53,7 @@ def test_pairwise(self, mock_loadtxt): ) type_map = ["foo", "bar"] - zbl_model = PairTabModel(tab_file=file_path, rcut=0.3, sel=2) + zbl_model = PairTabAtomicModel(tab_file=file_path, rcut=0.3, sel=2) dp_model = DPAtomicModel(ds, ft, type_map=type_map) wgt_model = DPZBLLinearAtomicModel( @@ -145,7 +145,7 @@ def setUp(self, mock_loadtxt): ) type_map = ["foo", "bar"] dp_model = DPAtomicModel(ds, ft, type_map=type_map) - zbl_model = PairTabModel(file_path, self.rcut, sum(self.sel)) + zbl_model = PairTabAtomicModel(file_path, self.rcut, sum(self.sel)) self.md0 = DPZBLLinearAtomicModel( dp_model, zbl_model, diff --git a/source/tests/common/dpmodel/test_pairtab_atomic_model.py b/source/tests/common/dpmodel/test_pairtab_atomic_model.py index f1e7bd257c..d004f9a37a 100644 --- a/source/tests/common/dpmodel/test_pairtab_atomic_model.py +++ b/source/tests/common/dpmodel/test_pairtab_atomic_model.py @@ -6,8 +6,8 @@ import numpy as np -from deepmd.dpmodel.model.pairtab_atomic_model import ( - PairTabModel, +from deepmd.dpmodel.atomic_model.pairtab_atomic_model import ( + PairTabAtomicModel, ) @@ -24,7 +24,7 @@ def setUp(self, mock_loadtxt) -> None: ] ) - self.model = PairTabModel(tab_file=file_path, rcut=0.02, sel=2) + self.model = PairTabAtomicModel(tab_file=file_path, rcut=0.02, sel=2) self.extended_coord = np.array( [ @@ -68,7 +68,7 @@ def test_with_mask(self): np.testing.assert_allclose(result["energy"], expected_result, 0.0001, 0.0001) def test_deserialize(self): - model1 = PairTabModel.deserialize(self.model.serialize()) + model1 = PairTabAtomicModel.deserialize(self.model.serialize()) np.testing.assert_allclose(self.model.tab_data, model1.tab_data) np.testing.assert_allclose(self.model.tab_info, model1.tab_info) @@ -166,7 +166,7 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: ] ) - model = PairTabModel(tab_file=file_path, rcut=rcut, sel=2) + model = PairTabAtomicModel(tab_file=file_path, rcut=rcut, sel=2) results.append( model.forward_atomic(extended_coord, extended_atype, nlist)["energy"] ) diff --git a/source/tests/pt/model/test_dp_atomic_model.py b/source/tests/pt/model/test_dp_atomic_model.py index fb7e684eaa..d56bec3bb2 100644 --- a/source/tests/pt/model/test_dp_atomic_model.py +++ b/source/tests/pt/model/test_dp_atomic_model.py @@ -4,15 +4,15 @@ import numpy as np import torch -from deepmd.dpmodel import DPAtomicModel as DPDPAtomicModel +from deepmd.dpmodel.atomic_model import DPAtomicModel as DPDPAtomicModel from deepmd.dpmodel.descriptor import DescrptSeA as DPDescrptSeA from deepmd.dpmodel.fitting import InvarFitting as DPInvarFitting +from deepmd.pt.model.atomic_model import ( + DPAtomicModel, +) from deepmd.pt.model.descriptor.se_a import ( DescrptSeA, ) -from deepmd.pt.model.model.dp_atomic_model import ( - DPAtomicModel, -) from deepmd.pt.model.task.ener import ( InvarFitting, ) diff --git a/source/tests/pt/model/test_dp_model.py b/source/tests/pt/model/test_dp_model.py index d970c8a542..fb936e59cb 100644 --- a/source/tests/pt/model/test_dp_model.py +++ b/source/tests/pt/model/test_dp_model.py @@ -10,7 +10,7 @@ from deepmd.pt.model.descriptor.se_a import ( DescrptSeA, ) -from deepmd.pt.model.model.ener import ( +from deepmd.pt.model.model import ( DPModel, EnergyModel, ) diff --git a/source/tests/pt/model/test_linear_atomic_model.py b/source/tests/pt/model/test_linear_atomic_model.py index 14fc6b386a..bb836d4e46 100644 --- a/source/tests/pt/model/test_linear_atomic_model.py +++ b/source/tests/pt/model/test_linear_atomic_model.py @@ -7,23 +7,19 @@ import numpy as np import torch -from deepmd.dpmodel.model.linear_atomic_model import ( +from deepmd.dpmodel.atomic_model import ( DPZBLLinearAtomicModel as DPDPZBLLinearAtomicModel, ) -from deepmd.pt.model.descriptor.se_a import ( - DescrptSeA, -) -from deepmd.pt.model.model.dp_atomic_model import ( +from deepmd.pt.model.atomic_model import ( DPAtomicModel, -) -from deepmd.pt.model.model.ener import ( - ZBLModel, -) -from deepmd.pt.model.model.linear_atomic_model import ( DPZBLLinearAtomicModel, + PairTabAtomicModel, +) +from deepmd.pt.model.descriptor.se_a import ( + DescrptSeA, ) -from deepmd.pt.model.model.pairtab_atomic_model import ( - PairTabModel, +from deepmd.pt.model.model import ( + DPZBLModel, ) from deepmd.pt.model.task.ener import ( InvarFitting, @@ -73,7 +69,7 @@ def test_pairwise(self, mock_loadtxt): ).to(env.DEVICE) type_map = ["foo", "bar"] - zbl_model = PairTabModel(tab_file=file_path, rcut=0.3, sel=2) + zbl_model = PairTabAtomicModel(tab_file=file_path, rcut=0.3, sel=2) dp_model = DPAtomicModel(ds, ft, type_map=type_map).to(env.DEVICE) wgt_model = DPZBLLinearAtomicModel( dp_model, @@ -141,7 +137,7 @@ def setUp(self, mock_loadtxt): ).to(env.DEVICE) type_map = ["foo", "bar"] dp_model = DPAtomicModel(ds, ft, type_map=type_map).to(env.DEVICE) - zbl_model = PairTabModel(file_path, self.rcut, sum(self.sel)) + zbl_model = PairTabAtomicModel(file_path, self.rcut, sum(self.sel)) self.md0 = DPZBLLinearAtomicModel( dp_model, zbl_model, @@ -152,7 +148,7 @@ def setUp(self, mock_loadtxt): env.DEVICE ) self.md2 = DPDPZBLLinearAtomicModel.deserialize(self.md0.serialize()) - self.md3 = ZBLModel(dp_model, zbl_model, sw_rmin=0.1, sw_rmax=0.25) + self.md3 = DPZBLModel(dp_model, zbl_model, sw_rmin=0.1, sw_rmax=0.25) def test_self_consistency(self): args = [ diff --git a/source/tests/pt/model/test_make_hessian_model.py b/source/tests/pt/model/test_make_hessian_model.py index 650d35e019..81aee758bf 100644 --- a/source/tests/pt/model/test_make_hessian_model.py +++ b/source/tests/pt/model/test_make_hessian_model.py @@ -13,7 +13,7 @@ from deepmd.pt.model.model import ( make_hessian_model, ) -from deepmd.pt.model.model.ener import ( +from deepmd.pt.model.model.dp_model import ( DPModel, ) from deepmd.pt.model.task.ener import ( diff --git a/source/tests/pt/model/test_pairtab_atomic_model.py b/source/tests/pt/model/test_pairtab_atomic_model.py index f58ac76211..69e305b40b 100644 --- a/source/tests/pt/model/test_pairtab_atomic_model.py +++ b/source/tests/pt/model/test_pairtab_atomic_model.py @@ -7,9 +7,9 @@ import numpy as np import torch -from deepmd.dpmodel.model.pairtab_atomic_model import PairTabModel as DPPairTabModel -from deepmd.pt.model.model.pairtab_atomic_model import ( - PairTabModel, +from deepmd.dpmodel.atomic_model import PairTabAtomicModel as DPPairTabAtomicModel +from deepmd.pt.model.atomic_model import ( + PairTabAtomicModel, ) from deepmd.pt.utils.utils import ( to_numpy_array, @@ -29,7 +29,7 @@ def setUp(self, mock_loadtxt) -> None: ] ) - self.model = PairTabModel(tab_file=file_path, rcut=0.02, sel=2) + self.model = PairTabAtomicModel(tab_file=file_path, rcut=0.02, sel=2) self.extended_coord = torch.tensor( [ @@ -87,7 +87,7 @@ def test_jit(self): self.assertEqual(model.get_type_map(), None) def test_deserialize(self): - model1 = PairTabModel.deserialize(self.model.serialize()) + model1 = PairTabAtomicModel.deserialize(self.model.serialize()) torch.testing.assert_close(self.model.tab_data, model1.tab_data) torch.testing.assert_close(self.model.tab_info, model1.tab_info) @@ -110,7 +110,7 @@ def test_deserialize(self): def test_cross_deserialize(self): model_dict = self.model.serialize() # pytorch model to dict - model1 = DPPairTabModel.deserialize(model_dict) # dict to numpy model + model1 = DPPairTabAtomicModel.deserialize(model_dict) # dict to numpy model np.testing.assert_allclose(self.model.tab_data, model1.tab_data) np.testing.assert_allclose(self.model.tab_info, model1.tab_info) @@ -209,7 +209,7 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: ] ) - model = PairTabModel(tab_file=file_path, rcut=rcut, sel=2) + model = PairTabAtomicModel(tab_file=file_path, rcut=rcut, sel=2) results.append( model.forward_atomic(extended_coord, extended_atype, nlist)["energy"] )