From fe14402b9ecc8f7a52bf7a6b072321c8d0bbce81 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 19 Feb 2024 22:07:53 -0500 Subject: [PATCH] fix neighbor stat mixed_types input (#3304) It was broken by #3289, which replaced distinguished types by mixed types but didn't change the input. Add tests for two types. --------- Signed-off-by: Jinzhe Zeng --- deepmd/dpmodel/utils/neighbor_stat.py | 8 ++-- deepmd/entrypoints/neighbor_stat.py | 6 +-- deepmd/main.py | 1 + deepmd/pt/utils/neighbor_stat.py | 14 ++++--- deepmd/tf/entrypoints/train.py | 14 +++---- deepmd/tf/utils/neighbor_stat.py | 8 ++-- deepmd/utils/neighbor_stat.py | 8 ++-- .../common/dpmodel/test_neighbor_stat.py | 13 +++--- source/tests/pt/test_neighbor_stat.py | 13 +++--- source/tests/tf/test_neighbor_stat.py | 41 +++++++++++-------- source/tests/tf/test_virtual_type.py | 2 +- 11 files changed, 71 insertions(+), 57 deletions(-) diff --git a/deepmd/dpmodel/utils/neighbor_stat.py b/deepmd/dpmodel/utils/neighbor_stat.py index bf0c30c153..96b39d20ad 100644 --- a/deepmd/dpmodel/utils/neighbor_stat.py +++ b/deepmd/dpmodel/utils/neighbor_stat.py @@ -115,7 +115,7 @@ class NeighborStat(BaseNeighborStat): The num of atom types rcut : float The cut-off radius - one_type : bool, optional, default=False + mixed_type : bool, optional, default=False Treat all types as a single type. """ @@ -123,10 +123,10 @@ def __init__( self, ntypes: int, rcut: float, - one_type: bool = False, + mixed_type: bool = False, ) -> None: - super().__init__(ntypes, rcut, one_type) - self.op = NeighborStatOP(ntypes, rcut, not one_type) + super().__init__(ntypes, rcut, mixed_type) + self.op = NeighborStatOP(ntypes, rcut, mixed_type) def iterator( self, data: DeepmdDataSystem diff --git a/deepmd/entrypoints/neighbor_stat.py b/deepmd/entrypoints/neighbor_stat.py index 8a496fb6f0..a68a3fd3bb 100644 --- a/deepmd/entrypoints/neighbor_stat.py +++ b/deepmd/entrypoints/neighbor_stat.py @@ -22,7 +22,7 @@ def neighbor_stat( system: str, rcut: float, type_map: List[str], - one_type: bool = False, + mixed_type: bool = False, backend: str = "tensorflow", **kwargs, ): @@ -36,7 +36,7 @@ def neighbor_stat( cutoff radius type_map : list[str] type map - one_type : bool, optional, default=False + mixed_type : bool, optional, default=False treat all types as a single type backend : str, optional, default="tensorflow" backend to use @@ -89,7 +89,7 @@ def neighbor_stat( type_map=type_map, ) data.get_batch() - nei = NeighborStat(data.get_ntypes(), rcut, one_type=one_type) + nei = NeighborStat(data.get_ntypes(), rcut, mixed_type=mixed_type) min_nbor_dist, max_nbor_size = nei.get_stat(data) log.info("min_nbor_dist: %f" % min_nbor_dist) log.info("max_nbor_size: %s" % str(max_nbor_size)) diff --git a/deepmd/main.py b/deepmd/main.py index d31cab30c2..2bde6376f2 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -640,6 +640,7 @@ def main_parser() -> argparse.ArgumentParser: help="type map", ) parser_neighbor_stat.add_argument( + "--mixed-type", "--one-type", action="store_true", default=False, diff --git a/deepmd/pt/utils/neighbor_stat.py b/deepmd/pt/utils/neighbor_stat.py index 6361e7b9c7..d5b5c74bdc 100644 --- a/deepmd/pt/utils/neighbor_stat.py +++ b/deepmd/pt/utils/neighbor_stat.py @@ -89,14 +89,16 @@ def forward( ) assert list(diff.shape) == [nframes, nloc, nall, 3] # remove the diagonal elements - mask = torch.eye(nloc, nall, dtype=torch.bool) + mask = torch.eye(nloc, nall, dtype=torch.bool, device=diff.device) diff[:, mask] = torch.inf rr2 = torch.sum(torch.square(diff), dim=-1) min_rr2, _ = torch.min(rr2, dim=-1) # count the number of neighbors if not self.mixed_types: mask = rr2 < self.rcut**2 - nnei = torch.zeros((nframes, nloc, self.ntypes), dtype=torch.int32) + nnei = torch.zeros( + (nframes, nloc, self.ntypes), dtype=torch.int32, device=mask.device + ) for ii in range(self.ntypes): nnei[:, :, ii] = torch.sum( mask & extend_atype.eq(ii)[:, None, :], dim=-1 @@ -120,7 +122,7 @@ class NeighborStat(BaseNeighborStat): The num of atom types rcut : float The cut-off radius - one_type : bool, optional, default=False + mixed_type : bool, optional, default=False Treat all types as a single type. """ @@ -128,10 +130,10 @@ def __init__( self, ntypes: int, rcut: float, - one_type: bool = False, + mixed_type: bool = False, ) -> None: - super().__init__(ntypes, rcut, one_type) - op = NeighborStatOP(ntypes, rcut, not one_type) + super().__init__(ntypes, rcut, mixed_type) + op = NeighborStatOP(ntypes, rcut, mixed_type) self.op = torch.jit.script(op) self.auto_batch_size = AutoBatchSize() diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index f0b9481d99..3759ff9331 100755 --- a/deepmd/tf/entrypoints/train.py +++ b/deepmd/tf/entrypoints/train.py @@ -370,7 +370,7 @@ def get_type_map(jdata): return jdata["model"].get("type_map", None) -def get_nbor_stat(jdata, rcut, one_type: bool = False): +def get_nbor_stat(jdata, rcut, mixed_type: bool = False): # it seems that DeepmdDataSystem does not need rcut # it's not clear why there is an argument... # max_rcut = get_rcut(jdata) @@ -414,7 +414,7 @@ def get_nbor_stat(jdata, rcut, one_type: bool = False): map_ntypes = data_ntypes ntypes = max([map_ntypes, data_ntypes]) - neistat = NeighborStat(ntypes, rcut, one_type=one_type) + neistat = NeighborStat(ntypes, rcut, mixed_type=mixed_type) min_nbor_dist, max_nbor_size = neistat.get_stat(train_data) @@ -430,8 +430,8 @@ def get_nbor_stat(jdata, rcut, one_type: bool = False): return min_nbor_dist, max_nbor_size -def get_sel(jdata, rcut, one_type: bool = False): - _, max_nbor_size = get_nbor_stat(jdata, rcut, one_type=one_type) +def get_sel(jdata, rcut, mixed_type: bool = False): + _, max_nbor_size = get_nbor_stat(jdata, rcut, mixed_type=mixed_type) return max_nbor_size @@ -468,12 +468,12 @@ def wrap_up_4(xx): return 4 * ((int(xx) + 3) // 4) -def update_one_sel(jdata, descriptor, one_type: bool = False): +def update_one_sel(jdata, descriptor, mixed_type: bool = False): rcut = descriptor["rcut"] tmp_sel = get_sel( jdata, rcut, - one_type=one_type, + mixed_type=mixed_type, ) sel = descriptor["sel"] if isinstance(sel, int): @@ -493,7 +493,7 @@ def update_one_sel(jdata, descriptor, one_type: bool = False): "not less than %d, but you set it to %d. The accuracy" " of your model may get worse." % (ii, tt, dd) ) - if one_type: + if mixed_type: descriptor["sel"] = sel = sum(sel) return descriptor diff --git a/deepmd/tf/utils/neighbor_stat.py b/deepmd/tf/utils/neighbor_stat.py index 61531e7857..f668d4a4da 100644 --- a/deepmd/tf/utils/neighbor_stat.py +++ b/deepmd/tf/utils/neighbor_stat.py @@ -168,7 +168,7 @@ class NeighborStat(BaseNeighborStat): The num of atom types rcut The cut-off radius - one_type : bool, optional, default=False + mixed_type : bool, optional, default=False Treat all types as a single type. """ @@ -176,12 +176,12 @@ def __init__( self, ntypes: int, rcut: float, - one_type: bool = False, + mixed_type: bool = False, ) -> None: """Constructor.""" - super().__init__(ntypes, rcut, one_type) + super().__init__(ntypes, rcut, mixed_type) self.auto_batch_size = AutoBatchSize() - self.neighbor_stat = NeighborStatOP(ntypes, rcut, not one_type) + self.neighbor_stat = NeighborStatOP(ntypes, rcut, mixed_type) self.place_holders = {} with tf.Graph().as_default() as sub_graph: self.op = self.build() diff --git a/deepmd/utils/neighbor_stat.py b/deepmd/utils/neighbor_stat.py index c6327a705e..34200df007 100644 --- a/deepmd/utils/neighbor_stat.py +++ b/deepmd/utils/neighbor_stat.py @@ -32,7 +32,7 @@ class NeighborStat(ABC): The num of atom types rcut : float The cut-off radius - one_type : bool, optional, default=False + mixed_type : bool, optional, default=False Treat all types as a single type. """ @@ -40,11 +40,11 @@ def __init__( self, ntypes: int, rcut: float, - one_type: bool = False, + mixed_type: bool = False, ) -> None: self.rcut = rcut self.ntypes = ntypes - self.one_type = one_type + self.mixed_type = mixed_type def get_stat(self, data: DeepmdDataSystem) -> Tuple[float, np.ndarray]: """Get the data statistics of the training data, including nearest nbor distance between atoms, max nbor size of atoms. @@ -62,7 +62,7 @@ def get_stat(self, data: DeepmdDataSystem) -> Tuple[float, np.ndarray]: An array with ntypes integers, denotes the actual achieved max sel """ min_nbor_dist = 100.0 - max_nbor_size = np.zeros(1 if self.one_type else self.ntypes, dtype=int) + max_nbor_size = np.zeros(1 if self.mixed_type else self.ntypes, dtype=int) for mn, dt, jj in self.iterator(data): if np.isinf(dt): diff --git a/source/tests/common/dpmodel/test_neighbor_stat.py b/source/tests/common/dpmodel/test_neighbor_stat.py index aa2ed5b3db..90764a049c 100644 --- a/source/tests/common/dpmodel/test_neighbor_stat.py +++ b/source/tests/common/dpmodel/test_neighbor_stat.py @@ -39,14 +39,14 @@ def tearDown(self): def test_neighbor_stat(self): for rcut in (0.0, 1.0, 2.0, 4.0): - for one_type in (True, False): - with self.subTest(rcut=rcut, one_type=one_type): + for mixed_type in (True, False): + with self.subTest(rcut=rcut, mixed_type=mixed_type): rcut += 1e-3 # prevent numerical errors min_nbor_dist, max_nbor_size = neighbor_stat( system="system_0", rcut=rcut, - type_map=["TYPE"], - one_type=one_type, + type_map=["TYPE", "NO_THIS_TYPE"], + mixed_type=mixed_type, backend="numpy", ) upper = np.ceil(rcut) + 1 @@ -58,4 +58,7 @@ def test_neighbor_stat(self): np.logical_and(distance > 0, distance <= rcut) ) self.assertAlmostEqual(min_nbor_dist, 1.0, 6) - self.assertEqual(max_nbor_size, [expected_neighbors]) + ret = [expected_neighbors] + if not mixed_type: + ret.append(0) + np.testing.assert_array_equal(max_nbor_size, ret) diff --git a/source/tests/pt/test_neighbor_stat.py b/source/tests/pt/test_neighbor_stat.py index 7e4be0909e..63ec684792 100644 --- a/source/tests/pt/test_neighbor_stat.py +++ b/source/tests/pt/test_neighbor_stat.py @@ -39,14 +39,14 @@ def tearDown(self): def test_neighbor_stat(self): for rcut in (0.0, 1.0, 2.0, 4.0): - for one_type in (True, False): - with self.subTest(rcut=rcut, one_type=one_type): + for mixed_type in (True, False): + with self.subTest(rcut=rcut, mixed_type=mixed_type): rcut += 1e-3 # prevent numerical errors min_nbor_dist, max_nbor_size = neighbor_stat( system="system_0", rcut=rcut, - type_map=["TYPE"], - one_type=one_type, + type_map=["TYPE", "NO_THIS_TYPE"], + mixed_type=mixed_type, backend="pytorch", ) upper = np.ceil(rcut) + 1 @@ -58,4 +58,7 @@ def test_neighbor_stat(self): np.logical_and(distance > 0, distance <= rcut) ) self.assertAlmostEqual(min_nbor_dist, 1.0, 6) - self.assertEqual(max_nbor_size, [expected_neighbors]) + ret = [expected_neighbors] + if not mixed_type: + ret.append(0) + np.testing.assert_array_equal(max_nbor_size, ret) diff --git a/source/tests/tf/test_neighbor_stat.py b/source/tests/tf/test_neighbor_stat.py index 9806e2053a..6fcba36914 100644 --- a/source/tests/tf/test_neighbor_stat.py +++ b/source/tests/tf/test_neighbor_stat.py @@ -38,21 +38,26 @@ def tearDown(self): shutil.rmtree("system_0") def test_neighbor_stat(self): - # set rcut to 0. will cause a core dumped - # TODO: check what is wrong - for rcut in (1.0, 2.0, 4.0): - with self.subTest(): - rcut += 1e-3 # prevent numerical errors - min_nbor_dist, max_nbor_size = neighbor_stat( - system="system_0", rcut=rcut, type_map=["TYPE"] - ) - upper = np.ceil(rcut) + 1 - X, Y, Z = np.mgrid[-upper:upper, -upper:upper, -upper:upper] - positions = np.vstack([X.ravel(), Y.ravel(), Z.ravel()]).T - # distance to (0,0,0) - distance = np.linalg.norm(positions, axis=1) - expected_neighbors = np.count_nonzero( - np.logical_and(distance > 0, distance <= rcut) - ) - self.assertAlmostEqual(min_nbor_dist, 1.0, 6) - self.assertEqual(max_nbor_size, [expected_neighbors]) + for rcut in (0.0, 1.0, 2.0, 4.0): + for mixed_type in (True, False): + with self.subTest(rcut=rcut, mixed_type=mixed_type): + rcut += 1e-3 # prevent numerical errors + min_nbor_dist, max_nbor_size = neighbor_stat( + system="system_0", + rcut=rcut, + type_map=["TYPE", "NO_THIS_TYPE"], + mixed_type=mixed_type, + ) + upper = np.ceil(rcut) + 1 + X, Y, Z = np.mgrid[-upper:upper, -upper:upper, -upper:upper] + positions = np.vstack([X.ravel(), Y.ravel(), Z.ravel()]).T + # distance to (0,0,0) + distance = np.linalg.norm(positions, axis=1) + expected_neighbors = np.count_nonzero( + np.logical_and(distance > 0, distance <= rcut) + ) + self.assertAlmostEqual(min_nbor_dist, 1.0, 6) + ret = [expected_neighbors] + if not mixed_type: + ret.append(0) + np.testing.assert_array_equal(max_nbor_size, ret) diff --git a/source/tests/tf/test_virtual_type.py b/source/tests/tf/test_virtual_type.py index 5ceb1c7637..e9c675fe3a 100644 --- a/source/tests/tf/test_virtual_type.py +++ b/source/tests/tf/test_virtual_type.py @@ -138,5 +138,5 @@ def test_data_mixed_type(self): data = DeepmdDataSystem(systems, batch_size, test_size, rcut, type_map=type_map) data.get_batch() # neighbor stat - nei_stat = NeighborStat(len(type_map), rcut, one_type=True) + nei_stat = NeighborStat(len(type_map), rcut, mixed_type=True) min_nbor_dist, max_nbor_size = nei_stat.get_stat(data)