Fix bugs in descriptor se_t

deepmd/descriptor/__init__.py

@@ -1,9 +1,9 @@
-from .hybrid    import DescrptHybrid
-from .se_a      import DescrptSeA
-from .se_r      import DescrptSeR
-from .se_ar     import DescrptSeAR
-from .se_a_t    import DescrptSeAT
-from .se_a_ebd  import DescrptSeAEbd
-from .se_a_ef   import DescrptSeAEf
-from .se_a_ef   import DescrptSeAEfLower
+from .hybrid import DescrptHybrid
+from .se_a import DescrptSeA
+from .se_r import DescrptSeR
+from .se_ar import DescrptSeAR
+from .se_t import DescrptSeT
+from .se_a_ebd import DescrptSeAEbd
+from .se_a_ef import DescrptSeAEf
+from .se_a_ef import DescrptSeAEfLower
 from .loc_frame import DescrptLocFrame

deepmd/descriptor/hybrid.py

@@ -8,16 +8,16 @@
 from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
 # from deepmd.descriptor import DescrptLocFrame
 # from deepmd.descriptor import DescrptSeA
-# from deepmd.descriptor import DescrptSeAT
+# from deepmd.descriptor import DescrptSeT
 # from deepmd.descriptor import DescrptSeAEbd
 # from deepmd.descriptor import DescrptSeAEf
 # from deepmd.descriptor import DescrptSeR
-from .se_a      import DescrptSeA
-from .se_r      import DescrptSeR
-from .se_ar     import DescrptSeAR
-from .se_a_t    import DescrptSeAT
-from .se_a_ebd  import DescrptSeAEbd
-from .se_a_ef   import DescrptSeAEf
+from .se_a import DescrptSeA
+from .se_r import DescrptSeR
+from .se_ar import DescrptSeAR
+from .se_t import DescrptSeT
+from .se_a_ebd import DescrptSeAEbd
+from .se_a_ef import DescrptSeAEf
 from .loc_frame import DescrptLocFrame
 
 class DescrptHybrid ():

deepmd/descriptor/se_a_t.py → deepmd/descriptor/se_t.py

@@ -10,7 +10,7 @@
 from deepmd.env import default_tf_session_config
 from deepmd.utils.network import embedding_net
 
