Merge pull request #149 from amcadmus/devel

scale and shift the polar predictions, scale the tensor loss; fix bug of lammps pppm interface

Author: amcadmus

source/CMakeLists.txt

@@ -133,6 +133,16 @@ if (USE_TTM)
   set(TTM_DEF "-DUSE_TTM")
 endif (USE_TTM)
 
+# old pppm interface
+if(NOT DEFINED OLD_LMP_PPPM)
+  set(OLD_LMP_PPPM FALSE)
+endif(NOT DEFINED OLD_LMP_PPPM)
+if (OLD_LMP_PPPM)
+  set(OLD_LMP_PPPM_DEF "-DOLD_LMP_PPPM")
+  message(STATUS "Use old lammps pppm interface")
+endif()
+add_definitions (${OLD_LMP_PPPM_DEF})
+
 # define build type
 if ((NOT DEFINED CMAKE_BUILD_TYPE) OR CMAKE_BUILD_TYPE STREQUAL "")
    set (CMAKE_BUILD_TYPE release)

source/lmp/env.sh.in

@@ -6,6 +6,6 @@ TF_INCLUDE_DIRS=`echo $TENSORFLOW_INCLUDE_DIRS | sed "s/;/ -I/g"`
 TF_LIBRARY_PATH=`echo $TENSORFLOW_LIBRARY_PATH | sed "s/;/ -L/g"`
 TF_RPATH=`echo $TENSORFLOW_LIBRARY_PATH | sed "s/;/ -Wl,-rpath=/g"`
 
-NNP_INC=" -std=c++11 @PREC_DEF@ @TTM_DEF@ -I$TF_INCLUDE_DIRS -I$DEEPMD_ROOT/include/deepmd "
+NNP_INC=" -std=c++11 @PREC_DEF@ @TTM_DEF@ @OLD_LMP_PPPM_DEF@ -I$TF_INCLUDE_DIRS -I$DEEPMD_ROOT/include/deepmd "
 NNP_PATH=" -L$TF_LIBRARY_PATH -L$DEEPMD_ROOT/lib"
 NNP_LIB=" -Wl,--no-as-needed -l@LIB_DEEPMD_OP@ -l@LIB_DEEPMD_OP_CUDA@ -l@LIB_DEEPMD@ -ltensorflow_cc -ltensorflow_framework -Wl,-rpath=$TF_RPATH -Wl,-rpath=$DEEPMD_ROOT/lib"

source/lmp/pppm_dplr.cpp

@@ -40,8 +40,13 @@ enum{FORWARD_IK,FORWARD_AD,FORWARD_IK_PERATOM,FORWARD_AD_PERATOM};
 
 /* ---------------------------------------------------------------------- */
 
+#ifdef OLD_LMP_PPPM
 PPPMDPLR::PPPMDPLR(LAMMPS *lmp, int narg, char **arg) :
   PPPM(lmp, narg, arg)
+#else
+PPPMDPLR::PPPMDPLR(LAMMPS *lmp) :
+  PPPM(lmp)
+#endif
 {
   triclinic_support = 1;
 }

source/lmp/pppm_dplr.h

@@ -35,7 +35,11 @@ namespace LAMMPS_NS {
 
   class PPPMDPLR : public PPPM {
 public:
+#ifdef OLD_LMP_PPPM
     PPPMDPLR(class LAMMPS *, int, char **);
+#else
+    PPPMDPLR(class LAMMPS *);
+#endif
     virtual ~PPPMDPLR () {};
     void init();
     const vector<double > & get_fele() const {return fele;};

source/train/Fitting.py

@@ -383,6 +383,8 @@ def __init__ (self, jdata, descrpt) :
                .add('neuron',           list,   default = [120,120,120], alias = 'n_neuron')\
                .add('resnet_dt',        bool,   default = True)\
                .add('fit_diag',         bool,   default = True)\
+               .add('diag_shift',       [list,float], default = [0.0 for ii in range(self.ntypes)])\
+               .add('scale',            [list,float], default = [1.0 for ii in range(self.ntypes)])\
                .add('sel_type',         [list,int],   default = [ii for ii in range(self.ntypes)], alias = 'pol_type')\
                .add('seed',             int)
         class_data = args.parse(jdata)
@@ -391,6 +393,14 @@ def __init__ (self, jdata, descrpt) :
         self.sel_type = class_data['sel_type']
         self.fit_diag = class_data['fit_diag']
         self.seed = class_data['seed']
+        self.diag_shift = class_data['diag_shift']
+        self.scale = class_data['scale']
+        if type(self.sel_type) is not list:
+            self.sel_type = [self.sel_type]
+        if type(self.diag_shift) is not list:
+            self.diag_shift = [self.diag_shift]
+        if type(self.scale) is not list:
+            self.scale = [self.scale]
         self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1()
         self.dim_rot_mat = self.dim_rot_mat_1 * 3
         self.useBN = False
@@ -477,6 +487,10 @@ def build (self,
             final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a = True)
             # nframes x natoms x 3 x 3
             final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3, 3])
+            # shift and scale
+            sel_type_idx = self.sel_type.index(type_i)
+            final_layer = final_layer * self.scale[sel_type_idx]
+            final_layer = final_layer + self.diag_shift[sel_type_idx] * tf.eye(3, batch_shape=[tf.shape(inputs)[0], natoms[2+type_i]], dtype = global_tf_float_precision)
 
             # concat the results
             if count == 0:

source/train/Loss.py

@@ -284,6 +284,10 @@ def __init__ (self, jdata, **kwarg) :
         self.tensor_size = kwarg['tensor_size']
         self.label_name = kwarg['label_name']
         self.atomic = kwarg.get('atomic', True)
+        if jdata is not None:
+            self.scale = jdata.get('scale', 1.0)
+        else:
+            self.scale = 1.0
         # data required
         add_data_requirement(self.label_name, 
                              self.tensor_size, 
@@ -300,7 +304,7 @@ def build (self,
                suffix):        
         polar_hat = label_dict[self.label_name]
         polar = model_dict[self.tensor_name]
-        l2_loss = tf.reduce_mean( tf.square(polar - polar_hat), name='l2_'+suffix)
+        l2_loss = tf.reduce_mean( tf.square(self.scale*(polar - polar_hat)), name='l2_'+suffix)
         if not self.atomic :
             atom_norm  = 1./ global_cvt_2_tf_float(natoms[0]) 
             l2_loss = l2_loss * atom_norm