so3krates implementation

libs/lib_criteria.py

@@ -395,6 +395,9 @@ def uncert_strconvter(value):
     """
     if value == '----          ':
         return value
+    if isinstance(value, np.float32):
+        value = float(value)
+
     return '{:.5e}'.format(Decimal(value))
 
 

libs/lib_input.py

@@ -49,6 +49,9 @@ def __init__(self, input_file='input.in'):
         self.output_format = 'trajectory.son'
 
         ##[Active learning setting]
+
+        self.MLIP = 'nequip'
+
         # nstep: int
         #     The number of subsampling sets
         self.nstep = 1
@@ -184,6 +187,41 @@ def __init__(self, input_file='input.in'):
         #     Recommend to use the ground state total energy.
         self.E_gs = 0.0
 
+        ##[so3krates setting]
+        self.r_cut = 5.0
+        self.l = 3
+        self.f = 132
+        self.l_min = 1
+        self.l_max = 3
+        self.max_body_order = 2
+        self.f_body_order = 1
+        self.we = 0.01 # weight-energy
+        self.wf = 1.00 # weight-forces
+        self.ws = 'None' # weight-stress
+        self.loss_variance_scaling = '--no-loss-variance-scaling'
+        self.epochs = 2000
+        self.eval_energy_t = 'None'
+        self.mic = '--mic'
+        self.float64 = '--no-float64'
+        self.lr = 0.001
+        self.lr_stop = 0.00001
+        self.lr_decay_exp_transition_steps = 100000
+        self.lr_decay_exp_decay_factor = 0.7
+        self.clip_by_global_norm = 'None'
+        self.shift_mean = '--shift-mean'
+        self.size_batch = 'None'
+        self.size_batch_training = 'None'
+        self.size_batch_validation = 'None'
+        # self.seed_model = 0
+        self.seed_data = 0
+        self.seed_training = 0
+        self.wandb_name = 'None'
+        self.wandb_group = 'None'
+        self.wandb_project = 'None'
+        # self.outfile_inputs = 'inputs.json'
+        self.overwrite_module = '--no-overwrite-module'
+        self.ace = '--no-ace'
+
         ##[Constants]
         # kB: float
         #     Boltzmann constant in units of eV/K
@@ -214,6 +252,8 @@ def __init__(self, input_file='input.in'):
         # ntotal: Total number of added training and valdiation data for all subsamplings for each iteractive step
         self.ntotal = (self.ntrain + self.nval) * self.nstep
 
+        self.train_split = self.ntrain/(self.ntrain + self.nval)
+
         ### Initization of parameters
         # index: int
         #     The index of AL interactive steps

libs/lib_load_model.py

@@ -3,6 +3,8 @@
 from nequip.ase import nequip_calculator
 from libs.lib_util     import mpi_print
 
+from glp import instantiate
+from glp.ase import Calculator
 import torch
 torch.set_default_dtype(torch.float64)
 
@@ -22,17 +24,23 @@ def load_model(inputs):
     for index_nmodel in range(inputs.nmodel):
         for index_nstep in range(inputs.nstep):
             if (index_nmodel * inputs.nstep + index_nstep) % inputs.size == inputs.rank:
-                dply_model = f'deployed-model_{index_nmodel}_{index_nstep}.pth'
-                if os.path.exists(f'{workpath}/{dply_model}'):
-                    mpi_print(f'\t\tFound the deployed model: {dply_model}', rank=0)
-                    inputs.calc_MLIP.append(
-                        nequip_calculator.NequIPCalculator.from_deployed_model(f'{workpath}/{dply_model}', device=inputs.device)
-                    )
-                else:
-                    # If there is no model, turn on the termination signal
-                    mpi_print(f'\t\tCannot find the model: {dply_model}', rank=0)
-                    signal = 1
-                    signal = inputs.comm.bcast(signal, root=inputs.rank)
+
+                if inputs.MLIP == 'nequip':
+                    dply_model = f'deployed-model_{index_nmodel}_{index_nstep}.pth'
+                    if os.path.exists(f'{workpath}/{dply_model}'):
+                        mpi_print(f'\t\tFound the deployed model: {dply_model}', rank=0)
+                        inputs.calc_MLIP.append(
+                            nequip_calculator.NequIPCalculator.from_deployed_model(f'{workpath}/{dply_model}', device=inputs.device)
+                        )
+                    else:
+                        # If there is no model, turn on the termination signal
+                        mpi_print(f'\t\tCannot find the model: {dply_model}', rank=0)
+                        signal = 1
+                        signal = inputs.comm.bcast(signal, root=inputs.rank)
+                elif inputs.MLIP == 'so3krates':
+                    potential_dict = {"mlff": {"folder": f"{workpath}/deployed-model_{index_nmodel}_{index_nstep}"}}
+                    get_calculator = instantiate.get_calculator(potential_dict, {"atom_pair": {"skin": 0.1}})
+                    inputs.calc_MLIP.append(Calculator(get_calculator))
 
     # Check the termination signal
     if signal == 1:

libs/lib_mainloop.py

@@ -2,7 +2,7 @@
 from ase.io.trajectory import TrajectoryWriter
 
 import os
-import son
+from vibes import son
 import copy
 import random
 import numpy as np

libs/lib_mainloop_new.py

@@ -3,7 +3,7 @@
 from ase.io.bundletrajectory import BundleTrajectory
 
 import os
-import son
+from vibes import son
 import copy
 import random
 import numpy as np

libs/lib_npz.py

@@ -103,11 +103,19 @@ def generate_npz_DFT_init(inputs, traj, workpath):
             'F': np.array(F_train_store),
             'R': np.array(R_train_store),
             'z': np.array(z_train_store),
-            'CELL': np.array(CELL_train_store),
-            'PBC': np.array(PBC_train_store),
             'sigma': np.array(sigma_train_store)
         }
 
+        if inputs.MLIP == 'nequip':
+            arrays_to_save['CELL'] = np.array(CELL_train_store)
+            arrays_to_save['PBC'] = np.array(PBC_train_store)
+        elif inputs.MLIP == 'so3krates':
+            arrays_to_save['unit_cell'] = np.array(CELL_train_store)
+            arrays_to_save['pbc'] = np.array(PBC_train_store)
+        else:
+            arrays_to_save['CELL'] = np.array(CELL_train_store)
+            arrays_to_save['PBC'] = np.array(PBC_train_store)
+
         if inputs.ensemble == 'NPTisoiso' or inputs.train_stress:
             arrays_to_save['stress'] = np.array(stress_train_store)
 
@@ -346,11 +354,19 @@ def generate_npz_DFT(inputs, workpath):
                 'F': np.array(F_train_store),
                 'R': np.array(R_train_store),
                 'z': np.array(z_train_store),
-                'CELL': np.array(CELL_train_store),
-                'PBC': np.array(PBC_train_store),
                 'sigma': np.array(sigma_train_store),
             }
 
+            if inputs.MLIP == 'nequip':
+                arrays_to_save['CELL'] = np.array(CELL_train_store)
+                arrays_to_save['PBC'] = np.array(PBC_train_store)
+            elif inputs.MLIP == 'so3krates':
+                arrays_to_save['unit_cell'] = np.array(CELL_train_store)
+                arrays_to_save['pbc'] = np.array(PBC_train_store)
+            else:
+                arrays_to_save['CELL'] = np.array(CELL_train_store)
+                arrays_to_save['PBC'] = np.array(PBC_train_store)
+
             if inputs.ensemble == 'NPTisoiso' or inputs.train_stress:
                 arrays_to_save['stress'] = np.array(stress_train_store)
 
@@ -468,11 +484,19 @@ def generate_npz_DFT_rand_init(inputs, traj, ntrain, nval, workpath):
             'F': np.array(F_train_store),
             'R': np.array(R_train_store),
             'z': np.array(z_train_store),
-            'CELL': np.array(CELL_train_store),
-            'PBC': np.array(PBC_train_store),
             'sigma': np.array(sigma_train_store)
         }
 
+        if inputs.MLIP == 'nequip':
+            arrays_to_save['CELL'] = np.array(CELL_train_store)
+            arrays_to_save['PBC'] = np.array(PBC_train_store)
+        elif inputs.MLIP == 'so3krates':
+            arrays_to_save['unit_cell'] = np.array(CELL_train_store)
+            arrays_to_save['pbc'] = np.array(PBC_train_store)
+        else:
+            arrays_to_save['CELL'] = np.array(CELL_train_store)
+            arrays_to_save['PBC'] = np.array(PBC_train_store)
+
         if inputs.ensemble == 'NPTisoiso' or inputs.train_stress:
             arrays_to_save['stress'] = np.array(stress_train_store)
 
@@ -626,11 +650,19 @@ def generate_npz_DFT_rand(
                 'F': np.array(F_train_store),
                 'R': np.array(R_train_store),
                 'z': np.array(z_train_store),
-                'CELL': np.array(CELL_train_store),
-                'PBC': np.array(PBC_train_store),
                 'sigma': np.array(sigma_train_store)
             }
 
+            if inputs.MLIP == 'nequip':
+                arrays_to_save['CELL'] = np.array(CELL_train_store)
+                arrays_to_save['PBC'] = np.array(PBC_train_store)
+            elif inputs.MLIP == 'so3krates':
+                arrays_to_save['unit_cell'] = np.array(CELL_train_store)
+                arrays_to_save['pbc'] = np.array(PBC_train_store)
+            else:
+                arrays_to_save['CELL'] = np.array(CELL_train_store)
+                arrays_to_save['PBC'] = np.array(PBC_train_store)
+
             if inputs.ensemble == 'NPTisoiso' or inputs.train_stress:
                 arrays_to_save['stress'] = np.array(stress_train_store)
 

libs/lib_progress.py

@@ -6,7 +6,6 @@
 
 from libs.lib_util        import check_mkdir, mpi_print, single_print, generate_msg
 from libs.lib_criteria    import get_result
-from libs.lib_train       import get_train_job
 from libs.lib_termination import get_testerror