-class DescrptSeAT ():
+class DescrptSeT ():
     @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys()))
     def __init__ (self, 
                   rcut: float,
@@ -95,7 +95,7 @@ def __init__ (self,
             self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms')
             self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh')
             self.stat_descrpt, descrpt_deriv, rij, nlist \
-                = op_module.descrpt_se_a(self.place_holders['coord'],
+                = op_module.prod_env_mat_a(self.place_holders['coord'],
                                          self.place_holders['type'],
                                          self.place_holders['natoms_vec'],
                                          self.place_holders['box'],
@@ -282,7 +282,7 @@ def build (self,
         atype = tf.reshape (atype_, [-1, natoms[1]])
 
         self.descrpt, self.descrpt_deriv, self.rij, self.nlist \
-            = op_module.descrpt_se_a (coord,
+            = op_module.prod_env_mat_a (coord,
                                        atype,
                                        natoms,
                                        box,

deepmd/train/trainer.py

@@ -11,7 +11,7 @@
 from deepmd.fit import EnerFitting, WFCFitting, PolarFittingLocFrame, PolarFittingSeA, GlobalPolarFittingSeA, DipoleFittingSeA
 from deepmd.descriptor import DescrptLocFrame
 from deepmd.descriptor import DescrptSeA
-from deepmd.descriptor import DescrptSeAT
+from deepmd.descriptor import DescrptSeT
 from deepmd.descriptor import DescrptSeAEbd
 from deepmd.descriptor import DescrptSeAEf
 from deepmd.descriptor import DescrptSeR
@@ -57,8 +57,8 @@ def _generate_descrpt_from_param_dict(descrpt_param):
         descrpt = DescrptLocFrame(**descrpt_param)
     elif descrpt_type == 'se_a' :            
         descrpt = DescrptSeA(**descrpt_param)
-    elif descrpt_type == 'se_a_3be' or descrpt_type == 'se_at' :
-        descrpt = DescrptSeAT(**descrpt_param)
+    elif descrpt_type == 'se_a_3be' or descrpt_type == 'se_at' or descrpt_type == 'se_t' :
+        descrpt = DescrptSeT(**descrpt_param)
     elif descrpt_type == 'se_a_tpe' or descrpt_type == 'se_a_ebd' :
         descrpt = DescrptSeAEbd(**descrpt_param)
     elif descrpt_type == 'se_a_ef' :
@@ -568,7 +568,7 @@ def valid_on_the_fly(self,
     @staticmethod
     def print_header(fp, train_results, valid_results):
         print_str = ''
-        print_str += "# %5s" % 'batch'
+        print_str += "# %5s" % 'step'
         if valid_results is not None:
             prop_fmt =  '   %11s %11s'
             for k in train_results.keys():

deepmd/utils/argcheck.py

@@ -71,7 +71,7 @@ def descrpt_se_a_args():
     ]
 
 
-def descrpt_se_a_3be_args():
+def descrpt_se_t_args():
     doc_sel = 'A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius.'
     doc_rcut = 'The cut-off radius.'
     doc_rcut_smth = 'Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`'
@@ -164,22 +164,22 @@ def descrpt_variant_type_args():
     link_lf = make_link('loc_frame', 'model/descriptor[loc_frame]')
     link_se_a = make_link('se_a', 'model/descriptor[se_a]')
     link_se_r = make_link('se_r', 'model/descriptor[se_r]')
-    link_se_a_3be = make_link('se_a_3be', 'model/descriptor[se_a_3be]')
+    link_se_t = make_link('se_t', 'model/descriptor[se_t]')
     link_se_a_tpe = make_link('se_a_tpe', 'model/descriptor[se_a_tpe]')
     link_hybrid = make_link('hybrid', 'model/descriptor[hybrid]')
     doc_descrpt_type = f'The type of the descritpor. See explanation below. \n\n\
 - `loc_frame`: Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame.\n\n\
 - `se_a`: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor.\n\n\
 - `se_r`: Used by the smooth edition of Deep Potential. Only the distance between atoms is used to construct the descriptor.\n\n\
-- `se_a_3be`: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Three-body embedding will be used by this descriptor.\n\n\
+- `se_t`: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Three-body embedding will be used by this descriptor.\n\n\
 - `se_a_tpe`: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Type embedding will be used by this descriptor.\n\n\
 - `hybrid`: Concatenate of a list of descriptors as a new descriptor.'
 
     return Variant("type", [
         Argument("loc_frame", dict, descrpt_local_frame_args()),
         Argument("se_a", dict, descrpt_se_a_args()),
         Argument("se_r", dict, descrpt_se_r_args()),
-        Argument("se_a_3be", dict, descrpt_se_a_3be_args(), alias = ['se_at']),
+        Argument("se_t", dict, descrpt_se_t_args(), alias = ['se_at', 'se_a_3be']),
         Argument("se_a_tpe", dict, descrpt_se_a_tpe_args(), alias = ['se_a_ebd']),
         Argument("hybrid", dict, descrpt_hybrid_args()),
     ], doc = doc_descrpt_type)

source/tests/test_model_se_t.py

@@ -0,0 +1,116 @@
+import dpdata,os,sys,unittest
+import numpy as np
+from deepmd.env import tf
+from common import Data,gen_data, j_loader
+
+from deepmd.utils.data_system import DataSystem
+from deepmd.descriptor import DescrptSeT
+from deepmd.fit import EnerFitting
+from deepmd.model import EnerModel
+from deepmd.common import j_must_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 = 'water_se_t.json'
+        jdata = j_loader(jfile)
+
+        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
+
+        jdata['model']['descriptor'].pop('type', None)        
+        descrpt = DescrptSeT(**jdata['model']['descriptor'])
+        jdata['model']['fitting_net']['descrpt'] = descrpt
+        fitting = EnerFitting(**jdata['model']['fitting_net'])
+        model = EnerModel(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']]
+        }
+        model._compute_input_stat(input_data)
+        model.descrpt.bias_atom_e = data.compute_energy_shift()
+
+        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 = "se_t", 
+                           reuse = False)
+        energy = model_pred['energy']
+        force  = model_pred['force']
+        virial = model_pred['virial']
+        atom_ener =  model_pred['atom_ener']
+
+        feed_dict_test = {t_prop_c:        test_data['prop_c'],
+                          t_energy:        test_data['energy']              [:numb_test],
+                          t_force:         np.reshape(test_data['force']    [:numb_test, :], [-1]),
+                          t_virial:        np.reshape(test_data['virial']   [:numb_test, :], [-1]),
+                          t_atom_ener:     np.reshape(test_data['atom_ener'][:numb_test, :], [-1]),
+                          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())
+        [e, f, v] = sess.run([energy, force, virial], 
+                             feed_dict = feed_dict_test)
+
+        e = e.reshape([-1])
+        f = f.reshape([-1])
+        v = v.reshape([-1])
+        np.savetxt('e.out', e.reshape([1, -1]))
+        np.savetxt('f.out', f.reshape([1, -1]), delimiter = ',')
+        np.savetxt('v.out', v.reshape([1, -1]), delimiter = ',')
+        refe = [4.826771866004193612e+01]
+        reff = [5.355088169393570574e+00,5.606772412401632266e+00,2.703270748296462966e-01,5.381408138049708967e+00,5.261355614357515975e+00,-4.079549918988090162e-01,-5.182324474551911919e+00,3.695481388907447262e-01,-5.238474288082559799e-02,1.665564584447352670e-01,-5.955401876564963892e+00,-2.217626865156164251e-01,-5.967343479332643419e+00,9.073821102416884665e-02,3.703103995504785639e-01,2.466151879965444438e-01,-5.373012500109097367e+00,4.146494691512622732e-02]
+        refv = [-1.336768232407933077e+01,4.818050125305787801e-01,3.589284283410607568e-01,4.818050125305786691e-01,-1.225345559839458964e+01,-1.701405121682751653e-01,3.589284283410607568e-01,-1.701405121682752486e-01,-3.428455515842296353e-02]
+        refe = np.reshape(refe, [-1])
+        reff = np.reshape(reff, [-1])
+        refv = np.reshape(refv, [-1])
+
+        places = 6
+        for ii in range(e.size) :
+            self.assertAlmostEqual(e[ii], refe[ii], places = places)
+        for ii in range(f.size) :
+            self.assertAlmostEqual(f[ii], reff[ii], places = places)
+        for ii in range(v.size) :
+            self.assertAlmostEqual(v[ii], refv[ii], places = places)

source/tests/water_se_t.json

@@ -0,0 +1,53 @@
+{
+    "_comment": " model parameters",
+    "model" : {
+	"descriptor" : {
+	    "type":		"se_3be",
+	    "sel":		[40, 80],
+	    "rcut_smth":	0.50,
+	    "rcut":		6.00,
+	    "neuron":		[2, 4, 8],
+	    "resnet_dt":	false,
+	    "seed":		1
+	},
+	"fitting_net" : {
+	    "neuron":		[240, 240, 240],
+	    "resnet_dt":	true,
+	    "seed":		1
+	}
+    },
+
+    "_comment": " traing controls",
+    "systems":		["system"],
+    "set_prefix":	"set",    
+    "stop_batch":	1000000,
+    "batch_size":	1,
+    "start_lr":		0.005,
+    "decay_steps":	5000,
+    "decay_rate":	0.95,
+
+    "start_pref_e":	0.02,
+    "limit_pref_e":	1,
+    "start_pref_f":	1000,
+    "limit_pref_f":	1,
+    "start_pref_v":	0,
+    "limit_pref_v":	0,
+
+    "seed":		1,
+
+    "_comment": " display and restart",
+    "_comment": " frequencies counted in batch",
+    "disp_file":	"lcurve.out",
+    "disp_freq":	100,
+    "numb_test":	1,
+    "save_freq":	1000,
+    "save_ckpt":	"model.ckpt",
+    "load_ckpt":	"model.ckpt",
+    "disp_training":	true,
+    "time_training":	true,
+    "profiling":	false,
+    "profiling_file":	"timeline.json",
+
+    "_comment":		"that's all"
+}
+