From 02e546c4289965dbfdd3838ef18b193b72e02aec Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 11 Apr 2024 14:11:54 +0800 Subject: [PATCH 01/43] chore: try remove dp model --- .../atomic_model/dipole_atomic_model.py | 9 ++ .../atomic_model/polar_atomic_model.py | 9 ++ deepmd/pt/model/atomic_model/__init__.py | 8 ++ deepmd/pt/model/model/dipole_model.py | 40 +++++- deepmd/pt/model/model/dos_model.py | 40 +++++- deepmd/pt/model/model/dp_model.py | 124 ------------------ deepmd/pt/model/model/polar_model.py | 40 +++++- 7 files changed, 132 insertions(+), 138 deletions(-) create mode 100644 deepmd/dpmodel/atomic_model/dipole_atomic_model.py create mode 100644 deepmd/dpmodel/atomic_model/polar_atomic_model.py delete mode 100644 deepmd/pt/model/model/dp_model.py diff --git a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py new file mode 100644 index 0000000000..47cf78c6a7 --- /dev/null +++ b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py @@ -0,0 +1,9 @@ +from .dp_atomic_model import ( + DPAtomicModel, +) + +class DPDipoleAtomicModel(DPAtomicModel): + + def apply_out_stat(self, ret, atype): + # dipole not applying bias + pass \ No newline at end of file diff --git a/deepmd/dpmodel/atomic_model/polar_atomic_model.py b/deepmd/dpmodel/atomic_model/polar_atomic_model.py new file mode 100644 index 0000000000..ca57c91d64 --- /dev/null +++ b/deepmd/dpmodel/atomic_model/polar_atomic_model.py @@ -0,0 +1,9 @@ +from .dp_atomic_model import ( + DPAtomicModel, +) + +class DPPolarAtomicModel(DPAtomicModel): + + def apply_out_stat(self, ret, atype): + # TODO: migrate bias + pass \ No newline at end of file diff --git a/deepmd/pt/model/atomic_model/__init__.py b/deepmd/pt/model/atomic_model/__init__.py index a747f28556..54263f7d0f 100644 --- a/deepmd/pt/model/atomic_model/__init__.py +++ b/deepmd/pt/model/atomic_model/__init__.py @@ -27,11 +27,19 @@ from .pairtab_atomic_model import ( PairTabAtomicModel, ) +from .polar_atomic_model import ( + DPPolarAtomicModel, +) +from .dipole_atomic_model import ( + DPDipoleAtomicModel, +) __all__ = [ "BaseAtomicModel", "DPAtomicModel", "PairTabAtomicModel", "LinearEnergyAtomicModel", + "DPPolarAtomicModel", + "DPDipoleAtomicModel", "DPZBLLinearEnergyAtomicModel", ] diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index 45b120771b..05b2da3ca2 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -6,12 +6,17 @@ import torch -from .dp_model import ( - DPModel, +from .make_model import ( + make_model, ) - - -class DipoleModel(DPModel): +from deepmd.pt.model.atomic_model import( + DPDipoleAtomicModel +) +from deepmd.pt.model.model.model import ( + BaseModel, +) +@BaseModel.register("standard") +class DipoleModel(make_model(DPDipoleAtomicModel)): model_type = "dipole" def __init__( @@ -57,6 +62,31 @@ def forward( model_predict["updated_coord"] += coord return model_predict + @classmethod + def update_sel(cls, global_jdata: dict, local_jdata: dict): + """Update the selection and perform neighbor statistics. + + Parameters + ---------- + global_jdata : dict + The global data, containing the training section + local_jdata : dict + The local data refer to the current class + """ + local_jdata_cpy = local_jdata.copy() + local_jdata_cpy["descriptor"] = cls.get_descriptor().update_sel( + 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 + @torch.jit.export def forward_lower( self, diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index e043700bee..66930b21d7 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -6,12 +6,17 @@ import torch -from .dp_model import ( - DPModel, +from .make_model import ( + make_model, ) - - -class DOSModel(DPModel): +from deepmd.pt.model.atomic_model import( + DPAtomicModel +) +from deepmd.pt.model.model.model import ( + BaseModel, +) +@BaseModel.register("standard") +class DOSModel(make_model(DPAtomicModel)): model_type = "dos" def __init__( @@ -50,6 +55,31 @@ def forward( model_predict["updated_coord"] += coord return model_predict + @classmethod + def update_sel(cls, global_jdata: dict, local_jdata: dict): + """Update the selection and perform neighbor statistics. + + Parameters + ---------- + global_jdata : dict + The global data, containing the training section + local_jdata : dict + The local data refer to the current class + """ + local_jdata_cpy = local_jdata.copy() + local_jdata_cpy["descriptor"] = cls.get_descriptor().update_sel( + 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 + @torch.jit.export def get_numb_dos(self) -> int: """Get the number of DOS for DOSFittingNet.""" diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py deleted file mode 100644 index d7b3c4f4e2..0000000000 --- a/deepmd/pt/model/model/dp_model.py +++ /dev/null @@ -1,124 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Dict, - Optional, -) - -import torch - -from deepmd.pt.model.atomic_model import ( - DPAtomicModel, -) -from deepmd.pt.model.descriptor.base_descriptor import ( - BaseDescriptor, -) -from deepmd.pt.model.model.model import ( - BaseModel, -) -from deepmd.pt.model.task.dipole import ( - DipoleFittingNet, -) -from deepmd.pt.model.task.dos import ( - DOSFittingNet, -) -from deepmd.pt.model.task.ener import ( - EnergyFittingNet, - EnergyFittingNetDirect, -) -from deepmd.pt.model.task.polarizability import ( - PolarFittingNet, -) - -from .make_model import ( - make_model, -) - - -@BaseModel.register("standard") -class DPModel(make_model(DPAtomicModel)): - def __new__( - cls, - descriptor=None, - fitting=None, - *args, - # disallow positional atomic_model_ - atomic_model_: Optional[DPAtomicModel] = None, - **kwargs, - ): - from deepmd.pt.model.model.dipole_model import ( - DipoleModel, - ) - from deepmd.pt.model.model.dos_model import ( - DOSModel, - ) - from deepmd.pt.model.model.ener_model import ( - EnergyModel, - ) - from deepmd.pt.model.model.polar_model import ( - PolarModel, - ) - - if atomic_model_ is not None: - fitting = atomic_model_.fitting_net - else: - assert fitting is not None, "fitting network is not provided" - - # according to the fitting network to decide the type of the model - if cls is DPModel: - # map fitting to model - if isinstance(fitting, EnergyFittingNet) or isinstance( - fitting, EnergyFittingNetDirect - ): - cls = EnergyModel - elif isinstance(fitting, DipoleFittingNet): - cls = DipoleModel - elif isinstance(fitting, PolarFittingNet): - cls = PolarModel - elif isinstance(fitting, DOSFittingNet): - cls = DOSModel - # else: unknown fitting type, fall back to DPModel - return super().__new__(cls) - - @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): - """Update the selection and perform neighbor statistics. - - Parameters - ---------- - global_jdata : dict - The global data, containing the training section - local_jdata : dict - The local data refer to the current class - """ - local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["descriptor"] = BaseDescriptor.update_sel( - 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 - - 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]: - # directly call the forward_common method when no specific transform rule - return self.forward_common( - coord, - atype, - box, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index 403058aa47..a25c5ca6cd 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -6,12 +6,19 @@ import torch -from .dp_model import ( - DPModel, -) +from .make_model import ( + make_model, +) +from deepmd.pt.model.atomic_model import( + DPPolarAtomicModel +) -class PolarModel(DPModel): +from deepmd.pt.model.model.model import ( + BaseModel, +) +@BaseModel.register("standard") +class PolarModel(make_model(DPPolarAtomicModel)): model_type = "polar" def __init__( @@ -49,6 +56,31 @@ def forward( model_predict["updated_coord"] += coord return model_predict + @classmethod + def update_sel(cls, global_jdata: dict, local_jdata: dict): + """Update the selection and perform neighbor statistics. + + Parameters + ---------- + global_jdata : dict + The global data, containing the training section + local_jdata : dict + The local data refer to the current class + """ + local_jdata_cpy = local_jdata.copy() + local_jdata_cpy["descriptor"] = cls.get_descriptor().update_sel( + 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 + @torch.jit.export def forward_lower( self, From 4c026fba295416ddb38614570ff98acf3540a472 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 11 Apr 2024 14:14:40 +0800 Subject: [PATCH 02/43] chore: try remove dp model --- deepmd/{dpmodel => pt/model}/atomic_model/dipole_atomic_model.py | 0 deepmd/{dpmodel => pt/model}/atomic_model/polar_atomic_model.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename deepmd/{dpmodel => pt/model}/atomic_model/dipole_atomic_model.py (100%) rename deepmd/{dpmodel => pt/model}/atomic_model/polar_atomic_model.py (100%) diff --git a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py similarity index 100% rename from deepmd/dpmodel/atomic_model/dipole_atomic_model.py rename to deepmd/pt/model/atomic_model/dipole_atomic_model.py diff --git a/deepmd/dpmodel/atomic_model/polar_atomic_model.py b/deepmd/pt/model/atomic_model/polar_atomic_model.py similarity index 100% rename from deepmd/dpmodel/atomic_model/polar_atomic_model.py rename to deepmd/pt/model/atomic_model/polar_atomic_model.py From 81573d50643be2c6de01009a17be75322070f2e0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 11 Apr 2024 06:16:18 +0000 Subject: [PATCH 03/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/atomic_model/__init__.py | 6 +++--- .../pt/model/atomic_model/dipole_atomic_model.py | 7 ++++--- .../pt/model/atomic_model/polar_atomic_model.py | 7 ++++--- deepmd/pt/model/model/dipole_model.py | 15 +++++++++------ deepmd/pt/model/model/dos_model.py | 15 +++++++++------ deepmd/pt/model/model/polar_model.py | 15 ++++++++------- 6 files changed, 37 insertions(+), 28 deletions(-) diff --git a/deepmd/pt/model/atomic_model/__init__.py b/deepmd/pt/model/atomic_model/__init__.py index 54263f7d0f..3adc60d8bd 100644 --- a/deepmd/pt/model/atomic_model/__init__.py +++ b/deepmd/pt/model/atomic_model/__init__.py @@ -17,6 +17,9 @@ from .base_atomic_model import ( BaseAtomicModel, ) +from .dipole_atomic_model import ( + DPDipoleAtomicModel, +) from .dp_atomic_model import ( DPAtomicModel, ) @@ -30,9 +33,6 @@ from .polar_atomic_model import ( DPPolarAtomicModel, ) -from .dipole_atomic_model import ( - DPDipoleAtomicModel, -) __all__ = [ "BaseAtomicModel", diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index 47cf78c6a7..3e077fa5f4 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -1,9 +1,10 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later from .dp_atomic_model import ( DPAtomicModel, ) + class DPDipoleAtomicModel(DPAtomicModel): - def apply_out_stat(self, ret, atype): - # dipole not applying bias - pass \ No newline at end of file + # dipole not applying bias + pass diff --git a/deepmd/pt/model/atomic_model/polar_atomic_model.py b/deepmd/pt/model/atomic_model/polar_atomic_model.py index ca57c91d64..2bbec9f9b3 100644 --- a/deepmd/pt/model/atomic_model/polar_atomic_model.py +++ b/deepmd/pt/model/atomic_model/polar_atomic_model.py @@ -1,9 +1,10 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later from .dp_atomic_model import ( DPAtomicModel, ) + class DPPolarAtomicModel(DPAtomicModel): - def apply_out_stat(self, ret, atype): - # TODO: migrate bias - pass \ No newline at end of file + # TODO: migrate bias + pass diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index 05b2da3ca2..198546f474 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -6,15 +6,18 @@ import torch -from .make_model import ( - make_model, -) -from deepmd.pt.model.atomic_model import( - DPDipoleAtomicModel +from deepmd.pt.model.atomic_model import ( + DPDipoleAtomicModel, ) from deepmd.pt.model.model.model import ( BaseModel, ) + +from .make_model import ( + make_model, +) + + @BaseModel.register("standard") class DipoleModel(make_model(DPDipoleAtomicModel)): model_type = "dipole" @@ -86,7 +89,7 @@ def get_fitting_net(self): def get_descriptor(self): """Get the descriptor.""" return self.atomic_model.descriptor - + @torch.jit.export def forward_lower( self, diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index 66930b21d7..a7a656d6bf 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -6,15 +6,18 @@ import torch -from .make_model import ( - make_model, -) -from deepmd.pt.model.atomic_model import( - DPAtomicModel +from deepmd.pt.model.atomic_model import ( + DPAtomicModel, ) from deepmd.pt.model.model.model import ( BaseModel, ) + +from .make_model import ( + make_model, +) + + @BaseModel.register("standard") class DOSModel(make_model(DPAtomicModel)): model_type = "dos" @@ -79,7 +82,7 @@ def get_fitting_net(self): 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.""" diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index a25c5ca6cd..4eff8992d1 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -6,17 +6,18 @@ import torch +from deepmd.pt.model.atomic_model import ( + DPPolarAtomicModel, +) +from deepmd.pt.model.model.model import ( + BaseModel, +) from .make_model import ( make_model, ) -from deepmd.pt.model.atomic_model import( - DPPolarAtomicModel -) -from deepmd.pt.model.model.model import ( - BaseModel, -) + @BaseModel.register("standard") class PolarModel(make_model(DPPolarAtomicModel)): model_type = "polar" @@ -80,7 +81,7 @@ def get_fitting_net(self): def get_descriptor(self): """Get the descriptor.""" return self.atomic_model.descriptor - + @torch.jit.export def forward_lower( self, From 3ad9637311705ff1a60d03664426ab7222d21c39 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 11 Apr 2024 14:50:00 +0800 Subject: [PATCH 04/43] fix: import --- deepmd/pt/model/model/dipole_model.py | 20 ++------------------ deepmd/pt/model/model/dos_model.py | 21 ++------------------- deepmd/pt/model/model/dp_model.py | 23 +++++++++++++++++++++++ deepmd/pt/model/model/dp_zbl_model.py | 20 ++------------------ deepmd/pt/model/model/ener_model.py | 13 +++++++++++-- deepmd/pt/model/model/polar_model.py | 21 ++------------------- deepmd/pt/model/model/spin_model.py | 9 ++++++--- 7 files changed, 48 insertions(+), 79 deletions(-) create mode 100644 deepmd/pt/model/model/dp_model.py diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index 198546f474..1eb118e575 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -13,13 +13,14 @@ BaseModel, ) +from .dp_model import DPModel from .make_model import ( make_model, ) @BaseModel.register("standard") -class DipoleModel(make_model(DPDipoleAtomicModel)): +class DipoleModel(make_model(DPDipoleAtomicModel), DPModel): model_type = "dipole" def __init__( @@ -65,23 +66,6 @@ def forward( model_predict["updated_coord"] += coord return model_predict - @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): - """Update the selection and perform neighbor statistics. - - Parameters - ---------- - global_jdata : dict - The global data, containing the training section - local_jdata : dict - The local data refer to the current class - """ - local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["descriptor"] = cls.get_descriptor().update_sel( - global_jdata, local_jdata["descriptor"] - ) - return local_jdata_cpy - def get_fitting_net(self): """Get the fitting network.""" return self.atomic_model.fitting_net diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index a7a656d6bf..0bfe711067 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -12,14 +12,14 @@ from deepmd.pt.model.model.model import ( BaseModel, ) - +from .dp_model import DPModel from .make_model import ( make_model, ) @BaseModel.register("standard") -class DOSModel(make_model(DPAtomicModel)): +class DOSModel(make_model(DPAtomicModel), DPModel): model_type = "dos" def __init__( @@ -58,23 +58,6 @@ def forward( model_predict["updated_coord"] += coord return model_predict - @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): - """Update the selection and perform neighbor statistics. - - Parameters - ---------- - global_jdata : dict - The global data, containing the training section - local_jdata : dict - The local data refer to the current class - """ - local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["descriptor"] = cls.get_descriptor().update_sel( - global_jdata, local_jdata["descriptor"] - ) - return local_jdata_cpy - def get_fitting_net(self): """Get the fitting network.""" return self.atomic_model.fitting_net diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py new file mode 100644 index 0000000000..19a9a74db5 --- /dev/null +++ b/deepmd/pt/model/model/dp_model.py @@ -0,0 +1,23 @@ +from deepmd.pt.model.descriptor.base_descriptor import ( + BaseDescriptor, +) + +class DPModel: + """A base class to implement common methods for all the Models. """ + + @classmethod + def update_sel(cls, global_jdata: dict, local_jdata: dict): + """Update the selection and perform neighbor statistics. + + Parameters + ---------- + global_jdata : dict + The global data, containing the training section + local_jdata : dict + The local data refer to the current class + """ + local_jdata_cpy = local_jdata.copy() + local_jdata_cpy["descriptor"] = BaseDescriptor.update_sel( + global_jdata, local_jdata["descriptor"] + ) + return local_jdata_cpy \ No newline at end of file diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index bbc82b8d77..800f222765 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -20,11 +20,12 @@ make_model, ) +from .dp_model import DPModel DPZBLModel_ = make_model(DPZBLLinearEnergyAtomicModel) @BaseModel.register("zbl") -class DPZBLModel(DPZBLModel_): +class DPZBLModel(DPZBLModel_, DPModel): model_type = "ener" def __init__( @@ -103,20 +104,3 @@ def forward_lower( assert model_ret["dforce"] is not None model_predict["dforce"] = model_ret["dforce"] return model_predict - - @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): - """Update the selection and perform neighbor statistics. - - Parameters - ---------- - global_jdata : dict - The global data, containing the training section - local_jdata : dict - The local data refer to the current class - """ - local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["dpmodel"] = DPModel.update_sel( - global_jdata, local_jdata["dpmodel"] - ) - return local_jdata_cpy diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index 5217293623..932b115f33 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -9,9 +9,18 @@ from .dp_model import ( DPModel, ) +from .make_model import ( + make_model, +) +from deepmd.pt.model.atomic_model import ( + DPAtomicModel, +) +from deepmd.pt.model.model.model import ( + BaseModel, +) - -class EnergyModel(DPModel): +@BaseModel.register("standard") +class EnergyModel(make_model(DPAtomicModel),DPModel): model_type = "ener" def __init__( diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index 4eff8992d1..cee497f764 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -12,14 +12,14 @@ from deepmd.pt.model.model.model import ( BaseModel, ) - +from .dp_model import DPModel from .make_model import ( make_model, ) @BaseModel.register("standard") -class PolarModel(make_model(DPPolarAtomicModel)): +class PolarModel(make_model(DPPolarAtomicModel), DPModel): model_type = "polar" def __init__( @@ -57,23 +57,6 @@ def forward( model_predict["updated_coord"] += coord return model_predict - @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): - """Update the selection and perform neighbor statistics. - - Parameters - ---------- - global_jdata : dict - The global data, containing the training section - local_jdata : dict - The local data refer to the current class - """ - local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["descriptor"] = cls.get_descriptor().update_sel( - global_jdata, local_jdata["descriptor"] - ) - return local_jdata_cpy - def get_fitting_net(self): """Get the fitting network.""" return self.atomic_model.fitting_net diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index df2f48e2e4..6cc523e7d3 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -18,8 +18,11 @@ Spin, ) -from .dp_model import ( - DPModel, +from .make_model import ( + make_model, +) +from deepmd.pt.model.atomic_model import ( + DPAtomicModel, ) @@ -474,7 +477,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data) -> "SpinModel": - backbone_model_obj = DPModel.deserialize(data["backbone_model"]) + backbone_model_obj = make_model(DPAtomicModel).deserialize(data["backbone_model"]) spin = Spin.deserialize(data["spin"]) return cls( backbone_model=backbone_model_obj, From c64e5208f0ff53f11893117624a6d494436e9b1f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 11 Apr 2024 06:50:43 +0000 Subject: [PATCH 05/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/model/dipole_model.py | 4 +++- deepmd/pt/model/model/dos_model.py | 5 ++++- deepmd/pt/model/model/dp_model.py | 6 ++++-- deepmd/pt/model/model/dp_zbl_model.py | 4 +++- deepmd/pt/model/model/ener_model.py | 16 +++++++++------- deepmd/pt/model/model/polar_model.py | 5 ++++- deepmd/pt/model/model/spin_model.py | 10 ++++++---- 7 files changed, 33 insertions(+), 17 deletions(-) diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index 1eb118e575..1d5023fb5f 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -13,7 +13,9 @@ BaseModel, ) -from .dp_model import DPModel +from .dp_model import ( + DPModel, +) from .make_model import ( make_model, ) diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index 0bfe711067..2e87d35a52 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -12,7 +12,10 @@ from deepmd.pt.model.model.model import ( BaseModel, ) -from .dp_model import DPModel + +from .dp_model import ( + DPModel, +) from .make_model import ( make_model, ) diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index 19a9a74db5..5f7970bde6 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -1,9 +1,11 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later from deepmd.pt.model.descriptor.base_descriptor import ( BaseDescriptor, ) + class DPModel: - """A base class to implement common methods for all the Models. """ + """A base class to implement common methods for all the Models.""" @classmethod def update_sel(cls, global_jdata: dict, local_jdata: dict): @@ -20,4 +22,4 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict): local_jdata_cpy["descriptor"] = BaseDescriptor.update_sel( global_jdata, local_jdata["descriptor"] ) - return local_jdata_cpy \ No newline at end of file + return local_jdata_cpy diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index 800f222765..9dbc7c3555 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -16,11 +16,13 @@ BaseModel, ) +from .dp_model import ( + DPModel, +) from .make_model import ( make_model, ) -from .dp_model import DPModel DPZBLModel_ = make_model(DPZBLLinearEnergyAtomicModel) diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index 932b115f33..f91f8733aa 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -6,12 +6,6 @@ import torch -from .dp_model import ( - DPModel, -) -from .make_model import ( - make_model, -) from deepmd.pt.model.atomic_model import ( DPAtomicModel, ) @@ -19,8 +13,16 @@ BaseModel, ) +from .dp_model import ( + DPModel, +) +from .make_model import ( + make_model, +) + + @BaseModel.register("standard") -class EnergyModel(make_model(DPAtomicModel),DPModel): +class EnergyModel(make_model(DPAtomicModel), DPModel): model_type = "ener" def __init__( diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index cee497f764..2009dd22a7 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -12,7 +12,10 @@ from deepmd.pt.model.model.model import ( BaseModel, ) -from .dp_model import DPModel + +from .dp_model import ( + DPModel, +) from .make_model import ( make_model, ) diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index 6cc523e7d3..4c71344c8a 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -8,6 +8,9 @@ import torch +from deepmd.pt.model.atomic_model import ( + DPAtomicModel, +) from deepmd.pt.utils.utils import ( to_torch_tensor, ) @@ -21,9 +24,6 @@ from .make_model import ( make_model, ) -from deepmd.pt.model.atomic_model import ( - DPAtomicModel, -) class SpinModel(torch.nn.Module): @@ -477,7 +477,9 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data) -> "SpinModel": - backbone_model_obj = make_model(DPAtomicModel).deserialize(data["backbone_model"]) + backbone_model_obj = make_model(DPAtomicModel).deserialize( + data["backbone_model"] + ) spin = Spin.deserialize(data["spin"]) return cls( backbone_model=backbone_model_obj, From acd1b993e3b40d64a4e63bf36bf16690070cc948 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 11 Apr 2024 14:54:56 +0800 Subject: [PATCH 06/43] fix: import --- deepmd/pt/model/model/dp_zbl_model.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index 9dbc7c3555..88a177924a 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -6,9 +6,6 @@ import torch -from deepmd.dpmodel.model.dp_model import ( - DPModel, -) from deepmd.pt.model.atomic_model import ( DPZBLLinearEnergyAtomicModel, ) From db2751c7f5315e0ba1e8d9b99a7bde0f654d6ad0 Mon Sep 17 00:00:00 2001 From: anyangml Date: Thu, 11 Apr 2024 07:39:23 +0000 Subject: [PATCH 07/43] fix: UTs --- deepmd/pt/model/model/__init__.py | 22 +++++++++++++++++++++- deepmd/pt/model/model/dipole_model.py | 2 +- deepmd/pt/model/model/dos_model.py | 2 +- deepmd/pt/model/model/dp_zbl_model.py | 2 +- deepmd/pt/model/model/ener_model.py | 10 +++++++++- deepmd/pt/model/model/polar_model.py | 2 +- 6 files changed, 34 insertions(+), 6 deletions(-) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 1675215d7b..d341ba2fd6 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -39,6 +39,15 @@ from .ener_model import ( EnergyModel, ) +from .dos_model import ( + DOSModel, +) +from .polar_model import ( + PolarModel, +) +from .dipole_model import ( + DipoleModel, +) from .frozen import ( FrozenModel, ) @@ -160,7 +169,18 @@ def get_standard_model(model_params): atom_exclude_types = model_params.get("atom_exclude_types", []) pair_exclude_types = model_params.get("pair_exclude_types", []) - model = DPModel( + if fitting_net["type"] == "dipole": + modelcls = DipoleModel + elif fitting_net["type"] == "polar": + modelcls = PolarModel + elif fitting_net["type"] == "dos": + modelcls = DOSModel + elif fitting_net["type"] == "ener": + modelcls = EnergyModel + else: + pass + + model = modelcls( descriptor=descriptor, fitting=fitting, type_map=model_params["type_map"], diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index 1d5023fb5f..68b758ddb9 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -22,7 +22,7 @@ @BaseModel.register("standard") -class DipoleModel(make_model(DPDipoleAtomicModel), DPModel): +class DipoleModel(DPModel, make_model(DPDipoleAtomicModel)): model_type = "dipole" def __init__( diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index 2e87d35a52..e5da6e66fe 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -22,7 +22,7 @@ @BaseModel.register("standard") -class DOSModel(make_model(DPAtomicModel), DPModel): +class DOSModel(DPModel, make_model(DPAtomicModel)): model_type = "dos" def __init__( diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index 88a177924a..06537313e5 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -24,7 +24,7 @@ @BaseModel.register("zbl") -class DPZBLModel(DPZBLModel_, DPModel): +class DPZBLModel(DPModel, DPZBLModel_ ): model_type = "ener" def __init__( diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index f91f8733aa..6a28a87d84 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -22,7 +22,7 @@ @BaseModel.register("standard") -class EnergyModel(make_model(DPAtomicModel), DPModel): +class EnergyModel(DPModel, make_model(DPAtomicModel)): model_type = "ener" def __init__( @@ -70,6 +70,14 @@ 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, diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index 2009dd22a7..3eaa9779ed 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -22,7 +22,7 @@ @BaseModel.register("standard") -class PolarModel(make_model(DPPolarAtomicModel), DPModel): +class PolarModel(DPModel, make_model(DPPolarAtomicModel)): model_type = "polar" def __init__( From 9612b02e09f59c5438c5547e3859fc2bd6decdda Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 11 Apr 2024 07:39:56 +0000 Subject: [PATCH 08/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/model/__init__.py | 18 +++++++++--------- deepmd/pt/model/model/dp_zbl_model.py | 2 +- deepmd/pt/model/model/ener_model.py | 2 +- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index d341ba2fd6..85d0adbf96 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -30,6 +30,12 @@ Spin, ) +from .dipole_model import ( + DipoleModel, +) +from .dos_model import ( + DOSModel, +) from .dp_model import ( DPModel, ) @@ -39,15 +45,6 @@ from .ener_model import ( EnergyModel, ) -from .dos_model import ( - DOSModel, -) -from .polar_model import ( - PolarModel, -) -from .dipole_model import ( - DipoleModel, -) from .frozen import ( FrozenModel, ) @@ -60,6 +57,9 @@ from .model import ( BaseModel, ) +from .polar_model import ( + PolarModel, +) from .spin_model import ( SpinEnergyModel, SpinModel, diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index 06537313e5..0c139ae7c1 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -24,7 +24,7 @@ @BaseModel.register("zbl") -class DPZBLModel(DPModel, DPZBLModel_ ): +class DPZBLModel(DPModel, DPZBLModel_): model_type = "ener" def __init__( diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index 6a28a87d84..262be4573b 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -77,7 +77,7 @@ def get_fitting_net(self): def get_descriptor(self): """Get the descriptor.""" return self.atomic_model.descriptor - + @torch.jit.export def forward_lower( self, From a7384d22e6e7f87948926557e9d467f0bd585315 Mon Sep 17 00:00:00 2001 From: anyangml Date: Thu, 11 Apr 2024 08:13:09 +0000 Subject: [PATCH 09/43] fix: UTs --- deepmd/pt/model/model/__init__.py | 4 ++-- source/tests/pt/model/test_dp_model.py | 23 +++++++++++------------ 2 files changed, 13 insertions(+), 14 deletions(-) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 85d0adbf96..624d170a39 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -175,10 +175,10 @@ def get_standard_model(model_params): modelcls = PolarModel elif fitting_net["type"] == "dos": modelcls = DOSModel - elif fitting_net["type"] == "ener": + elif fitting_net["type"] in ["ener", "direct_force_ener"]: modelcls = EnergyModel else: - pass + raise RuntimeError(f"Unknown fitting type: {fitting_net['type']}") model = modelcls( descriptor=descriptor, diff --git a/source/tests/pt/model/test_dp_model.py b/source/tests/pt/model/test_dp_model.py index 7470cf96d0..b0f2288b19 100644 --- a/source/tests/pt/model/test_dp_model.py +++ b/source/tests/pt/model/test_dp_model.py @@ -11,7 +11,6 @@ DescrptSeA, ) from deepmd.pt.model.model import ( - DPModel, EnergyModel, ) from deepmd.pt.model.task.ener import ( @@ -57,8 +56,8 @@ def test_self_consistency(self): mixed_types=ds.mixed_types(), ).to(env.DEVICE) type_map = ["foo", "bar"] - md0 = DPModel(ds, ft, type_map=type_map).to(env.DEVICE) - md1 = DPModel.deserialize(md0.serialize()).to(env.DEVICE) + md0 = EnergyModel(ds, ft, type_map=type_map).to(env.DEVICE) + md1 = EnergyModel.deserialize(md0.serialize()).to(env.DEVICE) args = [to_torch_tensor(ii) for ii in [self.coord, self.atype, self.cell]] ret0 = md0.forward_common(*args) ret1 = md1.forward_common(*args) @@ -133,7 +132,7 @@ def test_dp_consistency(self): ) type_map = ["foo", "bar"] md0 = DPDPModel(ds, ft, type_map=type_map) - md1 = DPModel.deserialize(md0.serialize()).to(env.DEVICE) + md1 = EnergyModel.deserialize(md0.serialize()).to(env.DEVICE) rng = np.random.default_rng() fparam = rng.normal(size=[self.nf, nfp]) @@ -174,7 +173,7 @@ def test_dp_consistency_nopbc(self): ) type_map = ["foo", "bar"] md0 = DPDPModel(ds, ft, type_map=type_map) - md1 = DPModel.deserialize(md0.serialize()).to(env.DEVICE) + md1 = EnergyModel.deserialize(md0.serialize()).to(env.DEVICE) rng = np.random.default_rng() fparam = rng.normal(size=[self.nf, nfp]) @@ -217,7 +216,7 @@ def test_prec_consistency(self): aparam = rng.normal(size=[self.nf, nloc, nap]) md0 = DPDPModel(ds, ft, type_map=type_map) - md1 = DPModel.deserialize(md0.serialize()).to(env.DEVICE) + md1 = EnergyModel.deserialize(md0.serialize()).to(env.DEVICE) args64 = [to_torch_tensor(ii) for ii in [self.coord, self.atype, self.cell]] args64[0] = args64[0].to(torch.float64) @@ -267,8 +266,8 @@ def test_self_consistency(self): mixed_types=ds.mixed_types(), ).to(env.DEVICE) type_map = ["foo", "bar"] - md0 = DPModel(ds, ft, type_map=type_map).to(env.DEVICE) - md1 = DPModel.deserialize(md0.serialize()).to(env.DEVICE) + md0 = EnergyModel(ds, ft, type_map=type_map).to(env.DEVICE) + md1 = EnergyModel.deserialize(md0.serialize()).to(env.DEVICE) args = [ to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist] ] @@ -319,7 +318,7 @@ def test_dp_consistency(self): ) type_map = ["foo", "bar"] md0 = DPDPModel(ds, ft, type_map=type_map) - md1 = DPModel.deserialize(md0.serialize()).to(env.DEVICE) + md1 = EnergyModel.deserialize(md0.serialize()).to(env.DEVICE) args0 = [self.coord_ext, self.atype_ext, self.nlist] args1 = [ to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist] @@ -358,7 +357,7 @@ def test_prec_consistency(self): aparam = rng.normal(size=[self.nf, nloc, nap]) md0 = DPDPModel(ds, ft, type_map=type_map) - md1 = DPModel.deserialize(md0.serialize()).to(env.DEVICE) + md1 = EnergyModel.deserialize(md0.serialize()).to(env.DEVICE) args64 = [ to_torch_tensor(ii) for ii in [self.coord_ext, self.atype_ext, self.nlist] @@ -405,7 +404,7 @@ def test_jit(self): mixed_types=ds.mixed_types(), ).to(env.DEVICE) type_map = ["foo", "bar"] - md0 = DPModel(ds, ft, type_map=type_map).to(env.DEVICE) + md0 = EnergyModel(ds, ft, type_map=type_map).to(env.DEVICE) md0 = torch.jit.script(md0) md0.get_rcut() md0.get_type_map() @@ -455,7 +454,7 @@ def setUp(self): mixed_types=ds.mixed_types(), ).to(env.DEVICE) type_map = ["foo", "bar"] - self.md = DPModel(ds, ft, type_map=type_map).to(env.DEVICE) + self.md = EnergyModel(ds, ft, type_map=type_map).to(env.DEVICE) def test_nlist_eq(self): # n_nnei == nnei From 30bd37847a8ca127efc9f3a5a41a4ae8d6ca2db3 Mon Sep 17 00:00:00 2001 From: anyangml Date: Thu, 11 Apr 2024 08:49:29 +0000 Subject: [PATCH 10/43] fix: UTs --- .../model/atomic_model/dipole_atomic_model.py | 2 +- .../model/atomic_model/polar_atomic_model.py | 2 +- deepmd/pt/model/model/make_model.py | 19 +++++++++++++++++++ .../tests/pt/model/test_make_hessian_model.py | 8 ++++---- 4 files changed, 25 insertions(+), 6 deletions(-) diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index 3e077fa5f4..df0ebabb87 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -7,4 +7,4 @@ class DPDipoleAtomicModel(DPAtomicModel): def apply_out_stat(self, ret, atype): # dipole not applying bias - pass + return ret diff --git a/deepmd/pt/model/atomic_model/polar_atomic_model.py b/deepmd/pt/model/atomic_model/polar_atomic_model.py index 2bbec9f9b3..caacc697de 100644 --- a/deepmd/pt/model/atomic_model/polar_atomic_model.py +++ b/deepmd/pt/model/atomic_model/polar_atomic_model.py @@ -7,4 +7,4 @@ class DPPolarAtomicModel(DPAtomicModel): def apply_out_stat(self, ret, atype): # TODO: migrate bias - pass + return ret diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 25a8ec9201..21272a5f7d 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -526,5 +526,24 @@ def mixed_types(self) -> bool: """ return self.atomic_model.mixed_types() + + 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]: + # directly call the forward_common method when no specific transform rule + return self.forward_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) return CM diff --git a/source/tests/pt/model/test_make_hessian_model.py b/source/tests/pt/model/test_make_hessian_model.py index 7d9ae2b810..b036f02831 100644 --- a/source/tests/pt/model/test_make_hessian_model.py +++ b/source/tests/pt/model/test_make_hessian_model.py @@ -13,8 +13,8 @@ from deepmd.pt.model.model import ( make_hessian_model, ) -from deepmd.pt.model.model.dp_model import ( - DPModel, +from deepmd.pt.model.model import ( + EnergyModel, ) from deepmd.pt.model.task.ener import ( InvarFitting, @@ -159,10 +159,10 @@ def setUp(self): neuron=[4, 4, 4], ).to(env.DEVICE) type_map = ["foo", "bar"] - self.model_hess = make_hessian_model(DPModel)(ds, ft0, type_map=type_map).to( + self.model_hess = make_hessian_model(EnergyModel)(ds, ft0, type_map=type_map).to( env.DEVICE ) - self.model_valu = DPModel.deserialize(self.model_hess.serialize()) + self.model_valu = EnergyModel.deserialize(self.model_hess.serialize()) self.model_hess.requires_hessian("energy") def test_output_def(self): From 6c9185e25c16ba8fc74cc41c772a725022c57612 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 11 Apr 2024 08:51:14 +0000 Subject: [PATCH 11/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/model/make_model.py | 2 +- source/tests/pt/model/test_make_hessian_model.py | 10 ++++------ 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 21272a5f7d..386b5e14f9 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -526,7 +526,7 @@ def mixed_types(self) -> bool: """ return self.atomic_model.mixed_types() - + def forward( self, coord, diff --git a/source/tests/pt/model/test_make_hessian_model.py b/source/tests/pt/model/test_make_hessian_model.py index b036f02831..b94e83bafc 100644 --- a/source/tests/pt/model/test_make_hessian_model.py +++ b/source/tests/pt/model/test_make_hessian_model.py @@ -10,11 +10,9 @@ from deepmd.pt.model.descriptor.se_a import ( DescrptSeA, ) -from deepmd.pt.model.model import ( - make_hessian_model, -) from deepmd.pt.model.model import ( EnergyModel, + make_hessian_model, ) from deepmd.pt.model.task.ener import ( InvarFitting, @@ -159,9 +157,9 @@ def setUp(self): neuron=[4, 4, 4], ).to(env.DEVICE) type_map = ["foo", "bar"] - self.model_hess = make_hessian_model(EnergyModel)(ds, ft0, type_map=type_map).to( - env.DEVICE - ) + self.model_hess = make_hessian_model(EnergyModel)( + ds, ft0, type_map=type_map + ).to(env.DEVICE) self.model_valu = EnergyModel.deserialize(self.model_hess.serialize()) self.model_hess.requires_hessian("energy") From 24859ac009799c3316f4d2302565f185eae1f149 Mon Sep 17 00:00:00 2001 From: anyangml Date: Thu, 11 Apr 2024 09:07:32 +0000 Subject: [PATCH 12/43] fix: UTs --- deepmd/pt/model/atomic_model/dipole_atomic_model.py | 8 +++++++- deepmd/pt/model/atomic_model/polar_atomic_model.py | 9 +++++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index df0ebabb87..d6bcd7a170 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -2,9 +2,15 @@ from .dp_atomic_model import ( DPAtomicModel, ) +from typing import Dict +import torch class DPDipoleAtomicModel(DPAtomicModel): - def apply_out_stat(self, ret, atype): + def apply_out_stat( + self, + ret: Dict[str, torch.Tensor], + atype: torch.Tensor, + ): # dipole not applying bias return ret diff --git a/deepmd/pt/model/atomic_model/polar_atomic_model.py b/deepmd/pt/model/atomic_model/polar_atomic_model.py index caacc697de..2def216162 100644 --- a/deepmd/pt/model/atomic_model/polar_atomic_model.py +++ b/deepmd/pt/model/atomic_model/polar_atomic_model.py @@ -2,9 +2,14 @@ from .dp_atomic_model import ( DPAtomicModel, ) - +from typing import Dict +import torch class DPPolarAtomicModel(DPAtomicModel): - def apply_out_stat(self, ret, atype): + def apply_out_stat( + self, + ret: Dict[str, torch.Tensor], + atype: torch.Tensor, + ): # TODO: migrate bias return ret From 754e18a971f86eb0abe566a63dbba4c6c6bb99f6 Mon Sep 17 00:00:00 2001 From: anyangml Date: Thu, 11 Apr 2024 09:20:57 +0000 Subject: [PATCH 13/43] fix: UTs --- deepmd/pt/model/model/dipole_model.py | 2 +- deepmd/pt/model/model/dos_model.py | 2 +- deepmd/pt/model/model/polar_model.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index 68b758ddb9..fcf204d173 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -21,7 +21,7 @@ ) -@BaseModel.register("standard") +@BaseModel.register("dipole") class DipoleModel(DPModel, make_model(DPDipoleAtomicModel)): model_type = "dipole" diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index e5da6e66fe..b600b1cd7b 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -21,7 +21,7 @@ ) -@BaseModel.register("standard") +@BaseModel.register("dos") class DOSModel(DPModel, make_model(DPAtomicModel)): model_type = "dos" diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index 3eaa9779ed..3d61a65b1d 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -21,7 +21,7 @@ ) -@BaseModel.register("standard") +@BaseModel.register("polar") class PolarModel(DPModel, make_model(DPPolarAtomicModel)): model_type = "polar" From 56accdfaa1779377d57913b1b05f77a011e69aee Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 11 Apr 2024 09:21:46 +0000 Subject: [PATCH 14/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/atomic_model/dipole_atomic_model.py | 8 ++++++-- deepmd/pt/model/atomic_model/polar_atomic_model.py | 9 +++++++-- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index d6bcd7a170..4152a3a358 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -1,9 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Dict, +) + +import torch + from .dp_atomic_model import ( DPAtomicModel, ) -from typing import Dict -import torch class DPDipoleAtomicModel(DPAtomicModel): diff --git a/deepmd/pt/model/atomic_model/polar_atomic_model.py b/deepmd/pt/model/atomic_model/polar_atomic_model.py index 2def216162..da19a31211 100644 --- a/deepmd/pt/model/atomic_model/polar_atomic_model.py +++ b/deepmd/pt/model/atomic_model/polar_atomic_model.py @@ -1,9 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Dict, +) + +import torch + from .dp_atomic_model import ( DPAtomicModel, ) -from typing import Dict -import torch + class DPPolarAtomicModel(DPAtomicModel): def apply_out_stat( From 9f437ed615d061d41ab7a74e2f5a2f5dad3d2afa Mon Sep 17 00:00:00 2001 From: anyangml Date: Fri, 12 Apr 2024 02:06:17 +0000 Subject: [PATCH 15/43] fix: address comments --- deepmd/pt/model/atomic_model/__init__.py | 9 +++++++++ deepmd/pt/model/atomic_model/dos_atomic_model.py | 7 +++++++ .../pt/model/atomic_model/energy_atomic_model.py | 7 +++++++ deepmd/pt/model/model/dipole_model.py | 12 ++---------- deepmd/pt/model/model/dos_model.py | 14 +++----------- deepmd/pt/model/model/dp_model.py | 11 ++++++++++- deepmd/pt/model/model/dp_zbl_model.py | 4 ++-- deepmd/pt/model/model/ener_model.py | 16 ++++------------ deepmd/pt/model/model/polar_model.py | 12 ++---------- 9 files changed, 46 insertions(+), 46 deletions(-) create mode 100644 deepmd/pt/model/atomic_model/dos_atomic_model.py create mode 100644 deepmd/pt/model/atomic_model/energy_atomic_model.py diff --git a/deepmd/pt/model/atomic_model/__init__.py b/deepmd/pt/model/atomic_model/__init__.py index 3adc60d8bd..b589bc484f 100644 --- a/deepmd/pt/model/atomic_model/__init__.py +++ b/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", diff --git a/deepmd/pt/model/atomic_model/dos_atomic_model.py b/deepmd/pt/model/atomic_model/dos_atomic_model.py new file mode 100644 index 0000000000..b979f93287 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/deepmd/pt/model/atomic_model/energy_atomic_model.py b/deepmd/pt/model/atomic_model/energy_atomic_model.py new file mode 100644 index 0000000000..bc843ad04a --- /dev/null +++ b/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 \ No newline at end of file diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index fcf204d173..e98aee3207 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -14,7 +14,7 @@ ) from .dp_model import ( - DPModel, + DPModelCommon, ) from .make_model import ( make_model, @@ -22,7 +22,7 @@ @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, diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index b600b1cd7b..490e256852 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -7,14 +7,14 @@ 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, @@ -22,7 +22,7 @@ @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.""" diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index 5f7970bde6..bc009585b8 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/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 diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index 0c139ae7c1..92e76da3d7 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/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__( diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index 262be4573b..b036616d33 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/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, diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index 3d61a65b1d..3349c164a6 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -14,7 +14,7 @@ ) from .dp_model import ( - DPModel, + DPModelCommon, ) from .make_model import ( make_model, @@ -22,7 +22,7 @@ @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, From ed34307c16f9618e4c7825fd6ae6a6dc6a6e6d80 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 12 Apr 2024 02:07:16 +0000 Subject: [PATCH 16/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/atomic_model/__init__.py | 13 ++++++------- deepmd/pt/model/atomic_model/dos_atomic_model.py | 3 ++- deepmd/pt/model/atomic_model/energy_atomic_model.py | 3 ++- deepmd/pt/model/model/dp_model.py | 1 - 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/deepmd/pt/model/atomic_model/__init__.py b/deepmd/pt/model/atomic_model/__init__.py index b589bc484f..3e94449057 100644 --- a/deepmd/pt/model/atomic_model/__init__.py +++ b/deepmd/pt/model/atomic_model/__init__.py @@ -20,9 +20,15 @@ from .dipole_atomic_model import ( DPDipoleAtomicModel, ) +from .dos_atomic_model import ( + DPDOSAtomicModel, +) from .dp_atomic_model import ( DPAtomicModel, ) +from .energy_atomic_model import ( + DPEnergyAtomicModel, +) from .linear_atomic_model import ( DPZBLLinearEnergyAtomicModel, LinearEnergyAtomicModel, @@ -33,13 +39,6 @@ from .polar_atomic_model import ( DPPolarAtomicModel, ) -from .dos_atomic_model import ( - DPDOSAtomicModel -) -from .energy_atomic_model import ( - DPEnergyAtomicModel -) - __all__ = [ "BaseAtomicModel", diff --git a/deepmd/pt/model/atomic_model/dos_atomic_model.py b/deepmd/pt/model/atomic_model/dos_atomic_model.py index b979f93287..aba391caef 100644 --- a/deepmd/pt/model/atomic_model/dos_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dos_atomic_model.py @@ -3,5 +3,6 @@ DPAtomicModel, ) + class DPDOSAtomicModel(DPAtomicModel): - pass \ No newline at end of file + pass diff --git a/deepmd/pt/model/atomic_model/energy_atomic_model.py b/deepmd/pt/model/atomic_model/energy_atomic_model.py index bc843ad04a..e2e9432d1e 100644 --- a/deepmd/pt/model/atomic_model/energy_atomic_model.py +++ b/deepmd/pt/model/atomic_model/energy_atomic_model.py @@ -3,5 +3,6 @@ DPAtomicModel, ) + class DPEnergyAtomicModel(DPAtomicModel): - pass \ No newline at end of file + pass diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index bc009585b8..fab1ff580f 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -23,7 +23,6 @@ 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.""" From 72a34fcf7870161ba093885db89d897769e133d4 Mon Sep 17 00:00:00 2001 From: anyangml Date: Fri, 12 Apr 2024 02:21:59 +0000 Subject: [PATCH 17/43] fix: import --- deepmd/pt/model/model/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 624d170a39..260ce92099 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -37,7 +37,7 @@ DOSModel, ) from .dp_model import ( - DPModel, + DPModelCommon, ) from .dp_zbl_model import ( DPZBLModel, @@ -203,7 +203,7 @@ def get_model(model_params): __all__ = [ "BaseModel", "get_model", - "DPModel", + "DPModelCommon", "EnergyModel", "FrozenModel", "SpinModel", From 53554289f982cb207f6892fc3e572a5089854814 Mon Sep 17 00:00:00 2001 From: anyangml Date: Fri, 12 Apr 2024 02:29:14 +0000 Subject: [PATCH 18/43] fix: UTs --- deepmd/pt/model/model/ener_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index b036616d33..893d6a684a 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -22,6 +22,7 @@ @BaseModel.register("energy") +@BaseModel.register("standard") class EnergyModel(DPModelCommon, make_model(DPEnergyAtomicModel)): model_type = "ener" From 257936df03fe3d943fd704536c24fbf6781120bf Mon Sep 17 00:00:00 2001 From: anyangml Date: Sat, 13 Apr 2024 11:34:27 +0000 Subject: [PATCH 19/43] feat: try remove standard --- deepmd/dpmodel/__init__.py | 4 ++-- deepmd/dpmodel/model/__init__.py | 4 ++-- deepmd/dpmodel/model/base_model.py | 4 ++-- deepmd/dpmodel/model/dp_model.py | 3 +-- deepmd/dpmodel/model/ener_model.py | 6 ++++++ deepmd/dpmodel/model/model.py | 10 +++++----- deepmd/dpmodel/model/spin_model.py | 9 ++++++--- .../model/atomic_model/dipole_atomic_model.py | 4 ++++ deepmd/pt/model/model/dp_zbl_model.py | 18 +++++++++++++++++- deepmd/pt/model/model/ener_model.py | 1 - source/tests/consistent/model/test_ener.py | 2 +- source/tests/pt/model/test_dp_model.py | 12 ++++++------ 12 files changed, 52 insertions(+), 25 deletions(-) create mode 100644 deepmd/dpmodel/model/ener_model.py diff --git a/deepmd/dpmodel/__init__.py b/deepmd/dpmodel/__init__.py index 6a7bdb3585..6f83f849a3 100644 --- a/deepmd/dpmodel/__init__.py +++ b/deepmd/dpmodel/__init__.py @@ -5,7 +5,7 @@ NativeOP, ) from .model import ( - DPModel, + DPModelCommon, ) from .output_def import ( FittingOutputDef, @@ -19,7 +19,7 @@ ) __all__ = [ - "DPModel", + "DPModelCommon", "PRECISION_DICT", "DEFAULT_PRECISION", "NativeOP", diff --git a/deepmd/dpmodel/model/__init__.py b/deepmd/dpmodel/model/__init__.py index c1ff15ab0d..7cd68dea60 100644 --- a/deepmd/dpmodel/model/__init__.py +++ b/deepmd/dpmodel/model/__init__.py @@ -13,7 +13,7 @@ """ from .dp_model import ( - DPModel, + DPModelCommon, ) from .make_model import ( make_model, @@ -23,7 +23,7 @@ ) __all__ = [ - "DPModel", + "DPModelCommon", "SpinModel", "make_model", ] diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 5169d1b5fe..814e93bc50 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -35,7 +35,7 @@ class BaseBaseModel(ABC, PluginVariant, make_plugin_registry("model")): def __new__(cls, *args, **kwargs): if inspect.isabstract(cls): - cls = cls.get_class_by_type(kwargs.get("type", "standard")) + cls = cls.get_class_by_type(kwargs.get("type", "ener")) return super().__new__(cls) @abstractmethod @@ -151,7 +151,7 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict): local_jdata : dict The local data refer to the current class """ - cls = cls.get_class_by_type(local_jdata.get("type", "standard")) + cls = cls.get_class_by_type(local_jdata.get("type", "ener")) return cls.update_sel(global_jdata, local_jdata) return BaseBaseModel diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 8d84c435b4..1e4dae6dd2 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -17,8 +17,7 @@ # use "class" to resolve "Variable not allowed in type expression" -@BaseModel.register("standard") -class DPModel(make_model(DPAtomicModel)): +class DPModelCommon: @classmethod def update_sel(cls, global_jdata: dict, local_jdata: dict): """Update the selection and perform neighbor statistics. diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py new file mode 100644 index 0000000000..ed22fc9f53 --- /dev/null +++ b/deepmd/dpmodel/model/ener_model.py @@ -0,0 +1,6 @@ +from .make_model import make_model +from .dp_model import DPModelCommon +from deepmd.dpmodel.atomic_model.dp_atomic_model import DPAtomicModel + +class EnergyModel(DPModelCommon, make_model(DPAtomicModel)): + pass \ No newline at end of file diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 3fdf5b802b..0df6e94f05 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -5,8 +5,8 @@ from deepmd.dpmodel.fitting.ener_fitting import ( EnergyFittingNet, ) -from deepmd.dpmodel.model.dp_model import ( - DPModel, +from deepmd.dpmodel.model.ener_model import ( + EnergyModel, ) from deepmd.dpmodel.model.spin_model import ( SpinModel, @@ -16,8 +16,8 @@ ) -def get_standard_model(data: dict) -> DPModel: - """Get a standard DPModel from a dictionary. +def get_standard_model(data: dict) -> EnergyModel: + """Get a EnergyModel from a dictionary. Parameters ---------- @@ -41,7 +41,7 @@ def get_standard_model(data: dict) -> DPModel: ) else: raise ValueError(f"Unknown fitting type {fitting_type}") - return DPModel( + return EnergyModel( descriptor=descriptor, fitting=fitting, type_map=data["type_map"], diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 5b31b64fdf..2388c601f9 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -7,8 +7,11 @@ import numpy as np -from deepmd.dpmodel.model.dp_model import ( - DPModel, +from deepmd.dpmodel.model.make_model import ( + make_model, +) +from deepmd.dpmodel.atomic_model.dp_atomic_model import( + DPAtomicModel ) from deepmd.utils.spin import ( Spin, @@ -259,7 +262,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data) -> "SpinModel": - backbone_model_obj = DPModel.deserialize(data["backbone_model"]) + backbone_model_obj = make_model(DPAtomicModel).deserialize(data["backbone_model"]) spin = Spin.deserialize(data["spin"]) return cls( backbone_model=backbone_model_obj, diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index 4152a3a358..f770dc00b8 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -11,6 +11,10 @@ class DPDipoleAtomicModel(DPAtomicModel): + def __init__(self, *args, **kwargs): + assert isinstance(fitting, DipoleFittingNet) + super().__init__(*args, **kwargs) + def apply_out_stat( self, ret: Dict[str, torch.Tensor], diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index 92e76da3d7..cc744fed91 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -24,7 +24,7 @@ @BaseModel.register("zbl") -class DPZBLModel(DPModelCommon, DPZBLModel_): +class DPZBLModel(DPZBLModel_): model_type = "ener" def __init__( @@ -103,3 +103,19 @@ def forward_lower( assert model_ret["dforce"] is not None model_predict["dforce"] = model_ret["dforce"] return model_predict + + @classmethod + def update_sel(cls, global_jdata: dict, local_jdata: dict): + """Update the selection and perform neighbor statistics. + Parameters + ---------- + global_jdata : dict + The global data, containing the training section + local_jdata : dict + The local data refer to the current class + """ + local_jdata_cpy = local_jdata.copy() + local_jdata_cpy["dpmodel"] = DPModelCommon.update_sel( + global_jdata, local_jdata["dpmodel"] + ) + return local_jdata_cpy \ No newline at end of file diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index 893d6a684a..b036616d33 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -22,7 +22,6 @@ @BaseModel.register("energy") -@BaseModel.register("standard") class EnergyModel(DPModelCommon, make_model(DPEnergyAtomicModel)): model_type = "ener" diff --git a/source/tests/consistent/model/test_ener.py b/source/tests/consistent/model/test_ener.py index da5033a3b6..c8ff9e4dcf 100644 --- a/source/tests/consistent/model/test_ener.py +++ b/source/tests/consistent/model/test_ener.py @@ -7,7 +7,7 @@ import numpy as np -from deepmd.dpmodel.model.dp_model import DPModel as EnergyModelDP +from deepmd.dpmodel.model.ener_model import EnergyModel as EnergyModelDP from deepmd.dpmodel.model.model import get_model as get_model_dp from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, diff --git a/source/tests/pt/model/test_dp_model.py b/source/tests/pt/model/test_dp_model.py index b0f2288b19..eb01674686 100644 --- a/source/tests/pt/model/test_dp_model.py +++ b/source/tests/pt/model/test_dp_model.py @@ -4,7 +4,7 @@ import numpy as np import torch -from deepmd.dpmodel import DPModel as DPDPModel +from deepmd.dpmodel.model.ener_model import EnergyModel as DPEnergyModel from deepmd.dpmodel.descriptor import DescrptSeA as DPDescrptSeA from deepmd.dpmodel.fitting import InvarFitting as DPInvarFitting from deepmd.pt.model.descriptor.se_a import ( @@ -131,7 +131,7 @@ def test_dp_consistency(self): numb_aparam=nap, ) type_map = ["foo", "bar"] - md0 = DPDPModel(ds, ft, type_map=type_map) + md0 = DPEnergyModel(ds, ft, type_map=type_map) md1 = EnergyModel.deserialize(md0.serialize()).to(env.DEVICE) rng = np.random.default_rng() @@ -172,7 +172,7 @@ def test_dp_consistency_nopbc(self): numb_aparam=nap, ) type_map = ["foo", "bar"] - md0 = DPDPModel(ds, ft, type_map=type_map) + md0 = DPEnergyModel(ds, ft, type_map=type_map) md1 = EnergyModel.deserialize(md0.serialize()).to(env.DEVICE) rng = np.random.default_rng() @@ -215,7 +215,7 @@ def test_prec_consistency(self): fparam = rng.normal(size=[self.nf, nfp]) aparam = rng.normal(size=[self.nf, nloc, nap]) - md0 = DPDPModel(ds, ft, type_map=type_map) + md0 = DPEnergyModel(ds, ft, type_map=type_map) md1 = EnergyModel.deserialize(md0.serialize()).to(env.DEVICE) args64 = [to_torch_tensor(ii) for ii in [self.coord, self.atype, self.cell]] @@ -317,7 +317,7 @@ def test_dp_consistency(self): mixed_types=ds.mixed_types(), ) type_map = ["foo", "bar"] - md0 = DPDPModel(ds, ft, type_map=type_map) + md0 = DPEnergyModel(ds, ft, type_map=type_map) md1 = EnergyModel.deserialize(md0.serialize()).to(env.DEVICE) args0 = [self.coord_ext, self.atype_ext, self.nlist] args1 = [ @@ -356,7 +356,7 @@ def test_prec_consistency(self): fparam = rng.normal(size=[self.nf, nfp]) aparam = rng.normal(size=[self.nf, nloc, nap]) - md0 = DPDPModel(ds, ft, type_map=type_map) + md0 = DPEnergyModel(ds, ft, type_map=type_map) md1 = EnergyModel.deserialize(md0.serialize()).to(env.DEVICE) args64 = [ From a51a792ab55bf955da72c376c01f2806f3c05cc3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 13 Apr 2024 11:35:00 +0000 Subject: [PATCH 20/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/model/dp_model.py | 10 ---------- deepmd/dpmodel/model/ener_model.py | 17 +++++++++++++---- deepmd/dpmodel/model/spin_model.py | 10 ++++++---- .../model/atomic_model/dipole_atomic_model.py | 2 +- deepmd/pt/model/model/dp_zbl_model.py | 3 ++- source/tests/pt/model/test_dp_model.py | 2 +- 6 files changed, 23 insertions(+), 21 deletions(-) diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 1e4dae6dd2..37cb426ab7 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -1,19 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from deepmd.dpmodel.atomic_model import ( - DPAtomicModel, -) from deepmd.dpmodel.descriptor.base_descriptor import ( BaseDescriptor, ) -from deepmd.dpmodel.model.base_model import ( - BaseModel, -) - -from .make_model import ( - make_model, -) # use "class" to resolve "Variable not allowed in type expression" diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index ed22fc9f53..6bf49b9906 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/deepmd/dpmodel/model/ener_model.py @@ -1,6 +1,15 @@ -from .make_model import make_model -from .dp_model import DPModelCommon -from deepmd.dpmodel.atomic_model.dp_atomic_model import DPAtomicModel +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.atomic_model.dp_atomic_model import ( + DPAtomicModel, +) + +from .dp_model import ( + DPModelCommon, +) +from .make_model import ( + make_model, +) + class EnergyModel(DPModelCommon, make_model(DPAtomicModel)): - pass \ No newline at end of file + pass diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 2388c601f9..90e2bb3fb4 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -7,12 +7,12 @@ import numpy as np +from deepmd.dpmodel.atomic_model.dp_atomic_model import ( + DPAtomicModel, +) from deepmd.dpmodel.model.make_model import ( make_model, ) -from deepmd.dpmodel.atomic_model.dp_atomic_model import( - DPAtomicModel -) from deepmd.utils.spin import ( Spin, ) @@ -262,7 +262,9 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data) -> "SpinModel": - backbone_model_obj = make_model(DPAtomicModel).deserialize(data["backbone_model"]) + backbone_model_obj = make_model(DPAtomicModel).deserialize( + data["backbone_model"] + ) spin = Spin.deserialize(data["spin"]) return cls( backbone_model=backbone_model_obj, diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index f770dc00b8..33b8d3f36b 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -14,7 +14,7 @@ class DPDipoleAtomicModel(DPAtomicModel): def __init__(self, *args, **kwargs): assert isinstance(fitting, DipoleFittingNet) super().__init__(*args, **kwargs) - + def apply_out_stat( self, ret: Dict[str, torch.Tensor], diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index cc744fed91..f18e1d097f 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -107,6 +107,7 @@ def forward_lower( @classmethod def update_sel(cls, global_jdata: dict, local_jdata: dict): """Update the selection and perform neighbor statistics. + Parameters ---------- global_jdata : dict @@ -118,4 +119,4 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict): local_jdata_cpy["dpmodel"] = DPModelCommon.update_sel( global_jdata, local_jdata["dpmodel"] ) - return local_jdata_cpy \ No newline at end of file + return local_jdata_cpy diff --git a/source/tests/pt/model/test_dp_model.py b/source/tests/pt/model/test_dp_model.py index eb01674686..c0ba07e754 100644 --- a/source/tests/pt/model/test_dp_model.py +++ b/source/tests/pt/model/test_dp_model.py @@ -4,9 +4,9 @@ import numpy as np import torch -from deepmd.dpmodel.model.ener_model import EnergyModel as DPEnergyModel from deepmd.dpmodel.descriptor import DescrptSeA as DPDescrptSeA from deepmd.dpmodel.fitting import InvarFitting as DPInvarFitting +from deepmd.dpmodel.model.ener_model import EnergyModel as DPEnergyModel from deepmd.pt.model.descriptor.se_a import ( DescrptSeA, ) From 4b14d549f80d136286856d82bb1b3210d1c66bb3 Mon Sep 17 00:00:00 2001 From: anyangml Date: Sat, 13 Apr 2024 11:48:10 +0000 Subject: [PATCH 21/43] fix: precommit --- deepmd/dpmodel/model/ener_model.py | 6 ++++-- deepmd/pt/model/atomic_model/dipole_atomic_model.py | 5 +++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index 6bf49b9906..bb15a27d81 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/deepmd/dpmodel/model/ener_model.py @@ -9,7 +9,9 @@ from .make_model import ( make_model, ) - - +from deepmd.dpmodel.model.base_model import ( + BaseModel, +) +@BaseModel.register("ener") class EnergyModel(DPModelCommon, make_model(DPAtomicModel)): pass diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index 33b8d3f36b..1a8febdbd9 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -8,12 +8,13 @@ from .dp_atomic_model import ( DPAtomicModel, ) +from deepmd.pt.model.task.dipole import DipoleFittingNet class DPDipoleAtomicModel(DPAtomicModel): - def __init__(self, *args, **kwargs): + def __init__(self, descriptor, fitting, type_map, **kwargs): assert isinstance(fitting, DipoleFittingNet) - super().__init__(*args, **kwargs) + super().__init__(descriptor, fitting, type_map, **kwargs) def apply_out_stat( self, From 548437590c3edcda32b34cb4b390058947527e54 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 13 Apr 2024 11:48:42 +0000 Subject: [PATCH 22/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/model/ener_model.py | 8 +++++--- deepmd/pt/model/atomic_model/dipole_atomic_model.py | 5 ++++- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index bb15a27d81..d786737c95 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/deepmd/dpmodel/model/ener_model.py @@ -2,6 +2,9 @@ from deepmd.dpmodel.atomic_model.dp_atomic_model import ( DPAtomicModel, ) +from deepmd.dpmodel.model.base_model import ( + BaseModel, +) from .dp_model import ( DPModelCommon, @@ -9,9 +12,8 @@ from .make_model import ( make_model, ) -from deepmd.dpmodel.model.base_model import ( - BaseModel, -) + + @BaseModel.register("ener") class EnergyModel(DPModelCommon, make_model(DPAtomicModel)): pass diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index 1a8febdbd9..1723a30f2d 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -5,10 +5,13 @@ import torch +from deepmd.pt.model.task.dipole import ( + DipoleFittingNet, +) + from .dp_atomic_model import ( DPAtomicModel, ) -from deepmd.pt.model.task.dipole import DipoleFittingNet class DPDipoleAtomicModel(DPAtomicModel): From 3c5c7951d0c29fc33dc87737a9dd86c53c4c33e5 Mon Sep 17 00:00:00 2001 From: anyangml Date: Sat, 13 Apr 2024 12:00:34 +0000 Subject: [PATCH 23/43] fix: import --- source/tests/common/dpmodel/test_dp_model.py | 12 ++++++------ source/tests/common/dpmodel/test_nlist.py | 6 +++--- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/source/tests/common/dpmodel/test_dp_model.py b/source/tests/common/dpmodel/test_dp_model.py index 9121c7cd07..67ab6a7c32 100644 --- a/source/tests/common/dpmodel/test_dp_model.py +++ b/source/tests/common/dpmodel/test_dp_model.py @@ -9,8 +9,8 @@ from deepmd.dpmodel.fitting import ( InvarFitting, ) -from deepmd.dpmodel.model import ( - DPModel, +from deepmd.dpmodel.model.ener_model import ( + EnergyModel, ) from .case_single_frame_with_nlist import ( @@ -40,8 +40,8 @@ def test_self_consistency( mixed_types=ds.mixed_types(), ) type_map = ["foo", "bar"] - md0 = DPModel(ds, ft, type_map=type_map) - md1 = DPModel.deserialize(md0.serialize()) + md0 = EnergyModel(ds, ft, type_map=type_map) + md1 = EnergyModel.deserialize(md0.serialize()) ret0 = md0.call_lower(self.coord_ext, self.atype_ext, self.nlist) ret1 = md1.call_lower(self.coord_ext, self.atype_ext, self.nlist) @@ -70,7 +70,7 @@ def test_prec_consistency(self): fparam = rng.normal(size=[self.nf, nfp]) aparam = rng.normal(size=[self.nf, nloc, nap]) - md1 = DPModel(ds, ft, type_map=type_map) + md1 = EnergyModel(ds, ft, type_map=type_map) args64 = [self.coord_ext, self.atype_ext, self.nlist] args64[0] = args64[0].astype(np.float64) @@ -122,7 +122,7 @@ def test_prec_consistency(self): fparam = rng.normal(size=[self.nf, nfp]) aparam = rng.normal(size=[self.nf, nloc, nap]) - md1 = DPModel(ds, ft, type_map=type_map) + md1 = EnergyModel(ds, ft, type_map=type_map) args64 = [self.coord, self.atype, self.cell] args64[0] = args64[0].astype(np.float64) diff --git a/source/tests/common/dpmodel/test_nlist.py b/source/tests/common/dpmodel/test_nlist.py index ee8a7139e7..404232013b 100644 --- a/source/tests/common/dpmodel/test_nlist.py +++ b/source/tests/common/dpmodel/test_nlist.py @@ -9,8 +9,8 @@ from deepmd.dpmodel.fitting import ( InvarFitting, ) -from deepmd.dpmodel.model import ( - DPModel, +from deepmd.dpmodel.model.ener_model import ( + EnergyModel, ) from deepmd.dpmodel.utils import ( build_multiple_neighbor_list, @@ -65,7 +65,7 @@ def setUp(self): mixed_types=ds.mixed_types(), ) type_map = ["foo", "bar"] - self.md = DPModel(ds, ft, type_map=type_map) + self.md = EnergyModel(ds, ft, type_map=type_map) def test_nlist_eq(self): # n_nnei == nnei From 224ef2751dc7a90d1a53724a13104259a1cbfe09 Mon Sep 17 00:00:00 2001 From: anyangml Date: Sat, 13 Apr 2024 12:37:18 +0000 Subject: [PATCH 24/43] fix: UTs --- deepmd/pt/model/atomic_model/dos_atomic_model.py | 7 ++++++- deepmd/pt/model/atomic_model/energy_atomic_model.py | 8 +++++++- deepmd/pt/model/atomic_model/polar_atomic_model.py | 7 +++++++ deepmd/pt/model/model/ener_model.py | 2 +- 4 files changed, 21 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/model/atomic_model/dos_atomic_model.py b/deepmd/pt/model/atomic_model/dos_atomic_model.py index aba391caef..556ba84480 100644 --- a/deepmd/pt/model/atomic_model/dos_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dos_atomic_model.py @@ -3,6 +3,11 @@ DPAtomicModel, ) +from deepmd.pt.model.task.dos import ( + DOSFittingNet, +) class DPDOSAtomicModel(DPAtomicModel): - pass + def __init__(self, descriptor, fitting, type_map, **kwargs): + assert isinstance(fitting, DOSFittingNet) + super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/pt/model/atomic_model/energy_atomic_model.py b/deepmd/pt/model/atomic_model/energy_atomic_model.py index e2e9432d1e..b81836203c 100644 --- a/deepmd/pt/model/atomic_model/energy_atomic_model.py +++ b/deepmd/pt/model/atomic_model/energy_atomic_model.py @@ -2,7 +2,13 @@ from .dp_atomic_model import ( DPAtomicModel, ) +from deepmd.pt.model.task.ener import ( + EnergyFittingNet, + EnergyFittingNetDirect, +) class DPEnergyAtomicModel(DPAtomicModel): - pass + def __init__(self, descriptor, fitting, type_map, **kwargs): + assert isinstance(fitting, EnergyFittingNet) or isinstance(fitting, EnergyFittingNetDirect) + super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/pt/model/atomic_model/polar_atomic_model.py b/deepmd/pt/model/atomic_model/polar_atomic_model.py index da19a31211..596b3a158e 100644 --- a/deepmd/pt/model/atomic_model/polar_atomic_model.py +++ b/deepmd/pt/model/atomic_model/polar_atomic_model.py @@ -8,9 +8,16 @@ from .dp_atomic_model import ( DPAtomicModel, ) +from deepmd.pt.model.task.polarizability import ( + PolarFittingNet, +) class DPPolarAtomicModel(DPAtomicModel): + def __init__(self, descriptor, fitting, type_map, **kwargs): + assert isinstance(fitting, PolarFittingNet) + super().__init__(descriptor, fitting, type_map, **kwargs) + def apply_out_stat( self, ret: Dict[str, torch.Tensor], diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index b036616d33..334c8a46ee 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -21,7 +21,7 @@ ) -@BaseModel.register("energy") +@BaseModel.register("ener") class EnergyModel(DPModelCommon, make_model(DPEnergyAtomicModel)): model_type = "ener" From bd079b96869a6a5db11c551b8177bd2864fde33f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 13 Apr 2024 12:37:50 +0000 Subject: [PATCH 25/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/atomic_model/dos_atomic_model.py | 7 ++++--- deepmd/pt/model/atomic_model/energy_atomic_model.py | 11 +++++++---- deepmd/pt/model/atomic_model/polar_atomic_model.py | 7 ++++--- 3 files changed, 15 insertions(+), 10 deletions(-) diff --git a/deepmd/pt/model/atomic_model/dos_atomic_model.py b/deepmd/pt/model/atomic_model/dos_atomic_model.py index 556ba84480..5e399f2aff 100644 --- a/deepmd/pt/model/atomic_model/dos_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dos_atomic_model.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.pt.model.task.dos import ( + DOSFittingNet, +) + from .dp_atomic_model import ( DPAtomicModel, ) -from deepmd.pt.model.task.dos import ( - DOSFittingNet, -) class DPDOSAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): diff --git a/deepmd/pt/model/atomic_model/energy_atomic_model.py b/deepmd/pt/model/atomic_model/energy_atomic_model.py index b81836203c..d1f159e7c2 100644 --- a/deepmd/pt/model/atomic_model/energy_atomic_model.py +++ b/deepmd/pt/model/atomic_model/energy_atomic_model.py @@ -1,14 +1,17 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from .dp_atomic_model import ( - DPAtomicModel, -) from deepmd.pt.model.task.ener import ( EnergyFittingNet, EnergyFittingNetDirect, ) +from .dp_atomic_model import ( + DPAtomicModel, +) + class DPEnergyAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - assert isinstance(fitting, EnergyFittingNet) or isinstance(fitting, EnergyFittingNetDirect) + assert isinstance(fitting, EnergyFittingNet) or isinstance( + fitting, EnergyFittingNetDirect + ) super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/pt/model/atomic_model/polar_atomic_model.py b/deepmd/pt/model/atomic_model/polar_atomic_model.py index 596b3a158e..3eb4136b6e 100644 --- a/deepmd/pt/model/atomic_model/polar_atomic_model.py +++ b/deepmd/pt/model/atomic_model/polar_atomic_model.py @@ -5,13 +5,14 @@ import torch -from .dp_atomic_model import ( - DPAtomicModel, -) from deepmd.pt.model.task.polarizability import ( PolarFittingNet, ) +from .dp_atomic_model import ( + DPAtomicModel, +) + class DPPolarAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): From 53dbe13c229d43eacb714fea10581099b8f02f8f Mon Sep 17 00:00:00 2001 From: anyangml Date: Sat, 13 Apr 2024 13:07:01 +0000 Subject: [PATCH 26/43] fix: UTs --- source/tests/pt/model/test_dp_model.py | 52 +++++++------------------- 1 file changed, 14 insertions(+), 38 deletions(-) diff --git a/source/tests/pt/model/test_dp_model.py b/source/tests/pt/model/test_dp_model.py index c0ba07e754..3d59a33ca2 100644 --- a/source/tests/pt/model/test_dp_model.py +++ b/source/tests/pt/model/test_dp_model.py @@ -5,7 +5,7 @@ import torch from deepmd.dpmodel.descriptor import DescrptSeA as DPDescrptSeA -from deepmd.dpmodel.fitting import InvarFitting as DPInvarFitting +from deepmd.dpmodel.fitting import EnergyFittingNet as DPEnergyFittingNet from deepmd.dpmodel.model.ener_model import EnergyModel as DPEnergyModel from deepmd.pt.model.descriptor.se_a import ( DescrptSeA, @@ -14,7 +14,7 @@ EnergyModel, ) from deepmd.pt.model.task.ener import ( - InvarFitting, + EnergyFittingNet, ) from deepmd.pt.utils import ( env, @@ -48,11 +48,9 @@ def test_self_consistency(self): self.rcut_smth, self.sel, ).to(env.DEVICE) - ft = InvarFitting( - "energy", + ft = EnergyFittingNet( self.nt, ds.get_dim_out(), - 1, mixed_types=ds.mixed_types(), ).to(env.DEVICE) type_map = ["foo", "bar"] @@ -121,11 +119,9 @@ def test_dp_consistency(self): self.rcut_smth, self.sel, ) - ft = DPInvarFitting( - "energy", + ft = DPEnergyFittingNet( self.nt, ds.get_dim_out(), - 1, mixed_types=ds.mixed_types(), numb_fparam=nfp, numb_aparam=nap, @@ -162,11 +158,9 @@ def test_dp_consistency_nopbc(self): self.rcut_smth, self.sel, ) - ft = DPInvarFitting( - "energy", + ft = DPEnergyFittingNet( self.nt, ds.get_dim_out(), - 1, mixed_types=ds.mixed_types(), numb_fparam=nfp, numb_aparam=nap, @@ -203,11 +197,9 @@ def test_prec_consistency(self): self.rcut_smth, self.sel, ) - ft = DPInvarFitting( - "energy", + ft = DPEnergyFittingNet( self.nt, ds.get_dim_out(), - 1, mixed_types=ds.mixed_types(), ) nfp, nap = 2, 3 @@ -258,11 +250,9 @@ def test_self_consistency(self): self.rcut_smth, self.sel, ).to(env.DEVICE) - ft = InvarFitting( - "energy", + ft = EnergyFittingNet( self.nt, ds.get_dim_out(), - 1, mixed_types=ds.mixed_types(), ).to(env.DEVICE) type_map = ["foo", "bar"] @@ -309,11 +299,9 @@ def test_dp_consistency(self): self.rcut_smth, self.sel, ) - ft = DPInvarFitting( - "energy", + ft = DPEnergyFittingNet( self.nt, ds.get_dim_out(), - 1, mixed_types=ds.mixed_types(), ) type_map = ["foo", "bar"] @@ -344,11 +332,9 @@ def test_prec_consistency(self): self.rcut_smth, self.sel, ) - ft = DPInvarFitting( - "energy", + ft = DPEnergyFittingNet( self.nt, ds.get_dim_out(), - 1, mixed_types=ds.mixed_types(), ) nfp, nap = 2, 3 @@ -396,11 +382,9 @@ def test_jit(self): self.rcut_smth, self.sel, ).to(env.DEVICE) - ft = InvarFitting( - "energy", + ft = EnergyFittingNet( self.nt, ds.get_dim_out(), - 1, mixed_types=ds.mixed_types(), ).to(env.DEVICE) type_map = ["foo", "bar"] @@ -446,11 +430,9 @@ def setUp(self): self.rcut_smth, self.sel, ).to(env.DEVICE) - ft = InvarFitting( - "energy", + ft = EnergyFittingNet( self.nt, ds.get_dim_out(), - 1, mixed_types=ds.mixed_types(), ).to(env.DEVICE) type_map = ["foo", "bar"] @@ -519,11 +501,9 @@ def test_self_consistency(self): self.rcut_smth, self.sel, ).to(env.DEVICE) - ft = InvarFitting( - "energy", + ft = EnergyFittingNet( self.nt, ds.get_dim_out(), - 1, mixed_types=ds.mixed_types(), ).to(env.DEVICE) type_map = ["foo", "bar"] @@ -589,11 +569,9 @@ def test_self_consistency(self): self.rcut_smth, self.sel, ).to(env.DEVICE) - ft = InvarFitting( - "energy", + ft = EnergyFittingNet( self.nt, ds.get_dim_out(), - 1, mixed_types=ds.mixed_types(), ).to(env.DEVICE) type_map = ["foo", "bar"] @@ -639,11 +617,9 @@ def test_jit(self): self.rcut_smth, self.sel, ).to(env.DEVICE) - ft = InvarFitting( - "energy", + ft = EnergyFittingNet( self.nt, ds.get_dim_out(), - 1, mixed_types=ds.mixed_types(), ).to(env.DEVICE) type_map = ["foo", "bar"] From 78ff90ddb89c6e7846666095902513a36868bcff Mon Sep 17 00:00:00 2001 From: anyangml Date: Sat, 13 Apr 2024 13:44:10 +0000 Subject: [PATCH 27/43] fix: UTs --- deepmd/pt/model/atomic_model/energy_atomic_model.py | 3 ++- deepmd/utils/argcheck.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/atomic_model/energy_atomic_model.py b/deepmd/pt/model/atomic_model/energy_atomic_model.py index d1f159e7c2..69eefc791b 100644 --- a/deepmd/pt/model/atomic_model/energy_atomic_model.py +++ b/deepmd/pt/model/atomic_model/energy_atomic_model.py @@ -2,6 +2,7 @@ from deepmd.pt.model.task.ener import ( EnergyFittingNet, EnergyFittingNetDirect, + InvarFitting, ) from .dp_atomic_model import ( @@ -13,5 +14,5 @@ class DPEnergyAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): assert isinstance(fitting, EnergyFittingNet) or isinstance( fitting, EnergyFittingNetDirect - ) + ) or isinstance(fitting, InvarFitting) super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 07add486c1..3dae7b07ef 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1372,7 +1372,7 @@ def model_args(exclude_hybrid=False): *hybrid_models, ], optional=True, - default_tag="standard", + default_tag="ener", ), ], ) From 7f9951223d75d283a57234c4e18d897e38759635 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 13 Apr 2024 13:44:42 +0000 Subject: [PATCH 28/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/atomic_model/energy_atomic_model.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/model/atomic_model/energy_atomic_model.py b/deepmd/pt/model/atomic_model/energy_atomic_model.py index 69eefc791b..7cedaa1ab3 100644 --- a/deepmd/pt/model/atomic_model/energy_atomic_model.py +++ b/deepmd/pt/model/atomic_model/energy_atomic_model.py @@ -12,7 +12,9 @@ class DPEnergyAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - assert isinstance(fitting, EnergyFittingNet) or isinstance( - fitting, EnergyFittingNetDirect - ) or isinstance(fitting, InvarFitting) + assert ( + isinstance(fitting, EnergyFittingNet) + or isinstance(fitting, EnergyFittingNetDirect) + or isinstance(fitting, InvarFitting) + ) super().__init__(descriptor, fitting, type_map, **kwargs) From 23fa34d92a17ba492ee4bbfbcf3928a8f9378c54 Mon Sep 17 00:00:00 2001 From: anyangml Date: Sat, 13 Apr 2024 14:02:38 +0000 Subject: [PATCH 29/43] fix: try modify argcheck --- deepmd/utils/argcheck.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 3dae7b07ef..a448334207 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1383,7 +1383,7 @@ def standard_model_args() -> Argument: doc_fitting = "The fitting of physical properties." ca = Argument( - "standard", + "ener", dict, [ Argument( From 16583decf0ef061408515ff429099598eee989e8 Mon Sep 17 00:00:00 2001 From: anyangml Date: Sat, 13 Apr 2024 14:14:07 +0000 Subject: [PATCH 30/43] fix: try modify argcheck --- deepmd/tf/model/dos.py | 2 +- deepmd/tf/model/ener.py | 2 +- deepmd/tf/model/model.py | 1 - deepmd/tf/model/tensor.py | 3 ++- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/deepmd/tf/model/dos.py b/deepmd/tf/model/dos.py index 265026b60a..e634d41631 100644 --- a/deepmd/tf/model/dos.py +++ b/deepmd/tf/model/dos.py @@ -22,7 +22,7 @@ merge_sys_stat, ) - +@StandardModel.register("dos") class DOSModel(StandardModel): """DOS model. diff --git a/deepmd/tf/model/ener.py b/deepmd/tf/model/ener.py index a493fe0517..e6b613a918 100644 --- a/deepmd/tf/model/ener.py +++ b/deepmd/tf/model/ener.py @@ -31,7 +31,7 @@ merge_sys_stat, ) - +@StandardModel.register("ener") class EnerModel(StandardModel): """Energy model. diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index fc8f862e3b..294797fdfb 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -589,7 +589,6 @@ def serialize(self, suffix: str = "") -> dict: raise NotImplementedError("Not implemented in class %s" % self.__name__) -@Model.register("standard") class StandardModel(Model): """Standard model, which must contain a descriptor and a fitting. diff --git a/deepmd/tf/model/tensor.py b/deepmd/tf/model/tensor.py index b232f40b13..445091fbf2 100644 --- a/deepmd/tf/model/tensor.py +++ b/deepmd/tf/model/tensor.py @@ -21,7 +21,8 @@ merge_sys_stat, ) - +@StandardModel.register("dipole") +@StandardModel.register("polar") class TensorModel(StandardModel): """Tensor model. From 23fd899de342c125c395eebe32f60295a2d52c0b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 13 Apr 2024 14:14:40 +0000 Subject: [PATCH 31/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/model/dos.py | 1 + deepmd/tf/model/ener.py | 1 + deepmd/tf/model/tensor.py | 1 + 3 files changed, 3 insertions(+) diff --git a/deepmd/tf/model/dos.py b/deepmd/tf/model/dos.py index e634d41631..2d244f5733 100644 --- a/deepmd/tf/model/dos.py +++ b/deepmd/tf/model/dos.py @@ -22,6 +22,7 @@ merge_sys_stat, ) + @StandardModel.register("dos") class DOSModel(StandardModel): """DOS model. diff --git a/deepmd/tf/model/ener.py b/deepmd/tf/model/ener.py index e6b613a918..4a0195b334 100644 --- a/deepmd/tf/model/ener.py +++ b/deepmd/tf/model/ener.py @@ -31,6 +31,7 @@ merge_sys_stat, ) + @StandardModel.register("ener") class EnerModel(StandardModel): """Energy model. diff --git a/deepmd/tf/model/tensor.py b/deepmd/tf/model/tensor.py index 445091fbf2..ba41879c50 100644 --- a/deepmd/tf/model/tensor.py +++ b/deepmd/tf/model/tensor.py @@ -21,6 +21,7 @@ merge_sys_stat, ) + @StandardModel.register("dipole") @StandardModel.register("polar") class TensorModel(StandardModel): From 4fe3661c9733c4a8621d4a832c8e30e43a8c0ef1 Mon Sep 17 00:00:00 2001 From: anyangml Date: Sat, 13 Apr 2024 14:22:46 +0000 Subject: [PATCH 32/43] fix: plugin --- deepmd/tf/model/dos.py | 2 ++ deepmd/tf/model/ener.py | 3 +++ deepmd/tf/model/tensor.py | 3 +++ 3 files changed, 8 insertions(+) diff --git a/deepmd/tf/model/dos.py b/deepmd/tf/model/dos.py index 2d244f5733..faca283e1f 100644 --- a/deepmd/tf/model/dos.py +++ b/deepmd/tf/model/dos.py @@ -199,6 +199,8 @@ def build( return model_dict + def serialize(self, suffix: str = "") -> dict: + return super().serialize(suffix= "").update({"type":"dos"}) def init_variables( self, graph: tf.Graph, diff --git a/deepmd/tf/model/ener.py b/deepmd/tf/model/ener.py index 4a0195b334..7555af5c1d 100644 --- a/deepmd/tf/model/ener.py +++ b/deepmd/tf/model/ener.py @@ -346,6 +346,9 @@ def build( return model_dict + def serialize(self, suffix: str = "") -> dict: + return super().serialize(suffix= "").update({"type":"ener"}) + def init_variables( self, graph: tf.Graph, diff --git a/deepmd/tf/model/tensor.py b/deepmd/tf/model/tensor.py index ba41879c50..741b25dc27 100644 --- a/deepmd/tf/model/tensor.py +++ b/deepmd/tf/model/tensor.py @@ -208,6 +208,9 @@ def build( return model_dict + def serialize(self, suffix: str = "") -> dict: + return super().serialize(suffix= "").update({"type":self.model_type}) + def init_variables( self, graph: tf.Graph, From 576e75d6e2093fd143cf756433888559649205fb Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 13 Apr 2024 14:23:18 +0000 Subject: [PATCH 33/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/model/dos.py | 3 ++- deepmd/tf/model/ener.py | 2 +- deepmd/tf/model/tensor.py | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/deepmd/tf/model/dos.py b/deepmd/tf/model/dos.py index faca283e1f..5fbea707a4 100644 --- a/deepmd/tf/model/dos.py +++ b/deepmd/tf/model/dos.py @@ -200,7 +200,8 @@ def build( return model_dict def serialize(self, suffix: str = "") -> dict: - return super().serialize(suffix= "").update({"type":"dos"}) + return super().serialize(suffix="").update({"type": "dos"}) + def init_variables( self, graph: tf.Graph, diff --git a/deepmd/tf/model/ener.py b/deepmd/tf/model/ener.py index 7555af5c1d..2ccb54f576 100644 --- a/deepmd/tf/model/ener.py +++ b/deepmd/tf/model/ener.py @@ -347,7 +347,7 @@ def build( return model_dict def serialize(self, suffix: str = "") -> dict: - return super().serialize(suffix= "").update({"type":"ener"}) + return super().serialize(suffix="").update({"type": "ener"}) def init_variables( self, diff --git a/deepmd/tf/model/tensor.py b/deepmd/tf/model/tensor.py index 741b25dc27..9c9d9d388c 100644 --- a/deepmd/tf/model/tensor.py +++ b/deepmd/tf/model/tensor.py @@ -209,7 +209,7 @@ def build( return model_dict def serialize(self, suffix: str = "") -> dict: - return super().serialize(suffix= "").update({"type":self.model_type}) + return super().serialize(suffix="").update({"type": self.model_type}) def init_variables( self, From 3173978ddfa5c49bd12af504adf1305dd001b47e Mon Sep 17 00:00:00 2001 From: anyangml Date: Sun, 14 Apr 2024 04:57:43 +0000 Subject: [PATCH 34/43] fix:inheritance --- deepmd/dpmodel/model/ener_model.py | 12 +++++++++--- deepmd/pt/model/model/dipole_model.py | 7 ++++--- deepmd/pt/model/model/dos_model.py | 6 ++++-- deepmd/pt/model/model/ener_model.py | 7 ++++--- deepmd/pt/model/model/polar_model.py | 7 ++++--- 5 files changed, 25 insertions(+), 14 deletions(-) diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index d786737c95..d0ef143115 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/deepmd/dpmodel/model/ener_model.py @@ -13,7 +13,13 @@ make_model, ) - +DPEnergyModel_ = make_model(DPAtomicModel) @BaseModel.register("ener") -class EnergyModel(DPModelCommon, make_model(DPAtomicModel)): - pass +class EnergyModel(DPModelCommon, DPEnergyModel_): + def __init__( + self, + *args, + **kwargs, + ): + DPModelCommon.__init__(self) + DPEnergyModel_.__init__(self, *args, **kwargs) \ No newline at end of file diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index e98aee3207..3879575ddb 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -20,9 +20,9 @@ make_model, ) - +DPDOSModel_ = make_model(DPDipoleAtomicModel) @BaseModel.register("dipole") -class DipoleModel(DPModelCommon, make_model(DPDipoleAtomicModel)): +class DipoleModel(DPModelCommon, DPDOSModel_): model_type = "dipole" def __init__( @@ -30,7 +30,8 @@ def __init__( *args, **kwargs, ): - super().__init__(*args, **kwargs) + DPModelCommon.__init__(self) + DPDOSModel_.__init__(self, *args, **kwargs) def forward( self, diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index 490e256852..949532e9e1 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -20,9 +20,10 @@ make_model, ) +DPDOSModel_ = make_model(DPDOSAtomicModel) @BaseModel.register("dos") -class DOSModel(DPModelCommon, make_model(DPDOSAtomicModel)): +class DOSModel(DPModelCommon, DPDOSModel_): model_type = "dos" def __init__( @@ -30,7 +31,8 @@ def __init__( *args, **kwargs, ): - super().__init__(*args, **kwargs) + DPModelCommon.__init__(self) + DPDOSModel_.__init__(self, *args, **kwargs) def forward( self, diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index 334c8a46ee..ba453e0d8b 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -20,9 +20,9 @@ make_model, ) - +DPEnergyModel_ = make_model(DPEnergyAtomicModel) @BaseModel.register("ener") -class EnergyModel(DPModelCommon, make_model(DPEnergyAtomicModel)): +class EnergyModel(DPModelCommon, DPEnergyModel_): model_type = "ener" def __init__( @@ -30,7 +30,8 @@ def __init__( *args, **kwargs, ): - super().__init__(*args, **kwargs) + DPModelCommon.__init__(self) + DPEnergyModel_.__init__(self, *args, **kwargs) def forward( self, diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index 3349c164a6..4d09168c34 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -20,9 +20,9 @@ make_model, ) - +DPDOSModel_ = make_model(DPPolarAtomicModel) @BaseModel.register("polar") -class PolarModel(DPModelCommon, make_model(DPPolarAtomicModel)): +class PolarModel(DPModelCommon, DPDOSModel_): model_type = "polar" def __init__( @@ -30,7 +30,8 @@ def __init__( *args, **kwargs, ): - super().__init__(*args, **kwargs) + DPModelCommon.__init__(self) + DPDOSModel_.__init__(self, *args, **kwargs) def forward( self, From 1e54270809ac64b63b9a70010451798116b907ad Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 14 Apr 2024 04:58:44 +0000 Subject: [PATCH 35/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/model/ener_model.py | 14 ++++++++------ deepmd/pt/model/model/dipole_model.py | 2 ++ deepmd/pt/model/model/dos_model.py | 1 + deepmd/pt/model/model/ener_model.py | 2 ++ deepmd/pt/model/model/polar_model.py | 2 ++ 5 files changed, 15 insertions(+), 6 deletions(-) diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index d0ef143115..5f21681830 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/deepmd/dpmodel/model/ener_model.py @@ -14,12 +14,14 @@ ) DPEnergyModel_ = make_model(DPAtomicModel) + + @BaseModel.register("ener") class EnergyModel(DPModelCommon, DPEnergyModel_): def __init__( - self, - *args, - **kwargs, - ): - DPModelCommon.__init__(self) - DPEnergyModel_.__init__(self, *args, **kwargs) \ No newline at end of file + self, + *args, + **kwargs, + ): + DPModelCommon.__init__(self) + DPEnergyModel_.__init__(self, *args, **kwargs) diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index 3879575ddb..b73976432b 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -21,6 +21,8 @@ ) DPDOSModel_ = make_model(DPDipoleAtomicModel) + + @BaseModel.register("dipole") class DipoleModel(DPModelCommon, DPDOSModel_): model_type = "dipole" diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index 949532e9e1..24095002d1 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -22,6 +22,7 @@ DPDOSModel_ = make_model(DPDOSAtomicModel) + @BaseModel.register("dos") class DOSModel(DPModelCommon, DPDOSModel_): model_type = "dos" diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index ba453e0d8b..4a0eb49945 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -21,6 +21,8 @@ ) DPEnergyModel_ = make_model(DPEnergyAtomicModel) + + @BaseModel.register("ener") class EnergyModel(DPModelCommon, DPEnergyModel_): model_type = "ener" diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index 4d09168c34..867afefc27 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -21,6 +21,8 @@ ) DPDOSModel_ = make_model(DPPolarAtomicModel) + + @BaseModel.register("polar") class PolarModel(DPModelCommon, DPDOSModel_): model_type = "polar" From 353ca6d421265389ca98d5444f8757d7e91e327e Mon Sep 17 00:00:00 2001 From: anyangml Date: Sun, 14 Apr 2024 11:31:39 +0000 Subject: [PATCH 36/43] fix: revert changes --- deepmd/dpmodel/model/base_model.py | 10 +++++++--- deepmd/tf/model/model.py | 2 +- deepmd/utils/argcheck.py | 4 ++-- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 814e93bc50..03612da6b2 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -35,7 +35,9 @@ class BaseBaseModel(ABC, PluginVariant, make_plugin_registry("model")): def __new__(cls, *args, **kwargs): if inspect.isabstract(cls): - cls = cls.get_class_by_type(kwargs.get("type", "ener")) + # getting model type based on fitting type + model_type = kwargs['fitting']['type'] + cls = cls.get_class_by_type(model_type) return super().__new__(cls) @abstractmethod @@ -118,7 +120,7 @@ def deserialize(cls, data: dict) -> "BaseBaseModel": The deserialized model """ if inspect.isabstract(cls): - return cls.get_class_by_type(data["type"]).deserialize(data) + return cls.get_class_by_type(data['fitting']["type"]).deserialize(data) raise NotImplementedError("Not implemented in class %s" % cls.__name__) model_def_script: str @@ -151,7 +153,9 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict): local_jdata : dict The local data refer to the current class """ - cls = cls.get_class_by_type(local_jdata.get("type", "ener")) + # getting model type based on fitting type + model_type = local_jdata['fitting']['type'] + cls = cls.get_class_by_type(model_type) return cls.update_sel(global_jdata, local_jdata) return BaseBaseModel diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 294797fdfb..09f6d40d57 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -588,7 +588,7 @@ def serialize(self, suffix: str = "") -> dict: """ raise NotImplementedError("Not implemented in class %s" % self.__name__) - +@Model.register("standard") class StandardModel(Model): """Standard model, which must contain a descriptor and a fitting. diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index a448334207..07add486c1 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1372,7 +1372,7 @@ def model_args(exclude_hybrid=False): *hybrid_models, ], optional=True, - default_tag="ener", + default_tag="standard", ), ], ) @@ -1383,7 +1383,7 @@ def standard_model_args() -> Argument: doc_fitting = "The fitting of physical properties." ca = Argument( - "ener", + "standard", dict, [ Argument( From 8de02ee8dc70c2aed3e34c58b08e1fbe0307f32c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 14 Apr 2024 11:33:17 +0000 Subject: [PATCH 37/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/model/base_model.py | 6 +++--- deepmd/tf/model/model.py | 1 + 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 03612da6b2..2cc55dd8d9 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -36,7 +36,7 @@ class BaseBaseModel(ABC, PluginVariant, make_plugin_registry("model")): def __new__(cls, *args, **kwargs): if inspect.isabstract(cls): # getting model type based on fitting type - model_type = kwargs['fitting']['type'] + model_type = kwargs["fitting"]["type"] cls = cls.get_class_by_type(model_type) return super().__new__(cls) @@ -120,7 +120,7 @@ def deserialize(cls, data: dict) -> "BaseBaseModel": The deserialized model """ if inspect.isabstract(cls): - return cls.get_class_by_type(data['fitting']["type"]).deserialize(data) + return cls.get_class_by_type(data["fitting"]["type"]).deserialize(data) raise NotImplementedError("Not implemented in class %s" % cls.__name__) model_def_script: str @@ -154,7 +154,7 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict): The local data refer to the current class """ # getting model type based on fitting type - model_type = local_jdata['fitting']['type'] + model_type = local_jdata["fitting"]["type"] cls = cls.get_class_by_type(model_type) return cls.update_sel(global_jdata, local_jdata) diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 09f6d40d57..fc8f862e3b 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -588,6 +588,7 @@ def serialize(self, suffix: str = "") -> dict: """ raise NotImplementedError("Not implemented in class %s" % self.__name__) + @Model.register("standard") class StandardModel(Model): """Standard model, which must contain a descriptor and a fitting. From e2a80b3ad507282d75f4cbcabde17011a1a9d9c6 Mon Sep 17 00:00:00 2001 From: anyangml Date: Sun, 14 Apr 2024 12:21:37 +0000 Subject: [PATCH 38/43] fix: UTs --- deepmd/dpmodel/model/base_model.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 2cc55dd8d9..34873bb417 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -36,7 +36,9 @@ class BaseBaseModel(ABC, PluginVariant, make_plugin_registry("model")): def __new__(cls, *args, **kwargs): if inspect.isabstract(cls): # getting model type based on fitting type - model_type = kwargs["fitting"]["type"] + model_type = kwargs.get("type","standard") + if model_type == "standard": + model_type = kwargs.get("fitting",{}).get("type","ener") cls = cls.get_class_by_type(model_type) return super().__new__(cls) @@ -120,7 +122,10 @@ def deserialize(cls, data: dict) -> "BaseBaseModel": The deserialized model """ if inspect.isabstract(cls): - return cls.get_class_by_type(data["fitting"]["type"]).deserialize(data) + model_type = data.get("type","standard") + if model_type == "standard": + model_type = data.get("fitting",{}).get("type","ener") + return cls.get_class_by_type(model_type).deserialize(data) raise NotImplementedError("Not implemented in class %s" % cls.__name__) model_def_script: str @@ -154,7 +159,9 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict): The local data refer to the current class """ # getting model type based on fitting type - model_type = local_jdata["fitting"]["type"] + model_type = local_jdata.get("type","standard") + if model_type == "standard": + model_type = local_jdata.get("fitting",{}).get("type","ener") cls = cls.get_class_by_type(model_type) return cls.update_sel(global_jdata, local_jdata) From efcf9514b8e40e10e738b516d51f9ab087271ed5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 14 Apr 2024 12:22:09 +0000 Subject: [PATCH 39/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/model/base_model.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 34873bb417..de69afcf6c 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -36,9 +36,9 @@ class BaseBaseModel(ABC, PluginVariant, make_plugin_registry("model")): def __new__(cls, *args, **kwargs): if inspect.isabstract(cls): # getting model type based on fitting type - model_type = kwargs.get("type","standard") + model_type = kwargs.get("type", "standard") if model_type == "standard": - model_type = kwargs.get("fitting",{}).get("type","ener") + model_type = kwargs.get("fitting", {}).get("type", "ener") cls = cls.get_class_by_type(model_type) return super().__new__(cls) @@ -122,9 +122,9 @@ def deserialize(cls, data: dict) -> "BaseBaseModel": The deserialized model """ if inspect.isabstract(cls): - model_type = data.get("type","standard") + model_type = data.get("type", "standard") if model_type == "standard": - model_type = data.get("fitting",{}).get("type","ener") + model_type = data.get("fitting", {}).get("type", "ener") return cls.get_class_by_type(model_type).deserialize(data) raise NotImplementedError("Not implemented in class %s" % cls.__name__) @@ -159,9 +159,9 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict): The local data refer to the current class """ # getting model type based on fitting type - model_type = local_jdata.get("type","standard") + model_type = local_jdata.get("type", "standard") if model_type == "standard": - model_type = local_jdata.get("fitting",{}).get("type","ener") + model_type = local_jdata.get("fitting", {}).get("type", "ener") cls = cls.get_class_by_type(model_type) return cls.update_sel(global_jdata, local_jdata) From 0c100fa173618eb5725c6ebdacc0d1beb96e4d97 Mon Sep 17 00:00:00 2001 From: anyangml Date: Sun, 14 Apr 2024 12:45:53 +0000 Subject: [PATCH 40/43] fix:UTs --- deepmd/tf/model/dos.py | 4 +++- deepmd/tf/model/ener.py | 4 +++- deepmd/tf/model/tensor.py | 4 +++- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/deepmd/tf/model/dos.py b/deepmd/tf/model/dos.py index 5fbea707a4..abfbdd262d 100644 --- a/deepmd/tf/model/dos.py +++ b/deepmd/tf/model/dos.py @@ -200,7 +200,9 @@ def build( return model_dict def serialize(self, suffix: str = "") -> dict: - return super().serialize(suffix="").update({"type": "dos"}) + dd = super().serialize(suffix) + dd.update({"type": "dos"}) + return dd def init_variables( self, diff --git a/deepmd/tf/model/ener.py b/deepmd/tf/model/ener.py index 2ccb54f576..c9e9a193ff 100644 --- a/deepmd/tf/model/ener.py +++ b/deepmd/tf/model/ener.py @@ -347,7 +347,9 @@ def build( return model_dict def serialize(self, suffix: str = "") -> dict: - return super().serialize(suffix="").update({"type": "ener"}) + dd = super().serialize(suffix) + dd.update({"type": "ener"}) + return dd def init_variables( self, diff --git a/deepmd/tf/model/tensor.py b/deepmd/tf/model/tensor.py index 9c9d9d388c..22f313bb8f 100644 --- a/deepmd/tf/model/tensor.py +++ b/deepmd/tf/model/tensor.py @@ -209,7 +209,9 @@ def build( return model_dict def serialize(self, suffix: str = "") -> dict: - return super().serialize(suffix="").update({"type": self.model_type}) + dd = super().serialize(suffix) + dd.update({"type": self.model_type}) + return dd def init_variables( self, From 752132d0a7aae9c5230bb4afd41d62947f1000d5 Mon Sep 17 00:00:00 2001 From: anyangml Date: Sun, 14 Apr 2024 13:27:24 +0000 Subject: [PATCH 41/43] fix: UTs --- deepmd/tf/model/dos.py | 5 ----- deepmd/tf/model/ener.py | 5 ----- deepmd/tf/model/tensor.py | 5 ----- 3 files changed, 15 deletions(-) diff --git a/deepmd/tf/model/dos.py b/deepmd/tf/model/dos.py index abfbdd262d..2d244f5733 100644 --- a/deepmd/tf/model/dos.py +++ b/deepmd/tf/model/dos.py @@ -199,11 +199,6 @@ def build( return model_dict - def serialize(self, suffix: str = "") -> dict: - dd = super().serialize(suffix) - dd.update({"type": "dos"}) - return dd - def init_variables( self, graph: tf.Graph, diff --git a/deepmd/tf/model/ener.py b/deepmd/tf/model/ener.py index c9e9a193ff..4a0195b334 100644 --- a/deepmd/tf/model/ener.py +++ b/deepmd/tf/model/ener.py @@ -346,11 +346,6 @@ def build( return model_dict - def serialize(self, suffix: str = "") -> dict: - dd = super().serialize(suffix) - dd.update({"type": "ener"}) - return dd - def init_variables( self, graph: tf.Graph, diff --git a/deepmd/tf/model/tensor.py b/deepmd/tf/model/tensor.py index 22f313bb8f..ba41879c50 100644 --- a/deepmd/tf/model/tensor.py +++ b/deepmd/tf/model/tensor.py @@ -208,11 +208,6 @@ def build( return model_dict - def serialize(self, suffix: str = "") -> dict: - dd = super().serialize(suffix) - dd.update({"type": self.model_type}) - return dd - def init_variables( self, graph: tf.Graph, From ff30173cddaa5d6f1a2ebdc72c8719fed874be29 Mon Sep 17 00:00:00 2001 From: anyangml Date: Mon, 15 Apr 2024 01:44:52 +0000 Subject: [PATCH 42/43] fix: tf register --- deepmd/tf/model/tensor.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/deepmd/tf/model/tensor.py b/deepmd/tf/model/tensor.py index ba41879c50..4a03e9f04e 100644 --- a/deepmd/tf/model/tensor.py +++ b/deepmd/tf/model/tensor.py @@ -22,8 +22,6 @@ ) -@StandardModel.register("dipole") -@StandardModel.register("polar") class TensorModel(StandardModel): """Tensor model. @@ -243,12 +241,12 @@ class WFCModel(TensorModel): def __init__(self, *args, **kwargs) -> None: TensorModel.__init__(self, "wfc", *args, **kwargs) - +@StandardModel.register("dipole") class DipoleModel(TensorModel): def __init__(self, *args, **kwargs) -> None: TensorModel.__init__(self, "dipole", *args, **kwargs) - +@StandardModel.register("polar") class PolarModel(TensorModel): def __init__(self, *args, **kwargs) -> None: TensorModel.__init__(self, "polar", *args, **kwargs) From 7399d3e78c199749f59958d332a4adae0e559138 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 15 Apr 2024 01:45:23 +0000 Subject: [PATCH 43/43] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/model/tensor.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/deepmd/tf/model/tensor.py b/deepmd/tf/model/tensor.py index 4a03e9f04e..ff20b771a7 100644 --- a/deepmd/tf/model/tensor.py +++ b/deepmd/tf/model/tensor.py @@ -241,11 +241,13 @@ class WFCModel(TensorModel): def __init__(self, *args, **kwargs) -> None: TensorModel.__init__(self, "wfc", *args, **kwargs) + @StandardModel.register("dipole") class DipoleModel(TensorModel): def __init__(self, *args, **kwargs) -> None: TensorModel.__init__(self, "dipole", *args, **kwargs) + @StandardModel.register("polar") class PolarModel(TensorModel): def __init__(self, *args, **kwargs) -> None: