From 0a691156e270cdd002059b811ebed3d0fcbe3aaf Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 13 Feb 2023 00:41:39 -0500 Subject: [PATCH 1/2] support exclude_types for se_atten Signed-off-by: Jinzhe Zeng --- deepmd/descriptor/se_atten.py | 121 ++++++++++++++++++++++------ source/tests/test_model_se_atten.py | 75 +++++++++++++++++ 2 files changed, 171 insertions(+), 25 deletions(-) diff --git a/deepmd/descriptor/se_atten.py b/deepmd/descriptor/se_atten.py index 912a4db74b..43c1cf1a2e 100644 --- a/deepmd/descriptor/se_atten.py +++ b/deepmd/descriptor/se_atten.py @@ -446,6 +446,18 @@ def _pass_filter( inputs_i = inputs inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) type_i = -1 + if len(self.exclude_types): + mask = self.build_type_exclude_mask( + self.exclude_types, + self.ntypes, + self.sel_a, + self.ndescrpt, + atype, + tf.shape(inputs_i)[0], + self.nei_type_vec, # extra input for atten + ) + inputs_i *= mask + layer, qmat = self._filter( inputs_i, type_i, @@ -854,10 +866,7 @@ def _filter_lower( ) # xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, outputs_size[-1])) else: - # we can safely return the final xyz_scatter filled with zero directly - return tf.cast( - tf.fill((natom, 4, outputs_size[-1]), 0.0), self.filter_precision - ) + raise RuntimeError("this should not be touched") # When using tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]) below # [588 24] -> [588 6 4] correct # but if sel is zero @@ -890,27 +899,6 @@ def _filter( shape = inputs.get_shape().as_list() outputs_size = [1] + self.filter_neuron outputs_size_2 = self.n_axis_neuron - all_excluded = all( - [ - (type_input, type_i) in self.exclude_types - for type_i in range(self.ntypes) - ] - ) - if all_excluded: - # all types are excluded so result and qmat should be zeros - # we can safaly return a zero matrix... - # See also https://stackoverflow.com/a/34725458/9567349 - # result: natom x outputs_size x outputs_size_2 - # qmat: natom x outputs_size x 3 - natom = tf.shape(inputs)[0] - result = tf.cast( - tf.fill((natom, outputs_size_2, outputs_size[-1]), 0.0), - GLOBAL_TF_FLOAT_PRECISION, - ) - qmat = tf.cast( - tf.fill((natom, outputs_size[-1], 3), 0.0), GLOBAL_TF_FLOAT_PRECISION - ) - return result, qmat start_index = 0 type_i = 0 @@ -1007,3 +995,86 @@ def init_variables( i, suffix, i ) ] + + def build_type_exclude_mask( + self, + exclude_types: List[Tuple[int, int]], + ntypes: int, + sel: List[int], + ndescrpt: int, + atype: tf.Tensor, + shape0: tf.Tensor, + nei_type_vec: tf.Tensor, + ) -> tf.Tensor: + r"""Build the type exclude mask for the attention descriptor. + + Notes + ----- + This method has the similiar way to build the type exclude mask as + :meth:`deepmd.descriptor.descriptor.Descriptor.build_type_exclude_mask`. + The mathmatical expression has been explained in that method. + The difference is that the attention descriptor has provided the type of + the neighbors (idx_j) that is not in order, so we use it from an extra + input. + + Parameters + ---------- + exclude_types : List[Tuple[int, int]] + The list of excluded types, e.g. [(0, 1), (1, 0)] means the interaction + between type 0 and type 1 is excluded. + ntypes : int + The number of types. + sel : List[int] + The list of the number of selected neighbors for each type. + ndescrpt : int + The number of descriptors for each atom. + atype : tf.Tensor + The type of atoms, with the size of shape0. + shape0 : tf.Tensor + The shape of the first dimension of the inputs, which is equal to + nsamples * natoms. + nei_type_vec : tf.Tensor + The type of neighbors, with the size of (shape0, nnei). + + Returns + ------- + tf.Tensor + The type exclude mask, with the shape of (shape0, ndescrpt), and the + precision of GLOBAL_TF_FLOAT_PRECISION. The mask has the value of 1 if the + interaction between two types is not excluded, and 0 otherwise. + + See Also + -------- + deepmd.descriptor.descriptor.Descriptor.build_type_exclude_mask + """ + # generate a mask + type_mask = np.array( + [ + [ + 1 if (tt_i, tt_j) not in exclude_types else 0 + for tt_i in range(ntypes) + ] + for tt_j in range(ntypes) + ], + dtype=bool, + ) + type_mask = tf.convert_to_tensor(type_mask, dtype=GLOBAL_TF_FLOAT_PRECISION) + type_mask = tf.reshape(type_mask, [-1]) + + # (nsamples * natoms, 1) + atype_expand = tf.reshape(atype, [-1, 1]) + # (nsamples * natoms, ndescrpt) + idx_i = tf.tile(atype_expand * ntypes, (1, ndescrpt)) + # idx_j has been provided by atten op + # (nsamples * natoms, nnei, 1) + idx_j = tf.reshape(nei_type_vec, [shape0, sel[0], 1]) + # (nsamples * natoms, nnei, ndescrpt // nnei) + idx_j = tf.tile(idx_j, (1, 1, ndescrpt // sel[0])) + # (nsamples * natoms, ndescrpt) + idx_j = tf.reshape(idx_j, [shape0, ndescrpt]) + idx = idx_i + idx_j + idx = tf.reshape(idx, [-1]) + mask = tf.nn.embedding_lookup(type_mask, idx) + # same as inputs_i, (nsamples * natoms, ndescrpt) + mask = tf.reshape(mask, [-1, ndescrpt]) + return mask diff --git a/source/tests/test_model_se_atten.py b/source/tests/test_model_se_atten.py index e9ea04d9ad..38f51ef6a3 100644 --- a/source/tests/test_model_se_atten.py +++ b/source/tests/test_model_se_atten.py @@ -192,3 +192,78 @@ def test_model(self): np.testing.assert_almost_equal(e, refe, places) np.testing.assert_almost_equal(f, reff, places) np.testing.assert_almost_equal(v, refv, places) + + def test_exclude_types(self): + """In this test, we make type 0 has no interaction with type 0 and type 1, + so the descriptor should be zero for type 0 atoms. + """ + jfile = "water_se_atten.json" + jdata = j_loader(jfile) + + systems = j_must_have(jdata, "systems") + 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 + test_size = 1 + rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + ntypes = 2 + + data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) + + test_data = data.get_test() + numb_test = 1 + + # set parameters + jdata["model"]["descriptor"]["exclude_types"] = [[0, 0], [0, 1]] + + t_prop_c = tf.placeholder(tf.float32, [5], name="t_prop_c") + t_coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="i_coord") + t_type = tf.placeholder(tf.int32, [None], name="i_type") + t_natoms = tf.placeholder(tf.int32, [ntypes + 2], name="i_natoms") + t_box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None, 9], name="i_box") + t_mesh = tf.placeholder(tf.int32, [None], name="i_mesh") + is_training = tf.placeholder(tf.bool) + + # successful + descrpt = DescrptSeAtten(ntypes=ntypes, **jdata["model"]["descriptor"]) + typeebd_param = jdata["model"]["type_embedding"] + typeebd = TypeEmbedNet( + neuron=typeebd_param["neuron"], + activation_function=None, + resnet_dt=typeebd_param["resnet_dt"], + seed=typeebd_param["seed"], + uniform_seed=True, + padding=True, + ) + type_embedding = typeebd.build( + ntypes, + ) + dout = descrpt.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + {"type_embedding": type_embedding}, + reuse=False, + suffix="_se_atten_exclude_types", + ) + + feed_dict_test1 = { + t_prop_c: test_data["prop_c"], + t_coord: np.reshape(test_data["coord"][:numb_test, :], [-1]), + t_box: test_data["box"][:numb_test, :], + t_type: np.reshape(test_data["type"][:numb_test, :], [-1]), + t_natoms: test_data["natoms_vec"], + t_mesh: test_data["default_mesh"], + is_training: False, + } + + with self.test_session() as sess: + sess.run(tf.global_variables_initializer()) + [des] = sess.run([dout], feed_dict=feed_dict_test1) + + np.testing.assert_almost_equal(des[:, 0:2], 0.0, 10) + with self.assertRaises(AssertionError): + np.testing.assert_almost_equal(des[:, 2:6], 0.0, 10) From 4d676761dfa8580d84fdf274bf106601a4435323 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 13 Feb 2023 01:35:33 -0500 Subject: [PATCH 2/2] add 1 for neighbor mask Signed-off-by: Jinzhe Zeng --- deepmd/descriptor/se_atten.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/deepmd/descriptor/se_atten.py b/deepmd/descriptor/se_atten.py index 43c1cf1a2e..da53d925d5 100644 --- a/deepmd/descriptor/se_atten.py +++ b/deepmd/descriptor/se_atten.py @@ -1048,11 +1048,12 @@ def build_type_exclude_mask( deepmd.descriptor.descriptor.Descriptor.build_type_exclude_mask """ # generate a mask + # op returns ntypes when the neighbor doesn't exist, so we need to add 1 type_mask = np.array( [ [ 1 if (tt_i, tt_j) not in exclude_types else 0 - for tt_i in range(ntypes) + for tt_i in range(ntypes + 1) ] for tt_j in range(ntypes) ], @@ -1064,7 +1065,7 @@ def build_type_exclude_mask( # (nsamples * natoms, 1) atype_expand = tf.reshape(atype, [-1, 1]) # (nsamples * natoms, ndescrpt) - idx_i = tf.tile(atype_expand * ntypes, (1, ndescrpt)) + idx_i = tf.tile(atype_expand * (ntypes + 1), (1, ndescrpt)) # idx_j has been provided by atten op # (nsamples * natoms, nnei, 1) idx_j = tf.reshape(nei_type_vec, [shape0, sel[0], 1])