Fix shape mismatch when type_embedding is enabled and type_one_side i… (#1074)

* Fix shape mismatch when type_embedding is enabled and type_one_side is disabled.

* Add UT to cover the case where batch_size is larger than 1.

* Fix random issue in unit tests.

Author: shishaochen

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
 
 

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
@@ -26,8 +25,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,14 +46,15 @@ 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)
-        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)
@@ -69,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)
 
@@ -128,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

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')