support no-pbc-training

Author: Han Wang

source/train/Data.py

@@ -25,6 +25,8 @@ def __init__ (self,
         self.type_map = self._load_type_map(sys_path)
         if self.type_map is not None:
             assert(len(self.type_map) >= max(self.atom_type)+1)
+        # check pbc
+        self.pbc = self._check_pbc(sys_path)
         # enforce type_map if necessary
         if type_map is not None and self.type_map is not None:
             atom_type_ = [type_map.index(self.type_map[ii]) for ii in self.atom_type]
@@ -348,6 +350,12 @@ def _load_type_map(self, sys_path) :
         else :
             return None
 
+    def _check_pbc(self, sys_path):
+        pbc = True
+        if os.path.isfile(os.path.join(sys_path, 'nopbc')) :
+            pbc = False
+        return pbc
+
 
 class DataSets (object):
     def __init__ (self, 

source/train/DataSystem.py

@@ -95,20 +95,24 @@ def _load_test(self, ntests = -1):
             for nn in test_system_data:
                 self.test_data[nn].append(test_system_data[nn])
 
+
     def _make_default_mesh(self):
         self.default_mesh = []
         cell_size = np.max (self.rcut)
         for ii in range(self.nsystems) :
-            test_system_data = self.data_systems[ii].get_batch(self.batch_size[ii])
-            self.data_systems[ii].reset_get_batch()
-            # test_system_data = self.data_systems[ii].get_test()
-            avg_box = np.average (test_system_data["box"], axis = 0)
-            avg_box = np.reshape (avg_box, [3,3])
-            ncell = (np.linalg.norm(avg_box, axis=1)/ cell_size).astype(np.int32)
-            ncell[ncell < 2] = 2
-            default_mesh = np.zeros (6, dtype = np.int32)
-            default_mesh[3:6] = ncell
-            self.default_mesh.append(default_mesh)
+            if self.data_systems[ii].pbc :
+                test_system_data = self.data_systems[ii].get_batch(self.batch_size[ii])
+                self.data_systems[ii].reset_get_batch()
+                # test_system_data = self.data_systems[ii].get_test()
+                avg_box = np.average (test_system_data["box"], axis = 0)
+                avg_box = np.reshape (avg_box, [3,3])
+                ncell = (np.linalg.norm(avg_box, axis=1)/ cell_size).astype(np.int32)
+                ncell[ncell < 2] = 2
+                default_mesh = np.zeros (6, dtype = np.int32)
+                default_mesh[3:6] = ncell
+                self.default_mesh.append(default_mesh)
+            else:
+                self.default_mesh.append(np.array([], dtype = np.int32))
 
 
     def compute_energy_shift(self, rcond = 1e-3, key = 'energy') :

source/train/DeepPot.py

@@ -53,6 +53,8 @@ def __init__(self,
             self.modifier_type = self.sess.run(t_modifier_type).decode('UTF-8')
         except ValueError:
             self.modifier_type = None
+        except KeyError:
+            self.modifier_type = None
         if self.modifier_type == 'dipole_charge':
             t_mdl_name = self.graph.get_tensor_by_name('load/modifier_attr/mdl_name:0')
             t_mdl_charge_map = self.graph.get_tensor_by_name('load/modifier_attr/mdl_charge_map:0')
@@ -108,9 +110,18 @@ def eval_inner(self,
              aparam = None, 
              atomic = False) :
         # standarize the shape of inputs
-        coords = np.array(coords)
-        cells = np.array(cells)
-        atom_types = np.array(atom_types, dtype = int)
+        atom_types = np.array(atom_types, dtype = int).reshape([-1])
+        natoms = atom_types.size
+        coords = np.reshape(np.array(coords), [-1, natoms * 3])
+        nframes = coords.shape[0]
+        if cells is None:
+            pbc = False
+            # make cells to work around the requirement of pbc
+            cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
+        else:
+            pbc = True
+            cells = np.array(cells).reshape([nframes, 9])
+        
         if self.has_fparam :
             assert(fparam is not None)
             fparam = np.array(fparam)
@@ -119,10 +130,6 @@ def eval_inner(self,
             aparam = np.array(aparam)
 
         # reshape the inputs 
-        cells = np.reshape(cells, [-1, 9])
-        nframes = cells.shape[0]
-        coords = np.reshape(coords, [nframes, -1])
-        natoms = coords.shape[1] // 3
         if self.has_fparam :
             fdim = self.get_dim_fparam()
             if fparam.size == nframes * fdim :
@@ -167,7 +174,10 @@ def eval_inner(self,
         for ii in range(nframes) :
             feed_dict_test[self.t_coord] = np.reshape(coords[ii:ii+1, :], [-1])
             feed_dict_test[self.t_box  ] = np.reshape(cells [ii:ii+1, :], [-1])
-            feed_dict_test[self.t_mesh ] = make_default_mesh(cells[ii:ii+1, :])
+            if pbc:
+                feed_dict_test[self.t_mesh ] = make_default_mesh(cells[ii:ii+1, :])
+            else:
+                feed_dict_test[self.t_mesh ] = np.array([], dtype = np.int32)
             if self.has_fparam:
                 feed_dict_test[self.t_fparam] = np.reshape(fparam[ii:ii+1, :], [-1])
             if self.has_aparam: