infer global polar

Author: Han Wang

deepmd/__init__.py

@@ -3,6 +3,7 @@
 from .DeepPot    import DeepPot
 from .DeepDipole import DeepDipole
 from .DeepPolar  import DeepPolar
+from .DeepPolar  import DeepGlobalPolar
 from .DeepWFC    import DeepWFC
 
 set_mkl()

source/scripts/freeze.py

@@ -44,6 +44,8 @@ def _make_node_names(model_type = None) :
         nodes = "o_dipole,descrpt_attr/rcut,descrpt_attr/ntypes,model_attr/tmap,model_attr/sel_type,model_attr/model_type"
     elif model_type == 'polar':
         nodes = "o_polar,descrpt_attr/rcut,descrpt_attr/ntypes,model_attr/tmap,model_attr/sel_type,model_attr/model_type"
+    elif model_type == 'global_polar':
+        nodes = "o_global_polar,descrpt_attr/rcut,descrpt_attr/ntypes,model_attr/tmap,model_attr/sel_type,model_attr/model_type"
     else:
         raise RuntimeError('unknow model type ' + model_type)
     return nodes

source/train/DeepEval.py

@@ -149,7 +149,8 @@ def get_sel_type(self):
     def eval(self,
              coords, 
              cells, 
-             atom_types) :
+             atom_types, 
+             atomic = True) :
         # standarize the shape of inputs
         coords = np.array(coords)
         cells = np.array(cells)
@@ -183,10 +184,13 @@ def eval(self,
             tensor.append(v_out[0])
 
         # reverse map of the outputs
-        tensor = np.array(tensor)
-        tensor = self.reverse_map(np.reshape(tensor, [nframes,-1,self.variable_dof]), sel_imap)
-
-        tensor = np.reshape(tensor, [nframes, len(sel_at), self.variable_dof])
+        if atomic:
+            tensor = np.array(tensor)
+            tensor = self.reverse_map(np.reshape(tensor, [nframes,-1,self.variable_dof]), sel_imap)
+            tensor = np.reshape(tensor, [nframes, len(sel_at), self.variable_dof])
+        else:
+            tensor = np.reshape(tensor, [nframes, self.variable_dof])
+        
         return tensor
 
 

source/train/DeepPolar.py

@@ -9,3 +9,14 @@ def __init__(self,
                  model_file) :
         DeepTensor.__init__(self, model_file, 'polar', 9)
 
+    
+class DeepGlobalPolar (DeepTensor) :
+    def __init__(self, 
+                 model_file) :
+        DeepTensor.__init__(self, model_file, 'global_polar', 9)
+
+    def eval(self,
+             coords, 
+             cells, 
+             atom_types) :
+        return DeepTensor.eval(self, coords, cells, atom_types, atomic = False)