diff --git a/src/analyze_all_data.py b/src/analyze_all_data.py new file mode 100644 index 0000000..9764fdb --- /dev/null +++ b/src/analyze_all_data.py @@ -0,0 +1,14 @@ +import numpy as np + +f = open('all_data.txt', 'r') +INF = f.readlines() +for line in INF: + # print(line.split()) + tmp_data_list = line.split() + index = tmp_data_list[0].strip(':') + # print(index) + ts_e_err = float(tmp_data_list[4]) + if abs(ts_e_err) > 0.15: + print(index, ts_e_err) +f.close() + diff --git a/src/bb.py b/src/bb.py new file mode 100644 index 0000000..dc2709a --- /dev/null +++ b/src/bb.py @@ -0,0 +1,112 @@ +import os +import sys +import json +import logging +import argparse +import shutil, errno +from ase.atoms import Atoms +from tabulate import tabulate +from ase.io import read, write +from ase.io.jsonio import decode +from maggma.stores import MongoStore +import numpy as np + + +def get_docs(launchpad_file, query, collections_name='fireworks'): + tasks_store = MongoStore.from_launchpad_file(launchpad_file, collections_name) + tasks_store.connect() + docs = list(tasks_store.query(query)) + return docs + + +def atomic_number_to_symbol(atomic_number): + periodic_table = { + 1: 'H', 2: 'He', 3: 'Li', 4: 'Be', 5: 'B', + 6: 'C', 7: 'N', 8: 'O', 9: 'F', 10: 'Ne', + } + return periodic_table.get(atomic_number, 'X') # 'X' for unknown element + + +def atoms_from_json(atoms_json): + atoms_dict = json.loads(atoms_json) + + atomic_numbers = atoms_dict['numbers']['__ndarray__'][2] + positions = atoms_dict['positions']['__ndarray__'][2] + + natoms = int(len(positions)/3) + coords = np.reshape(positions, (natoms, 3)) + + atomic_symbols = [atomic_number_to_symbol(i) for i in atomic_numbers] + return Atoms(symbols=atomic_symbols, positions=coords) + + +def download_data(docs, base_dir): + for doc in docs: + tag = doc['metadata']['tag'] + index = tag.split('-')[-1].split('_')[0] + # print('tag:', tag) + if tag.startswith('TS0'): + name = 'TS' + structure = atoms_from_json(doc['output']['atoms']['atoms_json']) + ''' + print( + name, + structure, + structure.get_positions() + ) + ''' + elif tag.startswith('irc'): + name = tag.split('-')[1] + structure = atoms_from_json(doc['output']['atoms']['atoms_json']) + ''' + print( + name, + structure, + structure.get_positions() + ) + ''' + else: + continue + + output_dir = os.path.join(base_dir, f"{index}") + + os.makedirs(output_dir, exist_ok=True) + + output_file = os.path.join(output_dir, f'{name}.xyz') + write(output_file, structure, format='xyz') + + +def main(): + tag = 'nov15' + noise_level = '00' + + query0 = { + "metadata.tags.class": tag, + "metadata.tag": { + '$regex': f'.*0.*_noise{noise_level}' + } + } + + launchPad_file = os.path.join(os.environ["HOME"],"fw_config/my_launchpad.yaml") + collections = 'quacc_results0' + + docs0 = get_docs(launchPad_file, query0, collections_name=collections) + download_data(docs0, "/global/scratch/users/kumaranu/saved_files/neb_nn_inputs/") + + ''' + for doc in docs0: + # print(doc.keys()) + print(doc['metadata']['tag']) + #a = 'irc-reverse0-021_noise00' + index = a.split('-')[-1].split('_')[0] + if a.split('-')[1] == 'reverse0': + print('reverse') + + #os.system(f'cp -rL {src} {dst}') + print(len(docs0)) + ''' + + +if __name__ == "__main__": + main() + diff --git a/src/calc_properties.py b/src/calc_properties.py new file mode 100644 index 0000000..109909b --- /dev/null +++ b/src/calc_properties.py @@ -0,0 +1,37 @@ +from ase.io import read, write +from calcs import calc + +# Read molecular configurations from XYZ file +atoms_list = read('/global/scratch/users/kumaranu/saved_files/neb_nn_inputs/263/output.xyz', index=':') + +# Create a calculator instance (e.g., EMT) +mlcalculator = calc() + +# Iterate over each molecular configuration +for atoms in atoms_list: + mlcalculator.calculate(atoms) + + # Calculate energies and forces + energy = mlcalculator.results['energy'] + forces = mlcalculator.results['forces'] + + # Store energies and forces in the atoms object + atoms.info['energy'] = energy + atoms.arrays['forces'] = forces + +# Write the updated atoms object with energies and forces to a new XYZ file +write('/global/scratch/users/kumaranu/saved_files/neb_nn_inputs/263/output1.xyz', atoms_list) + + +''' +h2 = Atoms(numbers=[1, 1], positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) +mlcalculator.calculate(h2) + +print(mlcalculator.results['energy']) # mean of calculated molecular energies, shape (1,) +print(mlcalculator.results['forces']) # mean of calculated atomic forces, shape (n_atom, 3) +print(mlcalculator.results['hessian']) # mean of calculated atomic Hessian, shape (n_atom, 3, n_atom, 3) +print(mlcalculator.results['energy_disagreement']) # disagreement among calculated molecular +''' + + + diff --git a/src/calcs.py b/src/calcs.py index 8846491..05bed12 100644 --- a/src/calcs.py +++ b/src/calcs.py @@ -18,13 +18,13 @@ def calc_mace(): def calc(): mlcalculator = MLAseCalculator( model_path=[ - '/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_52/models/best_model_state.tar', + '/global/u2/k/kumaranu/Documents/NewtonNet/example/predict/training_52/models/best_model_state.tar', #'/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_53/models/best_model_state.tar', #'/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_54/models/best_model_state.tar', #'/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_55/models/best_model_state.tar', ], # path to model file, str or list of str settings_path=[ - '/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_52/run_scripts/config0.yml', + '/global/u2/k/kumaranu/Documents/NewtonNet/example/predict/training_52/run_scripts/config0.yml', #'/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_53/run_scripts/config2.yml', #'/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_54/run_scripts/config1.yml', #'/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_55/run_scripts/config3.yml', diff --git a/src/calcs1.py b/src/calcs1.py new file mode 100644 index 0000000..03518b0 --- /dev/null +++ b/src/calcs1.py @@ -0,0 +1,43 @@ +from mace.calculators import MACECalculator +from newtonnet.utils.ase_interface import MLAseCalculator +from ase import Atoms + + +def calc_mace(): + + mlcalculator = MACECalculator( + model_paths='/global/home/users/kumaranu/Documents/gpu_jobs/MACE_model.model', + #model_paths='/global/home/users/kumaranu/Documents/mep_tests/src/MACE_model_cpu.model', + device='cuda', + #device='cpu', + default_dtype="float64", + ) + return mlcalculator + + +def calc(): + mlcalculator = MLAseCalculator( + model_path=[ + '/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_52/models/best_model_state.tar', + ], # path to model file, str or list of str + settings_path=[ + '/global/home/users/kumaranu/Documents/NewtonNet/example/predict/training_52/run_scripts/config0.yml', + ], # path to configuration file, str or list of str + hess_method=None, # method to calculate hessians. 'autograd', 'fwd_diff', 'cnt_diff', or None (default: 'autograd') + # hess_precision=1e-5, # hessian gradient calculation precision for 'fwd_diff' and 'cnt_diff', ignored otherwise (default: None) + disagreement='std', # method to calculate disagreement among models. 'std', 'std_outlierremoval', 'range':, 'values', or None (default: 'std') + #device='cuda' # 'cpu' or list of cuda + device='cpu' # 'cpu' or list of cuda + ) + return mlcalculator + +''' +h2 = Atoms(numbers=[1, 1], positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) +mlcalculator.calculate(h2) + +print(mlcalculator.results['energy']) # mean of calculated molecular energies, shape (1,) +print(mlcalculator.results['forces']) # mean of calculated atomic forces, shape (n_atom, 3) +print(mlcalculator.results['hessian']) # mean of calculated atomic Hessian, shape (n_atom, 3, n_atom, 3) +print(mlcalculator.results['energy_disagreement']) # disagreement among calculated molecular +''' + diff --git a/src/check.py b/src/check.py new file mode 100644 index 0000000..d08b0c0 --- /dev/null +++ b/src/check.py @@ -0,0 +1,18 @@ +import glob, os, sys + +logdir = '/global/scratch/users/kumaranu/saved_files/neb_nn_inputs' + +os.chdir(logdir) +for i in range(265): + if os.path.exists(f'{i:03}'): + os.chdir(f'{i:03}') + if os.path.exists('neb_band_aseneb_NEBOptimizer_None.txt'): + INF = open('neb_band_aseneb_NEBOptimizer_None.txt', 'r').readlines() + #if int(INF[-1].split()[1]) > 200: + print(f'{i:03}: {INF[-1].split()[1]}') + else: + print(f'{i:03}: No NEB file.') + os.chdir('../') + else: + print(f'{i:03} does not exist.') + diff --git a/src/check_graph_iso.py b/src/check_graph_iso.py new file mode 100644 index 0000000..460551e --- /dev/null +++ b/src/check_graph_iso.py @@ -0,0 +1,18 @@ +import glob, os +from ase.io import read + +ref_dir = '/global/cfs/cdirs/m2834/kumaranu/neb_nn_inputs' + +for i in range(265): + logdir = os.path.join(ref_dir, f'{i:03}') + if os.path.exists(logdir): + reactant_traj_file = os.path.join(logdir, 'reactant_opt.traj') + if os.path.exists(reactant_traj_file): + traj_atoms = read(reactant_traj_file, index=':', format="traj") + else: + print(f"Reactant's traj file could not be found for index {i:03}") + #for atoms in traj_atoms: + # print(i, atoms) + else: + print(f'Path could not be found for index {i:03}.') + diff --git a/src/check_neb_results.py b/src/check_neb_results.py new file mode 100644 index 0000000..b741b98 --- /dev/null +++ b/src/check_neb_results.py @@ -0,0 +1,82 @@ +import glob, os, sys +from ase.io import read +from calcs1 import calc +import numpy as np +from ase import Atoms +from multiprocessing import Pool + + +def align_geom(refgeom, geom): + """Find translation/rotation that moves a given geometry to maximally overlap + with a reference geometry. Implemented with Kabsch algorithm. + + Args: + refgeom: The reference geometry to be rotated to + geom: The geometry to be rotated and shifted + + Returns: + RMSD: Root-mean-squared difference between the rotated geometry + and the reference + new_geom: The rotated geometry that maximumally overal with the reference + """ + center = np.mean(refgeom, axis=0) # Find the geometric center + ref2 = refgeom - center + geom2 = geom - np.mean(geom, axis=0) + cov = np.dot(geom2.T, ref2) + v, sv, w = np.linalg.svd(cov) + + if np.linalg.det(v) * np.linalg.det(w) < 0: + sv[-1] = -sv[-1] + v[:, -1] = -v[:, -1] + u = np.dot(v, w) + new_geom = np.dot(geom2, u) + center + rmsd = np.sqrt(np.mean((new_geom - refgeom) ** 2)) + return rmsd, new_geom + + +def process_file(i): + try: + if os.path.exists(f'{i:03}'): + os.chdir(f'{i:03}') + if os.path.exists('neb_band_aseneb_NEBOptimizer_None.txt'): + ase_traj = read('optimized_path_aseneb_NEBOptimizer_None.xyz', ':') + energies = [i.get_potential_energy() for i in ase_traj] + ts_index = np.argmax(energies) + + mlcalculator = calc() + ts_atoms = read('TS.xyz') + mlcalculator.calculate(ts_atoms) + energy_ts = mlcalculator.results['energy'] + rmsd, new_geom = align_geom(ase_traj[ts_index].get_positions(), ts_atoms.get_positions()) + + fwd_barrier_neb = energies[ts_index] - energies[0] + rev_barrier_neb = energies[ts_index] - energies[-1] + ref_fwd_barrier = energy_ts - energies[0] + ref_rev_barrier = energy_ts - energies[-1] + + ts_e_err = energy_ts - energies[ts_index] + fwd_barrier_err = fwd_barrier_neb - ref_fwd_barrier + rev_barrier_err = rev_barrier_neb - ref_rev_barrier + os.chdir('../') + return f'{i:03}: TS energy err: {ts_e_err:.6f} fwd barrier err: {fwd_barrier_err:.6f} Rev barrier err: {rev_barrier_err:.6f} Geometry RMSD: {rmsd:.6f}' + else: + os.chdir('../') + except Exception as e: + return f'Error processing file {i}: {str(e)}' + +if __name__ == '__main__': + + logdir = '/global/scratch/users/kumaranu/saved_files/neb_nn_inputs' + os.chdir(logdir) + + with Pool(5) as pool: + results = pool.map(process_file, range(265)) + + missing_indices = [i for i, result in enumerate(results) if result is None] + if missing_indices: + print(f'Missing results for indices: {missing_indices}') + + for result in results: + if result is not None: + print(result) + diff --git a/src/initialize_path.py b/src/initialize_path.py index 546e458..4926215 100644 --- a/src/initialize_path.py +++ b/src/initialize_path.py @@ -1,32 +1,129 @@ import os -from ase.io import read +import copy +import numpy as np from calcs import calc +from ase.io import read, write +from ase.io import Trajectory from ase.optimize import BFGS, ODE12r +from sella import Sella +from typing import Any, Dict, Iterator, List, Sequence, Tuple, TypeVar, Union +import ase.units as units +from math import pi, sin, sqrt +from ase.units import kcal, mol, Ang +from sella_wrapper import sella_wrapper -def setup_images(): - N_intermediate = 40 - # xyz_r_p = '/global/home/users/kumaranu/trash/inputs/abcd/A_B.xyz' - # xyz_r_p = '/global/home/users/kumaranu/trash/inputs/abcd/A_C+D.xyz' - xyz_r_p = '/global/home/users/kumaranu/trash/inputs/abcd/B_C+D.xyz' - reactant = read(xyz_r_p, index='0') - reactant.calc = calc() - product = read(xyz_r_p, index='1') - product.calc = calc() - os.makedirs('/tmp/geodesic_calc_dir', exist_ok=True) - os.chdir('/tmp/geodesic_calc_dir') - os.system('geodesic_interpolate ' + str(xyz_r_p) + ' --output output.xyz --nimages ' + str(N_intermediate)) - intermediate = read('output.xyz', index=':') - - qn = ODE12r(reactant) - qn.run(fmax=1e-3, steps=1000) - - qn = ODE12r(product) - qn.run(fmax=1e-3, steps=1000) - - images = [reactant] - for image in intermediate: +def get_hessian( + ts, + hess_precision=1e-5, + method='fwd_diff' + ): + n_atoms = len(ts) + hessian = np.zeros((n_atoms, 3, n_atoms, 3)) + + # Calculate Hessian using finite difference + hess_precision = 1e-5 # Hessian precision + for A_ in range(n_atoms): + for X_ in range(3): + ts.positions[A_][X_] += hess_precision + ts.set_calculator(calc()) + forces_pos = ts.get_forces() + ts.positions[A_][X_] -= 2 * hess_precision + ts.set_calculator(calc()) + forces_neg = ts.get_forces() + hessian[:, :, A_, X_] = -(forces_pos - forces_neg) / (2 * hess_precision) * (kcal/mol/Ang) + return hessian + + +def get_hessian_2d(ts): + hessian = get_hessian(ts) + n_atoms = np.shape(hessian)[0] + return np.reshape(hessian, (3*n_atoms, 3*n_atoms)) + + +def get_eigs(hessian): + n_atoms = np.shape(hessian)[0] + eigvals, _ = np.linalg.eig(np.reshape(hessian, (3*n_atoms, 3*n_atoms))) + eigvals_real = np.real(eigvals) + formatted_eigvals = ["{:.3e}".format(val) for val in eigvals_real] + + print('eigvals') + for val in formatted_eigvals: + print(val) + + +def _energies_and_modes(ts) -> Tuple[np.ndarray, np.ndarray]: + n_atoms = len(ts) + masses = ts.get_masses() + + if not np.all(masses): + raise ValueError('Zero mass encountered in one or more of ' + 'the vibrated atoms. Use Atoms.set_masses()' + ' to set all masses to non-zero values.') + mass_weights = np.repeat(masses**-0.5, 3) + + positions = ts.get_positions() - ts.get_center_of_mass() + _, vectors_inertia = ts.get_moments_of_inertia(vectors=True) + vectors_transrot = np.zeros((6, n_atoms, 3)) + vectors_transrot[0, :, 0] = 1 + vectors_transrot[1, :, 1] = 1 + vectors_transrot[2, :, 2] = 1 + vectors_transrot[3] = positions @ vectors_inertia[[1]].T @ vectors_inertia[[2]] - positions @ vectors_inertia[[2]].T @ vectors_inertia[[1]] + vectors_transrot[4] = positions @ vectors_inertia[[2]].T @ vectors_inertia[[0]] - positions @ vectors_inertia[[0]].T @ vectors_inertia[[2]] + vectors_transrot[5] = positions @ vectors_inertia[[0]].T @ vectors_inertia[[1]] - positions @ vectors_inertia[[1]].T @ vectors_inertia[[0]] + vectors_transrot = vectors_transrot.reshape((6, n_atoms * 3)) + vectors_transrot = vectors_transrot / mass_weights + vectors_transrot, _ = np.linalg.qr(vectors_transrot.T) + vectors_transrot = vectors_transrot.T + proj = np.eye(n_atoms * 3) - vectors_transrot.T @ vectors_transrot + + omega2, vectors = np.linalg.eigh( + proj.T @ ( + mass_weights + * get_hessian_2d(ts) + * mass_weights[:, np.newaxis]) + @ proj) + + unit_conversion = units._hbar * units.m / sqrt(units._e * units._amu) + energies = unit_conversion * omega2.astype(complex)**0.5 + + modes = vectors.T.reshape(n_atoms * 3, n_atoms, 3) + modes = modes * masses[np.newaxis, :, np.newaxis]**-0.5 + + return (energies, modes) + + +def setup_analysis( + init_path, + ts, + highest_geodesic_ts_traj_file, + highest_geodesic_geom_file='highest_geodesic.xyz', + ): + ts = read(xyz_ts) + ts.calc = calc() + sella_wrapper(ts, traj_file, sella_order=1) + + ts_e = ts.get_potential_energy() + r_e = reactant.get_potential_energy() + p_e = product.get_potential_energy() + + energy_list = [r_e] + for image in images[1:-1]: image.calc = calc() - images.append(image) - images.append(product) - return images, None, None + energy_list.append(image.get_potential_energy()) + energy_list.append(p_e) + write(highest_geodesic_geom_file, images[np.argmax(energy_list)]) + + highest_geodesic = copy.deepcopy(images[np.argmax(energy_list)]) + highest_geodesic.calc = calc() + sella_wrapper(highest_geodesic, highest_geodesic_ts_traj_file, sella_order=1) + + eigs, _ = _energies_and_modes(ts) + print('\n', eigs) + + eigs, _ = _energies_and_modes(highest_geodesic) + print('\n', eigs) + + print(f'Forward_barrier: {ts_e - r_e}') + print(f'Reverse_barrier: {ts_e - p_e}') + diff --git a/src/neb_wrapper.py b/src/neb_wrapper.py index 0744e20..da4a7e2 100644 --- a/src/neb_wrapper.py +++ b/src/neb_wrapper.py @@ -1,14 +1,24 @@ - import os +from calcs import calc +from ase.io import write import matplotlib.pyplot as plt -from initialize_path import setup_images +from setup_images import setup_images from ase.neb import NEB from ase.neb import NEBTools -def run_neb_method(method, optimizer, precon, optmethod, testdir='trash/'): - images, _, _ = setup_images() - +def run_neb_method( + method, + optimizer, + precon, + optmethod, + logdir=None, + xyz_r_p=None, + xyz_ts=None, + N_intermediate=20, + ): + images = setup_images(logdir, xyz_r_p, xyz_ts, N_intermediate) + fmax_history = [] def save_fmax_history(mep): @@ -21,36 +31,51 @@ def save_fmax_history(mep): images, k=k, method=method, - climb=False, + climb=True, precon=precon, - remove_rotation_and_translation=False, + remove_rotation_and_translation=True, parallel=True, ) - os.makedirs(testdir, exist_ok=True) + os.makedirs(logdir, exist_ok=True) log_filename = f'neb_band_{method}_{optimizer.__name__}_{precon}.txt' - logfile_path = os.path.join(testdir, log_filename) + + logfile_path = os.path.join(logdir, log_filename) if optmethod is not None: - opt = optimizer(mep, method=optmethod, logfile=logfile_path) + opt = optimizer(mep, method=optmethod, logfile=logfile_path, verbose=2) else: - opt = optimizer(mep, logfile=logfile_path) + opt = optimizer(mep, logfile=logfile_path, verbose=2) opt.attach(save_fmax_history, 1, mep) - opt.run(fmax=1e-2, steps=1000) + opt.run(fmax=1e-2, steps=400) + # opt.run(fmax=1e-2, steps=1000) + + write(f'{logdir}/optimized_path_{method}_{optimizer.__name__}_{precon}.xyz', images) + + # nebtools = NEBTools(images) + + # Ef, dE = nebtools.get_barrier(fit=False) + # print(f'{method},{optimizer.__name__},{precon} ' + # f'=> Ef = {Ef:.3f}, dE = {dE:.3f}') + + # # Save the NEB band plot + # output_filename = f'neb_band_{method}_{optimizer}_{precon}.png' + # output_path = os.path.join(logdir, output_filename) - nebtools = NEBTools(images) + # # Plot and save the NEB band + # fig, ax = plt.subplots() + # # nebtools.plot_band(ax=ax) - Ef, dE = nebtools.get_barrier(fit=False) - print(f'{method},{optimizer.__name__},{precon} ' - f'=> Ef = {Ef:.3f}, dE = {dE:.3f}') + # for image in initial_images_copy: + # image.calc = calc() + # ax.plot(image.get_potential_energy(), label='Initial Image') + # for image in images: + # ax.plot(image.get_potential_energy(), label='Optimized Image') + # ax.set_xlabel('Image Index') + # ax.set_ylabel('Potential Energy (eV)') + # ax.legend() - # Save the NEB band plot - output_filename = f'neb_band_{method}_{optimizer.__name__}_{precon}.png' - output_path = os.path.join(testdir, output_filename) + # plt.savefig(output_path) + # plt.close(fig) - # Plot and save the NEB band - fig, ax = plt.subplots() - nebtools.plot_band(ax=ax) - plt.savefig(output_path) - plt.close(fig) diff --git a/src/plot.py b/src/plot.py new file mode 100644 index 0000000..1d94f6d --- /dev/null +++ b/src/plot.py @@ -0,0 +1,48 @@ +import os +from calcs import calc +from ase.io import read +from ase.neb import NEBTools +import matplotlib.pyplot as plt +from initialize_path import setup_images + +method = 'aseneb' +logdir = '/global/home/users/kumaranu/Documents/gpu_jobs' + +r_p = read(os.path.join(logdir, 'r_p.xyz'), ':') +intermediate_images = read(os.path.join(logdir, 'output.xyz'), ':') +optimized_images = read(os.path.join(logdir, 'optimized_path.xyz'), ':') + +reactant = r_p[0] +reactant.calc = calc() +product = r_p[1] +product.calc = calc() + +initial_images = [reactant] +for image in intermediate_images: + image.calc = calc() + initial_images.append(image) +initial_images.append(product) + +nebtools_i = NEBTools(initial_images) +Ef_i, dE_i = nebtools_i.get_barrier(fit=False) + +nebtools_f = NEBTools(optimized_images) +Ef_f, dE_f = nebtools_f.get_barrier(fit=False) + +output_filename = f'neb_band_{method}_NEBOptimizer_None.png' +output_path = os.path.join(logdir, output_filename) + +fig, ax = plt.subplots() +nebtools_i.plot_band(ax=ax) + +for image in initial_images: + ax.plot(image.get_potential_energy(), label='Initial Image', linestyle='--') +for image in optimized_images: + ax.plot(image.get_potential_energy(), label='Optimized Image', linestyle='-') +ax.set_xlabel('Image Index') +ax.set_ylabel('Potential Energy (eV)') +ax.legend() + +plt.savefig(output_path) +plt.close(fig) + diff --git a/src/plot1.py b/src/plot1.py new file mode 100644 index 0000000..3c7681b --- /dev/null +++ b/src/plot1.py @@ -0,0 +1,54 @@ +import os +from calcs import calc +from ase.io import read +from ase.neb import ElasticBand +import matplotlib.pyplot as plt +from initialize_path import setup_images +import numpy as np + +method = 'aseneb' +logdir = '/global/home/users/kumaranu/Documents/gpu_jobs' + +r_p = read(os.path.join(logdir, 'r_p.xyz'), ':') +intermediate_images = read(os.path.join(logdir, 'output.xyz'), ':') +optimized_images = read(os.path.join(logdir, 'optimized_path.xyz'), ':') + +reactant = r_p[0] +reactant.calc = calc() +product = r_p[1] +product.calc = calc() + +initial_images = [reactant] +for image in intermediate_images: + image.calc = calc() + initial_images.append(image) +initial_images.append(product) + +# Calculate RMSD values for initial images +initial_rmsd = [] +for image in initial_images: + initial_rmsd.append(image.get_all_distances(product).sum() / len(reactant)) + +# Calculate RMSD values for optimized images +optimized_rmsd = [] +for image in optimized_images: + optimized_rmsd.append(image.get_all_distances(product).sum() / len(reactant)) + +output_filename = f'neb_band_{method}_NEBOptimizer_None.png' +output_path = os.path.join(logdir, output_filename) + +fig, ax = plt.subplots() + +# Plot initial images +ax.plot(initial_rmsd, [image.get_potential_energy() for image in initial_images], marker='o', linestyle='-', label='Initial Image', color='blue') + +# Plot optimized images on top of initial images +ax.plot(optimized_rmsd, [image.get_potential_energy() for image in optimized_images], marker='o', linestyle='-', label='Optimized Image', color='red') + +ax.set_xlabel('RMSD from Reactant') +ax.set_ylabel('Potential Energy (eV)') +ax.legend() + +plt.savefig(output_path) +plt.close(fig) + diff --git a/src/plot2.py b/src/plot2.py new file mode 100644 index 0000000..518ac2c --- /dev/null +++ b/src/plot2.py @@ -0,0 +1,58 @@ +import os +import sys +import numpy as np +import matplotlib.pyplot as plt +from ase.io import read + +def plot_energy_from_xyz(xyz_file): + # Read the XYZ file + atoms = read(xyz_file, index=':') + + # Extract energies and geometry indices + energies = [] + geometry_indices = [] + + for i, atom in enumerate(atoms): + energy = atom.info['energy'] + energies.append(energy) + geometry_indices.append(i) + + # Convert lists to numpy arrays for plotting + energies = np.array(energies) + geometry_indices = np.array(geometry_indices) + + # Plot energies + plt.plot( + geometry_indices, + energies, + marker='o', + markersize=3, + linewidth=0.5, + linestyle='-', + label=os.path.basename(xyz_file)) + +# Check if two arguments are provided +if len(sys.argv) != 3: + print("Usage: python script.py ") + sys.exit(1) + +# Plot energies from the first XYZ file +plot_energy_from_xyz(sys.argv[1]) + +# Plot energies from the second XYZ file +plot_energy_from_xyz(sys.argv[2]) + +# Set plot labels and title +plt.xlabel('Geometry Index') +plt.ylabel('Energy (eV)') +plt.title('Energy vs. Geometry Index') +plt.legend() +plt.grid(True) + +# Save the plot in the same directory as the first XYZ file +output_file = os.path.splitext(sys.argv[1])[0] + '_energy_plot.png' +plt.savefig(output_file) + +# Show the plot +plt.show() + diff --git a/src/sella_wrapper.py b/src/sella_wrapper.py new file mode 100644 index 0000000..f55ad83 --- /dev/null +++ b/src/sella_wrapper.py @@ -0,0 +1,38 @@ +from ase.io import write +from ase.io import Trajectory +from sella import Sella +from typing import Any, Dict, Iterator, List, Sequence, Tuple, TypeVar, Union + + +def sella_wrapper( + atoms_object, + traj_file=None, + sella_order=0, + use_internal=True, + traj_log_interval=2, + fmax_cutoff=1e-3, + max_steps=1000 + ): + if traj_file: + traj = Trajectory( + traj_file, + 'w', + atoms_object, + ) + qn = Sella( + atoms_object, + order=sella_order, + internal=use_internal, + ) + if traj_file: + qn.attach( + traj.write, + interval=traj_log_interval, + ) + qn.run( + fmax=fmax_cutoff, + steps=max_steps, + ) + if traj_file: + traj.close() + diff --git a/src/setup_images.py b/src/setup_images.py new file mode 100644 index 0000000..58b97ed --- /dev/null +++ b/src/setup_images.py @@ -0,0 +1,43 @@ +import os +import copy +from calcs import calc +from ase.io import read, write +from typing import Any, Dict, Iterator, List, Sequence, Tuple, TypeVar, Union +from sella_wrapper import sella_wrapper + + +def setup_images( + logdir=None, + xyz_r_p=None, + xyz_ts=None, + N_intermediate=40, + ): + reactant = read(xyz_r_p, index='0') + reactant.calc = calc() + sella_wrapper(reactant, traj_file=logdir + '/reactant_opt.traj', sella_order=0) + + product = read(xyz_r_p, index='1') + product.calc = calc() + sella_wrapper(product, traj_file=logdir + '/product_opt.traj', sella_order=0) + + os.makedirs(logdir, exist_ok=True) + os.chdir(logdir) + write("r_p.xyz", [reactant.copy(), product.copy()]) + os.system('geodesic_interpolate r_p.xyz --output output.xyz --nimages ' + str(N_intermediate)) + images = read('output.xyz', index=':') + + for image in images: + image.calc = calc() + mlcalculator = calc() + mlcalculator.calculate(image) + + energy = mlcalculator.results['energy'] + forces = mlcalculator.results['forces'] + + image.info['energy'] = energy + #image.arrays['forces'] = forces + + write('geodesic_path.xyz', images) + + return images + diff --git a/src/test1.py b/src/test1.py index 56ee699..bbe4d20 100644 --- a/src/test1.py +++ b/src/test1.py @@ -1,32 +1,30 @@ - -from ase.optimize import BFGS -from ase.calculators.emt import EMT -from ase.mep.neb import NEBOptimizer +import sys +import ase +import toml from neb_wrapper import run_neb_method -if __name__ == "__main__": - # Define multiple sets of inputs - input_sets = [ - #('aseneb', BFGS, None, None), - #('improvedtangent', BFGS, None, None), - #('spline', BFGS, None, None), - #('string', BFGS, None, None), - ('aseneb', NEBOptimizer, None, None), - ('improvedtangent', NEBOptimizer, None, None), - ('spline', NEBOptimizer, None, None), - ('string', NEBOptimizer, None, None), - ] - # Set the test directory (you can modify this based on your needs) - testdir = '/global/home/users/kumaranu/Documents/gpu_jobs' +if __name__ == "__main__": + with open(sys.argv[1], 'r') as f: + inputs = toml.load(f) + + input_sets = inputs['input_sets'] + logdir = inputs['input_paths']['logdir'] - # Iterate over input sets and call run_neb_method for input_set in input_sets: + optimizer_name = input_set['optimizer'] + # ase_optimizer = getattr(ase.mep.neb, optimizer_name) + ase_optimizer = getattr(ase.neb, optimizer_name) + run_neb_method( - method=input_set[0], - optimizer=input_set[1], - precon=input_set[2], - optmethod=input_set[3], - testdir=testdir, + method=input_set['method'], + optimizer=ase_optimizer, + precon=None, + optmethod=None, + logdir=inputs['input_paths']['logdir'], + xyz_r_p=inputs['input_paths']['xyz_r_p'], + xyz_ts=inputs['input_paths']['xyz_ts'], + N_intermediate=inputs['geodesic_inputs']['N_intermediate'], ) +