From 5ee88c94fa13bc508b5e798e6d1a44f30a7489fc Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 19 Nov 2025 16:07:34 +0800 Subject: [PATCH 1/6] feat(pt): Enhance type embedding compression in se_t_tebd - Introduced a new method for type embedding compression, allowing precomputation of strip-mode type embeddings for all type pairs. - Added runtime checks to ensure compatibility with the strip input mode and the existence of necessary filter layers. - Updated the forward method to utilize precomputed type embeddings when available, improving performance during inference. These changes optimize the handling of type embeddings, enhancing the efficiency of the descriptor model. --- deepmd/pt/model/descriptor/se_t_tebd.py | 64 +++++++++++++++---------- 1 file changed, 40 insertions(+), 24 deletions(-) diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index e8ab0522ff..055f92eb5d 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -599,6 +599,7 @@ def enable_compression( self.se_ttebd.enable_compression( self.table.data, self.table_config, self.lower, self.upper ) + self.se_ttebd.type_embedding_compression(self.type_embedding) self.compress = True @@ -700,6 +701,7 @@ def __init__( self.compress_data = nn.ParameterList( [nn.Parameter(torch.zeros(0, dtype=self.prec, device=env.DEVICE))] ) + self.type_embd_data: Optional[torch.Tensor] = None def get_rcut(self) -> float: """Returns the cut-off radius.""" @@ -838,6 +840,27 @@ def reinit_exclude( self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) + def type_embedding_compression(self, type_embedding_net: TypeEmbedNet) -> None: + """Precompute strip-mode type embeddings for all type pairs.""" + if self.tebd_input_mode != "strip": + raise RuntimeError("Type embedding compression only works in strip mode") + if self.filter_layers_strip is None: + raise RuntimeError( + "filter_layers_strip must exist for type embedding compression" + ) + + with torch.no_grad(): + full_embd = type_embedding_net.get_full_embedding(env.DEVICE) + nt, t_dim = full_embd.shape + type_embedding_i = full_embd.view(nt, 1, t_dim).expand(nt, nt, t_dim) + type_embedding_j = full_embd.view(1, nt, t_dim).expand(nt, nt, t_dim) + two_side_type_embedding = torch.cat( + [type_embedding_i, type_embedding_j], dim=-1 + ).reshape(-1, t_dim * 2) + self.type_embd_data = self.filter_layers_strip.networks[0]( + two_side_type_embedding + ).detach() + def forward( self, nlist: torch.Tensor, @@ -986,31 +1009,24 @@ def forward( nei_type_j = nei_type.unsqueeze(1).expand([-1, nnei, -1]) idx_i = nei_type_i * ntypes_with_padding idx_j = nei_type_j - # (nf x nl x nt_i x nt_j) x ng - idx = ( - (idx_i + idx_j) - .view(-1, 1) - .expand(-1, ng) - .type(torch.long) - .to(torch.long) - ) - # ntypes * (ntypes) * nt - type_embedding_i = torch.tile( - type_embedding.view(ntypes_with_padding, 1, nt), - [1, ntypes_with_padding, 1], - ) - # (ntypes) * ntypes * nt - type_embedding_j = torch.tile( - type_embedding.view(1, ntypes_with_padding, nt), - [ntypes_with_padding, 1, 1], - ) - # (ntypes * ntypes) * (nt+nt) - two_side_type_embedding = torch.cat( - [type_embedding_i, type_embedding_j], -1 - ).reshape(-1, nt * 2) - tt_full = self.filter_layers_strip.networks[0](two_side_type_embedding) + idx = (idx_i + idx_j).reshape(-1).to(torch.long) + if self.type_embd_data is not None: + tt_full = self.type_embd_data + else: + type_embedding_i = torch.tile( + type_embedding.view(ntypes_with_padding, 1, nt), + [1, ntypes_with_padding, 1], + ) + type_embedding_j = torch.tile( + type_embedding.view(1, ntypes_with_padding, nt), + [ntypes_with_padding, 1, 1], + ) + two_side_type_embedding = torch.cat( + [type_embedding_i, type_embedding_j], -1 + ).reshape(-1, nt * 2) + tt_full = self.filter_layers_strip.networks[0](two_side_type_embedding) # (nfnl x nt_i x nt_j) x ng - gg_t = torch.gather(tt_full, dim=0, index=idx) + gg_t = tt_full[idx] # (nfnl x nt_i x nt_j) x ng gg_t = gg_t.reshape(nfnl, nnei, nnei, ng) if self.smooth: From 7753645dd6c03533bc829b57f808a1eaed8999e7 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 19 Nov 2025 16:19:21 +0800 Subject: [PATCH 2/6] refactor(pt): Update type embedding handling in se_t_tebd - Changed the assignment of type embedding data to use `register_buffer`, improving memory management. - Introduced a new variable `embd_tensor` for clarity and maintainability in the embedding computation process. These modifications enhance the structure of the code while maintaining existing functionality. --- deepmd/pt/model/descriptor/se_t_tebd.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 055f92eb5d..313999e1af 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -857,9 +857,10 @@ def type_embedding_compression(self, type_embedding_net: TypeEmbedNet) -> None: two_side_type_embedding = torch.cat( [type_embedding_i, type_embedding_j], dim=-1 ).reshape(-1, t_dim * 2) - self.type_embd_data = self.filter_layers_strip.networks[0]( + embd_tensor = self.filter_layers_strip.networks[0]( two_side_type_embedding ).detach() + self.register_buffer("type_embd_data", embd_tensor) def forward( self, From c81df99fb303850b11b9c889a36bc43d4d0fb85b Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 20 Nov 2025 15:24:52 +0800 Subject: [PATCH 3/6] fix(pt): Clear type embedding data before registering new buffer in DescrptBlockSeTTebd --- deepmd/pt/model/descriptor/se_t_tebd.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 313999e1af..79323b2a4f 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -860,6 +860,8 @@ def type_embedding_compression(self, type_embedding_net: TypeEmbedNet) -> None: embd_tensor = self.filter_layers_strip.networks[0]( two_side_type_embedding ).detach() + if hasattr(self, "type_embd_data"): + del self.type_embd_data self.register_buffer("type_embd_data", embd_tensor) def forward( From 6870dc4fd24fa1591abf5a48464819abe3c058ad Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 21 Nov 2025 15:25:45 +0800 Subject: [PATCH 4/6] feat(tests): add model compression tests for se_e3_tebd --- .../pt/test_model_compression_se_e3_tebd.py | 770 ++++++++++++++++++ 1 file changed, 770 insertions(+) create mode 100644 source/tests/pt/test_model_compression_se_e3_tebd.py diff --git a/source/tests/pt/test_model_compression_se_e3_tebd.py b/source/tests/pt/test_model_compression_se_e3_tebd.py new file mode 100644 index 0000000000..817abf93f9 --- /dev/null +++ b/source/tests/pt/test_model_compression_se_e3_tebd.py @@ -0,0 +1,770 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +# Set environment variables before any imports to avoid threading conflicts +import json +import os +import unittest + +import numpy as np + +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.infer.deep_eval import ( + DeepEval, +) + +from .common import ( + j_loader, + run_dp, + tests_path, +) + +if GLOBAL_NP_FLOAT_PRECISION == np.float32: + default_places = 4 +else: + default_places = 10 + + +def _file_delete(file) -> None: + if os.path.isdir(file): + os.rmdir(file) + elif os.path.isfile(file): + os.remove(file) + + +def _init_models(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-se-e3-tebd.pth") + compressed_model = str(tests_path / "dp-compressed-se-e3-tebd.pth") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + + # Configure se_e3_tebd descriptor + jdata["model"]["descriptor"] = { + "type": "se_e3_tebd", + "sel": 40, + "rcut_smth": 0.50, + "rcut": 4.00, + "neuron": [2, 4, 8], + "tebd_dim": 8, + "tebd_input_mode": "strip", # Use strip for compression compatibility + "activation_function": "tanh", + "resnet_dt": False, + "seed": 1, + } + + jdata["training"]["training_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp --pt train " + INPUT) + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp --pt freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp( + "dp --pt compress " + " -i " + frozen_model + " -o " + compressed_model + ) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +def _init_models_exclude_types(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-se-e3-tebd-exclude-types.pth") + compressed_model = str(tests_path / "dp-compressed-se-e3-tebd-exclude-types.pth") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + + # Configure se_e3_tebd descriptor with exclude_types + jdata["model"]["descriptor"] = { + "type": "se_e3_tebd", + "exclude_types": [[0, 1]], + "sel": 40, + "rcut_smth": 0.50, + "rcut": 4.00, + "neuron": [2, 4, 8], + "tebd_dim": 8, + "tebd_input_mode": "strip", # Use strip for compression compatibility + "activation_function": "tanh", + "resnet_dt": False, + "seed": 1, + } + + jdata["training"]["training_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp --pt train " + INPUT) + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp --pt freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp( + "dp --pt compress " + " -i " + frozen_model + " -o " + compressed_model + ) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +def _init_models_skip_neighbor_stat(): + suffix = "-skip-neighbor-stat" + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / f"dp-original-se-e3-tebd{suffix}.pth") + compressed_model = str(tests_path / f"dp-compressed-se-e3-tebd{suffix}.pth") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + + # Configure se_e3_tebd descriptor + jdata["model"]["descriptor"] = { + "type": "se_e3_tebd", + "sel": 40, + "rcut_smth": 0.50, + "rcut": 4.00, + "neuron": [2, 4, 8], + "tebd_dim": 8, + "tebd_input_mode": "strip", # Use strip for compression compatibility + "activation_function": "tanh", + "resnet_dt": False, + "seed": 1, + } + + jdata["training"]["training_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp --pt train " + INPUT + " --skip-neighbor-stat") + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp --pt freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp( + "dp --pt compress " + + " -i " + + frozen_model + + " -o " + + compressed_model + + " -t " + + INPUT + ) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +def setUpModule() -> None: + global \ + INPUT, \ + FROZEN_MODEL, \ + COMPRESSED_MODEL, \ + INPUT_ET, \ + FROZEN_MODEL_ET, \ + COMPRESSED_MODEL_ET, \ + FROZEN_MODEL_SKIP_NEIGHBOR_STAT, \ + COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT + INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() + _, FROZEN_MODEL_SKIP_NEIGHBOR_STAT, COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT = ( + _init_models_skip_neighbor_stat() + ) + INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types() + + +class TestDeepPotAPBC(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_attrs(self) -> None: + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 4.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 4.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self) -> None: + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + coords2, box2, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestDeepPotANoPBC(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = None + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self) -> None: + coords2 = np.concatenate((self.coords, self.coords)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + coords2, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestDeepPotALargeBoxNoPBC(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = np.array([19.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_ase(self) -> None: + from ase import ( + Atoms, + ) + + from deepmd.calculator import ( + DP, + ) + + water0 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(FROZEN_MODEL), + ) + water1 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(COMPRESSED_MODEL), + ) + ee0 = water0.get_potential_energy() + ff0 = water0.get_forces() + ee1 = water1.get_potential_energy() + ff1 = water1.get_forces() + # nframes = 1 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + + +class TestDeepPotAPBCExcludeTypes(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL_ET) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL_ET) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + @classmethod + def tearDownClass(cls) -> None: + _file_delete(INPUT_ET) + _file_delete(FROZEN_MODEL_ET) + _file_delete(COMPRESSED_MODEL_ET) + _file_delete("out.json") + _file_delete("compress.json") + _file_delete("checkpoint") + _file_delete("lcurve.out") + _file_delete("model.ckpt") + _file_delete("model-compression/checkpoint") + _file_delete("model-compression") + + def test_attrs(self) -> None: + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 4.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 4.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self) -> None: + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + coords2, box2, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestSkipNeighborStat(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL_SKIP_NEIGHBOR_STAT) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_attrs(self) -> None: + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 4.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 4.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self) -> None: + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + coords2, box2, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +if __name__ == "__main__": + unittest.main() From 878e3f8bbe843ac3de7838172330b4dc02f3fff3 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 21 Nov 2025 18:06:24 +0800 Subject: [PATCH 5/6] feat(tests): enhance cleanup in model compression tests --- .../pt/test_model_compression_se_e3_tebd.py | 38 ++++++++++++------- 1 file changed, 24 insertions(+), 14 deletions(-) diff --git a/source/tests/pt/test_model_compression_se_e3_tebd.py b/source/tests/pt/test_model_compression_se_e3_tebd.py index 817abf93f9..c4501db87c 100644 --- a/source/tests/pt/test_model_compression_se_e3_tebd.py +++ b/source/tests/pt/test_model_compression_se_e3_tebd.py @@ -2,6 +2,7 @@ # Set environment variables before any imports to avoid threading conflicts import json import os +import shutil import unittest import numpy as np @@ -27,7 +28,7 @@ def _file_delete(file) -> None: if os.path.isdir(file): - os.rmdir(file) + shutil.rmtree(file) elif os.path.isfile(file): os.remove(file) @@ -165,6 +166,28 @@ def setUpModule() -> None: INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types() +def tearDownModule() -> None: + # Clean up files created by _init_models + _file_delete(INPUT) + _file_delete(FROZEN_MODEL) + _file_delete(COMPRESSED_MODEL) + # Clean up files created by _init_models_skip_neighbor_stat + _file_delete(FROZEN_MODEL_SKIP_NEIGHBOR_STAT) + _file_delete(COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT) + # Clean up files created by _init_models_exclude_types + _file_delete(INPUT_ET) + _file_delete(FROZEN_MODEL_ET) + _file_delete(COMPRESSED_MODEL_ET) + # Clean up other artifacts + _file_delete("out.json") + _file_delete("compress.json") + _file_delete("checkpoint") + _file_delete("lcurve.out") + _file_delete("model.ckpt") + _file_delete("model-compression/checkpoint") + _file_delete("model-compression") + + class TestDeepPotAPBC(unittest.TestCase): @classmethod def setUpClass(cls) -> None: @@ -535,19 +558,6 @@ def setUpClass(cls) -> None: cls.atype = [0, 1, 1, 0, 1, 1] cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - @classmethod - def tearDownClass(cls) -> None: - _file_delete(INPUT_ET) - _file_delete(FROZEN_MODEL_ET) - _file_delete(COMPRESSED_MODEL_ET) - _file_delete("out.json") - _file_delete("compress.json") - _file_delete("checkpoint") - _file_delete("lcurve.out") - _file_delete("model.ckpt") - _file_delete("model-compression/checkpoint") - _file_delete("model-compression") - def test_attrs(self) -> None: self.assertEqual(self.dp_original.get_ntypes(), 2) self.assertAlmostEqual(self.dp_original.get_rcut(), 4.0, places=default_places) From 16b383f94a0c23d8a5c70db2365777359238aa37 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 23 Nov 2025 13:09:40 +0800 Subject: [PATCH 6/6] feat(pt): enforce strip mode for model compression and enhance type embedding handling --- deepmd/pt/model/descriptor/se_t_tebd.py | 70 +- .../pt/test_model_compression_se_e3_tebd.py | 780 ------------------ 2 files changed, 39 insertions(+), 811 deletions(-) delete mode 100644 source/tests/pt/test_model_compression_se_e3_tebd.py diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 79323b2a4f..8225d8e4af 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -553,6 +553,8 @@ def enable_compression( assert not self.se_ttebd.resnet_dt, ( "Model compression error: descriptor resnet_dt must be false!" ) + if self.tebd_input_mode != "strip": + raise RuntimeError("Cannot compress model when tebd_input_mode != 'strip'") for tt in self.se_ttebd.exclude_types: if (tt[0] not in range(self.se_ttebd.ntypes)) or ( tt[1] not in range(self.se_ttebd.ntypes) @@ -573,9 +575,6 @@ def enable_compression( "Empty embedding-nets are not supported in model compression!" ) - if self.tebd_input_mode != "strip": - raise RuntimeError("Cannot compress model when tebd_input_mode == 'concat'") - data = self.serialize() self.table = DPTabulate( self, @@ -597,9 +596,12 @@ def enable_compression( ) self.se_ttebd.enable_compression( - self.table.data, self.table_config, self.lower, self.upper + self.type_embedding, + self.table.data, + self.table_config, + self.lower, + self.upper, ) - self.se_ttebd.type_embedding_compression(self.type_embedding) self.compress = True @@ -695,13 +697,17 @@ def __init__( self.stats = None # compression related variables self.compress = False + # For geometric compression self.compress_info = nn.ParameterList( [nn.Parameter(torch.zeros(0, dtype=self.prec, device="cpu"))] ) self.compress_data = nn.ParameterList( [nn.Parameter(torch.zeros(0, dtype=self.prec, device=env.DEVICE))] ) - self.type_embd_data: Optional[torch.Tensor] = None + # For type embedding compression + self.register_buffer( + "type_embd_data", torch.zeros(0, dtype=self.prec, device=env.DEVICE) + ) def get_rcut(self) -> float: """Returns the cut-off radius.""" @@ -840,30 +846,6 @@ def reinit_exclude( self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) - def type_embedding_compression(self, type_embedding_net: TypeEmbedNet) -> None: - """Precompute strip-mode type embeddings for all type pairs.""" - if self.tebd_input_mode != "strip": - raise RuntimeError("Type embedding compression only works in strip mode") - if self.filter_layers_strip is None: - raise RuntimeError( - "filter_layers_strip must exist for type embedding compression" - ) - - with torch.no_grad(): - full_embd = type_embedding_net.get_full_embedding(env.DEVICE) - nt, t_dim = full_embd.shape - type_embedding_i = full_embd.view(nt, 1, t_dim).expand(nt, nt, t_dim) - type_embedding_j = full_embd.view(1, nt, t_dim).expand(nt, nt, t_dim) - two_side_type_embedding = torch.cat( - [type_embedding_i, type_embedding_j], dim=-1 - ).reshape(-1, t_dim * 2) - embd_tensor = self.filter_layers_strip.networks[0]( - two_side_type_embedding - ).detach() - if hasattr(self, "type_embd_data"): - del self.type_embd_data - self.register_buffer("type_embd_data", embd_tensor) - def forward( self, nlist: torch.Tensor, @@ -1013,7 +995,7 @@ def forward( idx_i = nei_type_i * ntypes_with_padding idx_j = nei_type_j idx = (idx_i + idx_j).reshape(-1).to(torch.long) - if self.type_embd_data is not None: + if self.compress: tt_full = self.type_embd_data else: type_embedding_i = torch.tile( @@ -1061,6 +1043,7 @@ def forward( def enable_compression( self, + type_embedding_net: TypeEmbedNet, table_data: dict, table_config: dict, lower: dict, @@ -1070,6 +1053,8 @@ def enable_compression( Parameters ---------- + type_embedding_net : TypeEmbedNet + The type embedding network table_data : dict The tabulated data from DPTabulate table_config : dict @@ -1079,6 +1064,13 @@ def enable_compression( upper : dict Upper bounds for compression """ + if self.tebd_input_mode != "strip": + raise RuntimeError("Type embedding compression only works in strip mode") + if self.filter_layers_strip is None: + raise RuntimeError( + "filter_layers_strip must exist for type embedding compression" + ) + # Compress the main geometric embedding network (self.filter_layers) net_key = "filter_net" self.compress_info[0] = torch.as_tensor( @@ -1097,6 +1089,22 @@ def enable_compression( device=env.DEVICE, dtype=self.prec ) + # Compress the type embedding network (self.filter_layers_strip) + with torch.no_grad(): + full_embd = type_embedding_net.get_full_embedding(env.DEVICE) + nt, t_dim = full_embd.shape + type_embedding_i = full_embd.view(nt, 1, t_dim).expand(nt, nt, t_dim) + type_embedding_j = full_embd.view(1, nt, t_dim).expand(nt, nt, t_dim) + two_side_type_embedding = torch.cat( + [type_embedding_i, type_embedding_j], dim=-1 + ).reshape(-1, t_dim * 2) + embd_tensor = self.filter_layers_strip.networks[0]( + two_side_type_embedding + ).detach() + if hasattr(self, "type_embd_data"): + del self.type_embd_data + self.register_buffer("type_embd_data", embd_tensor) + self.compress = True def has_message_passing(self) -> bool: diff --git a/source/tests/pt/test_model_compression_se_e3_tebd.py b/source/tests/pt/test_model_compression_se_e3_tebd.py deleted file mode 100644 index c4501db87c..0000000000 --- a/source/tests/pt/test_model_compression_se_e3_tebd.py +++ /dev/null @@ -1,780 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -# Set environment variables before any imports to avoid threading conflicts -import json -import os -import shutil -import unittest - -import numpy as np - -from deepmd.env import ( - GLOBAL_NP_FLOAT_PRECISION, -) -from deepmd.infer.deep_eval import ( - DeepEval, -) - -from .common import ( - j_loader, - run_dp, - tests_path, -) - -if GLOBAL_NP_FLOAT_PRECISION == np.float32: - default_places = 4 -else: - default_places = 10 - - -def _file_delete(file) -> None: - if os.path.isdir(file): - shutil.rmtree(file) - elif os.path.isfile(file): - os.remove(file) - - -def _init_models(): - data_file = str(tests_path / os.path.join("model_compression", "data")) - frozen_model = str(tests_path / "dp-original-se-e3-tebd.pth") - compressed_model = str(tests_path / "dp-compressed-se-e3-tebd.pth") - INPUT = str(tests_path / "input.json") - jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) - - # Configure se_e3_tebd descriptor - jdata["model"]["descriptor"] = { - "type": "se_e3_tebd", - "sel": 40, - "rcut_smth": 0.50, - "rcut": 4.00, - "neuron": [2, 4, 8], - "tebd_dim": 8, - "tebd_input_mode": "strip", # Use strip for compression compatibility - "activation_function": "tanh", - "resnet_dt": False, - "seed": 1, - } - - jdata["training"]["training_data"]["systems"] = data_file - with open(INPUT, "w") as fp: - json.dump(jdata, fp, indent=4) - - ret = run_dp("dp --pt train " + INPUT) - np.testing.assert_equal(ret, 0, "DP train failed!") - ret = run_dp("dp --pt freeze -o " + frozen_model) - np.testing.assert_equal(ret, 0, "DP freeze failed!") - ret = run_dp( - "dp --pt compress " + " -i " + frozen_model + " -o " + compressed_model - ) - np.testing.assert_equal(ret, 0, "DP model compression failed!") - return INPUT, frozen_model, compressed_model - - -def _init_models_exclude_types(): - data_file = str(tests_path / os.path.join("model_compression", "data")) - frozen_model = str(tests_path / "dp-original-se-e3-tebd-exclude-types.pth") - compressed_model = str(tests_path / "dp-compressed-se-e3-tebd-exclude-types.pth") - INPUT = str(tests_path / "input.json") - jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) - - # Configure se_e3_tebd descriptor with exclude_types - jdata["model"]["descriptor"] = { - "type": "se_e3_tebd", - "exclude_types": [[0, 1]], - "sel": 40, - "rcut_smth": 0.50, - "rcut": 4.00, - "neuron": [2, 4, 8], - "tebd_dim": 8, - "tebd_input_mode": "strip", # Use strip for compression compatibility - "activation_function": "tanh", - "resnet_dt": False, - "seed": 1, - } - - jdata["training"]["training_data"]["systems"] = data_file - with open(INPUT, "w") as fp: - json.dump(jdata, fp, indent=4) - - ret = run_dp("dp --pt train " + INPUT) - np.testing.assert_equal(ret, 0, "DP train failed!") - ret = run_dp("dp --pt freeze -o " + frozen_model) - np.testing.assert_equal(ret, 0, "DP freeze failed!") - ret = run_dp( - "dp --pt compress " + " -i " + frozen_model + " -o " + compressed_model - ) - np.testing.assert_equal(ret, 0, "DP model compression failed!") - return INPUT, frozen_model, compressed_model - - -def _init_models_skip_neighbor_stat(): - suffix = "-skip-neighbor-stat" - data_file = str(tests_path / os.path.join("model_compression", "data")) - frozen_model = str(tests_path / f"dp-original-se-e3-tebd{suffix}.pth") - compressed_model = str(tests_path / f"dp-compressed-se-e3-tebd{suffix}.pth") - INPUT = str(tests_path / "input.json") - jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) - - # Configure se_e3_tebd descriptor - jdata["model"]["descriptor"] = { - "type": "se_e3_tebd", - "sel": 40, - "rcut_smth": 0.50, - "rcut": 4.00, - "neuron": [2, 4, 8], - "tebd_dim": 8, - "tebd_input_mode": "strip", # Use strip for compression compatibility - "activation_function": "tanh", - "resnet_dt": False, - "seed": 1, - } - - jdata["training"]["training_data"]["systems"] = data_file - with open(INPUT, "w") as fp: - json.dump(jdata, fp, indent=4) - - ret = run_dp("dp --pt train " + INPUT + " --skip-neighbor-stat") - np.testing.assert_equal(ret, 0, "DP train failed!") - ret = run_dp("dp --pt freeze -o " + frozen_model) - np.testing.assert_equal(ret, 0, "DP freeze failed!") - ret = run_dp( - "dp --pt compress " - + " -i " - + frozen_model - + " -o " - + compressed_model - + " -t " - + INPUT - ) - np.testing.assert_equal(ret, 0, "DP model compression failed!") - return INPUT, frozen_model, compressed_model - - -def setUpModule() -> None: - global \ - INPUT, \ - FROZEN_MODEL, \ - COMPRESSED_MODEL, \ - INPUT_ET, \ - FROZEN_MODEL_ET, \ - COMPRESSED_MODEL_ET, \ - FROZEN_MODEL_SKIP_NEIGHBOR_STAT, \ - COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT - INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() - _, FROZEN_MODEL_SKIP_NEIGHBOR_STAT, COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT = ( - _init_models_skip_neighbor_stat() - ) - INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types() - - -def tearDownModule() -> None: - # Clean up files created by _init_models - _file_delete(INPUT) - _file_delete(FROZEN_MODEL) - _file_delete(COMPRESSED_MODEL) - # Clean up files created by _init_models_skip_neighbor_stat - _file_delete(FROZEN_MODEL_SKIP_NEIGHBOR_STAT) - _file_delete(COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT) - # Clean up files created by _init_models_exclude_types - _file_delete(INPUT_ET) - _file_delete(FROZEN_MODEL_ET) - _file_delete(COMPRESSED_MODEL_ET) - # Clean up other artifacts - _file_delete("out.json") - _file_delete("compress.json") - _file_delete("checkpoint") - _file_delete("lcurve.out") - _file_delete("model.ckpt") - _file_delete("model-compression/checkpoint") - _file_delete("model-compression") - - -class TestDeepPotAPBC(unittest.TestCase): - @classmethod - def setUpClass(cls) -> None: - cls.dp_original = DeepEval(FROZEN_MODEL) - cls.dp_compressed = DeepEval(COMPRESSED_MODEL) - cls.coords = np.array( - [ - 12.83, - 2.56, - 2.18, - 12.09, - 2.87, - 2.74, - 00.25, - 3.32, - 1.68, - 3.36, - 3.00, - 1.81, - 3.51, - 2.51, - 2.60, - 4.27, - 3.22, - 1.56, - ] - ) - cls.atype = [0, 1, 1, 0, 1, 1] - cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - - def test_attrs(self) -> None: - self.assertEqual(self.dp_original.get_ntypes(), 2) - self.assertAlmostEqual(self.dp_original.get_rcut(), 4.0, places=default_places) - self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) - self.assertEqual(self.dp_original.get_dim_fparam(), 0) - self.assertEqual(self.dp_original.get_dim_aparam(), 0) - - self.assertEqual(self.dp_compressed.get_ntypes(), 2) - self.assertAlmostEqual( - self.dp_compressed.get_rcut(), 4.0, places=default_places - ) - self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) - self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) - self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - - def test_1frame(self) -> None: - ee0, ff0, vv0 = self.dp_original.eval( - self.coords, self.box, self.atype, atomic=False - ) - ee1, ff1, vv1 = self.dp_compressed.eval( - self.coords, self.box, self.atype, atomic=False - ) - # check shape of the returns - nframes = 1 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - def test_1frame_atm(self) -> None: - ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( - self.coords, self.box, self.atype, atomic=True - ) - ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( - self.coords, self.box, self.atype, atomic=True - ) - # check shape of the returns - nframes = 1 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ae0.shape, (nframes, natoms, 1)) - self.assertEqual(av0.shape, (nframes, natoms, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - self.assertEqual(ae1.shape, (nframes, natoms, 1)) - self.assertEqual(av1.shape, (nframes, natoms, 9)) - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ae0, ae1, default_places) - np.testing.assert_almost_equal(av0, av1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - def test_2frame_atm(self) -> None: - coords2 = np.concatenate((self.coords, self.coords)) - box2 = np.concatenate((self.box, self.box)) - ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( - coords2, box2, self.atype, atomic=True - ) - ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( - coords2, box2, self.atype, atomic=True - ) - # check shape of the returns - nframes = 2 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ae0.shape, (nframes, natoms, 1)) - self.assertEqual(av0.shape, (nframes, natoms, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - self.assertEqual(ae1.shape, (nframes, natoms, 1)) - self.assertEqual(av1.shape, (nframes, natoms, 9)) - - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ae0, ae1, default_places) - np.testing.assert_almost_equal(av0, av1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - -class TestDeepPotANoPBC(unittest.TestCase): - @classmethod - def setUpClass(cls) -> None: - cls.dp_original = DeepEval(FROZEN_MODEL) - cls.dp_compressed = DeepEval(COMPRESSED_MODEL) - cls.coords = np.array( - [ - 12.83, - 2.56, - 2.18, - 12.09, - 2.87, - 2.74, - 00.25, - 3.32, - 1.68, - 3.36, - 3.00, - 1.81, - 3.51, - 2.51, - 2.60, - 4.27, - 3.22, - 1.56, - ] - ) - cls.atype = [0, 1, 1, 0, 1, 1] - cls.box = None - - def test_1frame(self) -> None: - ee0, ff0, vv0 = self.dp_original.eval( - self.coords, self.box, self.atype, atomic=False - ) - ee1, ff1, vv1 = self.dp_compressed.eval( - self.coords, self.box, self.atype, atomic=False - ) - # check shape of the returns - nframes = 1 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - def test_1frame_atm(self) -> None: - ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( - self.coords, self.box, self.atype, atomic=True - ) - ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( - self.coords, self.box, self.atype, atomic=True - ) - # check shape of the returns - nframes = 1 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ae0.shape, (nframes, natoms, 1)) - self.assertEqual(av0.shape, (nframes, natoms, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - self.assertEqual(ae1.shape, (nframes, natoms, 1)) - self.assertEqual(av1.shape, (nframes, natoms, 9)) - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ae0, ae1, default_places) - np.testing.assert_almost_equal(av0, av1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - def test_2frame_atm(self) -> None: - coords2 = np.concatenate((self.coords, self.coords)) - ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( - coords2, self.box, self.atype, atomic=True - ) - ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( - coords2, self.box, self.atype, atomic=True - ) - # check shape of the returns - nframes = 2 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ae0.shape, (nframes, natoms, 1)) - self.assertEqual(av0.shape, (nframes, natoms, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - self.assertEqual(ae1.shape, (nframes, natoms, 1)) - self.assertEqual(av1.shape, (nframes, natoms, 9)) - - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ae0, ae1, default_places) - np.testing.assert_almost_equal(av0, av1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - -class TestDeepPotALargeBoxNoPBC(unittest.TestCase): - @classmethod - def setUpClass(cls) -> None: - cls.dp_original = DeepEval(FROZEN_MODEL) - cls.dp_compressed = DeepEval(COMPRESSED_MODEL) - cls.coords = np.array( - [ - 12.83, - 2.56, - 2.18, - 12.09, - 2.87, - 2.74, - 00.25, - 3.32, - 1.68, - 3.36, - 3.00, - 1.81, - 3.51, - 2.51, - 2.60, - 4.27, - 3.22, - 1.56, - ] - ) - cls.atype = [0, 1, 1, 0, 1, 1] - cls.box = np.array([19.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - - def test_1frame(self) -> None: - ee0, ff0, vv0 = self.dp_original.eval( - self.coords, self.box, self.atype, atomic=False - ) - ee1, ff1, vv1 = self.dp_compressed.eval( - self.coords, self.box, self.atype, atomic=False - ) - # check shape of the returns - nframes = 1 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - def test_1frame_atm(self) -> None: - ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( - self.coords, self.box, self.atype, atomic=True - ) - ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( - self.coords, self.box, self.atype, atomic=True - ) - # check shape of the returns - nframes = 1 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ae0.shape, (nframes, natoms, 1)) - self.assertEqual(av0.shape, (nframes, natoms, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - self.assertEqual(ae1.shape, (nframes, natoms, 1)) - self.assertEqual(av1.shape, (nframes, natoms, 9)) - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ae0, ae1, default_places) - np.testing.assert_almost_equal(av0, av1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - def test_ase(self) -> None: - from ase import ( - Atoms, - ) - - from deepmd.calculator import ( - DP, - ) - - water0 = Atoms( - "OHHOHH", - positions=self.coords.reshape((-1, 3)), - cell=self.box.reshape((3, 3)), - calculator=DP(FROZEN_MODEL), - ) - water1 = Atoms( - "OHHOHH", - positions=self.coords.reshape((-1, 3)), - cell=self.box.reshape((3, 3)), - calculator=DP(COMPRESSED_MODEL), - ) - ee0 = water0.get_potential_energy() - ff0 = water0.get_forces() - ee1 = water1.get_potential_energy() - ff1 = water1.get_forces() - # nframes = 1 - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - - -class TestDeepPotAPBCExcludeTypes(unittest.TestCase): - @classmethod - def setUpClass(cls) -> None: - cls.dp_original = DeepEval(FROZEN_MODEL_ET) - cls.dp_compressed = DeepEval(COMPRESSED_MODEL_ET) - cls.coords = np.array( - [ - 12.83, - 2.56, - 2.18, - 12.09, - 2.87, - 2.74, - 00.25, - 3.32, - 1.68, - 3.36, - 3.00, - 1.81, - 3.51, - 2.51, - 2.60, - 4.27, - 3.22, - 1.56, - ] - ) - cls.atype = [0, 1, 1, 0, 1, 1] - cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - - def test_attrs(self) -> None: - self.assertEqual(self.dp_original.get_ntypes(), 2) - self.assertAlmostEqual(self.dp_original.get_rcut(), 4.0, places=default_places) - self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) - self.assertEqual(self.dp_original.get_dim_fparam(), 0) - self.assertEqual(self.dp_original.get_dim_aparam(), 0) - - self.assertEqual(self.dp_compressed.get_ntypes(), 2) - self.assertAlmostEqual( - self.dp_compressed.get_rcut(), 4.0, places=default_places - ) - self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) - self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) - self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - - def test_1frame(self) -> None: - ee0, ff0, vv0 = self.dp_original.eval( - self.coords, self.box, self.atype, atomic=False - ) - ee1, ff1, vv1 = self.dp_compressed.eval( - self.coords, self.box, self.atype, atomic=False - ) - # check shape of the returns - nframes = 1 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - def test_1frame_atm(self) -> None: - ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( - self.coords, self.box, self.atype, atomic=True - ) - ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( - self.coords, self.box, self.atype, atomic=True - ) - # check shape of the returns - nframes = 1 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ae0.shape, (nframes, natoms, 1)) - self.assertEqual(av0.shape, (nframes, natoms, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - self.assertEqual(ae1.shape, (nframes, natoms, 1)) - self.assertEqual(av1.shape, (nframes, natoms, 9)) - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ae0, ae1, default_places) - np.testing.assert_almost_equal(av0, av1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - def test_2frame_atm(self) -> None: - coords2 = np.concatenate((self.coords, self.coords)) - box2 = np.concatenate((self.box, self.box)) - ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( - coords2, box2, self.atype, atomic=True - ) - ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( - coords2, box2, self.atype, atomic=True - ) - # check shape of the returns - nframes = 2 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ae0.shape, (nframes, natoms, 1)) - self.assertEqual(av0.shape, (nframes, natoms, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - self.assertEqual(ae1.shape, (nframes, natoms, 1)) - self.assertEqual(av1.shape, (nframes, natoms, 9)) - - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ae0, ae1, default_places) - np.testing.assert_almost_equal(av0, av1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - -class TestSkipNeighborStat(unittest.TestCase): - @classmethod - def setUpClass(cls) -> None: - cls.dp_original = DeepEval(FROZEN_MODEL_SKIP_NEIGHBOR_STAT) - cls.dp_compressed = DeepEval(COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT) - cls.coords = np.array( - [ - 12.83, - 2.56, - 2.18, - 12.09, - 2.87, - 2.74, - 00.25, - 3.32, - 1.68, - 3.36, - 3.00, - 1.81, - 3.51, - 2.51, - 2.60, - 4.27, - 3.22, - 1.56, - ] - ) - cls.atype = [0, 1, 1, 0, 1, 1] - cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) - - def test_attrs(self) -> None: - self.assertEqual(self.dp_original.get_ntypes(), 2) - self.assertAlmostEqual(self.dp_original.get_rcut(), 4.0, places=default_places) - self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) - self.assertEqual(self.dp_original.get_dim_fparam(), 0) - self.assertEqual(self.dp_original.get_dim_aparam(), 0) - - self.assertEqual(self.dp_compressed.get_ntypes(), 2) - self.assertAlmostEqual( - self.dp_compressed.get_rcut(), 4.0, places=default_places - ) - self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) - self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) - self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) - - def test_1frame(self) -> None: - ee0, ff0, vv0 = self.dp_original.eval( - self.coords, self.box, self.atype, atomic=False - ) - ee1, ff1, vv1 = self.dp_compressed.eval( - self.coords, self.box, self.atype, atomic=False - ) - # check shape of the returns - nframes = 1 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - def test_1frame_atm(self) -> None: - ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( - self.coords, self.box, self.atype, atomic=True - ) - ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( - self.coords, self.box, self.atype, atomic=True - ) - # check shape of the returns - nframes = 1 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ae0.shape, (nframes, natoms, 1)) - self.assertEqual(av0.shape, (nframes, natoms, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - self.assertEqual(ae1.shape, (nframes, natoms, 1)) - self.assertEqual(av1.shape, (nframes, natoms, 9)) - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ae0, ae1, default_places) - np.testing.assert_almost_equal(av0, av1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - def test_2frame_atm(self) -> None: - coords2 = np.concatenate((self.coords, self.coords)) - box2 = np.concatenate((self.box, self.box)) - ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( - coords2, box2, self.atype, atomic=True - ) - ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( - coords2, box2, self.atype, atomic=True - ) - # check shape of the returns - nframes = 2 - natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes, 1)) - self.assertEqual(ff0.shape, (nframes, natoms, 3)) - self.assertEqual(vv0.shape, (nframes, 9)) - self.assertEqual(ae0.shape, (nframes, natoms, 1)) - self.assertEqual(av0.shape, (nframes, natoms, 9)) - self.assertEqual(ee1.shape, (nframes, 1)) - self.assertEqual(ff1.shape, (nframes, natoms, 3)) - self.assertEqual(vv1.shape, (nframes, 9)) - self.assertEqual(ae1.shape, (nframes, natoms, 1)) - self.assertEqual(av1.shape, (nframes, natoms, 9)) - - # check values - np.testing.assert_almost_equal(ff0, ff1, default_places) - np.testing.assert_almost_equal(ae0, ae1, default_places) - np.testing.assert_almost_equal(av0, av1, default_places) - np.testing.assert_almost_equal(ee0, ee1, default_places) - np.testing.assert_almost_equal(vv0, vv1, default_places) - - -if __name__ == "__main__": - unittest.main()