fix init

libs/lib_cont_nvtlangevin_bias.py

@@ -159,7 +159,7 @@ def cont_NVTLangevin_bias(
         )
 
     # Get averaged force from trained models
-    forces = get_forces_bias(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected, inputs.bias_A, inputs.bias_B)
+    forces = get_forces_bias(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected, inputs.bias_A, inputs.bias_B, inputs.idx_atom)
 
     # mpi_print(f'Step 3: {time.time()-time_init}', rank)
     # Go trough steps until the requested number of steps
@@ -172,7 +172,7 @@ def cont_NVTLangevin_bias(
         # mpi_print(f'Step 6: {time.time()-time_init}', rank)
         # Get averaged forces and velocities
         if forces is None:
-            forces = get_forces_bias(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected, inputs.bias_A, inputs.bias_B)
+            forces = get_forces_bias(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected, inputs.bias_A, inputs.bias_B, inputs.idx_atom)
         # Velocity is already calculated based on averaged forces
         # in the previous step
         velocity = struc.get_velocities()
@@ -225,7 +225,7 @@ def cont_NVTLangevin_bias(
         velocity = (struc.get_positions() - position - rnd_pos) / timestep
         comm.Barrier()
         # mpi_print(f'Step 10-1: {time.time()-time_init}', rank)
-        forces = get_forces_bias(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected, inputs.bias_A, inputs.bias_B)
+        forces = get_forces_bias(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected, inputs.bias_A, inputs.bias_B, inputs.idx_atom)
         comm.Barrier()
 
         # mpi_print(f'Step 10-2: {time.time()-time_init}', rank)
@@ -318,7 +318,7 @@ def cont_NVTLangevin_bias(
 
 
 def get_forces_bias(
-    struc, nstep, nmodel, calculator, harmonic_F, anharmonic_F, criteria, bias_A, bias_B
+    struc, nstep, nmodel, calculator, harmonic_F, anharmonic_F, criteria, bias_A, bias_B, idx_atom
 ):
     """Function [get_forces]
     Evalulate the average of forces from all different trained models.
@@ -406,7 +406,6 @@ def get_forces_bias(
         # print(f'Bias Ratio: {np.average(ratio_bias)}\n')
 
 
-        idx_atom = 0 # 120
         E_bias_deriv = E_step_std[idx_atom] / (bias_B ** 2) * bias_A * (np.exp((-1) * (E_step_std[idx_atom] ** 2) / (2 * bias_B ** 2)))
         F_bias_elem = np.array([0.0, 0.0, 0.0])
         for idx_E, idx_F in zip(np.array(E_step_filtered)[:,idx_atom], np.array(F_step_filtered)[:,idx_atom]):

libs/lib_cont_nvtlangevin_temp.py

@@ -159,7 +159,7 @@ def cont_NVTLangevin_temp(
         )
 
     # Get averaged force from trained models
-    forces, F_step_norm_std = get_forces_temp(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected)
+    forces, step_std = get_forces_temp(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected, inputs.al_type)
 
     # mpi_print(f'Step 3: {time.time()-time_init}', rank)
     # Go trough steps until the requested number of steps
@@ -172,7 +172,7 @@ def cont_NVTLangevin_temp(
         # mpi_print(f'Step 6: {time.time()-time_init}', rank)
         # Get averaged forces and velocities
         if forces is None:
-            forces, F_step_norm_std = get_forces_temp(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected)
+            forces, step_std = get_forces_temp(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected, inputs.al_type)
         # Velocity is already calculated based on averaged forces
         # in the previous step
         velocity = struc.get_velocities()
@@ -200,20 +200,28 @@ def cont_NVTLangevin_temp(
         c5 = timestep**1.5 * sigma / (2 * np.sqrt(3))
         c4 = inputs.friction / 2. * c5
 
-        uncert_max_idx = np.argmax(F_step_norm_std)
-
-        uncert_ratio = (F_step_norm_std[uncert_max_idx] - criteria_collected.Un_Abs_F_avg_i) / (criteria_collected.Un_Abs_F_std_i * 0.025)
-        temp_ratio = 1.0 if uncert_ratio < 0 else 300.0 if uncert_ratio > 299 else uncert_ratio + 1
-        if temp_ratio > 1.0:
-            print(f'Step {MD_step_index}; temp activate (atom {uncert_max_idx}):{temp_ratio}')
-        sigma_elem = np.sqrt(2 * (temperature * temp_ratio) * inputs.friction / masses[uncert_max_idx][0])
+        if inputs.al_type == 'force_max':
+            uncert_ratio = (step_std[inputs.idx_atom] - criteria_collected.Un_Abs_F_avg_i) / (criteria_collected.Un_Abs_F_std_i)
+            if step_std[inputs.idx_atom] - criteria_collected.Un_Abs_F_avg_i < 0:
+                temp_ratio = 1.0
+            else:
+                temp_ratio = np.exp( (-1/2) * (uncert_ratio)**2 )
+        elif inputs.al_type == 'energy_max':
+            uncert_ratio = (step_std[inputs.idx_atom] - criteria_collected.Un_Abs_E_avg_i) / (criteria_collected.Un_Abs_E_std_i)
+            if step_std[inputs.idx_atom] - criteria_collected.Un_Abs_E_avg_i < 0:
+                temp_ratio = 1.0
+            else:
+                temp_ratio = np.exp( (-1/2) * (uncert_ratio)**2 )
+
+        print(f'Step {MD_step_index}; temp activate (atom {inputs.idx_atom}):{temp_ratio}')
+        sigma_elem = np.sqrt(2 * (temperature + temp_ratio * (inputs.temp_factor * units.kB)) * inputs.friction / masses[inputs.idx_atom][0])
         c3_elem = np.sqrt(timestep) * sigma_elem / 2. - timestep**1.5 * inputs.friction * sigma_elem / 8.
         c5_elem = timestep**1.5 * sigma_elem / (2 * np.sqrt(3))
         c4_elem = inputs.friction / 2. * c5_elem
 
-        c3[uncert_max_idx] = c3_elem
-        c4[uncert_max_idx] = c4_elem
-        c5[uncert_max_idx] = c5_elem
+        c3[inputs.idx_atom] = c3_elem
+        c4[inputs.idx_atom] = c4_elem
+        c5[inputs.idx_atom] = c5_elem
 
         # c3, c4, c5 = [], [], []
 
@@ -260,7 +268,7 @@ def cont_NVTLangevin_temp(
         velocity = (struc.get_positions() - position - rnd_pos) / timestep
         comm.Barrier()
         # mpi_print(f'Step 10-1: {time.time()-time_init}', rank)
-        forces, F_step_norm_std = get_forces_temp(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected)
+        forces, step_std = get_forces_temp(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected, inputs.al_type)
         comm.Barrier()
 
         # mpi_print(f'Step 10-2: {time.time()-time_init}', rank)
@@ -353,7 +361,7 @@ def cont_NVTLangevin_temp(
 
 
 def get_forces_temp(
-    struc, nstep, nmodel, calculator, harmonic_F, anharmonic_F, criteria
+    struc, nstep, nmodel, calculator, harmonic_F, anharmonic_F, criteria, al_type
 ):
     """Function [get_forces]
     Evalulate the average of forces from all different trained models.
@@ -397,7 +405,10 @@ def get_forces_temp(
                     # mpi_print(f'Step 10-a1 second {rank}: {time.time()-time_init}', rank)
                     temp_force = struc.get_forces()
                     # mpi_print(f'Step 10-a1 third {rank}: {time.time()-time_init}', rank)
-                    energy.append(struc.get_potential_energy())
+                    if al_type == 'energy_max':
+                        energy.append(struc.get_potential_energies())
+                    else:
+                        energy.append(struc.get_potential_energy())
                     forces.append(temp_force)
                     # sigmas.append(eval_sigma(temp_force, struc.get_positions(), 'force_max'))
                     # mpi_print(f'Step 10-a1 last {rank}: {time.time()-time_init}', rank)
@@ -408,9 +419,10 @@ def get_forces_temp(
         # sigmas = comm.allgather(sigmas)
         # mpi_print(f'Step 10-a3: {time.time()-time_init}', rank)
         E_step_filtered = [jtem for item in energy if len(item) != 0 for jtem in item]
-        F_step_filtered = [jtem for item in forces if len(item) != 0 for jtem in item]
-        # sigma_filtered = [jtem for item in sigmas if len(item) != 0 for jtem in item]
+        E_step_avg = np.average(E_step_filtered, axis=0)
+        E_step_std = np.std(E_step_filtered, axis=0)
 
+        F_step_filtered = [jtem for item in forces if len(item) != 0 for jtem in item]
         F_step_avg = np.average(F_step_filtered, axis=0)
         F_step_norm = np.array([[np.linalg.norm(Fcomp) for Fcomp in Ftems] for Ftems in F_step_filtered - F_step_avg])
         F_step_norm_std = np.sqrt(np.average(F_step_norm ** 2, axis=0))
@@ -430,4 +442,7 @@ def get_forces_temp(
 
     # mpi_print(f'Step 10-d: {time.time()-time_init}', rank)
 
-    return force_avg, F_step_norm_std
+    if al_type == 'energy_max':
+        return force_avg, E_step_std
+    else:
+        return force_avg, F_step_norm_std

libs/lib_criteria.py

@@ -73,6 +73,20 @@ def eval_uncert(
 
         return (uncerts, Epot_step_avg, S_step_avg)
 
+    elif al_type == 'energy_max':
+        uncerts.UncertAbs_E = np.max(Epot_step_std)
+        uncerts.UncertRel_E = np.max(
+            np.array([std / np.absolute(avg) for avg, std in zip(Epot_step_avg, Epot_step_std)])
+            )
+        uncerts.UncertAbs_F = np.ndarray.max(F_step_norm_std)
+        uncerts.UncertRel_F = np.ndarray.max(
+            np.array([std / avg for avg, std in zip(F_step_norm_avg, F_step_norm_std)])
+            )
+        uncerts.UncertAbs_S = S_step_std
+        uncerts.UncertRel_S = S_step_std / S_step_avg
+
+        return (uncerts, np.sum(Epot_step_avg), S_step_avg)
+
     else:
         sys.exit("You need to set al_type.")
 
@@ -119,13 +133,19 @@ def eval_uncert_all(
     F_step = []
     prd_struc = []
     zndex = 0
+    natoms = len(struc_step)
 
     # Get predicted potential and total energies shifted by E_ref (ground state energy)
     for index_nmodel in range(nmodel):
         for index_nstep in range(nstep):
             if (index_nmodel*nstep + index_nstep) % size == rank:
                 struc_step.calc = calculator[zndex]
-                Epot_step.append(struc_step.get_potential_energy() - E_ref)
+
+                if al_type == 'energy_max':
+                    Epot_step.append(np.array(struc_step.get_potential_energies()) - E_ref/natoms)
+                else:
+                    Epot_step.append(struc_step.get_potential_energy() - E_ref)
+
                 F_step.append(struc_step.get_forces())
                 prd_struc.append(struc_step.get_positions())
                 zndex += 1
@@ -228,7 +248,7 @@ def get_criteria(
     criteria.Epotential_avg = result_data.loc[:, 'E_potent_avg_i'].to_numpy()[-1]
     criteria.Epotential_std = result_data.loc[:, 'E_potent_std_i'].to_numpy()[-1]
 
-    if al_type == 'energy':
+    if al_type == 'energy' or al_type == 'energy_max':
         criteria.Un_Abs_E_avg_i = result_data.loc[:, 'Un_Abs_E_avg_i'].to_numpy()[-1]
         criteria.Un_Abs_E_std_i = result_data.loc[:, 'Un_Abs_E_std_i'].to_numpy()[-1]
         criteria.Un_Rel_E_avg_i = result_data.loc[:, 'Un_Rel_E_avg_i'].to_numpy()[-1]
@@ -239,6 +259,17 @@ def get_criteria(
         criteria.Un_Rel_E_avg_i = 0.0
         criteria.Un_Rel_E_std_i = 0.0
 
+    if al_type == 'energy_max':
+        criteria.Un_Abs_Ea_avg_i = result_data.loc[:, 'Un_Abs_Ea_avg_i'].to_numpy()[-1]
+        criteria.Un_Abs_Ea_std_i = result_data.loc[:, 'Un_Abs_Ea_std_i'].to_numpy()[-1]
+        criteria.Un_Rel_Ea_avg_i = result_data.loc[:, 'Un_Rel_Ea_avg_i'].to_numpy()[-1]
+        criteria.Un_Rel_Ea_std_i = result_data.loc[:, 'Un_Rel_Ea_std_i'].to_numpy()[-1]
+    else:
+        criteria.Un_Abs_Ea_avg_i = 0.0
+        criteria.Un_Abs_Ea_std_i = 0.0
+        criteria.Un_Rel_Ea_avg_i = 0.0
+        criteria.Un_Rel_Ea_std_i = 0.0
+
     if al_type == 'force' or al_type == 'force_max':
         criteria.Un_Abs_F_avg_i = result_data.loc[:, 'Un_Abs_F_avg_i'].to_numpy()[-1]
         criteria.Un_Abs_F_std_i = result_data.loc[:, 'Un_Abs_F_std_i'].to_numpy()[-1]
@@ -294,7 +325,7 @@ def get_result(inputs, get_type):
 
     result_print = ''
     # Get their average and standard deviation
-    if inputs.al_type == 'energy':
+    if inputs.al_type == 'energy' or inputs.al_type == 'energy_max':
         UncerAbs_E_list = uncert_data.loc[:,'UncertAbs_E'].values
         UncerRel_E_list = uncert_data.loc[:,'UncertRel_E'].values
         criteria_UncertAbs_E_avg_all = uncert_average(UncerAbs_E_list[:])
@@ -442,7 +473,7 @@ def get_criteria_prob(inputs, Epot_step, uncerts, criteria):
     criteria_Uncert_S = 1
 
     # Calculate the probability based on energy, force, or both energy and force
-    if inputs.al_type == 'energy':
+    if inputs.al_type == 'energy' or inputs.al_type == 'energy_max':
         criteria_Uncert_E = get_criteria_uncert(
             inputs.uncert_type, inputs.uncert_shift, inputs.uncert_grad,
             uncerts.UncertAbs_E, criteria.Un_Abs_E_avg_i, criteria.Un_Abs_E_std_i,
@@ -469,7 +500,7 @@ def get_criteria_prob(inputs, Epot_step, uncerts, criteria):
 
     beta = inputs.kB * inputs.temperature
 
-    if inputs.ensemble == 'NVTLangevin_meta' or inputs.ensemble == 'NVTLangevin_bias' or inputs.ensemble == 'NPTisoiso' :
+    if inputs.ensemble == 'NVTLangevin_meta' or inputs.ensemble == 'NVTLangevin_bias' or inputs.ensemble == 'NVTLangevin_bias_temp' or inputs.ensemble == 'NPTisoiso' or inputs.criteria_energy == False:
         criteria_Prob = 1
     else:
         # Caculate the canonical ensemble propbability using the total energy

libs/lib_input.py

@@ -83,6 +83,8 @@ def __init__(self, input_file='input.in'):
         #     The number of iteration for DFT random sampling
         self.random_index = 1
 
+        self.train_stress = False
+
         ##[Molecular dynamics setting]
         # ensemble: str
         #     Type of MD ensembles; 'NVTLangevin'
@@ -114,6 +116,15 @@ def __init__(self, input_file='input.in'):
         self.meta_restart = False
         self.meta_r_crtria = 100.0
 
+        self.idx_atom = 0
+        self.bias_A = 0.0
+        self.bias_B = 999
+
+        self.temp_factor = 0.0
+
+        self.npz_test = 'test'
+        self.criteria_energy = True
+
         ##[runMD default inputs]
         # workpath: str
         #     A path to the directory containing trained models

libs/lib_mainloop_new.py

@@ -156,7 +156,7 @@ def MLMD_main(
             struc_init = atoms_read('start.in', format='aims')
             # Make it supercell
             struc_step = make_supercell(struc_init, inputs.supercell_init)
-            MaxwellBoltzmannDistribution(struc_step, temperature_K=inputs.temperature, force_temp=True)
+            MaxwellBoltzmannDistribution(struc_step, temperature_K=inputs.temperature*1.5, force_temp=True)
 
         elif os.path.exists('start.traj'):
             mpi_print(f'[runMD]\tFound the start.traj file. MD starts from this.', inputs.rank)
@@ -167,7 +167,7 @@ def MLMD_main(
             try:
                 struc_step.get_velocities()
             except AttributeError:
-                MaxwellBoltzmannDistribution(struc_step, temperature_K=inputs.temperature, force_temp=True)
+                MaxwellBoltzmannDistribution(struc_step, temperature_K=inputs.temperature*1.5, force_temp=True)
 
         elif os.path.exists('start.bundle'):
             from ase.io.bundletrajectory import BundleTrajectory
@@ -180,7 +180,7 @@ def MLMD_main(
             try:
                 struc_step.get_velocities()
             except AttributeError:
-                MaxwellBoltzmannDistribution(struc_step, temperature_K=inputs.temperature, force_temp=True)
+                MaxwellBoltzmannDistribution(struc_step, temperature_K=inputs.temperature*1.5, force_temp=True)
 
         else:
             mpi_print(f'[MLMD] Read a configuration from trajectory_train.son', inputs.rank)
@@ -207,7 +207,7 @@ def MLMD_main(
             traj_previous = Trajectory(traj_temp)
         struc_step = traj_previous[-1]; del traj_previous;
 
-    if os.path.exists('start.in') and (inputs.ensemble == 'NVTLangevin_meta' or inputs.ensemble == 'NVTLangevin_temp' or inputs.ensemble == 'NVTLangevin_bias') and MD_index == 0:
+    if os.path.exists('start.in') and (inputs.ensemble == 'NVTLangevin_meta' or inputs.ensemble == 'NVTLangevin_temp' or inputs.ensemble == 'NVTLangevin_bias' or inputs.ensemble == 'NVTLangevin_bias_temp') and MD_index == 0:
         mpi_print(f'[MLMD] Read a configuration from start.in', inputs.rank)
         # Read the ground state structure with the primitive cell
         struc_init = atoms_read('start.in', format='aims')
@@ -257,7 +257,7 @@ def traj_fromRealE(temperature, pressure, E_ref, E_gs, uncert_type, al_type, nto
     elif uncert_type == 'relative':
         uncert_piece = 'Rel'
 
-    if al_type == 'energy':
+    if al_type == 'energy' or 'energy_max':
         al_piece = 'E'
     elif al_type == 'force' or 'force_max':
         al_piece = 'F'

libs/lib_md.py

@@ -234,6 +234,20 @@ def cont_runMD(
             signal_uncert = signal_uncert,
             signal_append = signal_append
         )
+    elif inputs.ensemble == 'NVTLangevin_bias_temp':
+        from libs.lib_cont_nvtlangevin_bias_temp import cont_NVTLangevin_bias_temp
+        cont_NVTLangevin_bias_temp(
+            inputs = inputs,
+            struc = struc,
+            timestep = inputs.timestep * units.fs,
+            temperature = inputs.temperature * units.kB,
+            calculator = calculator,
+            E_ref = 0.0,
+            MD_index = MD_index,
+            MD_step_index = MD_step_index,
+            signal_uncert = signal_uncert,
+            signal_append = signal_append
+        )
     elif inputs.ensemble == 'NPTisoiso':
         from libs.lib_cont_nptisoiso import cont_NPTisoiso
         cont_NPTisoiso(