diff --git a/libs/lib_cont_nvtlangevin_bias.py b/libs/lib_cont_nvtlangevin_bias.py index a99dbac..d8dec65 100644 --- a/libs/lib_cont_nvtlangevin_bias.py +++ b/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]): diff --git a/libs/lib_cont_nvtlangevin_bias_temp.py b/libs/lib_cont_nvtlangevin_bias_temp.py new file mode 100644 index 0000000..c199087 --- /dev/null +++ b/libs/lib_cont_nvtlangevin_bias_temp.py @@ -0,0 +1,472 @@ +from ase.io.trajectory import Trajectory +from ase.io.trajectory import TrajectoryWriter +import ase.units as units +from ase.io.cif import write_cif + +import time +import os +import random +import numpy as np +import pandas as pd +from mpi4py import MPI +from decimal import Decimal +from ase.build import make_supercell +from ase.io import read as atoms_read +from ase.md.velocitydistribution import MaxwellBoltzmannDistribution + +from libs.lib_util import single_print, mpi_print +from libs.lib_MD_util import get_forces, get_MDinfo_temp, get_masses +from libs.lib_criteria import eval_uncert, uncert_strconvter, get_criteria, get_criteria_prob + +import torch +torch.set_default_dtype(torch.float64) + +def cont_NVTLangevin_bias_temp( + inputs, struc, timestep, temperature, calculator, E_ref, + MD_index, MD_step_index, signal_uncert=False, signal_append=True, fix_com=True, +): + """Function [NVTLangevin] + Evalulate the absolute and relative uncertainties of + predicted energies and forces. + This script is adopted from ASE Langevin function + and modified to use averaged forces from trained model. + + Parameters: + + struc: ASE atoms + A structral configuration of a starting point + timestep: float + The step interval for printing MD steps + temperature: float + The desired temperature in units of Kelvin (K) + + friction: float + Strength of the friction parameter in NVTLangevin ensemble + steps: int + The length of the Molecular Dynamics steps + loginterval: int + The step interval for printing MD steps + + nstep: int + The number of subsampling sets + nmodel: int + The number of ensemble model sets with different initialization + calculator: ASE calculator + Any calculator + E_ref: flaot + The energy of reference state (Here, ground state) + al_type: str + Type of active learning: 'energy', 'force', 'force_max' + trajectory: str + A name of MD trajectory file + + logfile: str (optional) + A name of MD logfile. With None, it will not print a log file. + signal_uncert: bool (optional) + + + fixcm: bool (optional) + If True, the position and momentum of the center of mass is + kept unperturbed. Default: True. + """ + + time_init = time.time() + + # Extract MPI infos + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + + # Initialization of index + condition = f'{inputs.temperature}K-{inputs.pressure}bar' + + trajectory = f'TEMPORARY/temp-{condition}_{inputs.index}.traj' + logfile = f'TEMPORARY/temp-{condition}_{inputs.index}.log' + + # Extract the criteria information from the initialization step + criteria_collected = get_criteria(inputs.temperature, inputs.pressure, inputs.index, inputs.steps_init, inputs.al_type) + + if os.path.exists(trajectory): + traj_temp = Trajectory(trajectory) + struc = traj_temp[-1] + MD_step_index = len(traj_temp) + del traj_temp + + # mpi_print(f'Step 1: {time.time()-time_init}', rank) + if MD_step_index == 0: # If appending and the file exists, + file_traj = TrajectoryWriter(filename=trajectory, mode='w') + # Add new configuration to the trajectory file + if rank == 0: + file_traj.write(atoms=struc) + + if isinstance(logfile, str): + if rank == 0: + file_log = open(logfile, 'w') + file_log.write( + 'Time[ps] \tEtot[eV] \tEpot[eV] \tEkin[eV] \t' + + 'Temperature[K]' + ) + if signal_uncert: + file_log.write( + '\tUncertAbs_E\tUncertRel_E\t' + + 'UncertAbs_F\tUncertRel_F\t' + + 'UncertAbs_S\tUncertRel_S\tS_average\n' + ) + else: + file_log.write('\n') + file_log.close() + + # Get MD information at the current step + info_TE, info_PE, info_KE, info_T = get_MDinfo_temp( + struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F + ) + + if signal_uncert: + # Get absolute and relative uncertainties of energy and force + # and also total energy + uncerts, Epot_step, S_step =\ + eval_uncert(struc, inputs.nstep, inputs.nmodel, E_ref, calculator, inputs.al_type, inputs.harmonic_F) + + # Log MD information at the current step in the log file + if rank == 0: + file_log = open(logfile, 'a') + file_log.write( + '{:.5f}'.format(Decimal(str(0.0))) + ' \t' + + '{:.5e}'.format(Decimal(str(info_TE))) + '\t' + + '{:.5e}'.format(Decimal(str(info_PE))) + '\t' + + '{:.5e}'.format(Decimal(str(info_KE))) + '\t' + + '{:.2f}'.format(Decimal(str(info_T))) + ) + if signal_uncert: + file_log.write( + ' \t' + + uncert_strconvter(uncerts.UncertAbs_E) + '\t' + + uncert_strconvter(uncerts.UncertRel_E) + '\t' + + uncert_strconvter(uncerts.UncertAbs_F) + '\t' + + uncert_strconvter(uncerts.UncertRel_F) + '\t' + + uncert_strconvter(uncerts.UncertAbs_S) + '\t' + + uncert_strconvter(uncerts.UncertRel_S) + '\t' + + uncert_strconvter(S_step) + '\n' + ) + else: + file_log.write('\n') + file_log.close() + else: + file_traj = TrajectoryWriter(filename=trajectory, mode='a') + + write_traj = TrajectoryWriter( + filename=f'TRAJ/traj-{condition}_{inputs.index+1}.traj', + mode='a' + ) + + # Get averaged force from trained models + forces, step_std = get_forces_bias_temp(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected, inputs.al_type, 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 + # If appending, it starts from Langevin_idx. Otherwise, Langevin_idx = 0 + while (MD_index < inputs.ntotal) or (inputs.calc_type == 'period' and MD_step_index < inputs.nperiod*inputs.loginterval): + + accept = '-- ' + natoms = len(struc) + + # mpi_print(f'Step 6: {time.time()-time_init}', rank) + # Get averaged forces and velocities + if forces is None: + forces, step_std = get_forces_bias_temp(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected, inputs.al_type, 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() + + # mpi_print(f'Step 7: {time.time()-time_init}', rank) + # Sample the random numbers for the temperature fluctuation + xi = np.empty(shape=(natoms, 3)) + eta = np.empty(shape=(natoms, 3)) + if rank == 0: + xi = np.random.standard_normal(size=(natoms, 3)) + eta = np.random.standard_normal(size=(natoms, 3)) + comm.Bcast(xi, root=0) + comm.Bcast(eta, root=0) + + # mpi_print(f'Step 4: {time.time()-time_init}', rank) + # Get essential properties + masses = get_masses(struc.get_masses(), natoms) + + # mpi_print(f'Step 5: {time.time()-time_init}', rank) + # Get Langevin coefficients + sigma = np.sqrt(2 * temperature * inputs.friction / masses) + c1 = timestep / 2. - timestep * timestep * inputs.friction / 8. + c2 = timestep * inputs.friction / 2 - timestep * timestep * inputs.friction * inputs.friction / 8. + c3 = np.sqrt(timestep) * sigma / 2. - timestep**1.5 * inputs.friction * sigma / 8. + c5 = timestep**1.5 * sigma / (2 * np.sqrt(3)) + c4 = inputs.friction / 2. * c5 + + 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) + uncert_ratio = (step_std[inputs.idx_atom] - criteria_collected.Un_Abs_Ea_avg_i) / (criteria_collected.Un_Abs_Ea_std_i) + # if step_std[inputs.idx_atom] - criteria_collected.Un_Abs_E_avg_i < 0: + if step_std[inputs.idx_atom] - criteria_collected.Un_Abs_Ea_avg_i < 0: + temp_ratio = 1.0 + else: + temp_ratio = np.exp( (-1/2) * (uncert_ratio)**2 ) + else: + temp_ratio = 0.0 + + 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[inputs.idx_atom] = c3_elem + c4[inputs.idx_atom] = c4_elem + c5[inputs.idx_atom] = c5_elem + + # mpi_print(f'Step 8: {time.time()-time_init}', rank) + # Get get changes of positions and velocities + rnd_pos = c5 * eta + rnd_vel = c3 * xi - c4 * eta + + # Check the center of mass + if fix_com: + rnd_pos -= rnd_pos.sum(axis=0) / natoms + rnd_vel -= (rnd_vel * masses).sum(axis=0) / (masses * natoms) + + # First halfstep in the velocity. + velocity += (c1 * forces / masses - c2 * velocity + rnd_vel) + + # mpi_print(f'Step 9: {time.time()-time_init}', rank) + # Full step in positions + position = struc.get_positions() + + # Step: x^n -> x^(n+1) - this applies constraints if any. + struc.set_positions(position + timestep * velocity + rnd_pos) + + # mpi_print(f'Step 10: {time.time()-time_init}', rank) + # recalc velocities after RATTLE constraints are applied + velocity = (struc.get_positions() - position - rnd_pos) / timestep + comm.Barrier() + # mpi_print(f'Step 10-1: {time.time()-time_init}', rank) + forces, step_std = get_forces_bias_temp(struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F, inputs.anharmonic_F, criteria_collected, inputs.al_type, inputs.bias_A, inputs.bias_B, inputs.idx_atom) + comm.Barrier() + + # mpi_print(f'Step 10-2: {time.time()-time_init}', rank) + # Update the velocities + velocity += (c1 * forces / masses - c2 * velocity + rnd_vel) + + # mpi_print(f'Step 10-3: {time.time()-time_init}', rank) + # Second part of RATTLE taken care of here + struc.set_momenta(velocity * masses) + + # mpi_print(f'Step 11: {time.time()-time_init}', rank) + # Log MD information at regular intervals + if (MD_step_index+1) % inputs.loginterval == 0: + + # Get absolute and relative uncertainties of energy and force + # and also total energy + uncerts, Epot_step, S_step =\ + eval_uncert(struc, inputs.nstep, inputs.nmodel, 0.0, calculator, inputs.al_type, inputs.harmonic_F) + + # Get a criteria probability from uncertainty and energy informations + criteria = get_criteria_prob(inputs, Epot_step, uncerts, criteria_collected) + + if inputs.rank == 0: + # Acceptance check with criteria + ##!! Epot_step should be rechecked. + if random.random() < criteria: # and Epot_step > 0.1: + accept = 'Accepted' + MD_index += 1 + write_traj.write(atoms=struc) + else: + accept = 'Vetoed' + + # Record the MD results at the current step + trajfile = open(f'UNCERT/uncertainty-{condition}_{inputs.index}.txt', 'a') + trajfile.write( + '{:.5e}'.format(Decimal(str(struc.get_temperature()))) + '\t' + + uncert_strconvter(uncerts.UncertAbs_E) + '\t' + + uncert_strconvter(uncerts.UncertRel_E) + '\t' + + uncert_strconvter(uncerts.UncertAbs_F) + '\t' + + uncert_strconvter(uncerts.UncertRel_F) + '\t' + + uncert_strconvter(uncerts.UncertAbs_S) + '\t' + + uncert_strconvter(uncerts.UncertRel_S) + '\t' + + uncert_strconvter(Epot_step) + '\t' + + uncert_strconvter(S_step) + '\t' + + str(MD_index) + ' \t' + + '{:.5e}'.format(Decimal(str(criteria))) + '\t' + + str(accept) + ' \n' + ) + trajfile.close() + + if isinstance(logfile, str): + # mpi_print(f'Step 12: {time.time()-time_init}', rank) + info_TE, info_PE, info_KE, info_T = get_MDinfo_temp( + struc, inputs.nstep, inputs.nmodel, calculator, inputs.harmonic_F + ) + + # mpi_print(f'Step 14: {time.time()-time_init}', rank) + if inputs.rank == 0: + file_log = open(logfile, 'a') + simtime = timestep*(MD_step_index+inputs.loginterval)/units.fs/1000 + file_log.write( + '{:.5f}'.format(Decimal(str(simtime))) + ' \t' + + '{:.5e}'.format(Decimal(str(info_TE))) + '\t' + + '{:.5e}'.format(Decimal(str(info_PE))) + '\t' + + '{:.5e}'.format(Decimal(str(info_KE))) + '\t' + + '{:.2f}'.format(Decimal(str(info_T))) + ) + if signal_uncert: + file_log.write( + ' \t' + + uncert_strconvter(uncerts.UncertAbs_E) + '\t' + + uncert_strconvter(uncerts.UncertRel_E) + '\t' + + uncert_strconvter(uncerts.UncertAbs_F) + '\t' + + uncert_strconvter(uncerts.UncertRel_F) + '\t' + + uncert_strconvter(uncerts.UncertAbs_S) + '\t' + + uncert_strconvter(uncerts.UncertRel_S) + '\t' + + uncert_strconvter(S_step) + '\n' + ) + else: + file_log.write('\n') + file_log.close() + # mpi_print(f'Step 15: {time.time()-time_init}', rank) + if rank == 0: + file_traj.write(atoms=struc) + MD_index = inputs.comm.bcast(MD_index, root=0) + + MD_step_index += 1 + MD_step_index = inputs.comm.bcast(MD_step_index, root=0) + # mpi_print(f'Step 16: {time.time()-time_init}', rank) + + +def get_forces_bias_temp( + struc, nstep, nmodel, calculator, harmonic_F, anharmonic_F, criteria, al_type, bias_A, bias_B, idx_atom +): + """Function [get_forces] + Evalulate the average of forces from all different trained models. + + Parameters: + + struc_step: ASE atoms + A structral configuration at the current step + nstep: int + The number of subsampling sets + nmodel: int + The number of ensemble model sets with different initialization + calculator: ASE calculator or list of ASE calculators + Calculators from trained models + + Returns: + + force_avg: np.array of float + Averaged forces across trained models + """ + + # time_init = time.time() + from libs.lib_util import eval_sigma + + # Extract MPI infos + comm = MPI.COMM_WORLD + size = comm.Get_size() + rank = comm.Get_rank() + + # mpi_print(f'Step 10-a: {time.time()-time_init}', rank) + if type(calculator) == list: + energies = [] + forces = [] + sigmas = [] + zndex = 0 + for index_nmodel in range(nmodel): + for index_nstep in range(nstep): + if (index_nmodel*nstep + index_nstep) % size == rank: + # mpi_print(f'Step 10-a1 first {rank}: {time.time()-time_init}', rank) + struc.calc = calculator[zndex] + # 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) + energies.append(struc.get_potential_energies()) + 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) + zndex += 1 + # mpi_print(f'Step 10-a2: {time.time()-time_init}', rank) + energies = comm.allgather(energies) + forces = comm.allgather(forces) + # sigmas = comm.allgather(sigmas) + + E_step_filtered = [jtem for item in energies 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)) + F_step_norm_avg = np.linalg.norm(F_step_avg, axis=1) + + # idx_E_std = np.argmax(F_step_norm_std) + + # !!!Entire!!! + # F_bias = [] + # natoms = len(struc) + # for idx in range(natoms): # mpi needs to be fixed. (coom. Sequence problem); Serial is fine. + # E_bias_deriv = E_step_std[idx] / (bias_B ** 2) * bias_A * (np.exp((-1) * (E_step_std[idx] ** 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], np.array(F_step_filtered)[:,idx]): + # F_bias_elem += (idx_E - E_step_avg[idx])*(idx_F - F_step_avg[idx]) + # # if idx == idx_E_std: print(f'idx_E:{idx_E - E_step_avg[idx]}, idx_F:{idx_F - F_step_avg[idx]}') + # if idx == idx_E_std: print(f'E_bias_deriv:{E_bias_deriv}, F_bias_elem:{F_bias_elem}') + # F_bias.append(E_bias_deriv * F_bias_elem) + + # force_avg = F_step_avg + np.array(F_bias) + + # norm_F_step = np.linalg.norm(F_step_avg, axis=1) + # norm_F_bias = np.linalg.norm(F_bias, axis=1) + + # ratio_bias = [] + # for item_norm_F_step, item_norm_F_bias in zip(norm_F_step, norm_F_bias): + # ratio_bias.append(item_norm_F_bias/item_norm_F_step) + # print(f'Bias Ratio: {np.average(ratio_bias)}\n') + + + 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]): + F_bias_elem += (idx_E - E_step_avg[idx_atom])*(idx_F - F_step_avg[idx_atom]) + # mpi_print(f'idx_E:{idx_E - E_step_avg[idx_atom]}, idx_F:{idx_F - F_step_avg[idx_atom]}', rank) + mpi_print(f'idx_atom:{idx_atom}| E_bias_deriv:{E_bias_deriv}, F_bias_elem:{F_bias_elem}', rank) + mpi_print(f'Uncert_E:{E_step_std[idx_atom]}', rank) + F_bias = E_bias_deriv * F_bias_elem + + norm_F_step = np.linalg.norm(F_step_avg[idx_atom]) + norm_F_bias = np.linalg.norm(F_bias) + + ratio_bias = norm_F_bias/norm_F_step + mpi_print(f'Bias Ratio: {np.average(ratio_bias)}\n', rank) + + force_avg = F_step_avg.copy() + + if ratio_bias > 10: + F_bias = F_bias / ratio_bias * 10 + force_avg[idx_atom] += F_bias + else: + force_avg[idx_atom] += F_bias + + else: + forces = None + if rank == 0: + struc.calc = calculator + forces = struc.get_forces(md=True) + force_avg = comm.bcast(forces, root=0) + + # mpi_print(f'Step 10-d: {time.time()-time_init}', rank) + + if al_type == 'energy_max': + return force_avg, E_step_std + else: + return force_avg, F_step_norm_std diff --git a/libs/lib_cont_nvtlangevin_temp.py b/libs/lib_cont_nvtlangevin_temp.py index 92011a2..7a5fd31 100644 --- a/libs/lib_cont_nvtlangevin_temp.py +++ b/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 diff --git a/libs/lib_criteria.py b/libs/lib_criteria.py index 26ac1e9..5385028 100644 --- a/libs/lib_criteria.py +++ b/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 diff --git a/libs/lib_input.py b/libs/lib_input.py index 464ef00..0256788 100644 --- a/libs/lib_input.py +++ b/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 diff --git a/libs/lib_mainloop_new.py b/libs/lib_mainloop_new.py index 8083b1d..aa3268b 100644 --- a/libs/lib_mainloop_new.py +++ b/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' diff --git a/libs/lib_md.py b/libs/lib_md.py index c445679..5f19b20 100644 --- a/libs/lib_md.py +++ b/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( diff --git a/libs/lib_npz.py b/libs/lib_npz.py index a8a575b..d6c7af3 100644 --- a/libs/lib_npz.py +++ b/libs/lib_npz.py @@ -6,7 +6,7 @@ from ase.io import read as atoms_read from libs.lib_util import read_aims, single_print -import son +from vibes import son def generate_npz_DFT_init(inputs, traj, workpath): """Function [generate_npz_DFT_init] @@ -39,17 +39,18 @@ def generate_npz_DFT_init(inputs, traj, workpath): z_train = [[] for i in range(inputs.nstep)] CELL_train = [[] for i in range(inputs.nstep)] PBC_train = [[] for i in range(inputs.nstep)] - if inputs.ensemble == 'NPTisoiso': + sigma_train = [[] for i in range(inputs.nstep)] + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train = [[] for i in range(inputs.nstep)] single_print(f'[npz]\tSample {inputs.nstep} different training data\n') # Random sampling for the structural configurations from trajectory for i, step in zip( - random.sample(range(0,len(traj)),inputs.ntotal), - tqdm(range(inputs.ntotal)) + random.sample(range(0,len(traj)),inputs.ntotal_init), + tqdm(range(inputs.ntotal_init)) ): for index_nstep in range(inputs.nstep): - if step < (inputs.ntrain+inputs.nval) * (index_nstep + 1) and step >= (inputs.ntrain+inputs.nval) * (index_nstep): + if step < (inputs.ntrain_init+inputs.nval_init) * (index_nstep + 1) and step >= (inputs.ntrain_init+inputs.nval_init) * (index_nstep): # Energy is shifted by the reference energy # to avoid the unsual weighting with forces in NequIP @@ -70,8 +71,17 @@ def generate_npz_DFT_init(inputs, traj, workpath): z_train[index_nstep].append([atomic_numbers[item[1]] for item in traj[i]['atoms']['symbols'] for jdx in range(item[0])]); CELL_train[index_nstep].append(traj[i]['atoms']['cell']); PBC_train[index_nstep].append(traj[i]['atoms']['pbc']) - if inputs.ensemble == 'NPTisoiso': + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train[index_nstep].append(traj[i]['calculator']['stress']) + + from libs.lib_util import eval_sigma + sigma_train[index_nstep].append( + eval_sigma( + struc_step_forces = traj[i]['calculator']['forces'], + struc_step_positions = traj[i]['atoms']['positions'], + al_type = 'sigma' + ) + ) break # Split the sampled data into individual files for each subsampling set @@ -82,7 +92,8 @@ def generate_npz_DFT_init(inputs, traj, workpath): z_train_store = z_train[index_nstep] CELL_train_store = CELL_train[index_nstep] PBC_train_store = PBC_train[index_nstep] - if inputs.ensemble == 'NPTisoiso': + sigma_train_store = sigma_train[index_nstep] + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train_store = stress_train[index_nstep] # Save each subsampling data @@ -93,10 +104,11 @@ def generate_npz_DFT_init(inputs, traj, workpath): 'R': np.array(R_train_store), 'z': np.array(z_train_store), 'CELL': np.array(CELL_train_store), - 'PBC': np.array(PBC_train_store) + 'PBC': np.array(PBC_train_store), + 'sigma': np.array(sigma_train_store) } - if inputs.ensemble == 'NPTisoiso': + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: arrays_to_save['stress'] = np.array(stress_train_store) np.savez(npz_name, **arrays_to_save) @@ -145,7 +157,8 @@ def generate_npz_DFT(inputs, workpath): z_train = [[] for i in range(inputs.nstep)] CELL_train = [[] for i in range(inputs.nstep)] PBC_train = [[] for i in range(inputs.nstep)] - if inputs.ensemble == 'NPTisoiso': + sigma_train = [[] for i in range(inputs.nstep)] + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train = [[] for i in range(inputs.nstep)] # Check the existence of NPZ files containing previous data @@ -190,7 +203,17 @@ def generate_npz_DFT(inputs, workpath): z_train[index_nstep].append(atoms.numbers); CELL_train[index_nstep].append(atoms.get_cell()); PBC_train[index_nstep].append(atoms.get_pbc()); - if inputs.ensemble == 'NPTisoiso': + + from libs.lib_util import eval_sigma + sigma_train[index_nstep].append( + eval_sigma( + struc_step_forces = F_step, + struc_step_positions = atoms.get_positions(), + al_type = 'sigma' + ) + ) + + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train[index_nstep].append(atom.get_stress(voigt=False)) break elif inputs.output_format == 'trajectory.son': @@ -222,7 +245,17 @@ def generate_npz_DFT(inputs, workpath): z_train[index_nstep].append(np.array(atom_numbers)); CELL_train[index_nstep].append(data[0]['atoms']['cell']); PBC_train[index_nstep].append(data[0]['atoms']['pbc']); - if inputs.ensemble == 'NPTisoiso': + + from libs.lib_util import eval_sigma + sigma_train[index_nstep].append( + eval_sigma( + struc_step_forces = data[0]['calculator']['forces'], + struc_step_positions = data[0]['atoms']['positions'], + al_type = 'sigma' + ) + ) + + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train[index_nstep].append(data[0]['calculator']['stress']) break elif inputs.output_format == 'nequip': @@ -251,7 +284,17 @@ def generate_npz_DFT(inputs, workpath): z_train[index_nstep].append(data[0].get_atomic_numbers()); CELL_train[index_nstep].append(data[0].get_cell()); PBC_train[index_nstep].append(data[0].get_pbc()); - if inputs.ensemble == 'NPTisoiso': + + from libs.lib_util import eval_sigma + sigma_train[index_nstep].append( + eval_sigma( + struc_step_forces = F_step, + struc_step_positions = data[0].get_positions(), + al_type = 'sigma' + ) + ) + + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train[index_nstep].append(data[0]['calculator']['stress']) break else: @@ -279,7 +322,9 @@ def generate_npz_DFT(inputs, workpath): ((data_train['CELL'], CELL_train[index_nstep]), axis=0) PBC_train_store = np.concatenate\ ((data_train['PBC'], PBC_train[index_nstep]), axis=0) - if inputs.ensemble == 'NPTisoiso': + sigma_train_store = np.concatenate\ + ((data_train['sigma'], sigma_train[index_nstep]), axis=0) + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train_store = np.concatenate\ ((data_train['stress'], stress_train[index_nstep]), axis=0) else: ##!! I don't think this part is no longer needed @@ -289,7 +334,8 @@ def generate_npz_DFT(inputs, workpath): z_train_store = z_train[index_nstep] CELL_train_store = CELL_train[index_nstep] PBC_train_store = PBC_train[index_nstep] - if inputs.ensemble == 'NPTisoiso': + sigma_train_store = sigma_train[index_nstep] + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train_store = stress_train[index_nstep] # Path to new data @@ -301,10 +347,11 @@ def generate_npz_DFT(inputs, workpath): 'R': np.array(R_train_store), 'z': np.array(z_train_store), 'CELL': np.array(CELL_train_store), - 'PBC': np.array(PBC_train_store) + 'PBC': np.array(PBC_train_store), + 'sigma': np.array(sigma_train_store), } - if inputs.ensemble == 'NPTisoiso': + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: arrays_to_save['stress'] = np.array(stress_train_store) np.savez(npz_name[:-4], **arrays_to_save) @@ -353,7 +400,8 @@ def generate_npz_DFT_rand_init(inputs, traj, ntrain, nval, workpath): z_train = [[] for i in range(inputs.nstep)]; CELL_train = [[] for i in range(inputs.nstep)]; PBC_train = [[] for i in range(inputs.nstep)]; - if inputs.ensemble == 'NPTisoiso': + sigma_train = [[] for i in range(inputs.nstep)]; + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train = [[] for i in range(inputs.nstep)] traj_idx = []; @@ -386,7 +434,17 @@ def generate_npz_DFT_rand_init(inputs, traj, ntrain, nval, workpath): z_train[index_nstep].append([atomic_numbers[item[1]] for item in traj[i]['atoms']['symbols'] for idx in range(item[0])]); CELL_train[index_nstep].append(traj[i]['atoms']['cell']); PBC_train[index_nstep].append(traj[i]['atoms']['pbc']) - if inputs.ensemble == 'NPTisoiso': + + from libs.lib_util import eval_sigma + sigma_train[index_nstep].append( + eval_sigma( + struc_step_forces = F_step, + struc_step_positions = traj[i]['atoms']['positions'], + al_type = 'sigma' + ) + ) + + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train[index_nstep].append(traj[i]['calculator']['stress']) traj_idx.append(i) break @@ -399,7 +457,8 @@ def generate_npz_DFT_rand_init(inputs, traj, ntrain, nval, workpath): z_train_store = z_train[index_nstep] CELL_train_store = CELL_train[index_nstep] PBC_train_store = PBC_train[index_nstep] - if inputs.ensemble == 'NPTisoiso': + sigma_train_store = sigma_train[index_nstep] + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train_store = stress_train[index_nstep] # Save each subsampling data @@ -410,10 +469,11 @@ def generate_npz_DFT_rand_init(inputs, traj, ntrain, nval, workpath): 'R': np.array(R_train_store), 'z': np.array(z_train_store), 'CELL': np.array(CELL_train_store), - 'PBC': np.array(PBC_train_store) + 'PBC': np.array(PBC_train_store), + 'sigma': np.array(sigma_train_store) } - if inputs.ensemble == 'NPTisoiso': + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: arrays_to_save['stress'] = np.array(stress_train_store) np.savez(npz_name[:-4], **arrays_to_save) @@ -477,7 +537,8 @@ def generate_npz_DFT_rand( z_train = [[] for i in range(inputs.nstep)]; CELL_train = [[] for i in range(inputs.nstep)]; PBC_train = [[] for i in range(inputs.nstep)]; - if inputs.ensemble == 'NPTisoiso': + sigma_train = [[] for i in range(inputs.nstep)]; + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train = [[] for i in range(inputs.nstep)] # Check the existence of NPZ files containing previous data @@ -515,7 +576,17 @@ def generate_npz_DFT_rand( z_train[index_nstep].append([atomic_numbers[item[1]] for item in traj[i]['atoms']['symbols'] for idx in range(item[0])]); CELL_train[index_nstep].append(traj[i]['atoms']['cell']); PBC_train[index_nstep].append(traj[i]['atoms']['pbc']) - if inputs.ensemble == 'NPTisoiso': + + from libs.lib_util import eval_sigma + sigma_train[index_nstep].append( + eval_sigma( + struc_step_forces = F_step, + struc_step_positions = traj[i]['atoms']['positions'], + al_type = 'sigma' + ) + ) + + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train[index_nstep].append(traj[i]['calculator']['stress']) traj_idx.append(i) break @@ -541,7 +612,9 @@ def generate_npz_DFT_rand( ((data_train['CELL'], CELL_train[index_nstep]), axis=0) PBC_train_store = np.concatenate\ ((data_train['PBC'], PBC_train[index_nstep]), axis=0) - if inputs.ensemble == 'NPTisoiso': + sigma_train_store = np.concatenate\ + ((data_train['sigma'], sigma_train[index_nstep]), axis=0) + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: stress_train_store = np.concatenate\ ((data_train['stress'], stress_train[index_nstep]), axis=0) @@ -554,10 +627,11 @@ def generate_npz_DFT_rand( 'R': np.array(R_train_store), 'z': np.array(z_train_store), 'CELL': np.array(CELL_train_store), - 'PBC': np.array(PBC_train_store) + 'PBC': np.array(PBC_train_store), + 'sigma': np.array(sigma_train_store) } - if inputs.ensemble == 'NPTisoiso': + if inputs.ensemble == 'NPTisoiso' or inputs.train_stress: arrays_to_save['stress'] = np.array(stress_train_store) np.savez(npz_name[:-4], **arrays_to_save) diff --git a/libs/lib_termination.py b/libs/lib_termination.py index 9d9a07c..5550d14 100644 --- a/libs/lib_termination.py +++ b/libs/lib_termination.py @@ -39,6 +39,8 @@ def termination(inputs): # Read the Uncertainty column if inputs.al_type == 'energy': result_uncert = result_data.loc[:,'TestError_E'] + elif inputs.al_type == 'energy_max': + result_uncert = result_data.loc[:,'Un_Abs_E_std_i'] elif inputs.al_type == 'force': result_uncert = result_data.loc[:,'TestError_F'] elif inputs.al_type == 'force_max': @@ -97,7 +99,10 @@ def get_testerror(inputs): from ase.io import read as atoms_read # Read testing data - npz_test = f'./MODEL/data-test.npz' # Path + if inputs.npz_test == 'train': + npz_test = f'./MODEL/{inputs.temperature}K-{inputs.pressure}bar_{inputs.index}/data-train_0.npz' # Path + else: + npz_test = f'./MODEL/data-test.npz' # Path data_test = np.load(npz_test) # Load ## Prepare the ground state structure @@ -123,6 +128,9 @@ def get_testerror(inputs): # Go through all configurations in the testing data prd_E_avg = [] prd_E_std = [] + if inputs.al_type == 'energy_max': + prd_E_atom_avg = [] + prd_E_atom_std = [] prd_F_avg = [] prd_F_std = [] prd_F_all_avg = [] @@ -131,7 +139,7 @@ def get_testerror(inputs): prd_sigma_max_avg = [] prd_sigma_max_std = [] real_E_list = [] - real_F_list = [] + real_F_list = [] for id_step, (id_R, id_z, id_CELL, id_PBC, id_E, id_F) in enumerate(zip( data_test['R'], data_test['z'], data_test['CELL'], data_test['PBC'], data_test['E'], data_test['F'] )): @@ -142,6 +150,7 @@ def get_testerror(inputs): prd_F = [] prd_R = [] zndex = 0 + natoms = len(struc) if id_step % inputs.printinterval == 0: mpi_print(f'[Termi]\tTesting: sample {id_step}', inputs.rank) @@ -152,7 +161,11 @@ def get_testerror(inputs): if (index_nmodel * inputs.nstep + index_nstep) % inputs.size == inputs.rank: struc.calc = inputs.calc_MLIP[zndex] # Predicted energy is shifted E_gs back, but the E_gs defualt is zero - prd_E.append(struc.get_potential_energy()) + + if inputs.al_type == 'energy_max': + prd_E.append(struc.get_potential_energies()) + else: + prd_E.append(struc.get_potential_energy()) prd_F.append(struc.get_forces(md=True)) prd_R.append(struc.get_positions()) zndex += 1 @@ -176,11 +189,21 @@ def get_testerror(inputs): prd_E = np.array(prd_E) + E_ha # id_E = id_E - E_ha - prd_E_avg.append(np.average(prd_E, axis=0)) + + if inputs.al_type == 'energy_max': + # mpi_print(prd_E, inputs.rank) + prd_E_atom = np.average(prd_E, axis=0) + prd_E_atom_avg.append(prd_E_atom) + prd_E_avg.append(np.sum(prd_E_atom, axis=0)) + else: + prd_E_avg.append(np.average(prd_E, axis=0)) + real_E_list.append(id_E) if inputs.al_type == 'energy': prd_E_std.append(np.std(prd_E, axis=0)) + elif inputs.al_type == 'energy_max': + prd_E_atom_std.append(np.std(prd_E, axis=0)) prd_F_step_avg = np.average(prd_F, axis=0) prd_F_all_avg.append(prd_F_step_avg) @@ -232,6 +255,23 @@ def get_testerror(inputs): UncertAbs_E_std = np.std(prd_E_std) UncertRel_E_avg = np.average(prd_E_std/prd_E_avg) UncertRel_E_std = np.std(prd_E_std/prd_E_avg) + elif inputs.al_type == 'energy_max': + prd_E_atom_std = np.array(prd_E_atom_std) + UncertAbs_E_max = [max(prd_E_std_step) for prd_E_std_step in prd_E_atom_std] + UncertAbs_E_avg = np.average(UncertAbs_E_max) + UncertAbs_E_std = np.std(UncertAbs_E_max) + + UncertRel_E_max = [max(prd_E_std_step/prd_E_avg_step) for prd_E_std_step, prd_E_avg_step in zip(prd_E_atom_std, prd_E_atom_avg)] + UncertRel_E_avg = np.average(UncertRel_E_max) + UncertRel_E_std = np.std(UncertRel_E_max) + + UncertAbs_E_atom = [prd_E_std_step[inputs.idx_atom] for prd_E_std_step in prd_E_atom_std] + UncertAbs_E_atom_avg = np.average(UncertAbs_E_atom) + UncertAbs_E_atom_std = np.std(UncertAbs_E_atom) + + UncertRel_E_atom = [prd_E_std_step[inputs.idx_atom]/np.absolute(prd_E_avg_step[inputs.idx_atom]) for prd_E_std_step, prd_E_avg_step in zip(prd_E_atom_std, prd_E_atom_avg)] + UncertRel_E_atom_avg = np.average(UncertRel_E_atom) + UncertRel_E_atom_std = np.std(UncertRel_E_atom) else: UncertAbs_E_avg = '---- ' UncertAbs_E_std = '---- ' @@ -291,12 +331,18 @@ def get_testerror(inputs): + '\t' + uncert_strconvter(np.average(prd_E_avg))\ + '\t' + uncert_strconvter(np.std(prd_E_avg)) - if inputs.al_type == 'energy': + if inputs.al_type == 'energy' or inputs.al_type == 'energy_max': result_print += '\t' + uncert_strconvter(UncertAbs_E_avg)\ + '\t' + uncert_strconvter(UncertAbs_E_std)\ + '\t' + uncert_strconvter(UncertRel_E_avg)\ + '\t' + uncert_strconvter(UncertRel_E_std) + if inputs.al_type == 'energy_max': + result_print += '\t' + uncert_strconvter(UncertAbs_E_atom_avg)\ + + '\t' + uncert_strconvter(UncertAbs_E_atom_std)\ + + '\t' + uncert_strconvter(UncertRel_E_atom_avg)\ + + '\t' + uncert_strconvter(UncertRel_E_atom_std) + if inputs.al_type == 'force' or inputs.al_type == 'force_max': result_print += '\t' + uncert_strconvter(UncertAbs_F_avg)\ + '\t' + uncert_strconvter(UncertAbs_F_std)\ diff --git a/libs/lib_util.py b/libs/lib_util.py index da1a3c0..b030ab0 100644 --- a/libs/lib_util.py +++ b/libs/lib_util.py @@ -300,10 +300,14 @@ def generate_msg(al_type): + 'TestError_E\tTestError_F\tTestError_S\t'\ + 'E_potent_avg_i\tE_potent_std_i' - if al_type == 'energy': + if al_type == 'energy' or al_type == 'energy_max': result_msg += '\tUn_Abs_E_avg_i\tUn_Abs_E_std_i'\ + '\tUn_Rel_E_avg_i\tUn_Rel_E_std_i' + if al_type == 'energy_max': + result_msg += '\tUn_Abs_Ea_avg_i\tUn_Abs_Ea_std_i'\ + + '\tUn_Rel_Ea_avg_i\tUn_Rel_Ea_std_i' + if al_type == 'force' or al_type == 'force_max': result_msg += '\tUn_Abs_F_avg_i\tUn_Abs_F_std_i'\ + '\tUn_Rel_F_avg_i\tUn_Rel_F_std_i' @@ -312,7 +316,7 @@ def generate_msg(al_type): result_msg += '\tUn_Abs_S_avg_i\tUn_Abs_S_std_i'\ + '\tUn_Rel_S_avg_i\tUn_Rel_S_std_i' - if al_type == 'energy': + if al_type == 'energy' or al_type == 'energy_max': result_msg += '\tUn_Abs_E_avg_a\tUn_Rel_E_avg_a' if al_type == 'force' or al_type == 'force_max': @@ -361,15 +365,15 @@ def read_input_file(file_path): value = value.strip() # Perform type conversions for specific variables - if name in ['supercell', 'supercell_init', 'mask', 'harmonic_F', 'anharmoic_F', 'meta_restart', 'signal_uncert']: + if name in ['supercell', 'supercell_init', 'mask', 'harmonic_F', 'anharmoic_F', 'meta_restart', 'signal_uncert', 'criteria_energy', 'train_stress']: value = eval(value) - elif name in ['crtria_cnvg', 'friction', 'compressibility', 'kB', 'E_gs', 'uncert_shift', 'uncert_grad', 'meta_Ediff', 'meta_r_crtria', 'ttime', 'pfactor', 'timestep', 'cell_factor', 'bias_A', 'bias_B']: + elif name in ['crtria_cnvg', 'friction', 'compressibility', 'kB', 'E_gs', 'uncert_shift', 'uncert_grad', 'meta_Ediff', 'meta_r_crtria', 'ttime', 'pfactor', 'timestep', 'cell_factor', 'bias_A', 'bias_B', 'temp_factor']: value = float(value) elif name in [ 'ntrain_init', 'ntrain', 'nstep', 'nmodel', 'nperiod', 'temperature', 'taut', 'pressure', 'taup', 'steps_ther', 'steps_init', 'steps_random', 'cutoff_ther', 'lmax', 'nfeatures', 'random_index', - 'wndex', 'steps', 'loginterval', 'num_calc', 'test_index', 'num_mdl_calc', 'printinterval' + 'wndex', 'steps', 'loginterval', 'num_calc', 'test_index', 'num_mdl_calc', 'printinterval', 'idx_atom' ]: value = int(value) else: diff --git a/scripts/almd.py b/scripts/almd.py index 1bfdc88..ed2c2f6 100644 --- a/scripts/almd.py +++ b/scripts/almd.py @@ -1,5 +1,5 @@ import os -import son +from vibes import son import sys import numpy as np diff --git a/scripts/lib_run_dft_init.py b/scripts/lib_run_dft_init.py index 2bddf12..490708b 100644 --- a/scripts/lib_run_dft_init.py +++ b/scripts/lib_run_dft_init.py @@ -1,4 +1,4 @@ -import son +from vibes import son from libs.lib_util import output_init, mpi_print, check_mkdir, job_dependency from libs.lib_npz import generate_npz_DFT_init diff --git a/scripts/lib_run_dft_rand.py b/scripts/lib_run_dft_rand.py index 5d145ea..30d3536 100644 --- a/scripts/lib_run_dft_rand.py +++ b/scripts/lib_run_dft_rand.py @@ -1,5 +1,5 @@ import os -import son +from vibes import son import numpy as np from libs.lib_util import output_init, mpi_print, check_mkdir diff --git a/scripts/utils.py b/scripts/utils.py index 234f962..440d1c2 100644 --- a/scripts/utils.py +++ b/scripts/utils.py @@ -1,7 +1,7 @@ import os import re import sys -import son +from vibes import son import random import argparse import collections @@ -250,6 +250,9 @@ def split_son(num_split, E_gs, harmonic_F=False): rm_file('trajectory_test.son') rm_file('trajectory_train.son') rm_file('MODEL/data-test.npz') + if metadata is not None: + son.dump(metadata, 'trajectory_test.son', is_metadata=True) + son.dump(metadata, 'trajectory_train.son', is_metadata=True) for test_item in test_data: son.dump(test_item, 'trajectory_test.son') # Save testing data into trajectory_test.son file for train_item in train_data: @@ -501,6 +504,7 @@ def harmonic2son(temperature, num_sample, output_format): single_print(f'[harmo2son]\tGo through all DFT results') for idx in range(num_sample): # Create a folder for each structral configuration + print(f'Sample {idx}') check_mkdir(f'{idx}') # Move to that folder os.chdir(f'{idx}')