From 438bc78fdb48004729fc5e76e5b1515600604d3d Mon Sep 17 00:00:00 2001 From: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> Date: Wed, 10 Jan 2024 10:00:47 +0800 Subject: [PATCH] Add dp model format sea (#3123) - add precision test for embedding net Limitations - only support `type_one_side` - does not support type embedding and `stripped_type_embedding` - does not support `exclude_types` - does not support spin --------- Co-authored-by: Han Wang --- deepmd_utils/model_format/__init__.py | 6 + deepmd_utils/model_format/network.py | 2 + deepmd_utils/model_format/se_e2_a.py | 195 ++++++++++++++++++++++++ source/tests/test_model_format_utils.py | 45 +++++- 4 files changed, 243 insertions(+), 5 deletions(-) create mode 100644 deepmd_utils/model_format/se_e2_a.py diff --git a/deepmd_utils/model_format/__init__.py b/deepmd_utils/model_format/__init__.py index 533dd9ffff..40769f187d 100644 --- a/deepmd_utils/model_format/__init__.py +++ b/deepmd_utils/model_format/__init__.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from .common import ( + DEFAULT_PRECISION, PRECISION_DICT, ) from .env_mat import ( @@ -13,8 +14,12 @@ save_dp_model, traverse_model_dict, ) +from .se_e2_a import ( + DescrptSeA, +) __all__ = [ + "DescrptSeA", "EnvMat", "EmbeddingNet", "NativeLayer", @@ -23,4 +28,5 @@ "save_dp_model", "traverse_model_dict", "PRECISION_DICT", + "DEFAULT_PRECISION", ] diff --git a/deepmd_utils/model_format/network.py b/deepmd_utils/model_format/network.py index 98c35636fa..682a349476 100644 --- a/deepmd_utils/model_format/network.py +++ b/deepmd_utils/model_format/network.py @@ -379,6 +379,7 @@ def __init__( self.neuron = neuron self.activation_function = activation_function self.resnet_dt = resnet_dt + self.precision = precision def serialize(self) -> dict: """Serialize the network to a dict. @@ -393,6 +394,7 @@ def serialize(self) -> dict: "neuron": self.neuron.copy(), "activation_function": self.activation_function, "resnet_dt": self.resnet_dt, + "precision": self.precision, "layers": [layer.serialize() for layer in self.layers], } diff --git a/deepmd_utils/model_format/se_e2_a.py b/deepmd_utils/model_format/se_e2_a.py new file mode 100644 index 0000000000..114f9df915 --- /dev/null +++ b/deepmd_utils/model_format/se_e2_a.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np + +try: + from deepmd_utils._version import version as __version__ +except ImportError: + __version__ = "unknown" + +from typing import ( + Any, + List, + Optional, +) + +from .common import ( + DEFAULT_PRECISION, + NativeOP, +) +from .env_mat import ( + EnvMat, +) +from .network import ( + EmbeddingNet, +) + + +class DescrptSeA(NativeOP): + def __init__( + self, + rcut: float, + rcut_smth: float, + sel: List[str], + neuron: List[int] = [24, 48, 96], + axis_neuron: int = 8, + resnet_dt: bool = False, + trainable: bool = True, + type_one_side: bool = True, + exclude_types: List[List[int]] = [], + set_davg_zero: bool = False, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + spin: Optional[Any] = None, + stripped_type_embedding: bool = False, + ) -> None: + ## seed, uniform_seed, multi_task, not included. + if not type_one_side: + raise NotImplementedError("type_one_side == False not implemented") + if stripped_type_embedding: + raise NotImplementedError("stripped_type_embedding is not implemented") + if exclude_types != []: + raise NotImplementedError("exclude_types is not implemented") + if spin is not None: + raise NotImplementedError("spin is not implemented") + + self.rcut = rcut + self.rcut_smth = rcut_smth + self.sel = sel + self.ntypes = len(self.sel) + self.neuron = neuron + self.axis_neuron = axis_neuron + self.resnet_dt = resnet_dt + self.trainable = trainable + self.type_one_side = type_one_side + self.exclude_types = exclude_types + self.set_davg_zero = set_davg_zero + self.activation_function = activation_function + self.precision = precision + self.spin = spin + self.stripped_type_embedding = stripped_type_embedding + + in_dim = 1 # not considiering type embedding + self.embeddings = [] + for ii in range(self.ntypes): + self.embeddings.append( + EmbeddingNet( + in_dim, + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + ) + ) + self.env_mat = EnvMat(self.rcut, self.rcut_smth) + self.nnei = np.sum(self.sel) + self.nneix4 = self.nnei * 4 + self.davg = np.zeros([self.ntypes, self.nneix4]) + self.dstd = np.ones([self.ntypes, self.nneix4]) + self.orig_sel = self.sel + + def __setitem__(self, key, value): + if key in ("avg", "data_avg", "davg"): + self.davg = value + elif key in ("std", "data_std", "dstd"): + self.dstd = value + else: + raise KeyError(key) + + def __getitem__(self, key): + if key in ("avg", "data_avg", "davg"): + return self.davg + elif key in ("std", "data_std", "dstd"): + return self.dstd + else: + raise KeyError(key) + + def cal_g( + self, + ss, + ll, + ): + nf, nloc, nnei = ss.shape[0:3] + ss = ss.reshape(nf, nloc, nnei, 1) + # nf x nloc x nnei x ng + gg = self.embeddings[ll].call(ss) + return gg + + def call( + self, + coord_ext, + atype_ext, + nlist, + ): + """Compute the environment matrix. + + 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 + + Returns + ------- + descriptor + The descriptor. shape: nf x nloc x ng x axis_neuron + """ + # nf x nloc x nnei x 4 + rr, ww = self.env_mat.call(nlist, coord_ext, atype_ext, self.davg, self.dstd) + nf, nloc, nnei, _ = rr.shape + sec = np.append([0], np.cumsum(self.sel)) + + ng = self.neuron[-1] + gr = np.zeros([nf, nloc, ng, 4]) + for tt in range(self.ntypes): + tr = rr[:, :, sec[tt] : sec[tt + 1], :] + ss = tr[..., 0:1] + gg = self.cal_g(ss, tt) + # nf x nloc x ng x 4 + gr += np.einsum("flni,flnj->flij", gg, tr) + gr /= self.nnei + gr1 = gr[:, :, : self.axis_neuron, :] + # nf x nloc x ng x ng1 + grrg = np.einsum("flid,fljd->flij", gr, gr1) + # nf x nloc x (ng x ng1) + grrg = grrg.reshape(nf, nloc, ng * self.axis_neuron) + return grrg + + def serialize(self) -> dict: + return { + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel, + "neuron": self.neuron, + "axis_neuron": self.axis_neuron, + "resnet_dt": self.resnet_dt, + "trainable": self.trainable, + "type_one_side": self.type_one_side, + "exclude_types": self.exclude_types, + "set_davg_zero": self.set_davg_zero, + "activation_function": self.activation_function, + "precision": self.precision, + "spin": self.spin, + "stripped_type_embedding": self.stripped_type_embedding, + "env_mat": self.env_mat.serialize(), + "embeddings": [ii.serialize() for ii in self.embeddings], + "@variables": { + "davg": self.davg, + "dstd": self.dstd, + }, + } + + @classmethod + def deserialize(cls, data: dict) -> "DescrptSeA": + variables = data.pop("@variables") + embeddings = data.pop("embeddings") + env_mat = data.pop("env_mat") + obj = cls(**data) + + obj["davg"] = variables["davg"] + obj["dstd"] = variables["dstd"] + obj.embeddings = [EmbeddingNet.deserialize(dd) for dd in embeddings] + obj.env_mat = EnvMat.deserialize(env_mat) + return obj diff --git a/source/tests/test_model_format_utils.py b/source/tests/test_model_format_utils.py index aeb717060d..7fd93c1366 100644 --- a/source/tests/test_model_format_utils.py +++ b/source/tests/test_model_format_utils.py @@ -9,6 +9,7 @@ import numpy as np from deepmd_utils.model_format import ( + DescrptSeA, EmbeddingNet, EnvMat, NativeLayer, @@ -97,12 +98,18 @@ def test_deserialize(self): np.testing.assert_array_equal(network[1]["resnet"], True) def test_embedding_net(self): - for ni, idt, act in itertools.product( + for ni, act, idt, prec in itertools.product( [1, 10], - [True, False], ["tanh", "none"], + [True, False], + ["double", "single"], ): - en0 = EmbeddingNet(ni) + en0 = EmbeddingNet( + ni, + activation_function=act, + precision=prec, + resnet_dt=idt, + ) en1 = EmbeddingNet.deserialize(en0.serialize()) inp = np.ones([ni]) np.testing.assert_allclose(en0.call(inp), en1.call(inp)) @@ -141,7 +148,7 @@ def tearDown(self) -> None: os.remove(self.filename) -class TestEnvMat(unittest.TestCase): +class TestCaseSingleFrameWithNlist: def setUp(self): # nloc == 3, nall == 4 self.nloc = 3 @@ -158,6 +165,7 @@ def setUp(self): ).reshape([1, self.nall * 3]) self.atype_ext = np.array([0, 0, 1, 0], dtype=int).reshape([1, self.nall]) # sel = [5, 2] + self.sel = [5, 2] self.nlist = np.array( [ [1, 3, -1, -1, -1, 2, -1], @@ -165,10 +173,15 @@ def setUp(self): [0, 1, -1, -1, -1, 0, -1], ], dtype=int, - ).reshape([1, self.nloc, 7]) + ).reshape([1, self.nloc, sum(self.sel)]) self.rcut = 0.4 self.rcut_smth = 2.2 + +class TestEnvMat(unittest.TestCase, TestCaseSingleFrameWithNlist): + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + def test_self_consistency( self, ): @@ -183,3 +196,25 @@ def test_self_consistency( mm1, ww1 = em1.call(self.nlist, self.coord_ext, self.atype_ext, davg, dstd) np.testing.assert_allclose(mm0, mm1) np.testing.assert_allclose(ww0, ww1) + + +class TestDescrptSeA(unittest.TestCase, TestCaseSingleFrameWithNlist): + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + + def test_self_consistency( + 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) + + em0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel) + em0.davg = davg + em0.dstd = dstd + em1 = DescrptSeA.deserialize(em0.serialize()) + mm0 = em0.call(self.coord_ext, self.atype_ext, self.nlist) + mm1 = em1.call(self.coord_ext, self.atype_ext, self.nlist) + np.testing.assert_allclose(mm0, mm1)