From c7c6e80378fd6cc01e262e67f6ead6fe289fa99a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 8 Jan 2024 08:11:16 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd_utils/model_format/__init__.py | 7 +-- deepmd_utils/model_format/common.py | 7 ++- deepmd_utils/model_format/env_mat.py | 78 ++++++++++++------------- deepmd_utils/model_format/network.py | 6 +- source/tests/test_model_format_utils.py | 78 +++++++++++++------------ 5 files changed, 90 insertions(+), 86 deletions(-) diff --git a/deepmd_utils/model_format/__init__.py b/deepmd_utils/model_format/__init__.py index 8e2e23806c..26aa1f4f04 100644 --- a/deepmd_utils/model_format/__init__.py +++ b/deepmd_utils/model_format/__init__.py @@ -2,6 +2,9 @@ from .common import ( PRECISION_DICT, ) +from .env_mat import ( + EnvMat, +) from .network import ( EmbeddingNet, NativeLayer, @@ -10,10 +13,6 @@ save_dp_model, traverse_model_dict, ) -from .env_mat import( - EnvMat, -) - __all__ = [ "EmbeddingNet", diff --git a/deepmd_utils/model_format/common.py b/deepmd_utils/model_format/common.py index 41387cecc8..2ca39321eb 100644 --- a/deepmd_utils/model_format/common.py +++ b/deepmd_utils/model_format/common.py @@ -1,4 +1,8 @@ -from abc import ABC +# SPDX-License-Identifier: LGPL-3.0-or-later +from abc import ( + ABC, +) + import numpy as np PRECISION_DICT = { @@ -8,6 +12,7 @@ "default": np.float64, } + class NativeOP(ABC): """The unit operation of a native model.""" diff --git a/deepmd_utils/model_format/env_mat.py b/deepmd_utils/model_format/env_mat.py index 1495a050b5..8ad6532c4f 100644 --- a/deepmd_utils/model_format/env_mat.py +++ b/deepmd_utils/model_format/env_mat.py @@ -1,9 +1,11 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later import numpy as np from .common import ( - NativeOP, + NativeOP, ) + def compute_smooth_weight( distance: np.ndarray, rmin: float, @@ -14,9 +16,10 @@ def compute_smooth_weight( max_mask = distance >= rmax mid_mask = np.logical_not(np.logical_or(min_mask, max_mask)) uu = (distance - rmin) / (rmax - rmin) - vv = uu * uu * uu * (-6. * uu * uu + 15. * uu - 10.) + 1. + vv = uu * uu * uu * (-6.0 * uu * uu + 15.0 * uu - 10.0) + 1.0 return vv * mid_mask + min_mask + def _make_env_mat( nlist, coord, @@ -27,15 +30,15 @@ def _make_env_mat( nf, nloc, nnei = nlist.shape # nf x nall x 3 coord = coord.reshape(nf, -1, 3) - mask = (nlist >= 0) - nlist = (nlist * mask) + mask = nlist >= 0 + nlist = nlist * mask # nf x (nloc x nnei) x 3 index = np.tile(nlist.reshape(nf, -1, 1), (1, 1, 3)) coord_r = np.take_along_axis(coord, index, 1) # nf x nloc x nnei x 3 coord_r = coord_r.reshape(nf, nloc, nnei, 3) # nf x nloc x 1 x 3 - coord_l = coord[:,:nloc].reshape(nf, -1, 1, 3) + coord_l = coord[:, :nloc].reshape(nf, -1, 1, 3) # nf x nloc x nnei x 3 diff = coord_r - coord_l # nf x nloc x nnei @@ -43,43 +46,40 @@ def _make_env_mat( # for index 0 nloc atom length = length + ~np.expand_dims(mask, -1) t0 = 1 / length - t1 = diff / length ** 2 + t1 = diff / length**2 weight = compute_smooth_weight(length, ruct_smth, rcut) - env_mat_se_a = ( - np.concatenate([t0, t1], axis=-1) * weight * np.expand_dims(mask, -1) - ) - return env_mat_se_a, diff * np.expand_dims(mask,-1), weight - + env_mat_se_a = np.concatenate([t0, t1], axis=-1) * weight * np.expand_dims(mask, -1) + return env_mat_se_a, diff * np.expand_dims(mask, -1), weight + class EnvMat(NativeOP): - def __init__( - self, - rcut, - rcut_smth, - ): - self.rcut = rcut - self.rcut_smth = rcut_smth + def __init__( + self, + rcut, + rcut_smth, + ): + self.rcut = rcut + self.rcut_smth = rcut_smth - def call( - self, - nlist, - coord_ext, - ): - em, diff, ww = _make_env_mat( - nlist, coord_ext, self.rcut, self.rcut_smth) - return em, ww + def call( + self, + nlist, + coord_ext, + ): + em, diff, ww = _make_env_mat(nlist, coord_ext, self.rcut, self.rcut_smth) + return em, ww - def serialize( - self, - )->dict: - return { - "rcut": self.rcut, - "rcut_smth": self.rcut_smth, - } + def serialize( + self, + ) -> dict: + return { + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + } - @classmethod - def deserialize( - cls, - data: dict, - )->"EnvMat": - return cls(**data) + @classmethod + def deserialize( + cls, + data: dict, + ) -> "EnvMat": + return cls(**data) diff --git a/deepmd_utils/model_format/network.py b/deepmd_utils/model_format/network.py index 3fe7bf4f6a..59d0bc1fdd 100644 --- a/deepmd_utils/model_format/network.py +++ b/deepmd_utils/model_format/network.py @@ -18,8 +18,8 @@ __version__ = "unknown" from .common import ( - NativeOP, - PRECISION_DICT, + PRECISION_DICT, + NativeOP, ) @@ -126,8 +126,6 @@ def load_dp_model(filename: str) -> dict: return model_dict - - class NativeLayer(NativeOP): """Native representation of a layer. diff --git a/source/tests/test_model_format_utils.py b/source/tests/test_model_format_utils.py index 81d5c68649..8c164aeb7d 100644 --- a/source/tests/test_model_format_utils.py +++ b/source/tests/test_model_format_utils.py @@ -9,8 +9,8 @@ import numpy as np from deepmd_utils.model_format import ( - EnvMat, EmbeddingNet, + EnvMat, NativeLayer, NativeNet, load_dp_model, @@ -20,7 +20,10 @@ class TestNativeLayer(unittest.TestCase): def test_serialize_deserize(self): - for (ni, no), bias, ut, activation_function, resnet, ashp, prec in itertools.product( + for ( + ni, + no, + ), bias, ut, activation_function, resnet, ashp, prec in itertools.product( [(5, 5), (5, 10), (5, 9), (9, 5)], [True, False], [True, False], @@ -139,39 +142,38 @@ def tearDown(self) -> None: class TestEnvMat(unittest.TestCase): - def setUp(self): - # nloc == 3, nall == 4 - self.nloc = 3 - self.nall = 4 - self.coord_ext = np.array( - [ [0, 0, 0], - [0, 1, 0], - [0, 0, 1], - [0, -2, 0], - ], - dtype = np.float64, - ).reshape([1, self.nall*3]) - self.atype_ext = np.array( - [0, 0, 1, 0], dtype=int - ).reshape([1, self.nall]) - # sel = [5, 2] - self.nlist = np.array( - [ - [1, 3, -1, -1, -1, 2, -1], - [0,-1, -1, -1, -1, 2, -1], - [0, 1, -1, -1, -1, 0, -1], - ], - dtype=int, - ).reshape([1, self.nloc, 7]) - self.rcut = .4 - self.rcut_smth = 2.2 - - def test_self_consistency( - self, - ): - em0 = EnvMat(self.rcut, self.rcut_smth) - em1 = EnvMat.deserialize(em0.serialize()) - mm0, ww0 = em0.call(self.nlist, self.coord_ext) - mm1, ww1 = em1.call(self.nlist, self.coord_ext) - np.testing.assert_allclose(mm0, mm1) - np.testing.assert_allclose(ww0, ww1) + def setUp(self): + # nloc == 3, nall == 4 + self.nloc = 3 + self.nall = 4 + self.coord_ext = np.array( + [ + [0, 0, 0], + [0, 1, 0], + [0, 0, 1], + [0, -2, 0], + ], + dtype=np.float64, + ).reshape([1, self.nall * 3]) + self.atype_ext = np.array([0, 0, 1, 0], dtype=int).reshape([1, self.nall]) + # sel = [5, 2] + self.nlist = np.array( + [ + [1, 3, -1, -1, -1, 2, -1], + [0, -1, -1, -1, -1, 2, -1], + [0, 1, -1, -1, -1, 0, -1], + ], + dtype=int, + ).reshape([1, self.nloc, 7]) + self.rcut = 0.4 + self.rcut_smth = 2.2 + + def test_self_consistency( + self, + ): + em0 = EnvMat(self.rcut, self.rcut_smth) + em1 = EnvMat.deserialize(em0.serialize()) + mm0, ww0 = em0.call(self.nlist, self.coord_ext) + mm1, ww1 = em1.call(self.nlist, self.coord_ext) + np.testing.assert_allclose(mm0, mm1) + np.testing.assert_allclose(ww0, ww1)