Merge pull request #51 from mir-group/auto_write_lmp

automatically write lmp file when mgp constructed #28

flare/mgp/mgp.py

@@ -6,7 +6,8 @@
 import subprocess
 import os
 
-from flare import gp, env, struc, kernels
+from flare import gp, struc, kernels
+from flare.env import AtomicEnvironment
 from flare.gp import GaussianProcess
 from flare.kernels import two_body, three_body, two_plus_three_body,\
     two_body_jit
@@ -20,27 +21,34 @@
 
 class MappedGaussianProcess:
     def __init__(self, GP: GaussianProcess, grid_params: dict,
-                 struc_params: dict, mean_only=False):
-
-        '''
-        :param: GP : gp model
-        :param: struc_params : {'species': [0, 1],
-                'cube_lat': cell, # should input the cell matrix
-                'mass_dict': {'0': 27 * unit, '1': 16 * unit}}
-        :param: grid_params : {'bounds_2': [[1.2], [3.5]], # [[lower_bound],
-                                                            [upper]]
-               'bounds_3': [[1.2, 1.2, 0], [3.5, 3.5, np.pi]],
-                    #[[lower,lower,0],[upper,upper,np.pi]]
-               'grid_num_2': 64,
-               'grid_num_3': [16, 16, 16],
-               'svd_rank_2': 64,
-               'svd_rank_3': 16**3,
-               'bodies': [2, 3],
-               'update': True,
-               'load_grid': None}
-        :param: mean_only : if True: only build mapping for mean (force)
+                 struc_params: dict, mean_only=False, lmp_file_name='lmp.mgp'):
 
         '''
+        Build MGP
+        :param: GP : trained GP model
+        :param: struc_params : information of training data
+        :param: grid_params : setting of grids for mapping
+        :param: mean_only : if True: only build mapping for mean (force)
+        :param: lmp_file_name : lammps coefficient file name
+        Examples:
+        
+        >>> struc_params = {'species': [0, 1],
+                            'cube_lat': cell, # should input the cell matrix
+                            'mass_dict': {'0': 27 * unit, '1': 16 * unit}}
+        >>> grid_params =  {'bounds_2': [[1.2], [3.5]], 
+                                        # [[lower_bound], [upper_bound]]
+                            'bounds_3': [[1.2, 1.2, 0], [3.5, 3.5, np.pi]],
+                                        # [[lower,lower,0],[upper,upper,np.pi]]
+                            'grid_num_2': 64,
+                            'grid_num_3': [16, 16, 16],
+                            'svd_rank_2': 64,
+                            'svd_rank_3': 16**3,
+                            'bodies': [2, 3],
+                            'update': True, # if True: accelerating grids 
+                                            # generating by saving intermediate 
+                                            # coeff when generating grids
+                            'load_grid': None}
+        '''
         self.GP = GP
         self.grid_params = grid_params
         self.struc_params = struc_params
@@ -62,20 +70,20 @@ def __init__(self, GP: GaussianProcess, grid_params: dict,
         if 2 in self.bodies:
             for b_struc in bond_struc[0]:
                 map_2 = Map2body(self.grid_num_2, self.bounds_2, self.GP,
-                                 b_struc, self.bodies,
-                                 grid_params['load_grid'],
-                                 self.svd_rank_2, self.mean_only)
+                                 b_struc, self.bodies, self.svd_rank_2, 
+                                 self.mean_only)
                 self.maps_2.append(map_2)
         if 3 in self.bodies:
             for b_struc in bond_struc[1]:
                 map_3 = Map3body(self.grid_num_3, self.bounds_3, self.GP,
-                                 b_struc, self.bodies,
-                                 grid_params['load_grid'],
-                                 self.svd_rank_3,
-                                 self.mean_only, self.update)
+                                 b_struc, self.bodies, self.svd_rank_3,
+                                 self.mean_only, grid_params['load_grid'],
+                                 self.update)
 
                 self.maps_3.append(map_3)
 
+        self.write_lmp_file(lmp_file_name)
+
     def build_bond_struc(self, struc_params):
 
         '''
@@ -144,7 +152,12 @@ def build_bond_struc(self, struc_params):
         spcs = [spc_2, spc_3]
         return bond_struc, spcs
 
-    def predict(self, atom_env, mean_only=False):
+    def predict(self, atom_env: AtomicEnvironment, mean_only: bool=False):
+        '''
+        predict force and variance for given atomic environment
+        :param atom_env: atomic environment (with a center atom and its neighbors)
+        :param mean_only: if True: only predict force (variance is always 0)
+        '''
         if self.mean_only:  # if not build mapping for var
             mean_only = True
 
