From 88f5079477ee05f2397ac4d618d83dde05ee567e Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Fri, 1 Nov 2024 16:55:54 +0800 Subject: [PATCH 01/19] chore: refactor dpmodel atomic model --- deepmd/dpmodel/atomic_model/__init__.py | 16 +++++ .../atomic_model/dipole_atomic_model.py | 23 ++++++++ .../dpmodel/atomic_model/dos_atomic_model.py | 14 +++++ .../atomic_model/energy_atomic_model.py | 18 ++++++ .../atomic_model/polar_atomic_model.py | 59 +++++++++++++++++++ deepmd/dpmodel/model/dos_model.py | 0 deepmd/dpmodel/model/ener_model.py | 6 +- 7 files changed, 133 insertions(+), 3 deletions(-) create mode 100644 deepmd/dpmodel/atomic_model/dipole_atomic_model.py create mode 100644 deepmd/dpmodel/atomic_model/dos_atomic_model.py create mode 100644 deepmd/dpmodel/atomic_model/energy_atomic_model.py create mode 100644 deepmd/dpmodel/atomic_model/polar_atomic_model.py create mode 100644 deepmd/dpmodel/model/dos_model.py diff --git a/deepmd/dpmodel/atomic_model/__init__.py b/deepmd/dpmodel/atomic_model/__init__.py index 37f6b8bf28..963a619b84 100644 --- a/deepmd/dpmodel/atomic_model/__init__.py +++ b/deepmd/dpmodel/atomic_model/__init__.py @@ -20,6 +20,18 @@ from .dp_atomic_model import ( DPAtomicModel, ) +from .energy_atomic_model import ( + DPEnergyAtomicModel, +) +from .dos_atomic_model import ( + DPDOSAtomicModel, +) +from .polar_atomic_model import ( + DPPolarAtomicModel, +) +from .dipole_atomic_model import ( + DPDipoleAtomicModel, +) from .linear_atomic_model import ( DPZBLLinearEnergyAtomicModel, LinearEnergyAtomicModel, @@ -35,7 +47,11 @@ "make_base_atomic_model", "BaseAtomicModel", "DPAtomicModel", + "DPEnergyAtomicModel", "PairTabAtomicModel", "LinearEnergyAtomicModel", "DPZBLLinearEnergyAtomicModel", + "DPDOSAtomicModel", + "DPPolarAtomicModel", + "DPDipoleAtomicModel", ] 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..8cd45044f4 --- /dev/null +++ b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +from deepmd.dpmodel.fitting.dipole_fitting import ( + DipoleFitting, +) + +from .dp_atomic_model import ( + DPAtomicModel, +) + + +class DPDipoleAtomicModel(DPAtomicModel): + def __init__(self, descriptor, fitting, type_map, **kwargs): + assert isinstance(fitting, DipoleFitting) + super().__init__(descriptor, fitting, type_map, **kwargs) + + def apply_out_stat( + self, + ret: dict[str, np.ndarray], + atype: np.ndarray, + ): + # dipole not applying bias + return ret diff --git a/deepmd/dpmodel/atomic_model/dos_atomic_model.py b/deepmd/dpmodel/atomic_model/dos_atomic_model.py new file mode 100644 index 0000000000..fc584bcb56 --- /dev/null +++ b/deepmd/dpmodel/atomic_model/dos_atomic_model.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.fitting.dos_fitting import ( + DOSFittingNet, +) + +from .dp_atomic_model import ( + DPAtomicModel, +) + + +class DPDOSAtomicModel(DPAtomicModel): + def __init__(self, descriptor, fitting, type_map, **kwargs): + assert isinstance(fitting, DOSFittingNet) + super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/dpmodel/atomic_model/energy_atomic_model.py b/deepmd/dpmodel/atomic_model/energy_atomic_model.py new file mode 100644 index 0000000000..88fe7ae7a7 --- /dev/null +++ b/deepmd/dpmodel/atomic_model/energy_atomic_model.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.fitting.ener_fitting import ( + EnergyFittingNet, + InvarFitting, +) + +from .dp_atomic_model import ( + DPAtomicModel, +) + + +class DPEnergyAtomicModel(DPAtomicModel): + def __init__(self, descriptor, fitting, type_map, **kwargs): + assert ( + isinstance(fitting, EnergyFittingNet) + or isinstance(fitting, InvarFitting) + ) + super().__init__(descriptor, fitting, type_map, **kwargs) 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..b01468e99e --- /dev/null +++ b/deepmd/dpmodel/atomic_model/polar_atomic_model.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +import numpy as np + +from deepmd.dpmodel.fitting.polarizability_fitting import ( + PolarFitting, +) + +from .dp_atomic_model import ( + DPAtomicModel, +) + + +class DPPolarAtomicModel(DPAtomicModel): + def __init__(self, descriptor, fitting, type_map, **kwargs): + assert isinstance(fitting, PolarFitting) + super().__init__(descriptor, fitting, type_map, **kwargs) + + def apply_out_stat( + self, + ret: dict[str, np.ndarray], + atype: np.ndarray, + ): + """Apply the stat to each atomic output. + + Parameters + ---------- + ret + The returned dict by the forward_atomic method + atype + The atom types. nf x nloc + + """ + out_bias, out_std = self._fetch_out_stat(self.bias_keys) + + if self.fitting_net.shift_diag: + nframes, nloc = atype.shape + dtype = out_bias[self.bias_keys[0]].dtype + for kk in self.bias_keys: + ntypes = out_bias[kk].shape[0] + temp = np.zeros(ntypes,dtype=dtype) + for i in range(ntypes): + temp[i] = np.mean(np.diagonal(out_bias[kk][i].reshape(3, 3))) + modified_bias = temp[atype] + + # (nframes, nloc, 1) + modified_bias = ( + modified_bias[..., np.newaxis] + * (self.fitting_net.scale[atype]) + ) + + eye = np.eye(3,dtype=dtype) + eye = np.tile(eye, (nframes, nloc, 1, 1)) + # (nframes, nloc, 3, 3) + modified_bias = modified_bias[..., np.newaxis] * eye + + # nf x nloc x odims, out_bias: ntypes x odims + ret[kk] = ret[kk] + modified_bias + return ret diff --git a/deepmd/dpmodel/model/dos_model.py b/deepmd/dpmodel/model/dos_model.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index 5f21681830..99ea239c9f 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/deepmd/dpmodel/model/ener_model.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from deepmd.dpmodel.atomic_model.dp_atomic_model import ( - DPAtomicModel, +from deepmd.dpmodel.atomic_model import ( + DPEnergyAtomicModel, ) from deepmd.dpmodel.model.base_model import ( BaseModel, @@ -13,7 +13,7 @@ make_model, ) -DPEnergyModel_ = make_model(DPAtomicModel) +DPEnergyModel_ = make_model(DPEnergyAtomicModel) @BaseModel.register("ener") From de96159581070af361c21bbc638ad9ba1fb118c7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 1 Nov 2024 08:57:43 +0000 Subject: [PATCH 02/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/atomic_model/__init__.py | 18 ++--- .../atomic_model/dipole_atomic_model.py | 1 + .../atomic_model/energy_atomic_model.py | 5 +- .../atomic_model/polar_atomic_model.py | 79 +++++++++---------- deepmd/dpmodel/model/dos_model.py | 1 + 5 files changed, 52 insertions(+), 52 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/__init__.py b/deepmd/dpmodel/atomic_model/__init__.py index 963a619b84..4f4ef32e03 100644 --- a/deepmd/dpmodel/atomic_model/__init__.py +++ b/deepmd/dpmodel/atomic_model/__init__.py @@ -17,20 +17,17 @@ from .base_atomic_model import ( BaseAtomicModel, ) -from .dp_atomic_model import ( - DPAtomicModel, -) -from .energy_atomic_model import ( - DPEnergyAtomicModel, +from .dipole_atomic_model import ( + DPDipoleAtomicModel, ) from .dos_atomic_model import ( DPDOSAtomicModel, ) -from .polar_atomic_model import ( - DPPolarAtomicModel, +from .dp_atomic_model import ( + DPAtomicModel, ) -from .dipole_atomic_model import ( - DPDipoleAtomicModel, +from .energy_atomic_model import ( + DPEnergyAtomicModel, ) from .linear_atomic_model import ( DPZBLLinearEnergyAtomicModel, @@ -42,6 +39,9 @@ from .pairtab_atomic_model import ( PairTabAtomicModel, ) +from .polar_atomic_model import ( + DPPolarAtomicModel, +) __all__ = [ "make_base_atomic_model", diff --git a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py index 8cd45044f4..6b7b9f470b 100644 --- a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np + from deepmd.dpmodel.fitting.dipole_fitting import ( DipoleFitting, ) diff --git a/deepmd/dpmodel/atomic_model/energy_atomic_model.py b/deepmd/dpmodel/atomic_model/energy_atomic_model.py index 88fe7ae7a7..ad00e1c2cb 100644 --- a/deepmd/dpmodel/atomic_model/energy_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/energy_atomic_model.py @@ -11,8 +11,7 @@ class DPEnergyAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - assert ( - isinstance(fitting, EnergyFittingNet) - or isinstance(fitting, InvarFitting) + assert isinstance(fitting, EnergyFittingNet) or isinstance( + fitting, InvarFitting ) super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/dpmodel/atomic_model/polar_atomic_model.py b/deepmd/dpmodel/atomic_model/polar_atomic_model.py index b01468e99e..a049491685 100644 --- a/deepmd/dpmodel/atomic_model/polar_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/polar_atomic_model.py @@ -17,43 +17,42 @@ def __init__(self, descriptor, fitting, type_map, **kwargs): super().__init__(descriptor, fitting, type_map, **kwargs) def apply_out_stat( - self, - ret: dict[str, np.ndarray], - atype: np.ndarray, - ): - """Apply the stat to each atomic output. - - Parameters - ---------- - ret - The returned dict by the forward_atomic method - atype - The atom types. nf x nloc - - """ - out_bias, out_std = self._fetch_out_stat(self.bias_keys) - - if self.fitting_net.shift_diag: - nframes, nloc = atype.shape - dtype = out_bias[self.bias_keys[0]].dtype - for kk in self.bias_keys: - ntypes = out_bias[kk].shape[0] - temp = np.zeros(ntypes,dtype=dtype) - for i in range(ntypes): - temp[i] = np.mean(np.diagonal(out_bias[kk][i].reshape(3, 3))) - modified_bias = temp[atype] - - # (nframes, nloc, 1) - modified_bias = ( - modified_bias[..., np.newaxis] - * (self.fitting_net.scale[atype]) - ) - - eye = np.eye(3,dtype=dtype) - eye = np.tile(eye, (nframes, nloc, 1, 1)) - # (nframes, nloc, 3, 3) - modified_bias = modified_bias[..., np.newaxis] * eye - - # nf x nloc x odims, out_bias: ntypes x odims - ret[kk] = ret[kk] + modified_bias - return ret + self, + ret: dict[str, np.ndarray], + atype: np.ndarray, + ): + """Apply the stat to each atomic output. + + Parameters + ---------- + ret + The returned dict by the forward_atomic method + atype + The atom types. nf x nloc + + """ + out_bias, out_std = self._fetch_out_stat(self.bias_keys) + + if self.fitting_net.shift_diag: + nframes, nloc = atype.shape + dtype = out_bias[self.bias_keys[0]].dtype + for kk in self.bias_keys: + ntypes = out_bias[kk].shape[0] + temp = np.zeros(ntypes, dtype=dtype) + for i in range(ntypes): + temp[i] = np.mean(np.diagonal(out_bias[kk][i].reshape(3, 3))) + modified_bias = temp[atype] + + # (nframes, nloc, 1) + modified_bias = ( + modified_bias[..., np.newaxis] * (self.fitting_net.scale[atype]) + ) + + eye = np.eye(3, dtype=dtype) + eye = np.tile(eye, (nframes, nloc, 1, 1)) + # (nframes, nloc, 3, 3) + modified_bias = modified_bias[..., np.newaxis] * eye + + # nf x nloc x odims, out_bias: ntypes x odims + ret[kk] = ret[kk] + modified_bias + return ret diff --git a/deepmd/dpmodel/model/dos_model.py b/deepmd/dpmodel/model/dos_model.py index e69de29bb2..6ceb116d85 100644 --- a/deepmd/dpmodel/model/dos_model.py +++ b/deepmd/dpmodel/model/dos_model.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later From ab3bd29789c3708b9b61ea09ece859e5a2c4358f Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Fri, 1 Nov 2024 17:10:05 +0800 Subject: [PATCH 03/19] feat: expose dos/polar/dipole model in dpmodel --- deepmd/dpmodel/model/dipole_model.py | 33 +++++++++++++++++++++++++++ deepmd/dpmodel/model/dos_model.py | 31 +++++++++++++++++++++++++ deepmd/dpmodel/model/polar_model.py | 30 ++++++++++++++++++++++++ deepmd/pt/model/model/dipole_model.py | 6 ++--- deepmd/pt/model/model/polar_model.py | 6 ++--- 5 files changed, 100 insertions(+), 6 deletions(-) create mode 100644 deepmd/dpmodel/model/dipole_model.py create mode 100644 deepmd/dpmodel/model/polar_model.py diff --git a/deepmd/dpmodel/model/dipole_model.py b/deepmd/dpmodel/model/dipole_model.py new file mode 100644 index 0000000000..4a799679ce --- /dev/null +++ b/deepmd/dpmodel/model/dipole_model.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + + +from deepmd.dpmodel.atomic_model import ( + DPDipoleAtomicModel, +) +from deepmd.dpmodel.model.model import ( + BaseModel, +) + +from .dp_model import ( + DPModelCommon, +) +from .make_model import ( + make_model, +) + +DPDipoleModel_ = make_model(DPDipoleAtomicModel) + + +@BaseModel.register("dipole") +class DipoleModel(DPModelCommon, DPDipoleModel_): + model_type = "dipole" + + def __init__( + self, + *args, + **kwargs, + ): + DPModelCommon.__init__(self) + DPDipoleModel_.__init__(self, *args, **kwargs) + + \ No newline at end of file diff --git a/deepmd/dpmodel/model/dos_model.py b/deepmd/dpmodel/model/dos_model.py index 6ceb116d85..4b50963eb0 100644 --- a/deepmd/dpmodel/model/dos_model.py +++ b/deepmd/dpmodel/model/dos_model.py @@ -1 +1,32 @@ # SPDX-License-Identifier: LGPL-3.0-or-later + +from deepmd.dpmodel.atomic_model import ( + DPDOSAtomicModel, +) +from deepmd.dpmodel.model.model import ( + BaseModel, +) + +from .dp_model import ( + DPModelCommon, +) +from .make_model import ( + make_model, +) + +DPDOSModel_ = make_model(DPDOSAtomicModel) + + +@BaseModel.register("dos") +class DOSModel(DPModelCommon, DPDOSModel_): + model_type = "dos" + + def __init__( + self, + *args, + **kwargs, + ): + DPModelCommon.__init__(self) + DPDOSModel_.__init__(self, *args, **kwargs) + + \ No newline at end of file diff --git a/deepmd/dpmodel/model/polar_model.py b/deepmd/dpmodel/model/polar_model.py new file mode 100644 index 0000000000..c7cce1c1fa --- /dev/null +++ b/deepmd/dpmodel/model/polar_model.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from deepmd.dpmodel.atomic_model import ( + DPPolarAtomicModel, +) +from deepmd.dpmodel.model.model import ( + BaseModel, +) + +from .dp_model import ( + DPModelCommon, +) +from .make_model import ( + make_model, +) + +DPPolarModel_ = make_model(DPPolarAtomicModel) + + +@BaseModel.register("polar") +class PolarModel(DPModelCommon, DPPolarModel_): + model_type = "polar" + + def __init__( + self, + *args, + **kwargs, + ): + DPModelCommon.__init__(self) + DPPolarModel_.__init__(self, *args, **kwargs) diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index c83d1f0bf7..53ae281306 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -22,11 +22,11 @@ make_model, ) -DPDOSModel_ = make_model(DPDipoleAtomicModel) +DPDipoleModel_ = make_model(DPDipoleAtomicModel) @BaseModel.register("dipole") -class DipoleModel(DPModelCommon, DPDOSModel_): +class DipoleModel(DPModelCommon, DPDipoleModel_): model_type = "dipole" def __init__( @@ -35,7 +35,7 @@ def __init__( **kwargs, ): DPModelCommon.__init__(self) - DPDOSModel_.__init__(self, *args, **kwargs) + DPDipoleModel_.__init__(self, *args, **kwargs) def translated_output_def(self): out_def_data = self.model_output_def().get_data() diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index 57379ba372..1cf935b09e 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -22,11 +22,11 @@ make_model, ) -DPDOSModel_ = make_model(DPPolarAtomicModel) +DPPolarModel_ = make_model(DPPolarAtomicModel) @BaseModel.register("polar") -class PolarModel(DPModelCommon, DPDOSModel_): +class PolarModel(DPModelCommon, DPPolarModel_): model_type = "polar" def __init__( @@ -35,7 +35,7 @@ def __init__( **kwargs, ): DPModelCommon.__init__(self) - DPDOSModel_.__init__(self, *args, **kwargs) + DPPolarModel_.__init__(self, *args, **kwargs) def translated_output_def(self): out_def_data = self.model_output_def().get_data() From 8e85e5962201686e0d99fedc722ee927317e7173 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 1 Nov 2024 09:11:23 +0000 Subject: [PATCH 04/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/model/dipole_model.py | 2 -- deepmd/dpmodel/model/dos_model.py | 2 -- 2 files changed, 4 deletions(-) diff --git a/deepmd/dpmodel/model/dipole_model.py b/deepmd/dpmodel/model/dipole_model.py index 4a799679ce..211debc2e8 100644 --- a/deepmd/dpmodel/model/dipole_model.py +++ b/deepmd/dpmodel/model/dipole_model.py @@ -29,5 +29,3 @@ def __init__( ): DPModelCommon.__init__(self) DPDipoleModel_.__init__(self, *args, **kwargs) - - \ No newline at end of file diff --git a/deepmd/dpmodel/model/dos_model.py b/deepmd/dpmodel/model/dos_model.py index 4b50963eb0..638bdb462d 100644 --- a/deepmd/dpmodel/model/dos_model.py +++ b/deepmd/dpmodel/model/dos_model.py @@ -28,5 +28,3 @@ def __init__( ): DPModelCommon.__init__(self) DPDOSModel_.__init__(self, *args, **kwargs) - - \ No newline at end of file From c5dab77e53f3fe6b0920b2a9e40124dd529d053a Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Fri, 1 Nov 2024 17:19:21 +0800 Subject: [PATCH 05/19] chore: refactor atomic model fitting assertion --- deepmd/dpmodel/atomic_model/dipole_atomic_model.py | 3 ++- deepmd/dpmodel/atomic_model/dos_atomic_model.py | 3 ++- deepmd/dpmodel/atomic_model/energy_atomic_model.py | 7 +++++-- deepmd/dpmodel/atomic_model/polar_atomic_model.py | 3 ++- deepmd/dpmodel/atomic_model/property_atomic_model.py | 3 ++- deepmd/pt/model/atomic_model/dipole_atomic_model.py | 3 ++- deepmd/pt/model/atomic_model/dos_atomic_model.py | 3 ++- deepmd/pt/model/atomic_model/energy_atomic_model.py | 7 +++++-- deepmd/pt/model/atomic_model/polar_atomic_model.py | 3 ++- deepmd/pt/model/atomic_model/property_atomic_model.py | 3 ++- 10 files changed, 26 insertions(+), 12 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py index 6b7b9f470b..263dba435e 100644 --- a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py @@ -12,7 +12,8 @@ class DPDipoleAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - assert isinstance(fitting, DipoleFitting) + if not isinstance(fitting, DipoleFitting): + raise TypeError("fitting must be an instance of DipoleFitting for DPDipoleAtomicModel") super().__init__(descriptor, fitting, type_map, **kwargs) def apply_out_stat( diff --git a/deepmd/dpmodel/atomic_model/dos_atomic_model.py b/deepmd/dpmodel/atomic_model/dos_atomic_model.py index fc584bcb56..8cb551313b 100644 --- a/deepmd/dpmodel/atomic_model/dos_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dos_atomic_model.py @@ -10,5 +10,6 @@ class DPDOSAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - assert isinstance(fitting, DOSFittingNet) + if not isinstance(fitting, DOSFittingNet): + raise TypeError("fitting must be an instance of DOSFittingNet for DPDOSAtomicModel") super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/dpmodel/atomic_model/energy_atomic_model.py b/deepmd/dpmodel/atomic_model/energy_atomic_model.py index ad00e1c2cb..a274b09ed4 100644 --- a/deepmd/dpmodel/atomic_model/energy_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/energy_atomic_model.py @@ -11,7 +11,10 @@ class DPEnergyAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - assert isinstance(fitting, EnergyFittingNet) or isinstance( + if not (isinstance(fitting, EnergyFittingNet) or isinstance( fitting, InvarFitting - ) + )): + raise TypeError( + "fitting must be an instance of EnergyFittingNet or InvarFitting for DPEnergyAtomicModel" + ) super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/dpmodel/atomic_model/polar_atomic_model.py b/deepmd/dpmodel/atomic_model/polar_atomic_model.py index a049491685..cea5ad2c54 100644 --- a/deepmd/dpmodel/atomic_model/polar_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/polar_atomic_model.py @@ -13,7 +13,8 @@ class DPPolarAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - assert isinstance(fitting, PolarFitting) + if not isinstance(fitting, PolarFitting): + raise TypeError("fitting must be an instance of PolarFitting for DPPolarAtomicModel") super().__init__(descriptor, fitting, type_map, **kwargs) def apply_out_stat( diff --git a/deepmd/dpmodel/atomic_model/property_atomic_model.py b/deepmd/dpmodel/atomic_model/property_atomic_model.py index ecc450bcd2..0cfd24ed91 100644 --- a/deepmd/dpmodel/atomic_model/property_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/property_atomic_model.py @@ -10,5 +10,6 @@ class DPPropertyAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - assert isinstance(fitting, PropertyFittingNet) + if not isinstance(fitting, PropertyFittingNet): + raise TypeError("fitting must be an instance of PropertyFittingNet for DPPropertyAtomicModel") super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index aa28294cc5..e980bd9c28 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -13,7 +13,8 @@ class DPDipoleAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - assert isinstance(fitting, DipoleFittingNet) + if not isinstance(fitting, DipoleFittingNet): + raise TypeError("fitting must be an instance of DipoleFittingNet for DPDipoleAtomicModel") super().__init__(descriptor, fitting, type_map, **kwargs) def apply_out_stat( diff --git a/deepmd/pt/model/atomic_model/dos_atomic_model.py b/deepmd/pt/model/atomic_model/dos_atomic_model.py index 5e399f2aff..8ec292f3ec 100644 --- a/deepmd/pt/model/atomic_model/dos_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dos_atomic_model.py @@ -10,5 +10,6 @@ class DPDOSAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - assert isinstance(fitting, DOSFittingNet) + if not isinstance(fitting, DOSFittingNet): + raise TypeError("fitting must be an instance of DOSFittingNet for DPDOSAtomicModel") 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 7cedaa1ab3..6d894b4aab 100644 --- a/deepmd/pt/model/atomic_model/energy_atomic_model.py +++ b/deepmd/pt/model/atomic_model/energy_atomic_model.py @@ -12,9 +12,12 @@ class DPEnergyAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - assert ( + if not ( isinstance(fitting, EnergyFittingNet) or isinstance(fitting, EnergyFittingNetDirect) or isinstance(fitting, InvarFitting) - ) + ): + raise TypeError( + "fitting must be an instance of EnergyFittingNet, EnergyFittingNetDirect or InvarFitting for DPEnergyAtomicModel" + ) 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 39cda2650d..e5854baa3f 100644 --- a/deepmd/pt/model/atomic_model/polar_atomic_model.py +++ b/deepmd/pt/model/atomic_model/polar_atomic_model.py @@ -13,7 +13,8 @@ class DPPolarAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - assert isinstance(fitting, PolarFittingNet) + if not isinstance(fitting, PolarFittingNet): + raise TypeError("fitting must be an instance of PolarFittingNet for DPPolarAtomicModel") super().__init__(descriptor, fitting, type_map, **kwargs) def apply_out_stat( diff --git a/deepmd/pt/model/atomic_model/property_atomic_model.py b/deepmd/pt/model/atomic_model/property_atomic_model.py index 2fac90100f..fb1cc498f8 100644 --- a/deepmd/pt/model/atomic_model/property_atomic_model.py +++ b/deepmd/pt/model/atomic_model/property_atomic_model.py @@ -13,7 +13,8 @@ class DPPropertyAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - assert isinstance(fitting, PropertyFittingNet) + if not isinstance(fitting, PropertyFittingNet): + raise TypeError("fitting must be an instance of PropertyFittingNet for DPPropertyAtomicModel") super().__init__(descriptor, fitting, type_map, **kwargs) def apply_out_stat( From 58b987e4c46e1eda2e3d0f27fde7babb297f7863 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 1 Nov 2024 09:21:45 +0000 Subject: [PATCH 06/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/atomic_model/dipole_atomic_model.py | 4 +++- deepmd/dpmodel/atomic_model/dos_atomic_model.py | 4 +++- deepmd/dpmodel/atomic_model/energy_atomic_model.py | 6 +++--- deepmd/dpmodel/atomic_model/polar_atomic_model.py | 4 +++- deepmd/dpmodel/atomic_model/property_atomic_model.py | 4 +++- deepmd/pt/model/atomic_model/dipole_atomic_model.py | 4 +++- deepmd/pt/model/atomic_model/dos_atomic_model.py | 4 +++- deepmd/pt/model/atomic_model/polar_atomic_model.py | 4 +++- deepmd/pt/model/atomic_model/property_atomic_model.py | 4 +++- 9 files changed, 27 insertions(+), 11 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py index 263dba435e..00428f4e95 100644 --- a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py @@ -13,7 +13,9 @@ class DPDipoleAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): if not isinstance(fitting, DipoleFitting): - raise TypeError("fitting must be an instance of DipoleFitting for DPDipoleAtomicModel") + raise TypeError( + "fitting must be an instance of DipoleFitting for DPDipoleAtomicModel" + ) super().__init__(descriptor, fitting, type_map, **kwargs) def apply_out_stat( diff --git a/deepmd/dpmodel/atomic_model/dos_atomic_model.py b/deepmd/dpmodel/atomic_model/dos_atomic_model.py index 8cb551313b..7ef6d10ebf 100644 --- a/deepmd/dpmodel/atomic_model/dos_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dos_atomic_model.py @@ -11,5 +11,7 @@ class DPDOSAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): if not isinstance(fitting, DOSFittingNet): - raise TypeError("fitting must be an instance of DOSFittingNet for DPDOSAtomicModel") + raise TypeError( + "fitting must be an instance of DOSFittingNet for DPDOSAtomicModel" + ) super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/dpmodel/atomic_model/energy_atomic_model.py b/deepmd/dpmodel/atomic_model/energy_atomic_model.py index a274b09ed4..4f9f8ec005 100644 --- a/deepmd/dpmodel/atomic_model/energy_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/energy_atomic_model.py @@ -11,9 +11,9 @@ class DPEnergyAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): - if not (isinstance(fitting, EnergyFittingNet) or isinstance( - fitting, InvarFitting - )): + if not ( + isinstance(fitting, EnergyFittingNet) or isinstance(fitting, InvarFitting) + ): raise TypeError( "fitting must be an instance of EnergyFittingNet or InvarFitting for DPEnergyAtomicModel" ) diff --git a/deepmd/dpmodel/atomic_model/polar_atomic_model.py b/deepmd/dpmodel/atomic_model/polar_atomic_model.py index cea5ad2c54..faff2444ea 100644 --- a/deepmd/dpmodel/atomic_model/polar_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/polar_atomic_model.py @@ -14,7 +14,9 @@ class DPPolarAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): if not isinstance(fitting, PolarFitting): - raise TypeError("fitting must be an instance of PolarFitting for DPPolarAtomicModel") + raise TypeError( + "fitting must be an instance of PolarFitting for DPPolarAtomicModel" + ) super().__init__(descriptor, fitting, type_map, **kwargs) def apply_out_stat( diff --git a/deepmd/dpmodel/atomic_model/property_atomic_model.py b/deepmd/dpmodel/atomic_model/property_atomic_model.py index 0cfd24ed91..6f69f8dfb6 100644 --- a/deepmd/dpmodel/atomic_model/property_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/property_atomic_model.py @@ -11,5 +11,7 @@ class DPPropertyAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): if not isinstance(fitting, PropertyFittingNet): - raise TypeError("fitting must be an instance of PropertyFittingNet for DPPropertyAtomicModel") + raise TypeError( + "fitting must be an instance of PropertyFittingNet for DPPropertyAtomicModel" + ) super().__init__(descriptor, fitting, type_map, **kwargs) diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index e980bd9c28..3796aa2e83 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -14,7 +14,9 @@ class DPDipoleAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): if not isinstance(fitting, DipoleFittingNet): - raise TypeError("fitting must be an instance of DipoleFittingNet for DPDipoleAtomicModel") + raise TypeError( + "fitting must be an instance of DipoleFittingNet for DPDipoleAtomicModel" + ) super().__init__(descriptor, fitting, type_map, **kwargs) def apply_out_stat( diff --git a/deepmd/pt/model/atomic_model/dos_atomic_model.py b/deepmd/pt/model/atomic_model/dos_atomic_model.py index 8ec292f3ec..2af1a4e052 100644 --- a/deepmd/pt/model/atomic_model/dos_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dos_atomic_model.py @@ -11,5 +11,7 @@ class DPDOSAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): if not isinstance(fitting, DOSFittingNet): - raise TypeError("fitting must be an instance of DOSFittingNet for DPDOSAtomicModel") + raise TypeError( + "fitting must be an instance of DOSFittingNet for DPDOSAtomicModel" + ) 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 e5854baa3f..fc6c351830 100644 --- a/deepmd/pt/model/atomic_model/polar_atomic_model.py +++ b/deepmd/pt/model/atomic_model/polar_atomic_model.py @@ -14,7 +14,9 @@ class DPPolarAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): if not isinstance(fitting, PolarFittingNet): - raise TypeError("fitting must be an instance of PolarFittingNet for DPPolarAtomicModel") + raise TypeError( + "fitting must be an instance of PolarFittingNet for DPPolarAtomicModel" + ) super().__init__(descriptor, fitting, type_map, **kwargs) def apply_out_stat( diff --git a/deepmd/pt/model/atomic_model/property_atomic_model.py b/deepmd/pt/model/atomic_model/property_atomic_model.py index fb1cc498f8..1fdc72b2b6 100644 --- a/deepmd/pt/model/atomic_model/property_atomic_model.py +++ b/deepmd/pt/model/atomic_model/property_atomic_model.py @@ -14,7 +14,9 @@ class DPPropertyAtomicModel(DPAtomicModel): def __init__(self, descriptor, fitting, type_map, **kwargs): if not isinstance(fitting, PropertyFittingNet): - raise TypeError("fitting must be an instance of PropertyFittingNet for DPPropertyAtomicModel") + raise TypeError( + "fitting must be an instance of PropertyFittingNet for DPPropertyAtomicModel" + ) super().__init__(descriptor, fitting, type_map, **kwargs) def apply_out_stat( From e2e9d60a9fbc8f7e7ca71edbaccb8c74a02674f3 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Fri, 15 Nov 2024 13:34:45 +0800 Subject: [PATCH 07/19] chore: fix typo --- deepmd/pt/model/atomic_model/dipole_atomic_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index e20d92c8fa..3796aa2e83 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -16,6 +16,7 @@ def __init__(self, descriptor, fitting, type_map, **kwargs): if not isinstance(fitting, DipoleFittingNet): raise TypeError( "fitting must be an instance of DipoleFittingNet for DPDipoleAtomicModel" + ) super().__init__(descriptor, fitting, type_map, **kwargs) def apply_out_stat( From 47114ecc388879fa1c110482c5c44ba18f9de506 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Fri, 15 Nov 2024 13:35:10 +0800 Subject: [PATCH 08/19] feat: try add dos consistent UT --- source/tests/consistent/model/test_dos.py | 245 ++++++++++++++++++++++ 1 file changed, 245 insertions(+) create mode 100644 source/tests/consistent/model/test_dos.py diff --git a/source/tests/consistent/model/test_dos.py b/source/tests/consistent/model/test_dos.py new file mode 100644 index 0000000000..c3734cee09 --- /dev/null +++ b/source/tests/consistent/model/test_dos.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.model.dos_model import DOSModel as DOSModelDP +from deepmd.dpmodel.model.model import get_model as get_model_dp +from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list, + extend_coord_with_ghosts, +) +from deepmd.dpmodel.utils.region import ( + normalize_coord, +) +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) + +from ..common import ( + INSTALLED_PT, + INSTALLED_TF, + SKIP_FLAG, + CommonTest, + parameterized, +) +from .common import ( + ModelTest, +) + +if INSTALLED_PT: + from deepmd.pt.model.model import get_model as get_model_pt + from deepmd.pt.model.model.dos_model import DOSModel as DOSModelPT + from deepmd.pt.utils.utils import to_numpy_array as torch_to_numpy + from deepmd.pt.utils.utils import to_torch_tensor as numpy_to_torch +else: + DOSModelPT = None +if INSTALLED_TF: + from deepmd.tf.model.dos import DOSModel as DOSModelTF +else: + DOSModelTF = None +from deepmd.utils.argcheck import ( + model_args, +) + + +@parameterized( + ( + [], + [[0, 1]], + ), + ( + [], + [1], + ), +) +class TestDOS(CommonTest, ModelTest, unittest.TestCase): + @property + def data(self) -> dict: + pair_exclude_types, atom_exclude_types = self.param + return { + "type_map": [ + "H" + ], + "descriptor": { + "type": "se_e2_a", + "sel": [ + 90 + ], + "rcut_smth": 1.8, + "rcut": 6.0, + "neuron": [ + 25, + 50, + 100 + ], + "resnet_dt": False, + "axis_neuron": 8, + "precision": "float64", + "seed": 1 + }, + "fitting_net": { + "type": "dos", + "numb_dos": 250, + "neuron": [ + 120, + 120, + 120 + ], + "resnet_dt": True, + "numb_fparam": 0, + "precision": "float64", + "seed": 1 + } + } + + tf_class = DOSModelTF + dp_class = DOSModelDP + pt_class = DOSModelPT + args = model_args() + + def get_reference_backend(self): + """Get the reference backend. + + We need a reference backend that can reproduce forces. + """ + if not self.skip_pt: + return self.RefBackend.PT + if not self.skip_tf: + return self.RefBackend.TF + if not self.skip_dp: + return self.RefBackend.DP + raise ValueError("No available reference") + + @property + def skip_tf(self): + return ( + self.data["pair_exclude_types"] != [] + or self.data["atom_exclude_types"] != [] + ) + + @property + def skip_jax(self) -> bool: + return True + + def pass_data_to_cls(self, cls, data) -> Any: + """Pass data to the class.""" + data = data.copy() + if cls is DOSModelDP: + return get_model_dp(data) + elif cls is DOSModelPT: + model = get_model_pt(data) + model.atomic_model.out_bias.uniform_() + return model + return cls(**data, **self.additional_data) + + def setUp(self) -> None: + CommonTest.setUp(self) + + self.ntypes = 2 + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + self.natoms = np.array([6, 6, 2, 4], dtype=np.int32) + + # TF requires the atype to be sort + idx_map = np.argsort(self.atype.ravel()) + self.atype = self.atype[:, idx_map] + self.coords = self.coords[:, idx_map] + + def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: + return self.build_tf_model( + obj, + self.natoms, + self.coords, + self.atype, + self.box, + suffix, + ) + + def eval_dp(self, dp_obj: Any) -> Any: + return self.eval_dp_model( + dp_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + + def eval_pt(self, pt_obj: Any) -> Any: + return self.eval_pt_model( + pt_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + + def eval_jax(self, jax_obj: Any) -> Any: + return self.eval_jax_model( + jax_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + + def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: + # shape not matched. ravel... + if backend is self.RefBackend.DP: + return ( + ret["energy_redu"].ravel(), + ret["energy"].ravel(), + SKIP_FLAG, + SKIP_FLAG, + SKIP_FLAG, + ) + elif backend is self.RefBackend.PT: + return ( + ret["energy"].ravel(), + ret["atom_energy"].ravel(), + ret["force"].ravel(), + ret["virial"].ravel(), + ret["atom_virial"].ravel(), + ) + elif backend is self.RefBackend.TF: + return ( + ret[0].ravel(), + ret[1].ravel(), + ret[2].ravel(), + ret[3].ravel(), + ret[4].ravel(), + ) + raise ValueError(f"Unknown backend: {backend}") + From 48d8d6c7ef84a92f15463c8337b6d23db6bc30c5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 15 Nov 2024 05:36:31 +0000 Subject: [PATCH 09/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/consistent/model/test_dos.py | 61 +++++++---------------- 1 file changed, 18 insertions(+), 43 deletions(-) diff --git a/source/tests/consistent/model/test_dos.py b/source/tests/consistent/model/test_dos.py index c3734cee09..87a83b7f9a 100644 --- a/source/tests/consistent/model/test_dos.py +++ b/source/tests/consistent/model/test_dos.py @@ -6,18 +6,8 @@ import numpy as np -from deepmd.dpmodel.common import ( - to_numpy_array, -) from deepmd.dpmodel.model.dos_model import DOSModel as DOSModelDP from deepmd.dpmodel.model.model import get_model as get_model_dp -from deepmd.dpmodel.utils.nlist import ( - build_neighbor_list, - extend_coord_with_ghosts, -) -from deepmd.dpmodel.utils.region import ( - normalize_coord, -) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) @@ -36,8 +26,6 @@ if INSTALLED_PT: from deepmd.pt.model.model import get_model as get_model_pt from deepmd.pt.model.model.dos_model import DOSModel as DOSModelPT - from deepmd.pt.utils.utils import to_numpy_array as torch_to_numpy - from deepmd.pt.utils.utils import to_torch_tensor as numpy_to_torch else: DOSModelPT = None if INSTALLED_TF: @@ -64,39 +52,27 @@ class TestDOS(CommonTest, ModelTest, unittest.TestCase): def data(self) -> dict: pair_exclude_types, atom_exclude_types = self.param return { - "type_map": [ - "H" - ], + "type_map": ["H"], "descriptor": { - "type": "se_e2_a", - "sel": [ - 90 - ], - "rcut_smth": 1.8, - "rcut": 6.0, - "neuron": [ - 25, - 50, - 100 - ], - "resnet_dt": False, - "axis_neuron": 8, - "precision": "float64", - "seed": 1 + "type": "se_e2_a", + "sel": [90], + "rcut_smth": 1.8, + "rcut": 6.0, + "neuron": [25, 50, 100], + "resnet_dt": False, + "axis_neuron": 8, + "precision": "float64", + "seed": 1, }, "fitting_net": { - "type": "dos", - "numb_dos": 250, - "neuron": [ - 120, - 120, - 120 - ], - "resnet_dt": True, - "numb_fparam": 0, - "precision": "float64", - "seed": 1 - } + "type": "dos", + "numb_dos": 250, + "neuron": [120, 120, 120], + "resnet_dt": True, + "numb_fparam": 0, + "precision": "float64", + "seed": 1, + }, } tf_class = DOSModelTF @@ -242,4 +218,3 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: ret[4].ravel(), ) raise ValueError(f"Unknown backend: {backend}") - From 81bd024c015cbe81814d596a41bef348a0e4f5f3 Mon Sep 17 00:00:00 2001 From: anyangml Date: Fri, 15 Nov 2024 16:31:20 +0800 Subject: [PATCH 10/19] fix: UT para --- source/tests/consistent/model/test_dos.py | 17 ++--------------- 1 file changed, 2 insertions(+), 15 deletions(-) diff --git a/source/tests/consistent/model/test_dos.py b/source/tests/consistent/model/test_dos.py index 87a83b7f9a..7cef5139be 100644 --- a/source/tests/consistent/model/test_dos.py +++ b/source/tests/consistent/model/test_dos.py @@ -37,20 +37,10 @@ ) -@parameterized( - ( - [], - [[0, 1]], - ), - ( - [], - [1], - ), -) + class TestDOS(CommonTest, ModelTest, unittest.TestCase): @property def data(self) -> dict: - pair_exclude_types, atom_exclude_types = self.param return { "type_map": ["H"], "descriptor": { @@ -95,10 +85,7 @@ def get_reference_backend(self): @property def skip_tf(self): - return ( - self.data["pair_exclude_types"] != [] - or self.data["atom_exclude_types"] != [] - ) + return False @property def skip_jax(self) -> bool: From 7490529a0a1d7c8e18c9671847af8b3261c7fd7f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 15 Nov 2024 08:32:46 +0000 Subject: [PATCH 11/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/consistent/model/test_dos.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/source/tests/consistent/model/test_dos.py b/source/tests/consistent/model/test_dos.py index 7cef5139be..8123d7d38d 100644 --- a/source/tests/consistent/model/test_dos.py +++ b/source/tests/consistent/model/test_dos.py @@ -17,7 +17,6 @@ INSTALLED_TF, SKIP_FLAG, CommonTest, - parameterized, ) from .common import ( ModelTest, @@ -37,7 +36,6 @@ ) - class TestDOS(CommonTest, ModelTest, unittest.TestCase): @property def data(self) -> dict: From c040df3912926be09f4efd6f494b461c806ea376 Mon Sep 17 00:00:00 2001 From: anyangml Date: Fri, 15 Nov 2024 17:26:48 +0800 Subject: [PATCH 12/19] fix: UT res extract --- deepmd/dpmodel/model/model.py | 2 +- source/tests/consistent/model/test_dos.py | 25 ++++++++--------------- 2 files changed, 9 insertions(+), 18 deletions(-) diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 19408e58c4..2b5890f42b 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -55,7 +55,7 @@ def get_standard_model(data: dict) -> EnergyModel: **data["fitting_net"], ) else: - raise ValueError(f"Unknown fitting type {fitting_type}") + raise ValueError(f"Unknown fitting type {fitting_type}") #fix return EnergyModel( descriptor=descriptor, fitting=fitting, diff --git a/source/tests/consistent/model/test_dos.py b/source/tests/consistent/model/test_dos.py index 8123d7d38d..c01e16bf47 100644 --- a/source/tests/consistent/model/test_dos.py +++ b/source/tests/consistent/model/test_dos.py @@ -40,13 +40,13 @@ class TestDOS(CommonTest, ModelTest, unittest.TestCase): @property def data(self) -> dict: return { - "type_map": ["H"], + "type_map": ["O", "H"], "descriptor": { "type": "se_e2_a", - "sel": [90], + "sel": [20, 20], "rcut_smth": 1.8, "rcut": 6.0, - "neuron": [25, 50, 100], + "neuron": [2,4,8], "resnet_dt": False, "axis_neuron": 8, "precision": "float64", @@ -55,7 +55,7 @@ def data(self) -> dict: "fitting_net": { "type": "dos", "numb_dos": 250, - "neuron": [120, 120, 120], + "neuron": [4,4,4], "resnet_dt": True, "numb_fparam": 0, "precision": "float64", @@ -180,26 +180,17 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: # shape not matched. ravel... if backend is self.RefBackend.DP: return ( - ret["energy_redu"].ravel(), - ret["energy"].ravel(), - SKIP_FLAG, - SKIP_FLAG, - SKIP_FLAG, + ret["dos_redu"].ravel(), + ret["dos"].ravel(), ) elif backend is self.RefBackend.PT: return ( - ret["energy"].ravel(), - ret["atom_energy"].ravel(), - ret["force"].ravel(), - ret["virial"].ravel(), - ret["atom_virial"].ravel(), + ret["dos"].ravel(), + ret["atom_dos"].ravel(), ) elif backend is self.RefBackend.TF: return ( ret[0].ravel(), ret[1].ravel(), - ret[2].ravel(), - ret[3].ravel(), - ret[4].ravel(), ) raise ValueError(f"Unknown backend: {backend}") From 8940dac66a3de56f2db6486d877c86207f96858f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 15 Nov 2024 09:28:32 +0000 Subject: [PATCH 13/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/model/model.py | 2 +- source/tests/consistent/model/test_dos.py | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 2b5890f42b..6db6368363 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -55,7 +55,7 @@ def get_standard_model(data: dict) -> EnergyModel: **data["fitting_net"], ) else: - raise ValueError(f"Unknown fitting type {fitting_type}") #fix + raise ValueError(f"Unknown fitting type {fitting_type}") # fix return EnergyModel( descriptor=descriptor, fitting=fitting, diff --git a/source/tests/consistent/model/test_dos.py b/source/tests/consistent/model/test_dos.py index c01e16bf47..5e83cc5f62 100644 --- a/source/tests/consistent/model/test_dos.py +++ b/source/tests/consistent/model/test_dos.py @@ -15,7 +15,6 @@ from ..common import ( INSTALLED_PT, INSTALLED_TF, - SKIP_FLAG, CommonTest, ) from .common import ( @@ -46,7 +45,7 @@ def data(self) -> dict: "sel": [20, 20], "rcut_smth": 1.8, "rcut": 6.0, - "neuron": [2,4,8], + "neuron": [2, 4, 8], "resnet_dt": False, "axis_neuron": 8, "precision": "float64", @@ -55,7 +54,7 @@ def data(self) -> dict: "fitting_net": { "type": "dos", "numb_dos": 250, - "neuron": [4,4,4], + "neuron": [4, 4, 4], "resnet_dt": True, "numb_fparam": 0, "precision": "float64", From 4802a9a3e13a0093f3610c9d7503325831c8ecbe Mon Sep 17 00:00:00 2001 From: anyangml Date: Mon, 18 Nov 2024 13:05:31 +0800 Subject: [PATCH 14/19] fix: add dos consistency UT --- deepmd/dpmodel/model/model.py | 10 +++++++++ source/tests/consistent/model/common.py | 25 +++++++++++++++-------- source/tests/consistent/model/test_dos.py | 5 +++-- 3 files changed, 30 insertions(+), 10 deletions(-) diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 6db6368363..2925b0ac32 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -11,6 +11,9 @@ from deepmd.dpmodel.fitting.ener_fitting import ( EnergyFittingNet, ) +from deepmd.dpmodel.fitting.dos_fitting import ( + DOSFittingNet +) from deepmd.dpmodel.model.base_model import ( BaseModel, ) @@ -54,6 +57,13 @@ def get_standard_model(data: dict) -> EnergyModel: mixed_types=descriptor.mixed_types(), **data["fitting_net"], ) + elif fitting_type == "dos": + fitting = DOSFittingNet( + ntypes=descriptor.get_ntypes(), + dim_descrpt=descriptor.get_dim_out(), + mixed_types=descriptor.mixed_types(), + **data["fitting_net"], + ) else: raise ValueError(f"Unknown fitting type {fitting_type}") # fix return EnergyModel( diff --git a/source/tests/consistent/model/common.py b/source/tests/consistent/model/common.py index 4eeb19b1f0..9dbc69ba91 100644 --- a/source/tests/consistent/model/common.py +++ b/source/tests/consistent/model/common.py @@ -34,7 +34,7 @@ class ModelTest: """Useful utilities for model tests.""" - def build_tf_model(self, obj, natoms, coords, atype, box, suffix): + def build_tf_model(self, obj, natoms, coords, atype, box, suffix, ret_key:str="energy"): t_coord = tf.placeholder( GLOBAL_TF_FLOAT_PRECISION, [None, None, None], name="i_coord" ) @@ -51,13 +51,22 @@ def build_tf_model(self, obj, natoms, coords, atype, box, suffix): {}, suffix=suffix, ) - return [ - ret["energy"], - ret["atom_ener"], - ret["force"], - ret["virial"], - ret["atom_virial"], - ], { + if ret_key == "energy": + ret_list = [ + ret["energy"], + ret["atom_ener"], + ret["force"], + ret["virial"], + ret["atom_virial"], + ] + elif ret_key == "dos": + ret_list = [ + ret["dos"], + ret["atom_dos"], + ] + else: + raise NotImplementedError + return ret_list, { t_coord: coords, t_type: atype, t_natoms: natoms, diff --git a/source/tests/consistent/model/test_dos.py b/source/tests/consistent/model/test_dos.py index 5e83cc5f62..5239c68ed7 100644 --- a/source/tests/consistent/model/test_dos.py +++ b/source/tests/consistent/model/test_dos.py @@ -53,7 +53,7 @@ def data(self) -> dict: }, "fitting_net": { "type": "dos", - "numb_dos": 250, + "numb_dos": 2, "neuron": [4, 4, 4], "resnet_dt": True, "numb_fparam": 0, @@ -82,7 +82,7 @@ def get_reference_backend(self): @property def skip_tf(self): - return False + return True # need to fix tf consistency @property def skip_jax(self) -> bool: @@ -146,6 +146,7 @@ def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: self.atype, self.box, suffix, + ret_key = "dos" ) def eval_dp(self, dp_obj: Any) -> Any: From 3355f502a8d7cc06f32a00071e7b400893096fa5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 18 Nov 2024 05:06:50 +0000 Subject: [PATCH 15/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/model/model.py | 6 +++--- source/tests/consistent/model/common.py | 4 +++- source/tests/consistent/model/test_dos.py | 10 ++-------- 3 files changed, 8 insertions(+), 12 deletions(-) diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 2925b0ac32..28e29cdcb7 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -8,12 +8,12 @@ from deepmd.dpmodel.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.dpmodel.fitting.dos_fitting import ( + DOSFittingNet, +) from deepmd.dpmodel.fitting.ener_fitting import ( EnergyFittingNet, ) -from deepmd.dpmodel.fitting.dos_fitting import ( - DOSFittingNet -) from deepmd.dpmodel.model.base_model import ( BaseModel, ) diff --git a/source/tests/consistent/model/common.py b/source/tests/consistent/model/common.py index 9dbc69ba91..ef1c7cf911 100644 --- a/source/tests/consistent/model/common.py +++ b/source/tests/consistent/model/common.py @@ -34,7 +34,9 @@ class ModelTest: """Useful utilities for model tests.""" - def build_tf_model(self, obj, natoms, coords, atype, box, suffix, ret_key:str="energy"): + def build_tf_model( + self, obj, natoms, coords, atype, box, suffix, ret_key: str = "energy" + ): t_coord = tf.placeholder( GLOBAL_TF_FLOAT_PRECISION, [None, None, None], name="i_coord" ) diff --git a/source/tests/consistent/model/test_dos.py b/source/tests/consistent/model/test_dos.py index 5239c68ed7..8f0b0309cc 100644 --- a/source/tests/consistent/model/test_dos.py +++ b/source/tests/consistent/model/test_dos.py @@ -82,7 +82,7 @@ def get_reference_backend(self): @property def skip_tf(self): - return True # need to fix tf consistency + return True # need to fix tf consistency @property def skip_jax(self) -> bool: @@ -140,13 +140,7 @@ def setUp(self) -> None: def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: return self.build_tf_model( - obj, - self.natoms, - self.coords, - self.atype, - self.box, - suffix, - ret_key = "dos" + obj, self.natoms, self.coords, self.atype, self.box, suffix, ret_key="dos" ) def eval_dp(self, dp_obj: Any) -> Any: From 5a11ff84002e46fe7218a331f23084f30a4d604d Mon Sep 17 00:00:00 2001 From: anyangml Date: Mon, 18 Nov 2024 14:52:52 +0800 Subject: [PATCH 16/19] fix: expose more dpmodel interface --- deepmd/dpmodel/model/dipole_model.py | 2 +- deepmd/dpmodel/model/dos_model.py | 2 +- deepmd/dpmodel/model/model.py | 79 +++++++++++++++++-------- deepmd/dpmodel/model/polar_model.py | 2 +- source/tests/consistent/model/common.py | 10 ++++ 5 files changed, 66 insertions(+), 29 deletions(-) diff --git a/deepmd/dpmodel/model/dipole_model.py b/deepmd/dpmodel/model/dipole_model.py index 211debc2e8..4ca523f79b 100644 --- a/deepmd/dpmodel/model/dipole_model.py +++ b/deepmd/dpmodel/model/dipole_model.py @@ -4,7 +4,7 @@ from deepmd.dpmodel.atomic_model import ( DPDipoleAtomicModel, ) -from deepmd.dpmodel.model.model import ( +from deepmd.dpmodel.model.base_model import ( BaseModel, ) diff --git a/deepmd/dpmodel/model/dos_model.py b/deepmd/dpmodel/model/dos_model.py index 638bdb462d..3df887b460 100644 --- a/deepmd/dpmodel/model/dos_model.py +++ b/deepmd/dpmodel/model/dos_model.py @@ -3,7 +3,7 @@ from deepmd.dpmodel.atomic_model import ( DPDOSAtomicModel, ) -from deepmd.dpmodel.model.model import ( +from deepmd.dpmodel.model.base_model import ( BaseModel, ) diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 28e29cdcb7..43eeb13898 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -8,8 +8,8 @@ from deepmd.dpmodel.descriptor.base_descriptor import ( BaseDescriptor, ) -from deepmd.dpmodel.fitting.dos_fitting import ( - DOSFittingNet, +from deepmd.dpmodel.fitting.base_fitting import ( + BaseFitting, ) from deepmd.dpmodel.fitting.ener_fitting import ( EnergyFittingNet, @@ -29,8 +29,35 @@ from deepmd.utils.spin import ( Spin, ) +import copy +from deepmd.dpmodel.model.dos_model import DOSModel +from deepmd.dpmodel.model.property_model import PropertyModel +from deepmd.dpmodel.model.dipole_model import DipoleModel +from deepmd.dpmodel.model.polar_model import PolarModel +def _get_standard_model_components(data, ntypes): + # descriptor + data["descriptor"]["ntypes"] = ntypes + data["descriptor"]["type_map"] = copy.deepcopy(data["type_map"]) + descriptor = BaseDescriptor(**data["descriptor"]) + # fitting + fitting_net = data.get("fitting_net", {}) + fitting_net["type"] = fitting_net.get("type", "ener") + fitting_net["ntypes"] = descriptor.get_ntypes() + fitting_net["type_map"] = copy.deepcopy(data["type_map"]) + fitting_net["mixed_types"] = descriptor.mixed_types() + if fitting_net["type"] in ["dipole", "polar"]: + fitting_net["embedding_width"] = descriptor.get_dim_emb() + fitting_net["dim_descrpt"] = descriptor.get_dim_out() + grad_force = "direct" not in fitting_net["type"] + if not grad_force: + fitting_net["out_dim"] = descriptor.get_dim_emb() + if "ener" in fitting_net["type"]: + fitting_net["return_energy"] = True + fitting = BaseFitting(**fitting_net) + return descriptor, fitting, fitting_net["type"] + def get_standard_model(data: dict) -> EnergyModel: """Get a EnergyModel from a dictionary. @@ -43,36 +70,36 @@ def get_standard_model(data: dict) -> EnergyModel: raise ValueError( "In the DP backend, type_embedding is not at the model level, but within the descriptor. See type embedding documentation for details." ) - data["descriptor"]["type_map"] = data["type_map"] - data["descriptor"]["ntypes"] = len(data["type_map"]) - fitting_type = data["fitting_net"].pop("type") - data["fitting_net"]["type_map"] = data["type_map"] - descriptor = BaseDescriptor( - **data["descriptor"], + data = copy.deepcopy(data) + ntypes = len(data["type_map"]) + descriptor, fitting, fitting_net_type = _get_standard_model_components( + data, ntypes ) - if fitting_type == "ener": - fitting = EnergyFittingNet( - ntypes=descriptor.get_ntypes(), - dim_descrpt=descriptor.get_dim_out(), - mixed_types=descriptor.mixed_types(), - **data["fitting_net"], - ) - elif fitting_type == "dos": - fitting = DOSFittingNet( - ntypes=descriptor.get_ntypes(), - dim_descrpt=descriptor.get_dim_out(), - mixed_types=descriptor.mixed_types(), - **data["fitting_net"], - ) + atom_exclude_types = data.get("atom_exclude_types", []) + pair_exclude_types = data.get("pair_exclude_types", []) + + + 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 in ["ener", "direct_force_ener"]: + modelcls = EnergyModel + elif fitting_net_type == "property": + modelcls = PropertyModel else: - raise ValueError(f"Unknown fitting type {fitting_type}") # fix - return EnergyModel( + raise RuntimeError(f"Unknown fitting type: {fitting_net_type}") + + model = modelcls( descriptor=descriptor, fitting=fitting, type_map=data["type_map"], - atom_exclude_types=data.get("atom_exclude_types", []), - pair_exclude_types=data.get("pair_exclude_types", []), + atom_exclude_types=atom_exclude_types, + pair_exclude_types=pair_exclude_types, ) + return model def get_zbl_model(data: dict) -> DPZBLModel: diff --git a/deepmd/dpmodel/model/polar_model.py b/deepmd/dpmodel/model/polar_model.py index c7cce1c1fa..994b3556c2 100644 --- a/deepmd/dpmodel/model/polar_model.py +++ b/deepmd/dpmodel/model/polar_model.py @@ -3,7 +3,7 @@ from deepmd.dpmodel.atomic_model import ( DPPolarAtomicModel, ) -from deepmd.dpmodel.model.model import ( +from deepmd.dpmodel.model.base_model import ( BaseModel, ) diff --git a/source/tests/consistent/model/common.py b/source/tests/consistent/model/common.py index ef1c7cf911..bb38abc5b6 100644 --- a/source/tests/consistent/model/common.py +++ b/source/tests/consistent/model/common.py @@ -66,6 +66,16 @@ def build_tf_model( ret["dos"], ret["atom_dos"], ] + elif ret_key == "dipole": + ret_list = [ + ret["global_dipole"], + ret["dipole"], + ] + elif ret_key == "polar": + ret_list = [ + ret["polar"], + ret["global_polar"], + ] else: raise NotImplementedError return ret_list, { From 9071e73fbd5470dd58dd1e13f54bc4c00fadf3c6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 18 Nov 2024 06:55:51 +0000 Subject: [PATCH 17/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/model/model.py | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 43eeb13898..1d18b70e8e 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -1,4 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import copy + from deepmd.dpmodel.atomic_model.dp_atomic_model import ( DPAtomicModel, ) @@ -17,23 +19,30 @@ from deepmd.dpmodel.model.base_model import ( BaseModel, ) +from deepmd.dpmodel.model.dipole_model import ( + DipoleModel, +) +from deepmd.dpmodel.model.dos_model import ( + DOSModel, +) from deepmd.dpmodel.model.dp_zbl_model import ( DPZBLModel, ) from deepmd.dpmodel.model.ener_model import ( EnergyModel, ) +from deepmd.dpmodel.model.polar_model import ( + PolarModel, +) +from deepmd.dpmodel.model.property_model import ( + PropertyModel, +) from deepmd.dpmodel.model.spin_model import ( SpinModel, ) from deepmd.utils.spin import ( Spin, ) -import copy -from deepmd.dpmodel.model.dos_model import DOSModel -from deepmd.dpmodel.model.property_model import PropertyModel -from deepmd.dpmodel.model.dipole_model import DipoleModel -from deepmd.dpmodel.model.polar_model import PolarModel def _get_standard_model_components(data, ntypes): @@ -57,7 +66,8 @@ def _get_standard_model_components(data, ntypes): fitting_net["return_energy"] = True fitting = BaseFitting(**fitting_net) return descriptor, fitting, fitting_net["type"] - + + def get_standard_model(data: dict) -> EnergyModel: """Get a EnergyModel from a dictionary. @@ -72,12 +82,9 @@ def get_standard_model(data: dict) -> EnergyModel: ) data = copy.deepcopy(data) ntypes = len(data["type_map"]) - descriptor, fitting, fitting_net_type = _get_standard_model_components( - data, ntypes - ) + descriptor, fitting, fitting_net_type = _get_standard_model_components(data, ntypes) atom_exclude_types = data.get("atom_exclude_types", []) pair_exclude_types = data.get("pair_exclude_types", []) - if fitting_net_type == "dipole": modelcls = DipoleModel From 1dc7fd95c35e7727672e7ef765fe87344a379553 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Tue, 19 Nov 2024 11:26:41 +0800 Subject: [PATCH 18/19] chore: vectorize for loop --- deepmd/dpmodel/atomic_model/polar_atomic_model.py | 3 +-- deepmd/pt/model/atomic_model/polar_atomic_model.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/polar_atomic_model.py b/deepmd/dpmodel/atomic_model/polar_atomic_model.py index faff2444ea..3dfd0dfc99 100644 --- a/deepmd/dpmodel/atomic_model/polar_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/polar_atomic_model.py @@ -42,8 +42,7 @@ def apply_out_stat( for kk in self.bias_keys: ntypes = out_bias[kk].shape[0] temp = np.zeros(ntypes, dtype=dtype) - for i in range(ntypes): - temp[i] = np.mean(np.diagonal(out_bias[kk][i].reshape(3, 3))) + temp = np.mean(np.diagonal(out_bias[kk].reshape(ntypes, 3, 3), axis1=1, axis2=2), axis=1) modified_bias = temp[atype] # (nframes, nloc, 1) diff --git a/deepmd/pt/model/atomic_model/polar_atomic_model.py b/deepmd/pt/model/atomic_model/polar_atomic_model.py index fc6c351830..7d5014d5a2 100644 --- a/deepmd/pt/model/atomic_model/polar_atomic_model.py +++ b/deepmd/pt/model/atomic_model/polar_atomic_model.py @@ -43,8 +43,7 @@ def apply_out_stat( for kk in self.bias_keys: ntypes = out_bias[kk].shape[0] temp = torch.zeros(ntypes, dtype=dtype, device=device) - for i in range(ntypes): - temp[i] = torch.mean(torch.diagonal(out_bias[kk][i].reshape(3, 3))) + temp = torch.mean(torch.diagonal(out_bias[kk].reshape(ntypes, 3, 3), dim1=-2, dim2=-1), dim=-1) modified_bias = temp[atype] # (nframes, nloc, 1) From 0afee21805af35d0ac6dc93487326d608c78279a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 19 Nov 2024 03:28:03 +0000 Subject: [PATCH 19/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/atomic_model/polar_atomic_model.py | 5 ++++- deepmd/pt/model/atomic_model/polar_atomic_model.py | 7 ++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/polar_atomic_model.py b/deepmd/dpmodel/atomic_model/polar_atomic_model.py index 3dfd0dfc99..6e1d32ff35 100644 --- a/deepmd/dpmodel/atomic_model/polar_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/polar_atomic_model.py @@ -42,7 +42,10 @@ def apply_out_stat( for kk in self.bias_keys: ntypes = out_bias[kk].shape[0] temp = np.zeros(ntypes, dtype=dtype) - temp = np.mean(np.diagonal(out_bias[kk].reshape(ntypes, 3, 3), axis1=1, axis2=2), axis=1) + temp = np.mean( + np.diagonal(out_bias[kk].reshape(ntypes, 3, 3), axis1=1, axis2=2), + axis=1, + ) modified_bias = temp[atype] # (nframes, nloc, 1) diff --git a/deepmd/pt/model/atomic_model/polar_atomic_model.py b/deepmd/pt/model/atomic_model/polar_atomic_model.py index 7d5014d5a2..6bd063591f 100644 --- a/deepmd/pt/model/atomic_model/polar_atomic_model.py +++ b/deepmd/pt/model/atomic_model/polar_atomic_model.py @@ -43,7 +43,12 @@ def apply_out_stat( for kk in self.bias_keys: ntypes = out_bias[kk].shape[0] temp = torch.zeros(ntypes, dtype=dtype, device=device) - temp = torch.mean(torch.diagonal(out_bias[kk].reshape(ntypes, 3, 3), dim1=-2, dim2=-1), dim=-1) + temp = torch.mean( + torch.diagonal( + out_bias[kk].reshape(ntypes, 3, 3), dim1=-2, dim2=-1 + ), + dim=-1, + ) modified_bias = temp[atype] # (nframes, nloc, 1)