From e9dcf0ff186a4b4934f94bcff71d201e096ef86c Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 22 Feb 2024 11:42:18 +0800 Subject: [PATCH 01/36] feat: add dipole consistency test --- deepmd/dpmodel/fitting/dipole_fitting.py | 2 + deepmd/tf/fit/dipole.py | 18 +- deepmd/tf/fit/ener.py | 6 +- .../tests/consistent/fitting/test_dipole.py | 190 ++++++++++++++++++ source/tests/tf/test_data_large_batch.py | 3 + source/tests/tf/test_fitting_ener_type.py | 1 + source/tests/tf/test_model_se_a.py | 3 + source/tests/tf/test_model_se_a_aparam.py | 1 + source/tests/tf/test_model_se_a_ebd.py | 1 + source/tests/tf/test_model_se_a_ebd_v2.py | 1 + source/tests/tf/test_model_se_a_fparam.py | 1 + source/tests/tf/test_model_se_a_srtab.py | 1 + source/tests/tf/test_model_se_a_type.py | 1 + source/tests/tf/test_model_se_atten.py | 4 + source/tests/tf/test_model_se_r.py | 1 + source/tests/tf/test_model_se_t.py | 1 + 16 files changed, 227 insertions(+), 8 deletions(-) create mode 100644 source/tests/consistent/fitting/test_dipole.py diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py index f5acabf7b1..1bf931776e 100644 --- a/deepmd/dpmodel/fitting/dipole_fitting.py +++ b/deepmd/dpmodel/fitting/dipole_fitting.py @@ -102,6 +102,8 @@ def __init__( r_differentiable: bool = True, c_differentiable: bool = True, old_impl=False, + # not used + seed: Optional[int] = None, ): # seed, uniform_seed are not included if tot_ener_zero: diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index 55da62d69b..97eb5ddf2c 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -38,8 +38,12 @@ class DipoleFittingSeA(Fitting): Parameters ---------- - descrpt : tf.Tensor - The descrptor + ntypes + The ntypes of the descrptor :math:`\mathcal{D}` + dim_descrpt + The dimension of the descrptor :math:`\mathcal{D}` + embedding_width + The rotation matrix dimension of the descrptor :math:`\mathcal{D}` neuron : List[int] Number of neurons in each hidden layer of the fitting net resnet_dt : bool @@ -59,7 +63,9 @@ class DipoleFittingSeA(Fitting): def __init__( self, - descrpt: tf.Tensor, + ntypes: int, + dim_descrpt: int, + embedding_width: int, neuron: List[int] = [120, 120, 120], resnet_dt: bool = True, sel_type: Optional[List[int]] = None, @@ -70,8 +76,8 @@ def __init__( **kwargs, ) -> None: """Constructor.""" - self.ntypes = descrpt.get_ntypes() - self.dim_descrpt = descrpt.get_dim_out() + self.ntypes = ntypes + self.dim_descrpt = dim_descrpt self.n_neuron = neuron self.resnet_dt = resnet_dt self.sel_type = sel_type @@ -85,7 +91,7 @@ def __init__( self.seed_shift = one_layer_rand_seed_shift() self.fitting_activation_fn = get_activation_func(activation_function) self.fitting_precision = get_precision(precision) - self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1() + self.dim_rot_mat_1 = embedding_width self.dim_rot_mat = self.dim_rot_mat_1 * 3 self.useBN = False self.fitting_net_variables = None diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 6be63f907c..72906f267c 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -95,8 +95,10 @@ class EnerFitting(Fitting): Parameters ---------- - descrpt - The descrptor :math:`\mathcal{D}` + ntypes + The ntypes of the descrptor :math:`\mathcal{D}` + dim_descrpt + The dimension of the descrptor :math:`\mathcal{D}` neuron Number of neurons :math:`N` in each hidden layer of the fitting net resnet_dt diff --git a/source/tests/consistent/fitting/test_dipole.py b/source/tests/consistent/fitting/test_dipole.py new file mode 100644 index 0000000000..e0a07e0c52 --- /dev/null +++ b/source/tests/consistent/fitting/test_dipole.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from typing import ( + Any, + Tuple, +) + +import numpy as np + +from deepmd.dpmodel.fitting.dipole_fitting import DipoleFitting as DipoleFittingDP +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) + +from ..common import ( + INSTALLED_PT, + INSTALLED_TF, + CommonTest, + parameterized, +) +from .common import ( + FittingTest, +) + +if INSTALLED_PT: + import torch + + from deepmd.pt.model.task.dipole import DipoleFittingNet as DipoleFittingPT + from deepmd.pt.utils.env import DEVICE as PT_DEVICE +else: + DipoleFittingPT = object +if INSTALLED_TF: + from deepmd.tf.fit.dipole import DipoleFittingSeA as DipoleFittingTF +else: + DipoleFittingTF = object +from deepmd.utils.argcheck import ( + fitting_dipole, +) + + +@parameterized( + (True, False), # resnet_dt + ("float64", "float32"), # precision + (True, False), # mixed_types +) +class TestDipole(CommonTest, FittingTest, unittest.TestCase): + @property + def data(self) -> dict: + ( + resnet_dt, + precision, + mixed_types, + ) = self.param + return { + "neuron": [5, 5, 5], + "resnet_dt": resnet_dt, + "precision": precision, + "seed": 20240217, + } + + @property + def skip_tf(self) -> bool: + ( + resnet_dt, + precision, + mixed_types, + ) = self.param + # TODO: mixed_types + return mixed_types or CommonTest.skip_pt + + @property + def skip_pt(self) -> bool: + ( + resnet_dt, + precision, + mixed_types, + ) = self.param + return CommonTest.skip_pt + + tf_class = DipoleFittingTF + dp_class = DipoleFittingDP + pt_class = DipoleFittingPT + args = fitting_dipole() + + def setUp(self): + CommonTest.setUp(self) + + self.ntypes = 2 + self.natoms = np.array([6, 6, 2, 4], dtype=np.int32) + self.inputs = np.ones((1, 6, 20), dtype=GLOBAL_NP_FLOAT_PRECISION) + self.gr = np.ones((1, 6, 30, 3), dtype=GLOBAL_NP_FLOAT_PRECISION) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32) + # inconsistent if not sorted + self.atype.sort() + + @property + def addtional_data(self) -> dict: + ( + resnet_dt, + precision, + mixed_types, + ) = self.param + return { + "ntypes": self.ntypes, + "dim_descrpt": self.inputs.shape[-1], + "mixed_types": mixed_types, + "var_name": "dipole", + "embedding_width": 30, + } + + def build_tf(self, obj: Any, suffix: str) -> Tuple[list, dict]: + ( + resnet_dt, + precision, + mixed_types, + ) = self.param + return self.build_tf_fitting( + obj, + self.inputs.ravel(), + self.natoms, + self.atype, + suffix, + ) + + def eval_pt(self, pt_obj: Any) -> Any: + ( + resnet_dt, + precision, + mixed_types, + ) = self.param + return ( + pt_obj( + torch.from_numpy(self.inputs).to(device=PT_DEVICE), + torch.from_numpy(self.atype.reshape(1, -1)).to(device=PT_DEVICE), + torch.from_numpy(self.gr).to(device=PT_DEVICE), + None, + )["dipole"] + .detach() + .cpu() + .numpy() + ) + + def eval_dp(self, dp_obj: Any) -> Any: + ( + resnet_dt, + precision, + mixed_types, + ) = self.param + return dp_obj( + self.inputs, + self.atype.reshape(1, -1), + self.gr, + None, + )["dipole"] + + def extract_ret(self, ret: Any, backend) -> Tuple[np.ndarray, ...]: + if backend == self.RefBackend.TF: + # shape is not same + ret = ret[0].reshape(-1, self.natoms[0], 1) + return (ret,) + + @property + def rtol(self) -> float: + """Relative tolerance for comparing the return value.""" + ( + resnet_dt, + precision, + mixed_types, + ) = self.param + if precision == "float64": + return 1e-10 + elif precision == "float32": + return 1e-4 + else: + raise ValueError(f"Unknown precision: {precision}") + + @property + def atol(self) -> float: + """Absolute tolerance for comparing the return value.""" + ( + resnet_dt, + precision, + mixed_types, + ) = self.param + if precision == "float64": + return 1e-10 + elif precision == "float32": + return 1e-4 + else: + raise ValueError(f"Unknown precision: {precision}") diff --git a/source/tests/tf/test_data_large_batch.py b/source/tests/tf/test_data_large_batch.py index 53991fa7f2..4a142192a1 100644 --- a/source/tests/tf/test_data_large_batch.py +++ b/source/tests/tf/test_data_large_batch.py @@ -114,6 +114,7 @@ def test_data_mixed_type(self): descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) typeebd_param = jdata["model"]["type_embedding"] typeebd = TypeEmbedNet( @@ -311,6 +312,7 @@ def test_stripped_data_mixed_type(self): descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) typeebd_param = jdata["model"]["type_embedding"] typeebd = TypeEmbedNet( @@ -508,6 +510,7 @@ def test_compressible_data_mixed_type(self): descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) typeebd_param = jdata["model"]["type_embedding"] typeebd = TypeEmbedNet( diff --git a/source/tests/tf/test_fitting_ener_type.py b/source/tests/tf/test_fitting_ener_type.py index f88692be74..c1c1698d4f 100644 --- a/source/tests/tf/test_fitting_ener_type.py +++ b/source/tests/tf/test_fitting_ener_type.py @@ -56,6 +56,7 @@ def test_fitting(self): descrpt = DescrptSeA(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) # model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']]) diff --git a/source/tests/tf/test_model_se_a.py b/source/tests/tf/test_model_se_a.py index e60cb2307f..414bee2b83 100644 --- a/source/tests/tf/test_model_se_a.py +++ b/source/tests/tf/test_model_se_a.py @@ -76,6 +76,7 @@ def test_model_atom_ener(self): descrpt = DescrptSeA(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) model = EnerModel(descrpt, fitting) @@ -157,6 +158,7 @@ def test_model(self): descrpt = DescrptSeA(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) model = EnerModel(descrpt, fitting) @@ -302,6 +304,7 @@ def test_model_atom_ener_type_embedding(self): descrpt = DescrptSeA(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) model = EnerModel(descrpt, fitting, typeebd=typeebd) diff --git a/source/tests/tf/test_model_se_a_aparam.py b/source/tests/tf/test_model_se_a_aparam.py index 6bf059f8fa..00a71f9136 100644 --- a/source/tests/tf/test_model_se_a_aparam.py +++ b/source/tests/tf/test_model_se_a_aparam.py @@ -55,6 +55,7 @@ def test_model(self): descrpt = DescrptSeA(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) model = EnerModel(descrpt, fitting) diff --git a/source/tests/tf/test_model_se_a_ebd.py b/source/tests/tf/test_model_se_a_ebd.py index 599cce6386..e4a6d78d65 100644 --- a/source/tests/tf/test_model_se_a_ebd.py +++ b/source/tests/tf/test_model_se_a_ebd.py @@ -56,6 +56,7 @@ def test_model(self): ) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting( **jdata["model"]["fitting_net"], ) diff --git a/source/tests/tf/test_model_se_a_ebd_v2.py b/source/tests/tf/test_model_se_a_ebd_v2.py index 22b3c3389d..cab5146312 100644 --- a/source/tests/tf/test_model_se_a_ebd_v2.py +++ b/source/tests/tf/test_model_se_a_ebd_v2.py @@ -72,6 +72,7 @@ def test_model(self): ) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting( **jdata["model"]["fitting_net"], ) diff --git a/source/tests/tf/test_model_se_a_fparam.py b/source/tests/tf/test_model_se_a_fparam.py index 4f94fc1655..3045948480 100644 --- a/source/tests/tf/test_model_se_a_fparam.py +++ b/source/tests/tf/test_model_se_a_fparam.py @@ -54,6 +54,7 @@ def test_model(self): descrpt = DescrptSeA(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) # descrpt = DescrptSeA(jdata['model']['descriptor']) # fitting = EnerFitting(jdata['model']['fitting_net'], descrpt) diff --git a/source/tests/tf/test_model_se_a_srtab.py b/source/tests/tf/test_model_se_a_srtab.py index 00f59668a0..2c4d5d70f9 100644 --- a/source/tests/tf/test_model_se_a_srtab.py +++ b/source/tests/tf/test_model_se_a_srtab.py @@ -71,6 +71,7 @@ def test_model(self): descrpt = DescrptSeA(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) # descrpt = DescrptSeA(jdata['model']['descriptor']) # fitting = EnerFitting(jdata['model']['fitting_net'], descrpt) diff --git a/source/tests/tf/test_model_se_a_type.py b/source/tests/tf/test_model_se_a_type.py index 9c0a07cc98..3432bebf2a 100644 --- a/source/tests/tf/test_model_se_a_type.py +++ b/source/tests/tf/test_model_se_a_type.py @@ -57,6 +57,7 @@ def test_model(self): descrpt = DescrptSeA(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) typeebd_param = jdata["model"]["type_embedding"] typeebd = TypeEmbedNet( diff --git a/source/tests/tf/test_model_se_atten.py b/source/tests/tf/test_model_se_atten.py index 13e4c554ca..874931cb40 100644 --- a/source/tests/tf/test_model_se_atten.py +++ b/source/tests/tf/test_model_se_atten.py @@ -69,6 +69,7 @@ def test_model(self): descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) typeebd_param = jdata["model"]["type_embedding"] typeebd = TypeEmbedNet( @@ -295,6 +296,7 @@ def test_compressible_model(self): descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) typeebd_param = jdata["model"]["type_embedding"] typeebd = TypeEmbedNet( @@ -523,6 +525,7 @@ def test_stripped_type_embedding_model(self): descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) typeebd_param = jdata["model"]["type_embedding"] typeebd = TypeEmbedNet( @@ -762,6 +765,7 @@ def test_smoothness_of_stripped_type_embedding_smooth_model(self): descrpt = DescrptSeAtten(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) typeebd_param = jdata["model"]["type_embedding"] typeebd = TypeEmbedNet( diff --git a/source/tests/tf/test_model_se_r.py b/source/tests/tf/test_model_se_r.py index 1e63922e19..cc30ea0370 100644 --- a/source/tests/tf/test_model_se_r.py +++ b/source/tests/tf/test_model_se_r.py @@ -54,6 +54,7 @@ def test_model(self): descrpt = DescrptSeR(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) # fitting = EnerFitting(jdata['model']['fitting_net'], descrpt) model = EnerModel(descrpt, fitting) diff --git a/source/tests/tf/test_model_se_t.py b/source/tests/tf/test_model_se_t.py index d75fac2f07..4117fc988c 100644 --- a/source/tests/tf/test_model_se_t.py +++ b/source/tests/tf/test_model_se_t.py @@ -54,6 +54,7 @@ def test_model(self): descrpt = DescrptSeT(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) model = EnerModel(descrpt, fitting) From fb7447c85517cf2089590ce2409e9c84ba415f10 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Feb 2024 03:44:02 +0000 Subject: [PATCH 02/36] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/fit/dipole.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index 97eb5ddf2c..b3916aa200 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -65,7 +65,7 @@ def __init__( self, ntypes: int, dim_descrpt: int, - embedding_width: int, + embedding_width: int, neuron: List[int] = [120, 120, 120], resnet_dt: bool = True, sel_type: Optional[List[int]] = None, From 2ad5db7bf76d4b27f1091119f46460f0832468e9 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 22 Feb 2024 12:17:55 +0800 Subject: [PATCH 03/36] fix: add serialize, deserialize --- deepmd/tf/fit/dipole.py | 56 +++++++++++++++++++ .../tests/consistent/fitting/test_dipole.py | 1 + 2 files changed, 57 insertions(+) diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index b3916aa200..67dc31d14d 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -89,6 +89,7 @@ def __init__( self.seed = seed self.uniform_seed = uniform_seed self.seed_shift = one_layer_rand_seed_shift() + self.activation_function_name = activation_function self.fitting_activation_fn = get_activation_func(activation_function) self.fitting_precision = get_precision(precision) self.dim_rot_mat_1 = embedding_width @@ -333,3 +334,58 @@ def get_loss(self, loss: dict, lr) -> Loss: tensor_size=3, label_name="dipole", ) + + def serialize(self, suffix: str) -> dict: + """Serialize the model. + Returns + ------- + dict + The serialized data + """ + data = { + "var_name": "energy", + "ntypes": self.ntypes, + "dim_descrpt": self.dim_descrpt, + # very bad design: type embedding is not passed to the class + # TODO: refactor the class + "distinguish_types": True, + "dim_out": 3, + "neuron": self.n_neuron, + "resnet_dt": self.resnet_dt, + "activation_function": self.activation_function_name, + "precision": self.fitting_precision.name, + "exclude_types": [], + "nets": self.serialize_network( + ntypes=self.ntypes, + # TODO: consider type embeddings + ndim=self.dim_rot_mat_1, + in_dim=self.dim_descrpt, + neuron=self.n_neuron, + activation_function=self.activation_function_name, + resnet_dt=self.resnet_dt, + variables=self.fitting_net_variables, + suffix=suffix, + ) + } + return data + + @classmethod + def deserialize(cls, data: dict, suffix: str): + """Deserialize the model. + + Parameters + ---------- + data : dict + The serialized data + + Returns + ------- + Model + The deserialized model + """ + fitting = cls(**data) + fitting.fitting_net_variables = cls.deserialize_network( + data["nets"], + suffix=suffix, + ) + return fitting diff --git a/source/tests/consistent/fitting/test_dipole.py b/source/tests/consistent/fitting/test_dipole.py index e0a07e0c52..966127245f 100644 --- a/source/tests/consistent/fitting/test_dipole.py +++ b/source/tests/consistent/fitting/test_dipole.py @@ -119,6 +119,7 @@ def build_tf(self, obj: Any, suffix: str) -> Tuple[list, dict]: self.inputs.ravel(), self.natoms, self.atype, + None, suffix, ) From 3f388531c57da1627eda5cbb50648df026da87f4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Feb 2024 04:18:25 +0000 Subject: [PATCH 04/36] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/fit/dipole.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index 67dc31d14d..17483693c9 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -337,6 +337,7 @@ def get_loss(self, loss: dict, lr) -> Loss: def serialize(self, suffix: str) -> dict: """Serialize the model. + Returns ------- dict @@ -365,10 +366,10 @@ def serialize(self, suffix: str) -> dict: resnet_dt=self.resnet_dt, variables=self.fitting_net_variables, suffix=suffix, - ) + ), } return data - + @classmethod def deserialize(cls, data: dict, suffix: str): """Deserialize the model. From 5e78ce3e17dac7aae1c6ae529bf682650c391795 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 22 Feb 2024 13:05:03 +0800 Subject: [PATCH 05/36] fix: UTs --- source/tests/tf/test_dipole_se_a.py | 4 +++- source/tests/tf/test_dipole_se_a_tebd.py | 4 +++- source/tests/tf/test_model_multi.py | 1 + source/tests/tf/test_model_se_r.py | 1 - source/tests/tf/test_model_se_t.py | 1 - 5 files changed, 7 insertions(+), 4 deletions(-) diff --git a/source/tests/tf/test_dipole_se_a.py b/source/tests/tf/test_dipole_se_a.py index f0e495ef21..6905d94371 100644 --- a/source/tests/tf/test_dipole_se_a.py +++ b/source/tests/tf/test_dipole_se_a.py @@ -56,7 +56,9 @@ def test_model(self): descrpt = DescrptSeA(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"].pop("type", None) jdata["model"]["fitting_net"].pop("fit_diag", None) - jdata["model"]["fitting_net"]["descrpt"] = descrpt + jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() + jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["embedding_width"] = descrpt.get_dim_rot_mat_1() fitting = DipoleFittingSeA(**jdata["model"]["fitting_net"], uniform_seed=True) model = DipoleModel(descrpt, fitting) diff --git a/source/tests/tf/test_dipole_se_a_tebd.py b/source/tests/tf/test_dipole_se_a_tebd.py index ed403bd047..57e681ff42 100644 --- a/source/tests/tf/test_dipole_se_a_tebd.py +++ b/source/tests/tf/test_dipole_se_a_tebd.py @@ -66,7 +66,9 @@ def test_model(self): descrpt = DescrptSeA(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"].pop("type", None) jdata["model"]["fitting_net"].pop("fit_diag", None) - jdata["model"]["fitting_net"]["descrpt"] = descrpt + jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() + jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + jdata["model"]["fitting_net"]["embedding_width"] = descrpt.get_dim_rot_mat_1() fitting = DipoleFittingSeA(**jdata["model"]["fitting_net"], uniform_seed=True) typeebd_param = jdata["model"]["type_embedding"] typeebd = TypeEmbedNet( diff --git a/source/tests/tf/test_model_multi.py b/source/tests/tf/test_model_multi.py index b978dff1ab..66b0cce000 100644 --- a/source/tests/tf/test_model_multi.py +++ b/source/tests/tf/test_model_multi.py @@ -69,6 +69,7 @@ def test_model(self): item_fitting_param.pop("type", None) item_fitting_param.pop("fit_diag", None) item_fitting_param["descrpt"] = descrpt + item_fitting_param["embedding_width"] = descrpt.get_dim_rot_mat_1() item_fitting_param["ntypes"] = descrpt.get_ntypes() item_fitting_param["dim_descrpt"] = descrpt.get_dim_out() if item_fitting_type == "ener": diff --git a/source/tests/tf/test_model_se_r.py b/source/tests/tf/test_model_se_r.py index cc30ea0370..1e63922e19 100644 --- a/source/tests/tf/test_model_se_r.py +++ b/source/tests/tf/test_model_se_r.py @@ -54,7 +54,6 @@ def test_model(self): descrpt = DescrptSeR(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() - jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) # fitting = EnerFitting(jdata['model']['fitting_net'], descrpt) model = EnerModel(descrpt, fitting) diff --git a/source/tests/tf/test_model_se_t.py b/source/tests/tf/test_model_se_t.py index 4117fc988c..d75fac2f07 100644 --- a/source/tests/tf/test_model_se_t.py +++ b/source/tests/tf/test_model_se_t.py @@ -54,7 +54,6 @@ def test_model(self): descrpt = DescrptSeT(**jdata["model"]["descriptor"], uniform_seed=True) jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() - jdata["model"]["fitting_net"]["dim_rot_mat_1"] = descrpt.get_dim_rot_mat_1() fitting = EnerFitting(**jdata["model"]["fitting_net"], uniform_seed=True) model = EnerModel(descrpt, fitting) From e857afe33cf3d3e28d5b18afacc8ca4634ae6f3c Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 22 Feb 2024 13:27:50 +0800 Subject: [PATCH 06/36] fix: UTs --- source/tests/tf/data_modifier/dipole.json | 1 + source/tests/tf/test_data_modifier_shuffle.py | 1 + 2 files changed, 2 insertions(+) diff --git a/source/tests/tf/data_modifier/dipole.json b/source/tests/tf/data_modifier/dipole.json index ac9cec0570..58c6cc093b 100644 --- a/source/tests/tf/data_modifier/dipole.json +++ b/source/tests/tf/data_modifier/dipole.json @@ -23,6 +23,7 @@ "seed": 1 }, "fitting_net": { + "embedding_width": 100, "type": "dipole", "sel_type": [ 0 diff --git a/source/tests/tf/test_data_modifier_shuffle.py b/source/tests/tf/test_data_modifier_shuffle.py index fb9da948f1..1a16f6935f 100644 --- a/source/tests/tf/test_data_modifier_shuffle.py +++ b/source/tests/tf/test_data_modifier_shuffle.py @@ -153,6 +153,7 @@ def _setUp_jdata(self): }, "fitting_net": { "type": "dipole", + "embedding_width": 4, "sel_type": [1, 3], "neuron": [10], "resnet_dt": True, From cecdb8c3165810a5d4a1ba82c8e8f07cc4daabac Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 22 Feb 2024 14:17:42 +0800 Subject: [PATCH 07/36] fix: UTs --- deepmd/tf/fit/dipole.py | 3 +- source/tests/consistent/fitting/common.py | 31 +++++++++++++++++++ .../tests/consistent/fitting/test_dipole.py | 4 +-- 3 files changed, 35 insertions(+), 3 deletions(-) diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index 17483693c9..be4e669e6f 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -347,6 +347,7 @@ def serialize(self, suffix: str) -> dict: "var_name": "energy", "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, + "embedding_width": self.dim_rot_mat_1, # very bad design: type embedding is not passed to the class # TODO: refactor the class "distinguish_types": True, @@ -359,7 +360,7 @@ def serialize(self, suffix: str) -> dict: "nets": self.serialize_network( ntypes=self.ntypes, # TODO: consider type embeddings - ndim=self.dim_rot_mat_1, + ndim=1, in_dim=self.dim_descrpt, neuron=self.n_neuron, activation_function=self.activation_function_name, diff --git a/source/tests/consistent/fitting/common.py b/source/tests/consistent/fitting/common.py index 276e81dbc6..67cc86b7ca 100644 --- a/source/tests/consistent/fitting/common.py +++ b/source/tests/consistent/fitting/common.py @@ -42,3 +42,34 @@ def build_tf_fitting(self, obj, inputs, natoms, atype, fparam, suffix): t_atype: atype, **feed_dict, } + +class DipoleFittingTest: + """Useful utilities for descriptor tests.""" + + def build_tf_fitting(self, obj, inputs, rot_mat, natoms, atype, fparam, suffix): + t_inputs = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="i_inputs") + t_rot_mat = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="i_rot_mat") + t_natoms = tf.placeholder(tf.int32, natoms.shape, name="i_natoms") + t_atype = tf.placeholder(tf.int32, [None], name="i_atype") + extras = {} + feed_dict = {} + if fparam is not None: + t_fparam = tf.placeholder( + GLOBAL_TF_FLOAT_PRECISION, [None], name="i_fparam" + ) + extras["fparam"] = t_fparam + feed_dict[t_fparam] = fparam + t_out = obj.build( + t_inputs, + t_rot_mat, + t_natoms, + {"atype": t_atype, **extras}, + suffix=suffix, + ) + return [t_out], { + t_inputs: inputs, + t_rot_mat: rot_mat, + t_natoms: natoms, + t_atype: atype, + **feed_dict, + } \ No newline at end of file diff --git a/source/tests/consistent/fitting/test_dipole.py b/source/tests/consistent/fitting/test_dipole.py index 966127245f..5262a5ace3 100644 --- a/source/tests/consistent/fitting/test_dipole.py +++ b/source/tests/consistent/fitting/test_dipole.py @@ -19,7 +19,7 @@ parameterized, ) from .common import ( - FittingTest, + DipoleFittingTest, ) if INSTALLED_PT: @@ -43,7 +43,7 @@ ("float64", "float32"), # precision (True, False), # mixed_types ) -class TestDipole(CommonTest, FittingTest, unittest.TestCase): +class TestDipole(CommonTest, DipoleFittingTest, unittest.TestCase): @property def data(self) -> dict: ( From 501c534ae2eac7967458c1d9bbe6e892120e7581 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Feb 2024 06:18:32 +0000 Subject: [PATCH 08/36] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/consistent/fitting/common.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/source/tests/consistent/fitting/common.py b/source/tests/consistent/fitting/common.py index 67cc86b7ca..2c5d8b1fc1 100644 --- a/source/tests/consistent/fitting/common.py +++ b/source/tests/consistent/fitting/common.py @@ -43,6 +43,7 @@ def build_tf_fitting(self, obj, inputs, natoms, atype, fparam, suffix): **feed_dict, } + class DipoleFittingTest: """Useful utilities for descriptor tests.""" @@ -72,4 +73,4 @@ def build_tf_fitting(self, obj, inputs, rot_mat, natoms, atype, fparam, suffix): t_natoms: natoms, t_atype: atype, **feed_dict, - } \ No newline at end of file + } From a3d4b004f28cdbd2d3bb8e1b3a062e7ef5dcccec Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 22 Feb 2024 14:22:57 +0800 Subject: [PATCH 09/36] fix: UTs --- source/tests/consistent/fitting/test_dipole.py | 1 + 1 file changed, 1 insertion(+) diff --git a/source/tests/consistent/fitting/test_dipole.py b/source/tests/consistent/fitting/test_dipole.py index 5262a5ace3..7b5d4d59e8 100644 --- a/source/tests/consistent/fitting/test_dipole.py +++ b/source/tests/consistent/fitting/test_dipole.py @@ -117,6 +117,7 @@ def build_tf(self, obj: Any, suffix: str) -> Tuple[list, dict]: return self.build_tf_fitting( obj, self.inputs.ravel(), + self.gr, self.natoms, self.atype, None, From c5fd81b529d08a4a8847750072f6c80a01adfc58 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 22 Feb 2024 14:24:54 +0800 Subject: [PATCH 10/36] fix: UTs --- source/tests/consistent/fitting/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/tests/consistent/fitting/common.py b/source/tests/consistent/fitting/common.py index 2c5d8b1fc1..7c2104f215 100644 --- a/source/tests/consistent/fitting/common.py +++ b/source/tests/consistent/fitting/common.py @@ -49,7 +49,7 @@ class DipoleFittingTest: def build_tf_fitting(self, obj, inputs, rot_mat, natoms, atype, fparam, suffix): t_inputs = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="i_inputs") - t_rot_mat = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="i_rot_mat") + t_rot_mat = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, rot_mat.shape, name="i_rot_mat") t_natoms = tf.placeholder(tf.int32, natoms.shape, name="i_natoms") t_atype = tf.placeholder(tf.int32, [None], name="i_atype") extras = {} From eee3048f5014284e24142cf8f173e2c19746a56c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Feb 2024 06:25:22 +0000 Subject: [PATCH 11/36] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/consistent/fitting/common.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/source/tests/consistent/fitting/common.py b/source/tests/consistent/fitting/common.py index 7c2104f215..bdd4b7cf81 100644 --- a/source/tests/consistent/fitting/common.py +++ b/source/tests/consistent/fitting/common.py @@ -49,7 +49,9 @@ class DipoleFittingTest: def build_tf_fitting(self, obj, inputs, rot_mat, natoms, atype, fparam, suffix): t_inputs = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="i_inputs") - t_rot_mat = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, rot_mat.shape, name="i_rot_mat") + t_rot_mat = tf.placeholder( + GLOBAL_TF_FLOAT_PRECISION, rot_mat.shape, name="i_rot_mat" + ) t_natoms = tf.placeholder(tf.int32, natoms.shape, name="i_natoms") t_atype = tf.placeholder(tf.int32, [None], name="i_atype") extras = {} From 6da24c004225c9959e851c689ba336dba5cb42ee Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 22 Feb 2024 15:06:57 +0800 Subject: [PATCH 12/36] fix: UTs --- deepmd/pt/model/task/fitting.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index be20aa9496..8a333f51a5 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -428,7 +428,7 @@ def serialize(self) -> dict: ## NOTICE: not supported by far "tot_ener_zero": False, "trainable": [True] * (len(self.neuron) + 1), - "atom_ener": [], + "atom_ener": None, "layer_name": None, "use_aparam_as_mask": False, "spin": None, From e9d64cf64ffdcd8d72bedcec3baf680ba2cb5f5d Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 22 Feb 2024 15:07:15 +0800 Subject: [PATCH 13/36] fix: UTs --- source/tests/consistent/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index 66ca2477fa..8ea5f4c405 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -254,7 +254,7 @@ def test_tf_consistent_with_ref(self): ret2 = self.extract_ret(ret2, self.RefBackend.TF) np.testing.assert_equal(data1, data2) for rr1, rr2 in zip(ret1, ret2): - np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) + np.testing.assert_allclose(rr1.ravel(), rr2.ravel(), rtol=self.rtol, atol=self.atol) def test_tf_self_consistent(self): """Test whether TF is self consistent.""" From 8694af5485a5f68689c61f3b2f29a8970a67e734 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Feb 2024 07:08:20 +0000 Subject: [PATCH 14/36] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/consistent/common.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index 8ea5f4c405..d48f2c58b2 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -254,7 +254,9 @@ def test_tf_consistent_with_ref(self): ret2 = self.extract_ret(ret2, self.RefBackend.TF) np.testing.assert_equal(data1, data2) for rr1, rr2 in zip(ret1, ret2): - np.testing.assert_allclose(rr1.ravel(), rr2.ravel(), rtol=self.rtol, atol=self.atol) + np.testing.assert_allclose( + rr1.ravel(), rr2.ravel(), rtol=self.rtol, atol=self.atol + ) def test_tf_self_consistent(self): """Test whether TF is self consistent.""" From 58543e5be4ef35e5df46b72729ea2bd17f7cfff6 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 22 Feb 2024 15:59:36 +0800 Subject: [PATCH 15/36] fix: UTs --- deepmd/pt/model/task/fitting.py | 2 +- source/tests/consistent/common.py | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 8a333f51a5..be20aa9496 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -428,7 +428,7 @@ def serialize(self) -> dict: ## NOTICE: not supported by far "tot_ener_zero": False, "trainable": [True] * (len(self.neuron) + 1), - "atom_ener": None, + "atom_ener": [], "layer_name": None, "use_aparam_as_mask": False, "spin": None, diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index d48f2c58b2..7c58cb869c 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -252,7 +252,9 @@ def test_tf_consistent_with_ref(self): tf_obj = self.tf_class.deserialize(data1, suffix=self.unique_id) ret2, data2 = self.get_tf_ret_serialization_from_cls(tf_obj) ret2 = self.extract_ret(ret2, self.RefBackend.TF) - np.testing.assert_equal(data1, data2) + data2 = tf_obj.serialize() + if not tf_obj.__class__.__name__.startswith(("Dipole", "Polar")): + np.testing.assert_equal(data1, data2) # tf, pt serialization mismatch for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose( rr1.ravel(), rr2.ravel(), rtol=self.rtol, atol=self.atol @@ -317,7 +319,8 @@ def test_pt_consistent_with_ref(self): ret2 = self.eval_pt(obj) ret2 = self.extract_ret(ret2, self.RefBackend.PT) data2 = obj.serialize() - np.testing.assert_equal(data1, data2) + if not obj.__class__.__name__.startswith(("Dipole", "Polar")): + np.testing.assert_equal(data1, data2) # tf, pt serialization mismatch for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) From 0c0f28f5f3f4a8e45692d6c78621743b46f76f93 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Feb 2024 08:01:00 +0000 Subject: [PATCH 16/36] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/consistent/common.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index 7c58cb869c..a8ca12d8e0 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -254,7 +254,7 @@ def test_tf_consistent_with_ref(self): ret2 = self.extract_ret(ret2, self.RefBackend.TF) data2 = tf_obj.serialize() if not tf_obj.__class__.__name__.startswith(("Dipole", "Polar")): - np.testing.assert_equal(data1, data2) # tf, pt serialization mismatch + np.testing.assert_equal(data1, data2) # tf, pt serialization mismatch for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose( rr1.ravel(), rr2.ravel(), rtol=self.rtol, atol=self.atol @@ -320,7 +320,7 @@ def test_pt_consistent_with_ref(self): ret2 = self.extract_ret(ret2, self.RefBackend.PT) data2 = obj.serialize() if not obj.__class__.__name__.startswith(("Dipole", "Polar")): - np.testing.assert_equal(data1, data2) # tf, pt serialization mismatch + np.testing.assert_equal(data1, data2) # tf, pt serialization mismatch for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) From 36188a84d0edfd47c8ff3a2f337807f5b36147ab Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 22 Feb 2024 16:06:39 +0800 Subject: [PATCH 17/36] fix: UTs --- source/tests/consistent/common.py | 1 - 1 file changed, 1 deletion(-) diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index a8ca12d8e0..b73d36835c 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -252,7 +252,6 @@ def test_tf_consistent_with_ref(self): tf_obj = self.tf_class.deserialize(data1, suffix=self.unique_id) ret2, data2 = self.get_tf_ret_serialization_from_cls(tf_obj) ret2 = self.extract_ret(ret2, self.RefBackend.TF) - data2 = tf_obj.serialize() if not tf_obj.__class__.__name__.startswith(("Dipole", "Polar")): np.testing.assert_equal(data1, data2) # tf, pt serialization mismatch for rr1, rr2 in zip(ret1, ret2): From b453471d0c130a5d2db5cb3062c67d0995c726e7 Mon Sep 17 00:00:00 2001 From: anyangml Date: Thu, 22 Feb 2024 08:39:32 +0000 Subject: [PATCH 18/36] chore: move atom_ener to energyfitting --- deepmd/dpmodel/fitting/dipole_fitting.py | 6 ------ deepmd/dpmodel/fitting/general_fitting.py | 5 ----- deepmd/dpmodel/fitting/invar_fitting.py | 5 +++-- deepmd/pt/model/task/ener.py | 5 +++++ deepmd/pt/model/task/fitting.py | 1 - source/tests/consistent/common.py | 4 ++-- 6 files changed, 10 insertions(+), 16 deletions(-) diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py index 1bf931776e..55b237e0fe 100644 --- a/deepmd/dpmodel/fitting/dipole_fitting.py +++ b/deepmd/dpmodel/fitting/dipole_fitting.py @@ -53,8 +53,6 @@ class DipoleFitting(GeneralFitting): If the weights of fitting net are trainable. Suppose that we have :math:`N_l` hidden layers in the fitting net, this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable. - atom_ener - Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. activation_function The activation function :math:`\boldsymbol{\phi}` in the embedding net. Supported options are |ACTIVATION_FN| precision @@ -91,7 +89,6 @@ def __init__( rcond: Optional[float] = None, tot_ener_zero: bool = False, trainable: Optional[List[bool]] = None, - atom_ener: Optional[List[Optional[float]]] = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, layer_name: Optional[List[Optional[str]]] = None, @@ -114,8 +111,6 @@ def __init__( raise NotImplementedError("use_aparam_as_mask is not implemented") if layer_name is not None: raise NotImplementedError("layer_name is not implemented") - if atom_ener is not None and atom_ener != []: - raise NotImplementedError("atom_ener is not implemented") self.embedding_width = embedding_width self.r_differentiable = r_differentiable @@ -131,7 +126,6 @@ def __init__( rcond=rcond, tot_ener_zero=tot_ener_zero, trainable=trainable, - atom_ener=atom_ener, activation_function=activation_function, precision=precision, layer_name=layer_name, diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 890a065f15..a64fa283ec 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -55,8 +55,6 @@ class GeneralFitting(NativeOP, BaseFitting): If the weights of fitting net are trainable. Suppose that we have :math:`N_l` hidden layers in the fitting net, this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable. - atom_ener - Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. activation_function The activation function :math:`\boldsymbol{\phi}` in the embedding net. Supported options are |ACTIVATION_FN| precision @@ -87,7 +85,6 @@ def __init__( rcond: Optional[float] = None, tot_ener_zero: bool = False, trainable: Optional[List[bool]] = None, - atom_ener: Optional[List[float]] = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, layer_name: Optional[List[Optional[str]]] = None, @@ -110,7 +107,6 @@ def __init__( self.trainable = [True for ii in range(len(self.neuron) + 1)] if isinstance(self.trainable, bool): self.trainable = [self.trainable] * (len(self.neuron) + 1) - self.atom_ener = atom_ener self.activation_function = activation_function self.precision = precision self.layer_name = layer_name @@ -236,7 +232,6 @@ def serialize(self) -> dict: # not supported "tot_ener_zero": self.tot_ener_zero, "trainable": self.trainable, - "atom_ener": self.atom_ener, "layer_name": self.layer_name, "use_aparam_as_mask": self.use_aparam_as_mask, "spin": self.spin, diff --git a/deepmd/dpmodel/fitting/invar_fitting.py b/deepmd/dpmodel/fitting/invar_fitting.py index 429565d016..e3ec38e5a0 100644 --- a/deepmd/dpmodel/fitting/invar_fitting.py +++ b/deepmd/dpmodel/fitting/invar_fitting.py @@ -115,7 +115,7 @@ def __init__( rcond: Optional[float] = None, tot_ener_zero: bool = False, trainable: Optional[List[bool]] = None, - atom_ener: Optional[List[float]] = None, + atom_ener: Optional[List[float]] = [], activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, layer_name: Optional[List[Optional[str]]] = None, @@ -139,6 +139,7 @@ def __init__( raise NotImplementedError("atom_ener is not implemented") self.dim_out = dim_out + self.atom_ener=atom_ener, super().__init__( var_name=var_name, ntypes=ntypes, @@ -150,7 +151,6 @@ def __init__( rcond=rcond, tot_ener_zero=tot_ener_zero, trainable=trainable, - atom_ener=atom_ener, activation_function=activation_function, precision=precision, layer_name=layer_name, @@ -163,6 +163,7 @@ def __init__( def serialize(self) -> dict: data = super().serialize() data["dim_out"] = self.dim_out + data["atom_ener"] = [] return data def _net_out_dim(self): diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index ed2dfbc02b..d0acc6fe2b 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -78,6 +78,8 @@ class InvarFitting(GeneralFitting): Random seed. exclude_types: List[int] Atomic contributions of the excluded atom types are set zero. + atom_ener + Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. """ @@ -98,9 +100,11 @@ def __init__( rcond: Optional[float] = None, seed: Optional[int] = None, exclude_types: List[int] = [], + atom_ener: Optional[List[float]] = None, **kwargs, ): self.dim_out = dim_out + self.atom_ener = atom_ener super().__init__( var_name=var_name, ntypes=ntypes, @@ -126,6 +130,7 @@ def _net_out_dim(self): def serialize(self) -> dict: data = super().serialize() data["dim_out"] = self.dim_out + data["atom_ener"] = self.atom_ener return data @property diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index be20aa9496..080bfb5172 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -428,7 +428,6 @@ def serialize(self) -> dict: ## NOTICE: not supported by far "tot_ener_zero": False, "trainable": [True] * (len(self.neuron) + 1), - "atom_ener": [], "layer_name": None, "use_aparam_as_mask": False, "spin": None, diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index b73d36835c..17c7229429 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -252,7 +252,7 @@ def test_tf_consistent_with_ref(self): tf_obj = self.tf_class.deserialize(data1, suffix=self.unique_id) ret2, data2 = self.get_tf_ret_serialization_from_cls(tf_obj) ret2 = self.extract_ret(ret2, self.RefBackend.TF) - if not tf_obj.__class__.__name__.startswith(("Dipole", "Polar")): + if not tf_obj.__class__.__name__.startswith(("Polar")): np.testing.assert_equal(data1, data2) # tf, pt serialization mismatch for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose( @@ -318,7 +318,7 @@ def test_pt_consistent_with_ref(self): ret2 = self.eval_pt(obj) ret2 = self.extract_ret(ret2, self.RefBackend.PT) data2 = obj.serialize() - if not obj.__class__.__name__.startswith(("Dipole", "Polar")): + if not obj.__class__.__name__.startswith(("Polar")): np.testing.assert_equal(data1, data2) # tf, pt serialization mismatch for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) From 13df1db7f2b1349f0f177167d6ff524e5e35d7ae Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Feb 2024 08:41:01 +0000 Subject: [PATCH 19/36] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/fitting/invar_fitting.py | 2 +- source/tests/consistent/common.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/fitting/invar_fitting.py b/deepmd/dpmodel/fitting/invar_fitting.py index e3ec38e5a0..2eed0f4e39 100644 --- a/deepmd/dpmodel/fitting/invar_fitting.py +++ b/deepmd/dpmodel/fitting/invar_fitting.py @@ -139,7 +139,7 @@ def __init__( raise NotImplementedError("atom_ener is not implemented") self.dim_out = dim_out - self.atom_ener=atom_ener, + self.atom_ener = (atom_ener,) super().__init__( var_name=var_name, ntypes=ntypes, diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index 17c7229429..8088497e99 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -252,7 +252,7 @@ def test_tf_consistent_with_ref(self): tf_obj = self.tf_class.deserialize(data1, suffix=self.unique_id) ret2, data2 = self.get_tf_ret_serialization_from_cls(tf_obj) ret2 = self.extract_ret(ret2, self.RefBackend.TF) - if not tf_obj.__class__.__name__.startswith(("Polar")): + if not tf_obj.__class__.__name__.startswith("Polar"): np.testing.assert_equal(data1, data2) # tf, pt serialization mismatch for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose( @@ -318,7 +318,7 @@ def test_pt_consistent_with_ref(self): ret2 = self.eval_pt(obj) ret2 = self.extract_ret(ret2, self.RefBackend.PT) data2 = obj.serialize() - if not obj.__class__.__name__.startswith(("Polar")): + if not obj.__class__.__name__.startswith("Polar"): np.testing.assert_equal(data1, data2) # tf, pt serialization mismatch for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) From 6052f7ebaac8a971fa6574b191f3a2d631d75011 Mon Sep 17 00:00:00 2001 From: anyangml Date: Thu, 22 Feb 2024 09:49:29 +0000 Subject: [PATCH 20/36] fix: UTs --- deepmd/tf/fit/dipole.py | 3 ++- deepmd/tf/fit/fitting.py | 6 +++++- source/tests/consistent/common.py | 16 ++++++++++++---- 3 files changed, 19 insertions(+), 6 deletions(-) diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index be4e669e6f..550305ca2f 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -344,7 +344,7 @@ def serialize(self, suffix: str) -> dict: The serialized data """ data = { - "var_name": "energy", + "var_name": "dipole", "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, "embedding_width": self.dim_rot_mat_1, @@ -362,6 +362,7 @@ def serialize(self, suffix: str) -> dict: # TODO: consider type embeddings ndim=1, in_dim=self.dim_descrpt, + out_dim = self.dim_rot_mat_1, neuron=self.n_neuron, activation_function=self.activation_function_name, resnet_dt=self.resnet_dt, diff --git a/deepmd/tf/fit/fitting.py b/deepmd/tf/fit/fitting.py index 458765f7c1..c3ecc1ac35 100644 --- a/deepmd/tf/fit/fitting.py +++ b/deepmd/tf/fit/fitting.py @@ -7,6 +7,7 @@ Callable, List, Type, + Optional, ) from deepmd.dpmodel.utils.network import ( @@ -175,6 +176,7 @@ def serialize_network( activation_function: str, resnet_dt: bool, variables: dict, + out_dim: Optional[int] = 1, suffix: str = "", ) -> dict: """Serialize network. @@ -197,6 +199,8 @@ def serialize_network( The input variables suffix : str, optional The suffix of the scope + out_dim : int, optional + The output dimension Returns ------- @@ -231,7 +235,7 @@ def serialize_network( # initialize the network if it is not initialized fittings[network_idx] = FittingNet( in_dim=in_dim, - out_dim=1, + out_dim=out_dim, neuron=neuron, activation_function=activation_function, resnet_dt=resnet_dt, diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index 8088497e99..9652922e1d 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -252,8 +252,12 @@ def test_tf_consistent_with_ref(self): tf_obj = self.tf_class.deserialize(data1, suffix=self.unique_id) ret2, data2 = self.get_tf_ret_serialization_from_cls(tf_obj) ret2 = self.extract_ret(ret2, self.RefBackend.TF) - if not tf_obj.__class__.__name__.startswith("Polar"): - np.testing.assert_equal(data1, data2) # tf, pt serialization mismatch + if tf_obj.__class__.__name__.startswith("Polar", "Dipole"): + # tf, pt serialization mismatch + common_keys = set(data1.keys()) & set(data2.keys()) + data1 = {k: data1[k] for k in common_keys} + data2 = {k: data2[k] for k in common_keys} + np.testing.assert_equal(data1, data2) for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose( rr1.ravel(), rr2.ravel(), rtol=self.rtol, atol=self.atol @@ -318,8 +322,12 @@ def test_pt_consistent_with_ref(self): ret2 = self.eval_pt(obj) ret2 = self.extract_ret(ret2, self.RefBackend.PT) data2 = obj.serialize() - if not obj.__class__.__name__.startswith("Polar"): - np.testing.assert_equal(data1, data2) # tf, pt serialization mismatch + if obj.__class__.__name__.startswith("Polar", "Dipole"): + # tf, pt serialization mismatch + common_keys = set(data1.keys()) & set(data2.keys()) + data1 = {k: data1[k] for k in common_keys} + data2 = {k: data2[k] for k in common_keys} + np.testing.assert_equal(data1, data2) for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) From 511e15efb5ea02dd0918be1c04ca34b1d5c5d33e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Feb 2024 09:51:10 +0000 Subject: [PATCH 21/36] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/fit/dipole.py | 2 +- deepmd/tf/fit/fitting.py | 2 +- source/tests/consistent/common.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index 550305ca2f..e808d407d4 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -362,7 +362,7 @@ def serialize(self, suffix: str) -> dict: # TODO: consider type embeddings ndim=1, in_dim=self.dim_descrpt, - out_dim = self.dim_rot_mat_1, + out_dim=self.dim_rot_mat_1, neuron=self.n_neuron, activation_function=self.activation_function_name, resnet_dt=self.resnet_dt, diff --git a/deepmd/tf/fit/fitting.py b/deepmd/tf/fit/fitting.py index c3ecc1ac35..598d020fe9 100644 --- a/deepmd/tf/fit/fitting.py +++ b/deepmd/tf/fit/fitting.py @@ -6,8 +6,8 @@ from typing import ( Callable, List, - Type, Optional, + Type, ) from deepmd.dpmodel.utils.network import ( diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index 9652922e1d..566f3bcb4e 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -257,7 +257,7 @@ def test_tf_consistent_with_ref(self): common_keys = set(data1.keys()) & set(data2.keys()) data1 = {k: data1[k] for k in common_keys} data2 = {k: data2[k] for k in common_keys} - np.testing.assert_equal(data1, data2) + np.testing.assert_equal(data1, data2) for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose( rr1.ravel(), rr2.ravel(), rtol=self.rtol, atol=self.atol @@ -327,7 +327,7 @@ def test_pt_consistent_with_ref(self): common_keys = set(data1.keys()) & set(data2.keys()) data1 = {k: data1[k] for k in common_keys} data2 = {k: data2[k] for k in common_keys} - np.testing.assert_equal(data1, data2) + np.testing.assert_equal(data1, data2) for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose(rr1, rr2, rtol=self.rtol, atol=self.atol) From 15b742582290b96e3e7a46eef2eb2ccb6cc80d2e Mon Sep 17 00:00:00 2001 From: anyangml Date: Thu, 22 Feb 2024 10:21:21 +0000 Subject: [PATCH 22/36] fix: UTs --- deepmd/dpmodel/fitting/polarizability_fitting.py | 6 ------ source/tests/consistent/common.py | 4 ++-- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index c3cbe7bd1a..9811e8e1c8 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/deepmd/dpmodel/fitting/polarizability_fitting.py @@ -56,8 +56,6 @@ class PolarFitting(GeneralFitting): If the weights of fitting net are trainable. Suppose that we have :math:`N_l` hidden layers in the fitting net, this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable. - atom_ener - Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. activation_function The activation function :math:`\boldsymbol{\phi}` in the embedding net. Supported options are |ACTIVATION_FN| precision @@ -93,7 +91,6 @@ def __init__( rcond: Optional[float] = None, tot_ener_zero: bool = False, trainable: Optional[List[bool]] = None, - atom_ener: Optional[List[Optional[float]]] = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, layer_name: Optional[List[Optional[str]]] = None, @@ -115,8 +112,6 @@ def __init__( raise NotImplementedError("use_aparam_as_mask is not implemented") if layer_name is not None: raise NotImplementedError("layer_name is not implemented") - if atom_ener is not None and atom_ener != []: - raise NotImplementedError("atom_ener is not implemented") self.embedding_width = embedding_width self.fit_diag = fit_diag @@ -142,7 +137,6 @@ def __init__( rcond=rcond, tot_ener_zero=tot_ener_zero, trainable=trainable, - atom_ener=atom_ener, activation_function=activation_function, precision=precision, layer_name=layer_name, diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index 566f3bcb4e..b8b671adf8 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -252,7 +252,7 @@ def test_tf_consistent_with_ref(self): tf_obj = self.tf_class.deserialize(data1, suffix=self.unique_id) ret2, data2 = self.get_tf_ret_serialization_from_cls(tf_obj) ret2 = self.extract_ret(ret2, self.RefBackend.TF) - if tf_obj.__class__.__name__.startswith("Polar", "Dipole"): + if tf_obj.__class__.__name__.startswith(("Polar", "Dipole")): # tf, pt serialization mismatch common_keys = set(data1.keys()) & set(data2.keys()) data1 = {k: data1[k] for k in common_keys} @@ -322,7 +322,7 @@ def test_pt_consistent_with_ref(self): ret2 = self.eval_pt(obj) ret2 = self.extract_ret(ret2, self.RefBackend.PT) data2 = obj.serialize() - if obj.__class__.__name__.startswith("Polar", "Dipole"): + if obj.__class__.__name__.startswith(("Polar","Dipole")): # tf, pt serialization mismatch common_keys = set(data1.keys()) & set(data2.keys()) data1 = {k: data1[k] for k in common_keys} From 737cf376f74cd1a17b9bbbcbf429994adec68fc9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 Feb 2024 10:21:53 +0000 Subject: [PATCH 23/36] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/consistent/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index b8b671adf8..759b78dd98 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -322,7 +322,7 @@ def test_pt_consistent_with_ref(self): ret2 = self.eval_pt(obj) ret2 = self.extract_ret(ret2, self.RefBackend.PT) data2 = obj.serialize() - if obj.__class__.__name__.startswith(("Polar","Dipole")): + if obj.__class__.__name__.startswith(("Polar", "Dipole")): # tf, pt serialization mismatch common_keys = set(data1.keys()) & set(data2.keys()) data1 = {k: data1[k] for k in common_keys} From 186a15ec2e6ef4d0203f51ae18a52881110d72bb Mon Sep 17 00:00:00 2001 From: anyangml Date: Fri, 23 Feb 2024 01:37:50 +0000 Subject: [PATCH 24/36] fix: UTs --- deepmd/dpmodel/fitting/invar_fitting.py | 4 ++-- deepmd/tf/model/model.py | 1 + deepmd/tf/model/multi.py | 1 + source/tests/tf/data_modifier/dipole.json | 1 - source/tests/tf/test_data_modifier.py | 2 +- source/tests/tf/test_data_modifier_shuffle.py | 3 +-- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/deepmd/dpmodel/fitting/invar_fitting.py b/deepmd/dpmodel/fitting/invar_fitting.py index 2eed0f4e39..3d958301ff 100644 --- a/deepmd/dpmodel/fitting/invar_fitting.py +++ b/deepmd/dpmodel/fitting/invar_fitting.py @@ -139,7 +139,7 @@ def __init__( raise NotImplementedError("atom_ener is not implemented") self.dim_out = dim_out - self.atom_ener = (atom_ener,) + self.atom_ener = atom_ener super().__init__( var_name=var_name, ntypes=ntypes, @@ -163,7 +163,7 @@ def __init__( def serialize(self) -> dict: data = super().serialize() data["dim_out"] = self.dim_out - data["atom_ener"] = [] + data["atom_ener"] = self.atom_ener return data def _net_out_dim(self): diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 7ca2b6f4ab..ab7cc5b0b1 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -693,6 +693,7 @@ def __init__( spin=self.spin, ntypes=self.descrpt.get_ntypes(), dim_descrpt=self.descrpt.get_dim_out(), + embedding_width=self.descrpt.get_dim_rot_mat_1(), ) self.rcut = self.descrpt.get_rcut() self.ntypes = self.descrpt.get_ntypes() diff --git a/deepmd/tf/model/multi.py b/deepmd/tf/model/multi.py index 833a700ebc..1ed74e35b1 100644 --- a/deepmd/tf/model/multi.py +++ b/deepmd/tf/model/multi.py @@ -139,6 +139,7 @@ def __init__( spin=self.spin, ntypes=self.descrpt.get_ntypes(), dim_descrpt=self.descrpt.get_dim_out(), + embedding_width=self.descrpt.get_dim_rot_mat_1(), ) # type embedding diff --git a/source/tests/tf/data_modifier/dipole.json b/source/tests/tf/data_modifier/dipole.json index 58c6cc093b..ac9cec0570 100644 --- a/source/tests/tf/data_modifier/dipole.json +++ b/source/tests/tf/data_modifier/dipole.json @@ -23,7 +23,6 @@ "seed": 1 }, "fitting_net": { - "embedding_width": 100, "type": "dipole", "sel_type": [ 0 diff --git a/source/tests/tf/test_data_modifier.py b/source/tests/tf/test_data_modifier.py index 98b6c41427..bb2af7179d 100644 --- a/source/tests/tf/test_data_modifier.py +++ b/source/tests/tf/test_data_modifier.py @@ -57,7 +57,7 @@ def _setUp(self): restart=None, init_model=None, log_path=None, log_level=30, mpi_log="master" ) jdata = j_loader(INPUT) - + jdata["model"]["fitting_net"]["embedding_width"] = jdata["model"]["descriptor"]["neuron"][-1] # init model model = DPTrainer(jdata, run_opt=run_opt) rcut = model.model.get_rcut() diff --git a/source/tests/tf/test_data_modifier_shuffle.py b/source/tests/tf/test_data_modifier_shuffle.py index 1a16f6935f..48b19672a3 100644 --- a/source/tests/tf/test_data_modifier_shuffle.py +++ b/source/tests/tf/test_data_modifier_shuffle.py @@ -59,7 +59,7 @@ def _setUp(self): ) jdata = self._setUp_jdata() self._setUp_data() - + jdata["model"]["fitting_net"]["embedding_width"] = jdata["model"]["descriptor"]["neuron"][-1] # init model model = DPTrainer(jdata, run_opt=run_opt) rcut = model.model.get_rcut() @@ -153,7 +153,6 @@ def _setUp_jdata(self): }, "fitting_net": { "type": "dipole", - "embedding_width": 4, "sel_type": [1, 3], "neuron": [10], "resnet_dt": True, From 0a84f3369491c571dc5712cf760f6790234b1e22 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 23 Feb 2024 01:38:23 +0000 Subject: [PATCH 25/36] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/tf/test_data_modifier.py | 4 +++- source/tests/tf/test_data_modifier_shuffle.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/source/tests/tf/test_data_modifier.py b/source/tests/tf/test_data_modifier.py index bb2af7179d..d44deff3d9 100644 --- a/source/tests/tf/test_data_modifier.py +++ b/source/tests/tf/test_data_modifier.py @@ -57,7 +57,9 @@ def _setUp(self): restart=None, init_model=None, log_path=None, log_level=30, mpi_log="master" ) jdata = j_loader(INPUT) - jdata["model"]["fitting_net"]["embedding_width"] = jdata["model"]["descriptor"]["neuron"][-1] + jdata["model"]["fitting_net"]["embedding_width"] = jdata["model"]["descriptor"][ + "neuron" + ][-1] # init model model = DPTrainer(jdata, run_opt=run_opt) rcut = model.model.get_rcut() diff --git a/source/tests/tf/test_data_modifier_shuffle.py b/source/tests/tf/test_data_modifier_shuffle.py index 48b19672a3..dc603998cc 100644 --- a/source/tests/tf/test_data_modifier_shuffle.py +++ b/source/tests/tf/test_data_modifier_shuffle.py @@ -59,7 +59,9 @@ def _setUp(self): ) jdata = self._setUp_jdata() self._setUp_data() - jdata["model"]["fitting_net"]["embedding_width"] = jdata["model"]["descriptor"]["neuron"][-1] + jdata["model"]["fitting_net"]["embedding_width"] = jdata["model"]["descriptor"][ + "neuron" + ][-1] # init model model = DPTrainer(jdata, run_opt=run_opt) rcut = model.model.get_rcut() From abad3dd76b55ef55f8a64ccbb6ba1b9eb319ec7e Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Fri, 23 Feb 2024 09:42:45 +0800 Subject: [PATCH 26/36] chore: refactor --- deepmd/tf/fit/dipole.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index e808d407d4..efb94a85c0 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -350,7 +350,7 @@ def serialize(self, suffix: str) -> dict: "embedding_width": self.dim_rot_mat_1, # very bad design: type embedding is not passed to the class # TODO: refactor the class - "distinguish_types": True, + "mixed_types": False, "dim_out": 3, "neuron": self.n_neuron, "resnet_dt": self.resnet_dt, From 3e0f3d8bf0db8ecf2a2887dcb273a93896045fe8 Mon Sep 17 00:00:00 2001 From: anyangml Date: Fri, 23 Feb 2024 02:16:48 +0000 Subject: [PATCH 27/36] fix: se_r, se_t --- deepmd/tf/descriptor/se_r.py | 4 ++++ deepmd/tf/descriptor/se_t.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py index f790d0a8fb..65eb490e55 100644 --- a/deepmd/tf/descriptor/se_r.py +++ b/deepmd/tf/descriptor/se_r.py @@ -203,6 +203,10 @@ def get_ntypes(self): def get_dim_out(self): """Returns the output dimension of this descriptor.""" return self.filter_neuron[-1] + + def get_dim_rot_mat_1(self) -> int: + """To accommodate dipole fitting, which needs embeeding_width.""" + return self.get_dim_out() def get_nlist(self): """Returns neighbor information. diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index c72296daa7..f2d40073c1 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -193,6 +193,10 @@ def get_ntypes(self) -> int: def get_dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.filter_neuron[-1] + + def get_dim_rot_mat_1(self) -> int: + """To accommodate dipole fitting, which needs embeeding_width.""" + return self.get_dim_out() def get_nlist(self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]: """Returns neighbor information. From d2cbe468ce8bb5715364a860601ce4d3bd8d4d5d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 23 Feb 2024 02:17:58 +0000 Subject: [PATCH 28/36] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/descriptor/se_r.py | 2 +- deepmd/tf/descriptor/se_t.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py index 65eb490e55..fd36c65a66 100644 --- a/deepmd/tf/descriptor/se_r.py +++ b/deepmd/tf/descriptor/se_r.py @@ -203,7 +203,7 @@ def get_ntypes(self): def get_dim_out(self): """Returns the output dimension of this descriptor.""" return self.filter_neuron[-1] - + def get_dim_rot_mat_1(self) -> int: """To accommodate dipole fitting, which needs embeeding_width.""" return self.get_dim_out() diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index f2d40073c1..ad5f2562ac 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -193,7 +193,7 @@ def get_ntypes(self) -> int: def get_dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.filter_neuron[-1] - + def get_dim_rot_mat_1(self) -> int: """To accommodate dipole fitting, which needs embeeding_width.""" return self.get_dim_out() From 8bbaf253baf0b08e50bb017eb9f84fc6547d12df Mon Sep 17 00:00:00 2001 From: anyangml Date: Fri, 23 Feb 2024 02:20:45 +0000 Subject: [PATCH 29/36] chore: refactor --- source/tests/tf/test_data_modifier.py | 3 --- source/tests/tf/test_data_modifier_shuffle.py | 4 +--- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/source/tests/tf/test_data_modifier.py b/source/tests/tf/test_data_modifier.py index d44deff3d9..cf2c50b761 100644 --- a/source/tests/tf/test_data_modifier.py +++ b/source/tests/tf/test_data_modifier.py @@ -57,9 +57,6 @@ def _setUp(self): restart=None, init_model=None, log_path=None, log_level=30, mpi_log="master" ) jdata = j_loader(INPUT) - jdata["model"]["fitting_net"]["embedding_width"] = jdata["model"]["descriptor"][ - "neuron" - ][-1] # init model model = DPTrainer(jdata, run_opt=run_opt) rcut = model.model.get_rcut() diff --git a/source/tests/tf/test_data_modifier_shuffle.py b/source/tests/tf/test_data_modifier_shuffle.py index dc603998cc..43941f3e55 100644 --- a/source/tests/tf/test_data_modifier_shuffle.py +++ b/source/tests/tf/test_data_modifier_shuffle.py @@ -59,9 +59,7 @@ def _setUp(self): ) jdata = self._setUp_jdata() self._setUp_data() - jdata["model"]["fitting_net"]["embedding_width"] = jdata["model"]["descriptor"][ - "neuron" - ][-1] + # init model model = DPTrainer(jdata, run_opt=run_opt) rcut = model.model.get_rcut() From 88d81a08e5c5176d7b269960916b0847f887294b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 23 Feb 2024 02:22:43 +0000 Subject: [PATCH 30/36] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/tf/test_data_modifier_shuffle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/tests/tf/test_data_modifier_shuffle.py b/source/tests/tf/test_data_modifier_shuffle.py index 43941f3e55..fb9da948f1 100644 --- a/source/tests/tf/test_data_modifier_shuffle.py +++ b/source/tests/tf/test_data_modifier_shuffle.py @@ -59,7 +59,7 @@ def _setUp(self): ) jdata = self._setUp_jdata() self._setUp_data() - + # init model model = DPTrainer(jdata, run_opt=run_opt) rcut = model.model.get_rcut() From 4c20c103a720148394006cbfe613764d3c477bda Mon Sep 17 00:00:00 2001 From: anyangml Date: Fri, 23 Feb 2024 02:30:57 +0000 Subject: [PATCH 31/36] chore: typo --- deepmd/tf/descriptor/se_r.py | 2 +- deepmd/tf/descriptor/se_t.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py index fd36c65a66..403eaff5eb 100644 --- a/deepmd/tf/descriptor/se_r.py +++ b/deepmd/tf/descriptor/se_r.py @@ -205,7 +205,7 @@ def get_dim_out(self): return self.filter_neuron[-1] def get_dim_rot_mat_1(self) -> int: - """To accommodate dipole fitting, which needs embeeding_width.""" + """To accommodate dipole fitting, which needs embedding_width.""" return self.get_dim_out() def get_nlist(self): diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index ad5f2562ac..fec0e6d9ba 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -195,7 +195,7 @@ def get_dim_out(self) -> int: return self.filter_neuron[-1] def get_dim_rot_mat_1(self) -> int: - """To accommodate dipole fitting, which needs embeeding_width.""" + """To accommodate dipole fitting, which needs embedding_width.""" return self.get_dim_out() def get_nlist(self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]: From 3e2d1f9ccb91bd469b9fbc7bd1bddba24e0c09ed Mon Sep 17 00:00:00 2001 From: anyangml Date: Fri, 23 Feb 2024 03:45:59 +0000 Subject: [PATCH 32/36] chore: refactor --- deepmd/tf/descriptor/se_r.py | 4 ---- deepmd/tf/descriptor/se_t.py | 4 ---- deepmd/tf/model/model.py | 3 ++- deepmd/tf/model/multi.py | 3 ++- 4 files changed, 4 insertions(+), 10 deletions(-) diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py index 403eaff5eb..f790d0a8fb 100644 --- a/deepmd/tf/descriptor/se_r.py +++ b/deepmd/tf/descriptor/se_r.py @@ -204,10 +204,6 @@ def get_dim_out(self): """Returns the output dimension of this descriptor.""" return self.filter_neuron[-1] - def get_dim_rot_mat_1(self) -> int: - """To accommodate dipole fitting, which needs embedding_width.""" - return self.get_dim_out() - def get_nlist(self): """Returns neighbor information. diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index fec0e6d9ba..c72296daa7 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -194,10 +194,6 @@ def get_dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.filter_neuron[-1] - def get_dim_rot_mat_1(self) -> int: - """To accommodate dipole fitting, which needs embedding_width.""" - return self.get_dim_out() - def get_nlist(self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]: """Returns neighbor information. diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index ab7cc5b0b1..b8fc39a918 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -687,13 +687,14 @@ def __init__( if isinstance(fitting_net, Fitting): self.fitting = fitting_net else: + if fitting_net["type"] in ["dipole", "polar"]: + fitting_net["embedding_width"] = self.descrpt.get_dim_rot_mat_1() self.fitting = Fitting( **fitting_net, descrpt=self.descrpt, spin=self.spin, ntypes=self.descrpt.get_ntypes(), dim_descrpt=self.descrpt.get_dim_out(), - embedding_width=self.descrpt.get_dim_rot_mat_1(), ) self.rcut = self.descrpt.get_rcut() self.ntypes = self.descrpt.get_ntypes() diff --git a/deepmd/tf/model/multi.py b/deepmd/tf/model/multi.py index 1ed74e35b1..278af4d509 100644 --- a/deepmd/tf/model/multi.py +++ b/deepmd/tf/model/multi.py @@ -133,13 +133,14 @@ def __init__( if isinstance(item_fitting_param, Fitting): fitting_dict[item] = item_fitting_param else: + if fitting_net["type"] in ["dipole", "polar"]: + fitting_net["embedding_width"] = self.descrpt.get_dim_rot_mat_1() fitting_dict[item] = Fitting( **item_fitting_param, descrpt=self.descrpt, spin=self.spin, ntypes=self.descrpt.get_ntypes(), dim_descrpt=self.descrpt.get_dim_out(), - embedding_width=self.descrpt.get_dim_rot_mat_1(), ) # type embedding From 9762912dcb1ab276452d52bd73d7592f7b454f9b Mon Sep 17 00:00:00 2001 From: anyangml Date: Fri, 23 Feb 2024 03:47:33 +0000 Subject: [PATCH 33/36] chore: typo --- deepmd/tf/model/multi.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/tf/model/multi.py b/deepmd/tf/model/multi.py index 278af4d509..f41bae0d73 100644 --- a/deepmd/tf/model/multi.py +++ b/deepmd/tf/model/multi.py @@ -133,8 +133,8 @@ def __init__( if isinstance(item_fitting_param, Fitting): fitting_dict[item] = item_fitting_param else: - if fitting_net["type"] in ["dipole", "polar"]: - fitting_net["embedding_width"] = self.descrpt.get_dim_rot_mat_1() + if item_fitting_param["type"] in ["dipole", "polar"]: + item_fitting_param["embedding_width"] = self.descrpt.get_dim_rot_mat_1() fitting_dict[item] = Fitting( **item_fitting_param, descrpt=self.descrpt, From 344910a61f6025e5e016b0e8f6d2336c57f4f99f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 23 Feb 2024 03:48:03 +0000 Subject: [PATCH 34/36] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/tf/model/multi.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/deepmd/tf/model/multi.py b/deepmd/tf/model/multi.py index f41bae0d73..2acf00fd52 100644 --- a/deepmd/tf/model/multi.py +++ b/deepmd/tf/model/multi.py @@ -134,7 +134,9 @@ def __init__( fitting_dict[item] = item_fitting_param else: if item_fitting_param["type"] in ["dipole", "polar"]: - item_fitting_param["embedding_width"] = self.descrpt.get_dim_rot_mat_1() + item_fitting_param[ + "embedding_width" + ] = self.descrpt.get_dim_rot_mat_1() fitting_dict[item] = Fitting( **item_fitting_param, descrpt=self.descrpt, From ea621ea2befc41051f3e475c7e929a8d0c50cfe4 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Mon, 26 Feb 2024 10:28:27 +0800 Subject: [PATCH 35/36] feat: add se_r descriptor --- deepmd/pt/model/descriptor/__init__.py | 2 + deepmd/pt/model/descriptor/env_mat.py | 50 +++- deepmd/pt/model/descriptor/se_r.py | 400 +++++++++++++++++++++++++ deepmd/pt/utils/env_mat_stat.py | 136 +++++++++ 4 files changed, 587 insertions(+), 1 deletion(-) create mode 100644 deepmd/pt/model/descriptor/se_r.py diff --git a/deepmd/pt/model/descriptor/__init__.py b/deepmd/pt/model/descriptor/__init__.py index 1c2e943369..0cdb35c55f 100644 --- a/deepmd/pt/model/descriptor/__init__.py +++ b/deepmd/pt/model/descriptor/__init__.py @@ -13,6 +13,7 @@ ) from .env_mat import ( prod_env_mat_se_a, + prod_env_mat_se_r, ) from .gaussian_lcc import ( DescrptGaussianLcc, @@ -38,6 +39,7 @@ "DescrptDPA1", "DescrptDPA2", "prod_env_mat_se_a", + "prod_env_mat_se_r", "DescrptGaussianLcc", "DescrptBlockHybrid", "DescrptBlockRepformers", diff --git a/deepmd/pt/model/descriptor/env_mat.py b/deepmd/pt/model/descriptor/env_mat.py index b3235de175..7fe5f29354 100644 --- a/deepmd/pt/model/descriptor/env_mat.py +++ b/deepmd/pt/model/descriptor/env_mat.py @@ -29,6 +29,27 @@ def _make_env_mat_se_a(nlist, coord, rcut: float, ruct_smth: float): env_mat_se_a = torch.cat([t0, t1], dim=-1) * weight return env_mat_se_a, diff * mask.unsqueeze(-1), weight +def _make_env_mat_se_r(nlist, coord, rcut: float, ruct_smth: float): + """Make smooth environment matrix.""" + bsz, natoms, nnei = nlist.shape + coord = coord.view(bsz, -1, 3) + nall = coord.shape[1] + mask = nlist >= 0 + # nlist = nlist * mask ## this impl will contribute nans in Hessian calculation. + nlist = torch.where(mask, nlist, nall - 1) + coord_l = coord[:, :natoms].view(bsz, -1, 1, 3) + index = nlist.view(bsz, -1).unsqueeze(-1).expand(-1, -1, 3) + coord_r = torch.gather(coord, 1, index) + coord_r = coord_r.view(bsz, natoms, nnei, 3) + diff = coord_r - coord_l + length = torch.linalg.norm(diff, dim=-1, keepdim=True) + # for index 0 nloc atom + length = length + ~mask.unsqueeze(-1) + t0 = 1 / length + weight = compute_smooth_weight(length, ruct_smth, rcut) + weight = weight * mask.unsqueeze(-1) + env_mat_se_r = t0 * weight + return env_mat_se_r, diff * mask.unsqueeze(-1), weight def prod_env_mat_se_a( extended_coord, nlist, atype, mean, stddev, rcut: float, rcut_smth: float @@ -50,7 +71,6 @@ def prod_env_mat_se_a( ------- - env_mat_se_a: Shape is [nframes, natoms[1]*nnei*4]. """ - nframes = extended_coord.shape[0] _env_mat_se_a, diff, switch = _make_env_mat_se_a( nlist, extended_coord, rcut, rcut_smth ) # shape [n_atom, dim, 4] @@ -58,3 +78,31 @@ def prod_env_mat_se_a( t_std = stddev[atype] # [n_atom, dim, 4] env_mat_se_a = (_env_mat_se_a - t_avg) / t_std return env_mat_se_a, diff, switch + +def prod_env_mat_se_r( + extended_coord, nlist, atype, mean, stddev, rcut: float, rcut_smth: float +): + """Generate smooth environment matrix from atom coordinates and other context. + + Args: + - extended_coord: Copied atom coordinates with shape [nframes, nall*3]. + - atype: Atom types with shape [nframes, nloc]. + - natoms: Batched atom statisics with shape [len(sec)+2]. + - box: Batched simulation box with shape [nframes, 9]. + - mean: Average value of descriptor per element type with shape [len(sec), nnei, 4]. + - stddev: Standard deviation of descriptor per element type with shape [len(sec), nnei, 4]. + - deriv_stddev: StdDev of descriptor derivative per element type with shape [len(sec), nnei, 4, 3]. + - rcut: Cut-off radius. + - rcut_smth: Smooth hyper-parameter for pair force & energy. + + Returns + ------- + - env_mat_se_a: Shape is [nframes, natoms[1]*nnei*4]. + """ + _env_mat_se_a, diff, switch = _make_env_mat_se_r( + nlist, extended_coord, rcut, rcut_smth + ) # shape [n_atom, dim, 1] + t_avg = mean[atype] # [n_atom, dim, 1] + t_std = stddev[atype] # [n_atom, dim, 1] + env_mat_se_a = (_env_mat_se_a - t_avg) / t_std + return env_mat_se_a, diff, switch \ No newline at end of file diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py new file mode 100644 index 0000000000..bb1c4e5ddd --- /dev/null +++ b/deepmd/pt/model/descriptor/se_r.py @@ -0,0 +1,400 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + ClassVar, + Dict, + List, + Optional, + Tuple, +) + +import numpy as np +import torch + +from deepmd.pt.model.descriptor import ( + Descriptor, + DescriptorBlock, + prod_env_mat_se_r, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + PRECISION_DICT, + RESERVED_PRECISON_DICT, +) +from deepmd.pt.utils.env_mat_stat import ( + EnvMatStatSeR, +) +from deepmd.utils.env_mat_stat import ( + StatItem, +) +from deepmd.utils.path import ( + DPPath, +) + +try: + from typing import ( + Final, + ) +except ImportError: + from torch.jit import Final + +from deepmd.dpmodel.utils import EnvMat as DPEnvMat +from deepmd.pt.model.network.mlp import ( + EmbeddingNet, + NetworkCollection, +) +from deepmd.pt.model.network.network import ( + TypeFilter, +) +from deepmd.pt.utils.exclude_mask import ( + PairExcludeMask, +) + + +@Descriptor.register("se_e2_r") +class DescrptSeR(Descriptor): + def __init__( + self, + rcut, + rcut_smth, + sel, + neuron=[25, 50, 100], + set_davg_zero: bool = False, + activation_function: str = "tanh", + precision: str = "float64", + resnet_dt: bool = False, + exclude_types: List[Tuple[int, int]] = [], + old_impl: bool = False, + **kwargs, + ): + super().__init__() + self.rcut = rcut + self.rcut_smth = rcut_smth + self.neuron = neuron + self.filter_neuron = self.neuron + self.set_davg_zero = set_davg_zero + self.activation_function = activation_function + self.precision = precision + self.prec = PRECISION_DICT[self.precision] + self.resnet_dt = resnet_dt + self.old_impl = old_impl + self.exclude_types = exclude_types + self.ntypes = len(sel) + self.emask = PairExcludeMask(len(sel), exclude_types=exclude_types) + + self.sel = sel + self.sec = torch.tensor( + np.append([0], np.cumsum(self.sel)), dtype=int, device=env.DEVICE + ) + self.split_sel = self.sel + self.nnei = sum(sel) + self.ndescrpt = self.nnei * 4 + + wanted_shape = (self.ntypes, self.nnei, 4) + mean = torch.zeros(wanted_shape, dtype=self.prec, device=env.DEVICE) + stddev = torch.ones(wanted_shape, dtype=self.prec, device=env.DEVICE) + self.register_buffer("mean", mean) + self.register_buffer("stddev", stddev) + self.filter_layers_old = None + self.filter_layers = None + + if self.old_impl: + filter_layers = [] + # TODO: remove + start_index = 0 + for type_i in range(self.ntypes): + one = TypeFilter(start_index, sel[type_i], self.filter_neuron) + filter_layers.append(one) + start_index += sel[type_i] + self.filter_layers_old = torch.nn.ModuleList(filter_layers) + else: + filter_layers = NetworkCollection( + ndim=1, ntypes=len(sel), network_type="embedding_network" + ) + # TODO: ndim=2 if type_one_side=False + for ii in range(self.ntypes): + filter_layers[(ii,)] = EmbeddingNet( + 1, + self.filter_neuron, + activation_function=self.activation_function, + precision=self.precision, + resnet_dt=self.resnet_dt, + ) + self.filter_layers = filter_layers + self.stats = None + + def get_rcut(self) -> float: + """Returns the cut-off radius.""" + return self.rcut + + def get_nsel(self) -> int: + """Returns the number of selected atoms in the cut-off radius.""" + return sum(self.sel) + + def get_sel(self) -> List[int]: + """Returns the number of selected atoms for each type.""" + return self.sel + + def get_ntypes(self) -> int: + """Returns the number of element types.""" + return self.ntypes + + def get_dim_out(self) -> int: + """Returns the output dimension.""" + return self.neuron[-1] + + def get_dim_emb(self) -> int: + """Returns the output dimension.""" + raise NotImplementedError + + def get_dim_in(self) -> int: + """Returns the input dimension.""" + return 0 + + def mixed_types(self) -> bool: + """If true, the discriptor + 1. assumes total number of atoms aligned across frames; + 2. requires a neighbor list that does not distinguish different atomic types. + + If false, the discriptor + 1. assumes total number of atoms of each atom type aligned across frames; + 2. requires a neighbor list that distinguishes different atomic types. + + """ + return False + + def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None): + """Update mean and stddev for descriptor elements.""" + env_mat_stat = EnvMatStatSeR(self) + if path is not None: + path = path / env_mat_stat.get_hash() + env_mat_stat.load_or_compute_stats(merged, path) + self.stats = env_mat_stat.stats + mean, stddev = env_mat_stat() + if not self.set_davg_zero: + self.mean.copy_(torch.tensor(mean, device=env.DEVICE)) + self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE)) + if not self.set_davg_zero: + self.mean.copy_(torch.tensor(mean, device=env.DEVICE)) + self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE)) + + def get_stats(self) -> Dict[str, StatItem]: + """Get the statistics of the descriptor.""" + if self.stats is None: + raise RuntimeError( + "The statistics of the descriptor has not been computed." + ) + return self.stats + + def __setitem__(self, key, value): + if key in ("avg", "data_avg", "davg"): + self.mean = value + elif key in ("std", "data_std", "dstd"): + self.stddev = value + else: + raise KeyError(key) + + def __getitem__(self, key): + if key in ("avg", "data_avg", "davg"): + return self.mean + elif key in ("std", "data_std", "dstd"): + return self.stddev + else: + raise KeyError(key) + + @classmethod + def get_data_process_key(cls, config): + """ + Get the keys for the data preprocess. + Usually need the information of rcut and sel. + TODO Need to be deprecated when the dataloader has been cleaned up. + """ + descrpt_type = config["type"] + assert descrpt_type in ["se_e2_r"] + return {"sel": config["sel"], "rcut": config["rcut"]} + + @property + def data_stat_key(self): + """ + Get the keys for the data statistic of the descriptor. + Return a list of statistic names needed, such as "sumr", "sumr2" or "sumn". + """ + return ["sumr", "sumn", "sumr2"] + + def forward( + self, + coord_ext: torch.Tensor, + atype_ext: torch.Tensor, + nlist: torch.Tensor, + mapping: Optional[torch.Tensor] = None, + ): + """Compute the descriptor. + + Parameters + ---------- + coord_ext + The extended coordinates of atoms. shape: nf x (nallx3) + atype_ext + The extended aotm types. shape: nf x nall + nlist + The neighbor list. shape: nf x nloc x nnei + mapping + The index mapping, not required by this descriptor. + + Returns + ------- + descriptor + The descriptor. shape: nf x nloc x (ng x axis_neuron) + gr + The rotationally equivariant and permutationally invariant single particle + representation. shape: nf x nloc x ng x 3 + g2 + The rotationally invariant pair-partical representation. + this descriptor returns None + h2 + The rotationally equivariant pair-partical representation. + this descriptor returns None + sw + The smooth switch function. + + """ + + del mapping + nloc = nlist.shape[1] + atype = atype_ext[:, :nloc] + dmatrix, diff, sw = prod_env_mat_se_r( + coord_ext, + nlist, + atype, + self.mean, + self.stddev, + self.rcut, + self.rcut_smth, + ) + + if self.old_impl: + assert self.filter_layers_old is not None + dmatrix = dmatrix.view( + -1, self.ndescrpt + ) # shape is [nframes*nall, self.ndescrpt] + xyz_scatter = torch.empty( + 1, + device=env.DEVICE, + ) + ret = self.filter_layers_old[0](dmatrix) + xyz_scatter = ret + for ii, transform in enumerate(self.filter_layers_old[1:]): + # shape is [nframes*nall, 1, self.filter_neuron[-1]] + ret = transform.forward(dmatrix) + xyz_scatter = xyz_scatter + ret + else: + assert self.filter_layers is not None + dmatrix = dmatrix.view(-1, self.nnei, 1) + dmatrix = dmatrix.to(dtype=self.prec) + nfnl = dmatrix.shape[0] + # pre-allocate a shape to pass jit + xyz_scatter = torch.zeros( + [nfnl, 1, self.filter_neuron[-1]], dtype=self.prec, device=env.DEVICE + ) + # nfnl x nnei + exclude_mask = self.emask(nlist, atype_ext).view(nfnl, -1) + for ii, ll in enumerate(self.filter_layers.networks): + # nfnl x nt + mm = exclude_mask[:, self.sec[ii] : self.sec[ii + 1]] + # nfnl x nt x 1 + rr = dmatrix[:, self.sec[ii] : self.sec[ii + 1], :] + rr = rr * mm[:, :, None] + ss = rr[:, :, :1] + # nfnl x nt x ng + gg = ll.forward(ss) + # nfnl x 1 x ng + gr = torch.matmul(rr.permute(0, 2, 1), gg) + xyz_scatter += gr + + xyz_scatter /= self.nnei + xyz_scatter_1 = xyz_scatter.permute(0, 2, 1) + xyz_scatter_2 = xyz_scatter[:, :, 0 : self.axis_neuron] + result = torch.matmul( + xyz_scatter_1, xyz_scatter_2 + ) # shape is [nframes*nall, self.filter_neuron[-1], self.axis_neuron] + result = result.view(-1, nloc, self.filter_neuron[-1] * self.axis_neuron) + return ( + result.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), + None, + None, + None, + sw, + ) + + def set_stat_mean_and_stddev( + self, + mean: torch.Tensor, + stddev: torch.Tensor, + ) -> None: + self.sea.mean = mean + self.sea.stddev = stddev + + def serialize(self) -> dict: + obj = self.sea + return { + "rcut": obj.rcut, + "rcut_smth": obj.rcut_smth, + "sel": obj.sel, + "neuron": obj.neuron, + "resnet_dt": obj.resnet_dt, + "set_davg_zero": obj.set_davg_zero, + "activation_function": obj.activation_function, + # make deterministic + "precision": RESERVED_PRECISON_DICT[obj.prec], + "embeddings": obj.filter_layers.serialize(), + "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(), + "exclude_types": obj.exclude_types, + "@variables": { + "davg": obj["davg"].detach().cpu().numpy(), + "dstd": obj["dstd"].detach().cpu().numpy(), + }, + ## to be updated when the options are supported. + "trainable": True, + "type_one_side": True, + "spin": None, + } + + @classmethod + def deserialize(cls, data: dict) -> "DescrptSeR": + data = data.copy() + variables = data.pop("@variables") + embeddings = data.pop("embeddings") + env_mat = data.pop("env_mat") + obj = cls(**data) + + def t_cvt(xx): + return torch.tensor(xx, dtype=obj.sea.prec, device=env.DEVICE) + + obj.sea["davg"] = t_cvt(variables["davg"]) + obj.sea["dstd"] = t_cvt(variables["dstd"]) + obj.sea.filter_layers = NetworkCollection.deserialize(embeddings) + return obj + +def analyze_descrpt(matrix, ndescrpt, natoms): + """Collect avg, square avg and count of descriptors in a batch.""" + ntypes = natoms.shape[1] - 2 + start_index = 0 + sysr = [] + sysn = [] + sysr2 = [] + for type_i in range(ntypes): + end_index = start_index + natoms[0, 2 + type_i] + dd = matrix[:, start_index:end_index] # all descriptors for this element + start_index = end_index + dd = np.reshape( + dd, [-1, 1] + ) # Shape is [nframes*natoms[2+type_id]*self.nnei, 1] + ddr = dd[:, :1] + sumr = np.sum(ddr) + sumn = dd.shape[0] # Value is nframes*natoms[2+type_id]*self.nnei + sumr2 = np.sum(np.multiply(ddr, ddr)) + sysr.append(sumr) + sysn.append(sumn) + sysr2.append(sumr2) + return sysr, sysr2, sysn diff --git a/deepmd/pt/utils/env_mat_stat.py b/deepmd/pt/utils/env_mat_stat.py index 2f3c728c99..fee4d5d071 100644 --- a/deepmd/pt/utils/env_mat_stat.py +++ b/deepmd/pt/utils/env_mat_stat.py @@ -14,6 +14,7 @@ ) from deepmd.pt.model.descriptor.env_mat import ( prod_env_mat_se_a, + prod_env_mat_se_r, ) from deepmd.pt.utils import ( env, @@ -193,3 +194,138 @@ def __call__(self): mean = np.stack(all_davg) stddev = np.stack(all_dstd) return mean, stddev + +class EnvMatStatSeR(EnvMatStat): + """Environmental matrix statistics for the se_r environemntal matrix. + + Parameters + ---------- + descriptor : DescriptorBlock + The descriptor of the model. + """ + + def __init__(self, descriptor: "DescriptorBlock"): + super().__init__() + self.descriptor = descriptor + + def iter( + self, data: List[Dict[str, torch.Tensor]] + ) -> Iterator[Dict[str, StatItem]]: + """Get the iterator of the environment matrix. + + Parameters + ---------- + data : List[Dict[str, torch.Tensor]] + The environment matrix. + + Yields + ------ + Dict[str, StatItem] + The statistics of the environment matrix. + """ + zero_mean = torch.zeros( + self.descriptor.get_ntypes(), + self.descriptor.get_nsel(), + 1, + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=env.DEVICE, + ) + one_stddev = torch.ones( + self.descriptor.get_ntypes(), + self.descriptor.get_nsel(), + 1, + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=env.DEVICE, + ) + for system in data: + coord, atype, box, natoms = ( + system["coord"], + system["atype"], + system["box"], + system["natoms"], + ) + ( + extended_coord, + extended_atype, + mapping, + nlist, + ) = extend_input_and_build_neighbor_list( + coord, + atype, + self.descriptor.get_rcut(), + self.descriptor.get_sel(), + mixed_types=self.descriptor.mixed_types(), + box=box, + ) + env_mat, _, _ = prod_env_mat_se_r( + extended_coord, + nlist, + atype, + zero_mean, + one_stddev, + self.descriptor.get_rcut(), + # TODO: export rcut_smth from DescriptorBlock + self.descriptor.rcut_smth, + ) + # reshape to nframes * nloc at the atom level, + # so nframes/mixed_type do not matter + env_mat = env_mat.view( + coord.shape[0] * coord.shape[1], self.descriptor.get_nsel(), 1 + ) + atype = atype.view(coord.shape[0] * coord.shape[1]) + # (1, nloc) eq (ntypes, 1), so broadcast is possible + # shape: (ntypes, nloc) + type_idx = torch.eq( + atype.view(1, -1), + torch.arange( + self.descriptor.get_ntypes(), device=env.DEVICE, dtype=torch.int32 + ).view(-1, 1), + ) + for type_i in range(self.descriptor.get_ntypes()): + dd = env_mat[type_idx[type_i]] + dd = dd.reshape([-1, 1]) # typen_atoms * nnei, 1 + env_mats = {} + env_mats[f"r_{type_i}"] = dd[:, :1] + yield self.compute_stat(env_mats) + + def get_hash(self) -> str: + """Get the hash of the environment matrix. + + Returns + ------- + str + The hash of the environment matrix. + """ + return get_hash( + { + "type": "se_r", + "ntypes": self.descriptor.get_ntypes(), + "rcut": round(self.descriptor.get_rcut(), 2), + "rcut_smth": round(self.descriptor.rcut_smth, 2), + "nsel": self.descriptor.get_nsel(), + "sel": self.descriptor.get_sel(), + "mixed_types": self.descriptor.mixed_types(), + } + ) + + def __call__(self): + avgs = self.get_avg() + stds = self.get_std() + + all_davg = [] + all_dstd = [] + for type_i in range(self.descriptor.get_ntypes()): + davgunit = [[avgs[f"r_{type_i}"]]] + dstdunit = [ + [ + stds[f"r_{type_i}"], + + ] + ] + davg = np.tile(davgunit, [self.descriptor.get_nsel(), 1]) + dstd = np.tile(dstdunit, [self.descriptor.get_nsel(), 1]) + all_davg.append(davg) + all_dstd.append(dstd) + mean = np.stack(all_davg) + stddev = np.stack(all_dstd) + return mean, stddev From fc911d208a5f03774dca8d12ad981100f26aeb3d Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Mon, 26 Feb 2024 11:43:10 +0800 Subject: [PATCH 36/36] feat: add UTs --- deepmd/pt/model/descriptor/__init__.py | 4 + deepmd/pt/model/descriptor/env_mat.py | 15 +- deepmd/pt/model/descriptor/se_r.py | 55 +++--- source/tests/pt/model/test_descriptor_se_r.py | 174 ++++++++++++++++++ 4 files changed, 213 insertions(+), 35 deletions(-) create mode 100644 source/tests/pt/model/test_descriptor_se_r.py diff --git a/deepmd/pt/model/descriptor/__init__.py b/deepmd/pt/model/descriptor/__init__.py index 0cdb35c55f..b69e5a82b1 100644 --- a/deepmd/pt/model/descriptor/__init__.py +++ b/deepmd/pt/model/descriptor/__init__.py @@ -28,6 +28,9 @@ DescrptBlockSeA, DescrptSeA, ) +from .se_r import ( + DescrptSeR, +) __all__ = [ "Descriptor", @@ -36,6 +39,7 @@ "DescrptBlockSeA", "DescrptBlockSeAtten", "DescrptSeA", + "DescrptSeR", "DescrptDPA1", "DescrptDPA2", "prod_env_mat_se_a", diff --git a/deepmd/pt/model/descriptor/env_mat.py b/deepmd/pt/model/descriptor/env_mat.py index 7fe5f29354..249b1e12c6 100644 --- a/deepmd/pt/model/descriptor/env_mat.py +++ b/deepmd/pt/model/descriptor/env_mat.py @@ -89,20 +89,21 @@ def prod_env_mat_se_r( - atype: Atom types with shape [nframes, nloc]. - natoms: Batched atom statisics with shape [len(sec)+2]. - box: Batched simulation box with shape [nframes, 9]. - - mean: Average value of descriptor per element type with shape [len(sec), nnei, 4]. - - stddev: Standard deviation of descriptor per element type with shape [len(sec), nnei, 4]. - - deriv_stddev: StdDev of descriptor derivative per element type with shape [len(sec), nnei, 4, 3]. + - mean: Average value of descriptor per element type with shape [len(sec), nnei, 1]. + - stddev: Standard deviation of descriptor per element type with shape [len(sec), nnei, 1]. + - deriv_stddev: StdDev of descriptor derivative per element type with shape [len(sec), nnei, 1, 3]. - rcut: Cut-off radius. - rcut_smth: Smooth hyper-parameter for pair force & energy. Returns ------- - - env_mat_se_a: Shape is [nframes, natoms[1]*nnei*4]. + - env_mat_se_a: Shape is [nframes, natoms[1]*nnei*1]. """ - _env_mat_se_a, diff, switch = _make_env_mat_se_r( + _env_mat_se_r, diff, switch = _make_env_mat_se_r( nlist, extended_coord, rcut, rcut_smth ) # shape [n_atom, dim, 1] + t_avg = mean[atype] # [n_atom, dim, 1] t_std = stddev[atype] # [n_atom, dim, 1] - env_mat_se_a = (_env_mat_se_a - t_avg) / t_std - return env_mat_se_a, diff, switch \ No newline at end of file + env_mat_se_r = (_env_mat_se_r - t_avg) / t_std + return env_mat_se_r, diff, switch \ No newline at end of file diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index bb1c4e5ddd..461667364b 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - ClassVar, Dict, List, Optional, @@ -12,7 +11,6 @@ from deepmd.pt.model.descriptor import ( Descriptor, - DescriptorBlock, prod_env_mat_se_r, ) from deepmd.pt.utils import ( @@ -89,9 +87,9 @@ def __init__( ) self.split_sel = self.sel self.nnei = sum(sel) - self.ndescrpt = self.nnei * 4 + self.ndescrpt = self.nnei * 1 - wanted_shape = (self.ntypes, self.nnei, 4) + wanted_shape = (self.ntypes, self.nnei, 1) mean = torch.zeros(wanted_shape, dtype=self.prec, device=env.DEVICE) stddev = torch.ones(wanted_shape, dtype=self.prec, device=env.DEVICE) self.register_buffer("mean", mean) @@ -272,6 +270,7 @@ def forward( self.rcut, self.rcut_smth, ) + assert dmatrix.shape == (2,3,7,1) if self.old_impl: assert self.filter_layers_old is not None @@ -297,6 +296,7 @@ def forward( xyz_scatter = torch.zeros( [nfnl, 1, self.filter_neuron[-1]], dtype=self.prec, device=env.DEVICE ) + # nfnl x nnei exclude_mask = self.emask(nlist, atype_ext).view(nfnl, -1) for ii, ll in enumerate(self.filter_layers.networks): @@ -314,11 +314,11 @@ def forward( xyz_scatter /= self.nnei xyz_scatter_1 = xyz_scatter.permute(0, 2, 1) - xyz_scatter_2 = xyz_scatter[:, :, 0 : self.axis_neuron] + result = torch.matmul( - xyz_scatter_1, xyz_scatter_2 - ) # shape is [nframes*nall, self.filter_neuron[-1], self.axis_neuron] - result = result.view(-1, nloc, self.filter_neuron[-1] * self.axis_neuron) + xyz_scatter_1, xyz_scatter + ) # shape is [nframes*nall, self.filter_neuron[-1], 1] + result = result.view(-1, nloc, self.filter_neuron[-1] * 1) return ( result.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), None, @@ -332,27 +332,26 @@ def set_stat_mean_and_stddev( mean: torch.Tensor, stddev: torch.Tensor, ) -> None: - self.sea.mean = mean - self.sea.stddev = stddev + self.mean = mean + self.stddev = stddev def serialize(self) -> dict: - obj = self.sea return { - "rcut": obj.rcut, - "rcut_smth": obj.rcut_smth, - "sel": obj.sel, - "neuron": obj.neuron, - "resnet_dt": obj.resnet_dt, - "set_davg_zero": obj.set_davg_zero, - "activation_function": obj.activation_function, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel, + "neuron": self.neuron, + "resnet_dt": self.resnet_dt, + "set_davg_zero": self.set_davg_zero, + "activation_function": self.activation_function, # make deterministic - "precision": RESERVED_PRECISON_DICT[obj.prec], - "embeddings": obj.filter_layers.serialize(), - "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(), - "exclude_types": obj.exclude_types, + "precision": RESERVED_PRECISON_DICT[self.prec], + "embeddings": self.filter_layers.serialize(), + "env_mat": DPEnvMat(self.rcut, self.rcut_smth).serialize(), + "exclude_types": self.exclude_types, "@variables": { - "davg": obj["davg"].detach().cpu().numpy(), - "dstd": obj["dstd"].detach().cpu().numpy(), + "davg": self["davg"].detach().cpu().numpy(), + "dstd": self["dstd"].detach().cpu().numpy(), }, ## to be updated when the options are supported. "trainable": True, @@ -369,11 +368,11 @@ def deserialize(cls, data: dict) -> "DescrptSeR": obj = cls(**data) def t_cvt(xx): - return torch.tensor(xx, dtype=obj.sea.prec, device=env.DEVICE) + return torch.tensor(xx, dtype=obj.prec, device=env.DEVICE) - obj.sea["davg"] = t_cvt(variables["davg"]) - obj.sea["dstd"] = t_cvt(variables["dstd"]) - obj.sea.filter_layers = NetworkCollection.deserialize(embeddings) + obj["davg"] = t_cvt(variables["davg"]) + obj["dstd"] = t_cvt(variables["dstd"]) + obj.filter_layers = NetworkCollection.deserialize(embeddings) return obj def analyze_descrpt(matrix, ndescrpt, natoms): diff --git a/source/tests/pt/model/test_descriptor_se_r.py b/source/tests/pt/model/test_descriptor_se_r.py new file mode 100644 index 0000000000..abf3a050d3 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_se_r.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import itertools +import unittest + +import numpy as np +import torch + +# from deepmd.dpmodel.descriptor import DescrptSeR as DPDescrptSeR +from deepmd.pt.model.descriptor.se_r import ( + DescrptSeR, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + PRECISION_DICT, +) + +from .test_env_mat import ( + TestCaseSingleFrameWithNlist, +) +from .test_mlp import ( + get_tols, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION + + +# to be merged with the tf test case +class TestDescrptSeR(unittest.TestCase, TestCaseSingleFrameWithNlist): + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + + def test_consistency( + self, + ): + rng = np.random.default_rng() + nf, nloc, nnei = self.nlist.shape + davg = rng.normal(size=(self.nt, nnei, 1)) + dstd = rng.normal(size=(self.nt, nnei, 1)) + dstd = 0.1 + np.abs(dstd) + + for idt, prec, em in itertools.product( + [False, True], + ["float64", "float32"], + [[], [[0, 1]], [[1, 1]]], + ): + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + err_msg = f"idt={idt} prec={prec}" + # sea new impl + dd0 = DescrptSeR( + self.rcut, + self.rcut_smth, + self.sel, + precision=prec, + resnet_dt=idt, + old_impl=False, + exclude_mask=em, + ).to(env.DEVICE) + dd0.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE) + dd0.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) + + rd0, _, _, _, _ = dd0( + torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), + torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), + torch.tensor(self.nlist, dtype=int, device=env.DEVICE), + ) + # serialization + dd1 = DescrptSeR.deserialize(dd0.serialize()) + rd1, gr1, _, _, sw1 = dd1( + torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), + torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), + torch.tensor(self.nlist, dtype=int, device=env.DEVICE), + ) + np.testing.assert_allclose( + rd0.detach().cpu().numpy(), + rd1.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + err_msg=err_msg, + ) + np.testing.assert_allclose( + rd0.detach().cpu().numpy()[0][self.perm[: self.nloc]], + rd0.detach().cpu().numpy()[1], + rtol=rtol, + atol=atol, + err_msg=err_msg, + ) + # dp impl + # dd2 = DPDescrptSeR.deserialize(dd0.serialize()) + # rd2, gr2, _, _, sw2 = dd2.call( + # self.coord_ext, + # self.atype_ext, + # self.nlist, + # ) + # for aa, bb in zip([rd1, gr1, sw1], [rd2, gr2, sw2]): + # np.testing.assert_allclose( + # aa.detach().cpu().numpy(), + # bb, + # rtol=rtol, + # atol=atol, + # err_msg=err_msg, + # ) + # old impl + if idt is False and prec == "float64": + dd3 = DescrptSeR( + self.rcut, + self.rcut_smth, + self.sel, + precision=prec, + resnet_dt=idt, + old_impl=True, + ).to(env.DEVICE) + dd0_state_dict = dd0.state_dict() + dd3_state_dict = dd3.state_dict() + for i in dd3_state_dict: + dd3_state_dict[i] = ( + dd0_state_dict[ + i.replace(".deep_layers.", ".layers.").replace( + "filter_layers_old.", "filter_layers.networks." + ) + ] + .detach() + .clone() + ) + if ".bias" in i: + dd3_state_dict[i] = dd3_state_dict[i].unsqueeze(0) + dd3.load_state_dict(dd3_state_dict) + + rd3, gr3, _, _, sw3 = dd3( + torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), + torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), + torch.tensor(self.nlist, dtype=int, device=env.DEVICE), + ) + for aa, bb in zip([rd1, gr1, sw1], [rd3, gr3, sw3]): + np.testing.assert_allclose( + aa.detach().cpu().numpy(), + bb.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + err_msg=err_msg, + ) + + def test_jit( + self, + ): + rng = np.random.default_rng() + nf, nloc, nnei = self.nlist.shape + davg = rng.normal(size=(self.nt, nnei, 4)) + dstd = rng.normal(size=(self.nt, nnei, 4)) + dstd = 0.1 + np.abs(dstd) + + for idt, prec in itertools.product( + [False, True], + ["float64", "float32"], + ): + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + err_msg = f"idt={idt} prec={prec}" + # sea new impl + dd0 = DescrptSeR( + self.rcut, + self.rcut_smth, + self.sel, + precision=prec, + resnet_dt=idt, + old_impl=False, + ) + dd0.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE) + dd0.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) + dd1 = DescrptSeR.deserialize(dd0.serialize()) + model = torch.jit.script(dd0) + model = torch.jit.script(dd1)