From dba8704b552b88246bbb43a2a6f5178f3ebc1a88 Mon Sep 17 00:00:00 2001 From: Shaochen Shi Date: Tue, 31 Aug 2021 13:19:58 +0800 Subject: [PATCH 1/3] Fix shape mismatch when type_embedding is enabled and type_one_side is disabled. --- deepmd/descriptor/se_a.py | 35 +++++++++++++++++++++++++++-------- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index fbc9a77b56..3777368c0f 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -667,17 +667,36 @@ def _concat_type_embedding( nframes, natoms, type_embedding, - ): + ): + '''Concatenate `type_embedding` of neighbors and `xyz_scatter`. + If not self.type_one_side, concatenate `type_embedding` of center atoms as well. + + Parameters + ---------- + xyz_scatter: + shape is [nframes*natoms[0]*self.nnei, 1] + nframes: + shape is [] + natoms: + shape is [1+1+self.ntypes] + type_embedding: + shape is [self.ntypes, Y] where Y=jdata['type_embedding']['neuron'][-1] + + Returns + ------- + embedding: + environment of each atom represented by embedding. + ''' te_out_dim = type_embedding.get_shape().as_list()[-1] - nei_embed = tf.nn.embedding_lookup(type_embedding,tf.cast(self.nei_type,dtype=tf.int32)) #nnei*nchnl - nei_embed = tf.tile(nei_embed,(nframes*natoms[0],1)) + nei_embed = tf.nn.embedding_lookup(type_embedding,tf.cast(self.nei_type,dtype=tf.int32)) # shape is [self.nnei, 1+te_out_dim] + nei_embed = tf.tile(nei_embed,(nframes*natoms[0],1)) # shape is [nframes*natoms[0]*self.nnei, te_out_dim] nei_embed = tf.reshape(nei_embed,[-1,te_out_dim]) - embedding_input = tf.concat([xyz_scatter,nei_embed],1) + embedding_input = tf.concat([xyz_scatter,nei_embed],1) # shape is [nframes*natoms[0]*self.nnei, 1+te_out_dim] if not self.type_one_side: - atm_embed = embed_atom_type(self.ntypes, natoms, type_embedding) - atm_embed = tf.tile(atm_embed,(1,self.nnei)) - atm_embed = tf.reshape(atm_embed,[-1,te_out_dim]) - embedding_input = tf.concat([embedding_input,atm_embed],1) + atm_embed = embed_atom_type(self.ntypes, natoms, type_embedding) # shape is [natoms[0], te_out_dim] + atm_embed = tf.tile(atm_embed,(nframes,self.nnei)) # shape is [nframes*natoms[0], self.nnei*te_out_dim] + atm_embed = tf.reshape(atm_embed,[-1,te_out_dim]) # shape is [nframes*natoms[0]*self.nnei, te_out_dim] + embedding_input = tf.concat([embedding_input,atm_embed],1) # shape is [nframes*natoms[0]*self.nnei, 1+te_out_dim+te_out_dim] return embedding_input From a1b1e84ce49f832b4f53aeaa39190372f4ae556b Mon Sep 17 00:00:00 2001 From: Shaochen Shi Date: Tue, 31 Aug 2021 14:01:21 +0800 Subject: [PATCH 2/3] Add UT to cover the case where batch_size is larger than 1. --- source/tests/common.py | 9 +++++---- source/tests/test_descrpt_se_a_type.py | 6 +++--- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/source/tests/common.py b/source/tests/common.py index 87dd5d6d70..ea083c565d 100644 --- a/source/tests/common.py +++ b/source/tests/common.py @@ -26,8 +26,8 @@ def del_data(): if os.path.isdir('system'): shutil.rmtree('system') -def gen_data() : - tmpdata = Data(rand_pert = 0.1, seed = 1) +def gen_data(nframes = 1) : + tmpdata = Data(rand_pert = 0.1, seed = 1, nframes = nframes) sys = dpdata.LabeledSystem() sys.data['atom_names'] = ['foo', 'bar'] sys.data['coords'] = tmpdata.coord @@ -47,10 +47,11 @@ class Data(): def __init__ (self, rand_pert = 0.1, seed = 1, - box_scale = 20) : + box_scale = 20, + nframes = 1): coord = [[0.0, 0.0, 0.1], [1.1, 0.0, 0.1], [0.0, 1.1, 0.1], [4.0, 0.0, 0.0], [5.1, 0.0, 0.0], [4.0, 1.1, 0.0]] - self.nframes = 1 + self.nframes = nframes self.coord = np.array(coord) self.coord = self._copy_nframes(self.coord) np.random.seed(seed) diff --git a/source/tests/test_descrpt_se_a_type.py b/source/tests/test_descrpt_se_a_type.py index fded177b66..f65338d0a4 100644 --- a/source/tests/test_descrpt_se_a_type.py +++ b/source/tests/test_descrpt_se_a_type.py @@ -17,8 +17,8 @@ GLOBAL_NP_FLOAT_PRECISION = np.float64 class TestModel(tf.test.TestCase): - def setUp(self) : - gen_data() + def setUp(self): + gen_data(nframes=2) def test_descriptor_two_sides(self): jfile = 'water_se_a_type.json' @@ -28,7 +28,7 @@ def test_descriptor_two_sides(self): set_pfx = j_must_have(jdata, 'set_prefix') batch_size = j_must_have(jdata, 'batch_size') test_size = j_must_have(jdata, 'numb_test') - batch_size = 1 + batch_size = 2 test_size = 1 stop_batch = j_must_have(jdata, 'stop_batch') rcut = j_must_have (jdata['model']['descriptor'], 'rcut') From f2e176de5d0fd429c41c0bab2882a60f14a66b7d Mon Sep 17 00:00:00 2001 From: Shaochen Shi Date: Tue, 31 Aug 2021 15:52:51 +0800 Subject: [PATCH 3/3] Fix random issue in unit tests. --- source/tests/common.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/source/tests/common.py b/source/tests/common.py index ea083c565d..e767f8a3eb 100644 --- a/source/tests/common.py +++ b/source/tests/common.py @@ -3,10 +3,9 @@ import pathlib from deepmd.env import tf -from deepmd.env import GLOBAL_TF_FLOAT_PRECISION from deepmd.env import GLOBAL_NP_FLOAT_PRECISION -from deepmd.env import GLOBAL_ENER_FLOAT_PRECISION from deepmd.common import j_loader as dp_j_loader +from deepmd.utils import random as dp_random if GLOBAL_NP_FLOAT_PRECISION == np.float32 : global_default_fv_hh = 1e-2 @@ -54,8 +53,8 @@ def __init__ (self, self.nframes = nframes self.coord = np.array(coord) self.coord = self._copy_nframes(self.coord) - np.random.seed(seed) - self.coord += rand_pert * np.random.random(self.coord.shape) + dp_random.seed(seed) + self.coord += rand_pert * dp_random.random(self.coord.shape) self.fparam = np.array([[0.1, 0.2]]) self.aparam = np.tile(self.fparam, [1, 6]) self.fparam = self._copy_nframes(self.fparam) @@ -70,7 +69,7 @@ def __init__ (self, self.coord = self.coord.reshape([self.nframes, -1, 3]) self.coord = self.coord[:,self.idx_map,:] self.coord = self.coord.reshape([self.nframes, -1]) - self.efield = np.random.random(self.coord.shape) + self.efield = dp_random.random(self.coord.shape) self.atype = self.atype[self.idx_map] self.datype = self._copy_nframes(self.atype) @@ -129,7 +128,7 @@ def get_test_box_data (self, coord0_, box0_, type0_ = self.get_data() coord = coord0_[0] box = box0_[0] - box += rand_pert * np.random.random(box.shape) + box += rand_pert * dp_random.random(box.shape) atype = type0_[0] nframes = 1 natoms = coord.size // 3