@@ -175,6 +188,9 @@ def predict(self, atom_env, mean_only=False):
         return f, v
 
     def get_2body_comp(self, atom_env, sig, ls, r_cut):
+        '''
+        get bonds grouped by species
+        '''
         bond_array_2 = atom_env.bond_array_2
         ctype = atom_env.ctype
         etypes = atom_env.etypes
@@ -189,6 +205,9 @@ def get_2body_comp(self, atom_env, sig, ls, r_cut):
         return kern2, spcs, comp_r, comp_xyz
 
     def get_3body_comp(self, atom_env, sig, ls, r_cut):
+        '''
+        get triplets and grouped by species 
+        '''
         bond_array_3 = atom_env.bond_array_3
         cross_bond_inds = atom_env.cross_bond_inds
         cross_bond_dists = atom_env.cross_bond_dists
@@ -206,7 +225,6 @@ def get_3body_comp(self, atom_env, sig, ls, r_cut):
         for d in range(3):
             kern3_gp[d] = three_body_mc(atom_env, atom_env, d+1, d+1,
                                         self.GP.hyps[-3:], self.GP.cutoffs)
-        # print(kern3, kern3_gp)
 
         spcs, comp_r, comp_xyz = \
             get_triplets(ctype, etypes, bond_array_3,
@@ -215,6 +233,10 @@ def get_3body_comp(self, atom_env, sig, ls, r_cut):
 
     def predict_multicomponent(self, atom_env, sig, ls, r_cut, get_comp,
                                spcs_list, mappings, mean_only):
+        '''
+        Add up results from `predict_component` to get the total contribution 
+        of all species
+        '''
         f_spcs = 0
         v_spcs = 0
 
@@ -234,9 +256,7 @@ def predict_multicomponent(self, atom_env, sig, ls, r_cut, get_comp,
 
     def predict_component(self, lengths, xyzs, mapping, mean_only):
         '''
-        predict for an atom environment
-        param: atom_env: ChemicalEnvironment
-        return force on an atom with its variance
+        predict force and variance contribution of one component
         '''
         lengths = np.array(lengths)
         xyzs = np.array(xyzs)
@@ -254,33 +274,11 @@ def predict_component(self, lengths, xyzs, mapping, mean_only):
             v = mapping.var.V @ v_d
         return f, v
 
-    def write_two_plus_three(self, lammps_name):
-        # write header
-        f = open(lammps_name, 'w')
-
-        header_comment = '''# #2bodyarray #3bodyarray
-        # elem1 elem2 a b order
-        '''
-        f.write(header_comment)
-
-        twobodyarray = len(self.spcs[0])
-        threebodyarray = len(self.spcs[1])
-        lower_cut = self.bounds_2[0][0]
-        two_cut = self.bounds_2[1][0]
-        three_cut = self.bounds_3[1][0]
-        grid_num_2 = self.grid_num_2
-        grid_num_3 = self.grid_num_3[0]
-        angle_lower = self.bounds_3[0][2]
-        angle_upper = self.bounds_3[1][2]
-
-        header = '\n{} {}\n'.format(twobodyarray, threebodyarray)
-        f.write(header)
-
-        a = lower_cut
-        b = two_cut
-        order = grid_num_2
+    def write_two_body(self, f):
+        a = self.bounds_2[0][0]
+        b = self.bounds_2[1][0]
+        order = self.grid_num_2
 
-        # write two body
         for ind, spc in enumerate(self.spcs[0]):
             coefs_2 = self.maps_2[ind].mean.model.__coeffs__
 
@@ -297,10 +295,10 @@ def write_two_plus_three(self, lammps_name):
 
             f.write('\n')
 
-        # write three body
-        a = [lower_cut, lower_cut, angle_lower]
-        b = [three_cut, three_cut, angle_upper]
-        order = [grid_num_3, grid_num_3, grid_num_3]
+    def write_three_body(self, f):
+        a = self.bounds_3[0]
+        b = self.bounds_3[1] 
+        order = self.grid_num_3 
 
         for ind, spc in enumerate(self.spcs[1]):
             coefs_3 = self.maps_3[ind].mean.model.__coeffs__
@@ -330,12 +328,40 @@ def write_two_plus_three(self, lammps_name):
             f.write('\n')
 
 
+    def write_lmp_file(self, lammps_name):
+        '''
+        write the coefficients to a file that can be used by lammps pair style
+        '''
+
+        # write header
+        f = open(lammps_name, 'w')
+
+        header_comment = '''# #2bodyarray #3bodyarray
+        # elem1 elem2 a b order
+        '''
+        f.write(header_comment)
+
+        twobodyarray = len(self.spcs[0])
+        threebodyarray = len(self.spcs[1])
+        header = '\n{} {}\n'.format(twobodyarray, threebodyarray)
+        f.write(header)
+
+        # write two body
+        if twobodyarray > 0:
+            self.write_two_body(f)
+
+        # write three body
+        if threebodyarray > 0:
+            self.write_three_body(f)
+
+        f.close()
+
+
 class Map2body:
     def __init__(self, grid_num, bounds, GP, bond_struc, bodies='2',
-                 load_prefix=None, svd_rank=0, mean_only=False):
+                 svd_rank=0, mean_only=False):
         '''
-        param grids: the 1st element is the number of grids for mean
-        prediction, the 2nd is for var
+        Build 2-body MGP
         '''
 
         self.grid_num = grid_num
@@ -370,7 +396,7 @@ def GenGrid(self, GP, bond_struc, processes=mp.cpu_count()):
         bond_lengths = np.linspace(self.l_bound[0], self.u_bound[0], nop)
         bond_means = np.zeros([nop])
         bond_vars = np.zeros([nop, len(GP.alpha)])
-        env12 = env.AtomicEnvironment(bond_struc, 0, self.cutoffs)
+        env12 = AtomicEnvironment(bond_struc, 0, self.cutoffs)
 
         pool_list = [(i, bond_lengths, GP, env12)
                      for i in range(nop)]
@@ -432,10 +458,9 @@ def build_map(self, y_mean, y_var):
 class Map3body:
 
     def __init__(self, grid_num, bounds, GP, bond_struc, bodies='3',
-                 load_grid=None, svd_rank=0, mean_only=False, update=True):
+                 svd_rank=0, mean_only=False, load_grid=None, update=True):
         '''
-        param grids: the 1st element is the number of grids for mean
-        prediction, the 2nd is for var
+        Build 3-body MGP
         '''
 
         self.grid_num = grid_num
