simplifies the implementation of Model (WFC and Polar) and WFCLoss

Author: Han Wang

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/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/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,

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,