Merge pull request #105 from amcadmus/devel

Implement dipole

Author: amcadmus

deepmd/__init__.py

@@ -1,8 +1,9 @@
 from .env import set_mkl
-from .DeepEval  import DeepEval
-from .DeepPot   import DeepPot
-from .DeepPolar import DeepPolar
-from .DeepWFC   import DeepWFC
+from .DeepEval   import DeepEval
+from .DeepPot    import DeepPot
+from .DeepDipole import DeepDipole
+from .DeepPolar  import DeepPolar
+from .DeepWFC    import DeepWFC
 
 set_mkl()
 

examples/water/train/polar.json

@@ -3,6 +3,7 @@
     "_comment": " model parameters",
     "model":{
 	"type_map":		["O", "H"],
+	"data_stat_nbatch":	1,
 	"descriptor": {
 	    "type":		"loc_frame",
 	    "sel_a":		[16, 32],

examples/water/train/polar_se_a.json

@@ -3,6 +3,7 @@
     "_comment": " model parameters",
     "model":{
 	"type_map":		["O", "H"],
+	"data_stat_nbatch":	1,
 	"descriptor" :{
 	    "type":		"se_a",
 	    "sel":		[46, 92],

examples/water/train/wannier.json

@@ -3,6 +3,7 @@
     "_comment": " model parameters",
     "model":{
 	"type_map":		["O", "H"],
+	"data_stat_nbatch":	1,
 	"descriptor": {
 	    "type":		"loc_frame",
 	    "sel_a":		[16, 32],

source/scripts/freeze.py

@@ -40,6 +40,8 @@ def _make_node_names(model_type = None) :
         nodes = "o_energy,o_force,o_virial,o_atom_energy,o_atom_virial,descrpt_attr/rcut,descrpt_attr/ntypes,fitting_attr/dfparam,fitting_attr/daparam,model_attr/tmap,model_attr/model_type"
     elif model_type == 'wfc':
         nodes = "o_wfc,descrpt_attr/rcut,descrpt_attr/ntypes,model_attr/tmap,model_attr/sel_type,model_attr/model_type"
+    elif model_type == 'dipole':
+        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"
     else:

source/tests/test_wfc.py

@@ -0,0 +1,97 @@
+import dpdata,os,sys,json,unittest
+import numpy as np
+from deepmd.env import tf
+from common import Data,gen_data
+
+from deepmd.RunOptions import RunOptions
+from deepmd.DataSystem import DataSystem
+from deepmd.DescrptLocFrame import DescrptLocFrame
+from deepmd.Fitting import WFCFitting
+from deepmd.Model import WFCModel
+from deepmd.common import j_must_have, j_must_have_d, j_have
+
+global_ener_float_precision = tf.float64
+global_tf_float_precision = tf.float64
+global_np_float_precision = np.float64
+
+class TestModel(unittest.TestCase):
+    def setUp(self) :
+        gen_data()
+
+    def test_model(self):
+        jfile = 'wfc.json'
+        with open(jfile) as fp:
+            jdata = json.load (fp)
+        run_opt = RunOptions(None) 
+        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
+        stop_batch = j_must_have(jdata, 'stop_batch')
+        rcut = j_must_have (jdata['model']['descriptor'], 'rcut')
+        
+        data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt = None)
+        
+        test_data = data.get_test ()
+        numb_test = 1
+        
+        descrpt = DescrptLocFrame(jdata['model']['descriptor'])
+        fitting = WFCFitting(jdata['model']['fitting_net'], descrpt)
+        model = WFCModel(jdata['model'], descrpt, fitting)
+
+        input_data = {'coord' : [test_data['coord']], 
+                      'box': [test_data['box']], 
+                      'type': [test_data['type']],
+                      'natoms_vec' : [test_data['natoms_vec']],
+                      'default_mesh' : [test_data['default_mesh']],
+                      'fparam': [test_data['fparam']],
+        }
+        model._compute_dstats(input_data)
+
+        t_prop_c           = tf.placeholder(tf.float32, [5],    name='t_prop_c')
+        t_energy           = tf.placeholder(global_ener_float_precision, [None], name='t_energy')
+        t_force            = tf.placeholder(global_tf_float_precision, [None], name='t_force')
+        t_virial           = tf.placeholder(global_tf_float_precision, [None], name='t_virial')
+        t_atom_ener        = tf.placeholder(global_tf_float_precision, [None], name='t_atom_ener')
+        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,   [model.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)
+        t_fparam = None
+
+        model_pred \
+            = model.build (t_coord, 
+                           t_type, 
+                           t_natoms, 
+                           t_box, 
+                           t_mesh,
+                           t_fparam,
+                           suffix = "wfc", 
+                           reuse = False)
+        wfc = model_pred['wfc']
+
+        feed_dict_test = {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}
+
+        sess = tf.Session()
+        sess.run(tf.global_variables_initializer())
+        [p] = sess.run([wfc], feed_dict = feed_dict_test)
+
+        p = p.reshape([-1])
+        refp = [-9.105016838228578990e-01,7.196284362034099935e-01,-9.548516928185298014e-02,2.764615027095288724e+00,2.661319598995644520e-01,7.579512949131941846e-02,-2.107409067376114997e+00,-1.299080016614967414e-01,-5.962778584850070285e-01,2.913899917663253514e-01,-1.226917174638697094e+00,1.829523069930876655e+00,1.015704024959750873e+00,-1.792333611099589386e-01,5.032898080485321834e-01,1.808561721292949453e-01,2.468863482075112081e+00,-2.566442546384765100e-01,-1.467453783795173994e-01,-1.822963931552128658e+00,5.843600156865462747e-01,-1.493875280832117403e+00,1.693322352814763398e-01,-1.877325443995481624e+00]
+
+        places = 6
+        for ii in range(p.size) :
+            self.assertAlmostEqual(p[ii], refp[ii], places = places)
+
+
+        

source/tests/wfc.json

@@ -0,0 +1,62 @@
+{
+    "with_distrib":	false,
+    "_comment": " model parameters",
+    "model":{
+	"type":			"polar",
+	"type_map":		["O", "H"],
+	"descriptor" :{
+	    "type":		"loc_frame",
+	    "sel_a":		[16, 32],
+	    "sel_r":		[30, 60],
+	    "rcut":		6.00,
+	    "axis_rule":	[0, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0],
+	    "_comment":	" default rule: []",
+	    "_comment":	" user defined rule: for each type provides two axes, ",
+	    "_comment":	"                    for each axis: (a_or_r, type, idx)",
+	    "_comment":	"                    if type < 0, exclude type -(type+1)",
+	    "_comment":	"                    for water (O:0, H:1) it can be",
+	    "_comment":	"                    [0, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0]",
+	    "_comment": " that's all"
+	},
+	"fitting_net": {
+	    "type":		"wfc",
+	    "wfc_numb":		4,
+	    "sel_type":		[0],
+	    "neuron":		[100, 100, 100],
+	    "resnet_dt":	true,
+	    "seed":		1,
+	    "_comment":		" that's all"
+	},
+	"_comment":	" that's all"
+    },
+    
+    "learning_rate" :{
+	"type":		"exp",
+	"start_lr":	0.001,
+	"decay_steps":	5000,
+	"decay_rate":	0.95,
+	"_comment":	"that's all"
+    },
+
+    "_comment": " traing controls",
+    "systems":		["system"],
+    "set_prefix":	"set",    
+    "stop_batch":	1000000,
+    "batch_size":	[1],
+
+    "seed":		1,
+
+    "_comment": " display and restart",
+    "_comment": " frequencies counted in batch",
+    "disp_file":	"lcurve.out",
+    "disp_freq":	100,
+    "numb_test":	10,
+    "save_freq":	1000,
+    "save_ckpt":	"model.ckpt",
+    "load_ckpt":	"model.ckpt",
+    "disp_training":true,
+    "time_training":true,
+
+    "_comment":		"that's all"
+}
+

source/train/DeepDipole.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+
+import os,sys
+import numpy as np
+from deepmd.DeepEval import DeepTensor
+
+class DeepDipole (DeepTensor) :
+    def __init__(self, 
+                 model_file) :
+        DeepTensor.__init__(self, model_file, 'dipole', 3)
+

source/train/Fitting.py

@@ -211,6 +211,9 @@ def get_sel_type(self):
     def get_wfc_numb(self):
         return self.wfc_numb
 
+    def get_out_size(self):
+        return self.wfc_numb * 3
+
     def build (self, 
                input_d,
                rot_mat,
@@ -283,6 +286,9 @@ def __init__ (self, jdata, descrpt) :
     def get_sel_type(self):
         return self.sel_type
 
+    def get_out_size(self):
+        return 9
+
     def build (self, 
                input_d,
                rot_mat,
@@ -360,6 +366,9 @@ def __init__ (self, jdata, descrpt) :
     def get_sel_type(self):
         return self.sel_type
 
+    def get_out_size(self):
+        return 9
+
     def build (self, 
                input_d,
                rot_mat,
@@ -412,3 +421,76 @@ def build (self,
             count += 1
 
         return tf.reshape(outs, [-1])
+
+
+class DipoleFittingSeA () :
+    def __init__ (self, jdata, descrpt) :
+        if not isinstance(descrpt, DescrptSeA) :
+            raise RuntimeError('DipoleFittingSeA only supports DescrptSeA')
+        self.ntypes = descrpt.get_ntypes()
+        self.dim_descrpt = descrpt.get_dim_out()
+        args = ClassArg()\
+               .add('neuron',           list,   default = [120,120,120], alias = 'n_neuron')\
+               .add('resnet_dt',        bool,   default = True)\
+               .add('sel_type',         [list,int],   default = [ii for ii in range(self.ntypes)], alias = 'dipole_type')\
+               .add('seed',             int)
+        class_data = args.parse(jdata)
+        self.n_neuron = class_data['neuron']
+        self.resnet_dt = class_data['resnet_dt']
+        self.sel_type = class_data['sel_type']
+        self.seed = class_data['seed']
+        self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1()
+        self.dim_rot_mat = self.dim_rot_mat_1 * 3
+        self.useBN = False
+
+    def get_sel_type(self):
+        return self.sel_type
+
+    def build (self, 
+               input_d,
+               rot_mat,
+               natoms,
+               reuse = None,
+               suffix = '') :
+        start_index = 0
+        inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]])
+        rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]])
+        shape = inputs.get_shape().as_list()
+
+        count = 0
+        for type_i in range(self.ntypes):
+            # cut-out inputs
+            inputs_i = tf.slice (inputs,
+                                 [ 0, start_index*      self.dim_descrpt],
+                                 [-1, natoms[2+type_i]* self.dim_descrpt] )
+            inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
+            rot_mat_i = tf.slice (rot_mat,
+                                  [ 0, start_index*      self.dim_rot_mat],
+                                  [-1, natoms[2+type_i]* self.dim_rot_mat] )
+            rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3])
+            start_index += natoms[2+type_i]
+            if not type_i in self.sel_type :
+                continue
+            layer = inputs_i
+            for ii in range(0,len(self.n_neuron)) :
+                if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
+                    layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt)
+                else :
+                    layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed)
+            # (nframes x natoms) x naxis
+            final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed)
+            # (nframes x natoms) x 1 * naxis
+            final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], 1, self.dim_rot_mat_1])
+            # (nframes x natoms) x 1 x 3(coord)
+            final_layer = tf.matmul(final_layer, rot_mat_i)
+            # nframes x natoms x 3
+            final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3])
+
+            # concat the results
+            if count == 0:
+                outs = final_layer
+            else:
+                outs = tf.concat([outs, final_layer], axis = 1)
+            count += 1
+
+        return tf.reshape(outs, [-1])