@@ -472,7 +497,7 @@ def GenGrid(self, GP, bond_struc, update, processes=mp.cpu_count()):
         angles = np.linspace(self.l_bound[2], self.u_bound[2], noa)
         bond_means = np.zeros([nop, nop, noa])
         bond_vars = np.zeros([nop, nop, noa, len(GP.alpha)])
-        env12 = env.AtomicEnvironment(bond_struc, 0, self.cutoffs)
+        env12 = AtomicEnvironment(bond_struc, 0, self.cutoffs)
 
         pool_list = [(i, angles[i], bond_lengths, GP, env12, update)\
                      for i in range(noa)]

flare/mgp/otf.py

@@ -33,6 +33,15 @@ def __init__(self, qe_input: str, dt: float, number_of_steps: int,
                  l_bound=None, two_d=False,
                  grid_params: dict={}, struc_params: dict={}):
 
+        '''
+        On-the-fly training with MGP inserted, added parameters:
+        :param non_mapping_steps: steps that not to use MGP
+        :param l_bound: initial minimal interatomic distance
+        :param two_d: if the system is a 2-D system
+        :param grid_params: see class `MappedGaussianProcess`
+        :param struc_params: see class `MappedGaussianProcess`
+        '''
+
         super().__init__(qe_input, dt, number_of_steps,
                          gp_model, pw_loc, 
                          std_tolerance_factor, 

flare/mgp/splines_methods.py

@@ -1,7 +1,5 @@
 import numpy as np
 import numpy
-import sys
-sys.path.append('../../flare')
 import time
 from memory_profiler import profile
 

flare/mgp/utils.py

@@ -1,9 +1,11 @@
 import numpy as np
 from numpy import array
 from numba import njit
-import io
-import sys
-import os
+import io, os, sys, time, random, logging
+import multiprocessing as mp
+import concurrent.futures
+import cProfile
+
 import flare.gp as gp
 import flare.env as env
 import flare.struc as struc
@@ -12,12 +14,6 @@
 from flare.kernels import triplet_kernel, three_body_helper_1, \
     three_body_helper_2, force_helper
 from flare.cutoffs import quadratic_cutoff
-import time
-import random
-import logging
-import multiprocessing as mp
-import concurrent.futures
-import cProfile
 
 
 def save_GP(GP, prefix):