source/train/Loss.py

@@ -166,57 +166,8 @@ def print_on_training(self,
         if self.has_pf:
             print_str += prop_fmt % (np.sqrt(error_pf_test) / natoms[0], np.sqrt(error_pf_train) / natoms[0])
 
-        return print_str
-        
-        
-
-class WFCLoss () :
-    def __init__ (self, jdata, **kwarg) :
-        model = kwarg['model']
-        # data required
-        add_data_requirement('wfc', 
-                             model.get_wfc_numb() * 3, 
-                             atomic=True,  
-                             must=True, 
-                             high_prec=False, 
-                             type_sel = model.get_sel_type())
-
-    def build (self, 
-               learning_rate,
-               natoms,
-               model_dict,
-               label_dict,
-               suffix):        
-        wfc_hat = label_dict['wfc']
-        wfc = model_dict['wfc']
-        l2_loss = tf.reduce_mean( tf.square(wfc - wfc_hat), name='l2_'+suffix)
-        self.l2_l = l2_loss
-        more_loss = {}
-
-        return l2_loss, more_loss
-
-    def print_header(self) :
-        prop_fmt = '   %9s %9s'
-        print_str = ''
-        print_str += prop_fmt % ('l2_tst', 'l2_trn')
-        return print_str
-
-    def print_on_training(self, 
-                          sess, 
-                          natoms,
-                          feed_dict_test,
-                          feed_dict_batch) :
-        error_test\
-            = sess.run([self.l2_l], \
-                       feed_dict=feed_dict_test)
-        error_train\
-            = sess.run([self.l2_l], \
-                       feed_dict=feed_dict_batch)
-        print_str = ""
-        prop_fmt = "   %9.2e %9.2e"
-        print_str += prop_fmt % (np.sqrt(error_test), np.sqrt(error_train))
+        return print_str        
 
-        return print_str
 
 
 class TensorLoss () :
@@ -235,7 +186,7 @@ def __init__ (self, jdata, **kwarg) :
                              atomic=True,  
                              must=True, 
                              high_prec=False, 
-                             type_sel = model.get_sel_type())
+                             type_sel = type_sel)
 
     def build (self, 
                learning_rate,