nebscape 0.1.0__py3-none-any.whl

nebscape/__init__.py

@@ -0,0 +1,3 @@
+from .nebscape import NEBscape, get_molecules_reactive_states
+
+__all__ = ["NEBscape", "get_molecules_reactive_states"]

nebscape/analysis.py

@@ -0,0 +1,597 @@
+"""
+Analyzing the result of NEB by looking at chemical fidelity. Also plotting multiple NEB results.
+"""
+
+from .utils import graphfromase, get_active_bonds_from_images, get_neb_path
+from pathlib import Path
+import matplotlib.pyplot as plt
+import numpy as np
+from ase.io import read
+from ase.geometry import find_mic
+from scipy.ndimage import label
+from scipy.signal import argrelextrema
+
+
+def determine_whether_reverse(
+    initial_ref, initial_trial, final_trial, slab_indices, mult=0.8
+):  # numpydoc ignore=GL08
+    # slab_indices = [at_i for at_i, tag in enumerate(initial_ref.arrays['tags']) if tag != 2]
+    initial_graph_ref = graphfromase(initial_ref, slab_indices=slab_indices, mult=mult)
+
+    initial_graph2 = graphfromase(initial_trial, slab_indices=slab_indices, mult=mult)
+    final_graph2 = graphfromase(final_trial, slab_indices=slab_indices, mult=mult)
+
+    isomorphic1 = initial_graph_ref.isomorphic_vf2(
+        initial_graph2,
+        color1=initial_graph_ref.vs["AtomicNums"],
+        color2=initial_graph2.vs["AtomicNums"],
+    )
+    if isomorphic1:
+        reverse = False
+
+    isomorphic2 = initial_graph_ref.isomorphic_vf2(
+        final_graph2,
+        color1=initial_graph_ref.vs["AtomicNums"],
+        color2=final_graph2.vs["AtomicNums"],
+    )
+    if isomorphic2:
+        reverse = True
+
+    return reverse
+
+
+def partition_neb_profile(images, slab_indices, mult=0.8):
+    """
+    Partition NEB profile based on molecular graph.
+
+    When there is single value in one partition, it can be merged to succeeding partition with merge=True
+    It often happens with transition state
+
+    Parameters
+    ----------
+    images : list(ASE atoms)
+        list of image as NEB trajectory
+    slab_indices : list
+        list of slab indices
+    mult : float, optional
+        mult value for ase neighborlist, by default 0.8
+
+    Returns
+    -------
+    list
+       Nested list of indices with image index with same chemical connectivity as a group.
+    """
+    list_of_graphs = []
+    for i, at in enumerate(images):
+        list_of_graphs.append(graphfromase(at, slab_indices=slab_indices, mult=mult))
+
+    corr = np.zeros([len(list_of_graphs), len(list_of_graphs)])
+    for i, graphs1 in enumerate(list_of_graphs):
+        for j, graphs2 in enumerate(list_of_graphs):
+            isomorphic = graphs1.isomorphic_vf2(
+                graphs2,
+                color1=graphs1.vs["AtomicNums"],
+                color2=graphs2.vs["AtomicNums"],
+            )
+            corr[i, j] = isomorphic
+
+    labeled_array, num_features = label(corr)
+    partitions = [
+        np.unique(np.where(labeled_array == i)[0]) for i in np.unique(np.diagonal(labeled_array))
+    ]
+
+    groups_dict = {i: np.where(corr[i, :])[0] for i in range(len(images))}
+    IS = groups_dict[0]
+    FS_index = len(images) - 1
+    FS = groups_dict[FS_index]
+    TS = np.arange(IS[-1] + 1, FS[0])
+    # print(f"{TS = }")
+    if len(TS) != 0:
+        TS_index = TS[0]
+    IS_indices = []
+    FS_indices = []
+    TS_indices = []
+    for idx, array in groups_dict.items():
+        if idx != 0 and len(array) == len(IS) and np.all(array == IS):
+            IS_indices.append(idx)
+        elif idx != FS_index and len(array) == len(FS) and np.all(array == FS):
+            FS_indices.append(idx)
+        elif len(TS) != 0 and idx != TS_index and len(array) == len(TS) and np.all(array == TS):
+            TS_indices.append(idx)
+    for i in IS_indices + FS_indices + TS_indices:
+        del groups_dict[i]
+
+    if len(TS) != 0:
+        for name, idx in zip(["IS", "TS", "FS"], [0, TS_index, FS_index]):
+            groups_dict[name] = groups_dict.pop(idx)
+    else:
+        for name, idx in zip(["IS", "FS"], [0, FS_index]):
+            groups_dict[name] = groups_dict.pop(idx)
+
+    return groups_dict, partitions
+
+
+def get_chemical_transition_state(energy, transition, maxima, minima, num_images):
+    """
+    Get the chemical transition state from the energy profile.
+
+    Parameters
+    ----------
+    energy : np.ndarray
+        The energy profile as a numpy array.
+    transition : float
+        The transition state position. (e.g 7.5 if transition happens between image #7 and #8)
+    maxima : list[int]
+        The indices of the maxima in the energy profile.
+    minima : list[int]
+        The indices of the minima in the energy profile.
+    num_images : int
+        The total number of images in the NEB calculation.
+
+    Returns
+    -------
+    int
+        The index of the chemical transition state.
+    """
+    search_maxima = True
+    index = np.ceil(transition)
+    while search_maxima:
+        if index in maxima:
+            search_maxima = False
+            upper_maxima = index
+        elif index in minima:
+            search_maxima = False
+            upper_maxima = None
+        if index == num_images - 1:
+            upper_maxima = index
+            break
+        index += 1
+
+    search_maxima = True
+    index = np.floor(transition)
+    while search_maxima:
+        if index in maxima:
+            search_maxima = False
+            lower_maxima = index
+        elif index in minima:
+            search_maxima = False
+            lower_maxima = None
+        if index == 0:
+            lower_maxima = 0
+            break
+        index -= 1
+
+    if lower_maxima is None:
+        TS_chemical = upper_maxima
+        return int(TS_chemical)
+    elif upper_maxima is None:
+        TS_chemical = lower_maxima
+        return int(TS_chemical)
+
+    if energy[int(lower_maxima)] < energy[int(upper_maxima)]:
+        TS_chemical = upper_maxima
+    else:
+        TS_chemical = lower_maxima
+    return int(TS_chemical)
+
+
+def check_chemical_fidelity(images, slab_indices, energy_key="last_op__neb_energy", verbose=True):
+    """
+    Check chemical fidelity of NEB results.
+
+    This function checks the followings.
+    (1) Check whether a NEB profile is single step process (Also identify "chemical tranition state")
+        - This is determined by looking at both chemical connectivity and energy profile.
+        - This should be ensured if NEB profile is indeed the a priori intended one
+        - "chemical tranition state" can be identified among energy maxima that corresponds to the intended chemical change.
+    (2) Check whether the highest point coincides with chemical transition state.
+        - If the highest point is not the same as "chemical tranition state", (e.g. diffusion barrier is the highest)
+          then obtained CI-NEB results are not meaningful.
+        - Therefore, additional refinement that changes initial and final that is close to "chemical tranition state" can be attempted.
+
+    Parameters
+    ----------
+    images : list[ase.Atoms]
+        List of NEB images
+    slab_indices : list
+        list of slab atom indices
+    energy_key : str, optional
+        energy key saved in ASE.atoms.info, by default "last_op__neb_energy"
+    verbose : bool, optional
+        verbosity of the output, by default True
+
+    Returns
+    -------
+    single_step_criterion : bool
+        True if NEB images shows single step process.
+    energy_criterion : bool
+        True if the hightest point is the same as "chemical tranition state".
+    barrier_less : bool
+        True if there's no apparent barrier for chemical transition.
+    indices : tuple
+        Tuple(new initial, TS_chemical, final indices) : indices of minima and TS_chemical.
+    """
+    energy = np.array([at.info[energy_key] for at in images])
+    groups_dict, partitions = partition_neb_profile(images, slab_indices, mult=0.8)
+    maxima = argrelextrema(energy, np.greater_equal, order=1)[0]
+    assert np.argmax(energy) in maxima, "Current list of maxima is incorrect"
+
+    minima = argrelextrema(energy, np.less_equal, order=1)[0]
+    assert np.argmin(energy) in minima, "Current list of minima is incorrect"
+    if verbose:
+        print(f"{partitions = }")
+        print(f"{maxima = }")
+        print(f"{minima = }")
+    single_step_criterion = False
+    barrier_less = False
+    if len(groups_dict) == 2:
+        # print("case 1")
+        transition = (partitions[0][-1] + partitions[1][0]) / 2
+        TS_chemical = get_chemical_transition_state(energy, transition, maxima, minima, len(images))
+
+        if TS_chemical == 0 or TS_chemical == len(images) - 1:
+            barrier_less = True
+            return True, False, barrier_less, None
+
+        single_step_criterion = True
+    elif len(groups_dict) == 3:
+        if "TS" in groups_dict.keys():
+            # print("case 2-1")
+            TS = groups_dict["TS"]
+            # check if Transition part of indices belong to the same barrier segment
+            for i in range(len(minima) - 1):
+                # barrier_range is defined as interval between each minima: concave down area of profile
+                barrier_range = np.arange(minima[i], minima[i + 1] + 1)
+                # TS region should be inside of barrier_range
+                # But also there should be no minima included in TS region
+                if np.all(np.isin(TS, barrier_range)) and np.all(~np.isin(TS, minima)):
+                    TS_chemical = barrier_range[np.argmax(energy[barrier_range])]
+                    single_step_criterion = True
+                    break
+        else:
+            # print("case 2-2")
+            transition = (groups_dict["IS"][-1] + groups_dict["FS"][0]) / 2
+            TS_chemical = get_chemical_transition_state(
+                energy, transition, maxima, minima, len(images)
+            )
+            single_step_criterion = True
+
+    if single_step_criterion and verbose:
+        print(f"{TS_chemical = }")
+    if single_step_criterion and TS_chemical == np.argmax(energy):
+        energy_criterion = True
+        indices = None
+    elif single_step_criterion and TS_chemical != np.argmax(energy):
+        smaller = minima[minima < TS_chemical].max()  # nearest smaller
+        larger = minima[minima > TS_chemical].min()  # nearest larger
+        energy_criterion = False
+        indices = smaller, TS_chemical, larger
+    else:
+        # There's no meaningful chemical TS
+        energy_criterion = False
+        indices = None
+
+    return single_step_criterion, energy_criterion, barrier_less, indices
+
+
+def get_same_with_dft_traj(dfttraj, trajs_dict, slab, verbose=False):  # numpydoc ignore=GL08
+    slab_indices = np.arange(len(slab))
+    same_with_dft = None
+    differences_init = {}
+    differences_fin = {}
+
+    active_bonds = get_active_bonds_from_images(dfttraj, slab_indices)
+    # If transfer reaction
+    if "Form" in active_bonds.keys() and "Break" in active_bonds.keys():
+        a, b = active_bonds["Break"]
+        c, d = active_bonds["Form"]
+        if verbose:
+            print(active_bonds)
+
+        v, l_init_dft = find_mic(
+            dfttraj[-1][a].position - dfttraj[-1][b].position,
+            cell=slab.cell,
+            pbc=slab.pbc,
+        )
+        v, l_fin_dft = find_mic(
+            dfttraj[0][c].position - dfttraj[0][d].position,
+            cell=dfttraj[0].cell,
+            pbc=dfttraj[0].pbc,
+        )
+        if verbose:
+            print("dft_l", l_init_dft, l_fin_dft)
+
+        # active_bonds = get_active_bonds_from_images(trajs_dict[min(trajs_dict.keys())], slab_indices)
+        # a,b = active_bonds["Break"]
+        # c,d = active_bonds["Form"]
+        # if verbose:
+        #     print(active_bonds)
+
+        for key, images in trajs_dict.items():
+            active_bonds = get_active_bonds_from_images(images, slab_indices)
+            a, b = active_bonds["Break"]
+            c, d = active_bonds["Form"]
+            if verbose:
+                print(active_bonds)
+
+            v, l_init = find_mic(
+                images[-1][a].position - images[-1][b].position,
+                cell=images[-1].cell,
+                pbc=images[-1].pbc,
+            )
+            v, l_fin = find_mic(
+                images[0][c].position - images[0][d].position,
+                cell=images[0].cell,
+                pbc=images[0].pbc,
+            )
+            if verbose:
+                print("l_init, l_fin", l_init, l_fin)
+
+            differences_init[key] = np.abs(l_init_dft - l_init)
+            differences_fin[key] = np.abs(l_fin_dft - l_fin)
+        same_with_dft_init = min(differences_init, key=differences_init.get)
+        same_with_dft_fin = min(differences_fin, key=differences_fin.get)
+
+        if same_with_dft_init == same_with_dft_fin:
+            same_with_dft = same_with_dft_init
+        else:
+            # print(f"same_with_dft conflict. ")
+            # print(f"{differences_init = }")
+            # print(f"{differences_fin = }")
+            for key in trajs_dict.keys():
+                if differences_init[key] < 1e-4 and differences_fin[key] < 1e-4:
+                    # print(f"{key} were selected.")
+                    return key
+
+            # if np.isclose(l_init_dft ,l_init, atol=atol) and np.isclose(l_fin_dft ,l_fin, atol=atol):
+            #     same_with_dft = key
+            #     break
+
+    # Dissociation reaction
+    elif "Form" not in active_bonds.keys() and "Break" in active_bonds.keys():
+        a, b = active_bonds["Break"]
+        if verbose:
+            print(active_bonds)
+
+        v, l_init_dft = find_mic(
+            dfttraj[0][a].position - dfttraj[0][b].position,
+            cell=slab.cell,
+            pbc=slab.pbc,
+        )
+
+        if verbose:
+            print(f"{l_init_dft=:0.3f}")
+
+        active_bonds = get_active_bonds_from_images(
+            trajs_dict[min(trajs_dict.keys())], slab_indices
+        )
+        a, b = active_bonds["Break"]
+        if verbose:
+            print(active_bonds)
+
+        for key, images in trajs_dict.items():
+            v, l_init = find_mic(
+                images[0][a].position - images[0][b].position,
+                cell=images[-1].cell,
+                pbc=images[-1].pbc,
+            )
+            if verbose:
+                print(f"{l_init=:0.3f}")
+
+            differences_init[key] = np.abs(l_init_dft - l_init)
+
+        same_with_dft = min(differences_init, key=differences_init.get)
+
+        # if np.isclose(l_init_dft ,l_init, atol=atol):
+        #     same_with_dft = key
+        #     break
+
+    return same_with_dft
+
+
+def plot_neb(
+    rootdir,
+    base_path="/work/Calculation/7_neb_protocol/benchmark_var_k",
+    dft_trajs_path="/work/Calculation/dft_trajs_for_release",
+    ylim=None,
+    annotate=False,
+):  # numpydoc ignore=GL08
+    markers = ["o", "s", "*"]
+    # rootdir = Path(f'/work/Calculation/7_neb_protocol/benchmark_var_k/transfer_6/transfer_ood_284_4627_6_111-8')
+
+    if len(str(rootdir).split("/")) == 1:
+        # print('hey')
+        reaction_type = rootdir.split("_")[0]
+        transfer_idx = int(rootdir.split("_")[4])
+        folder = rootdir
+        rootdir = Path(f"{base_path}/{reaction_type}_{transfer_idx}/{rootdir}")
+    else:
+        folder = str(rootdir).split("/")[-1]
+        reaction_type = folder.split("_")[0]
+
+    dfttraj = read(Path(f"{dft_trajs_path}/{reaction_type}s") / f"{folder}_neb1.0.traj", "-10:")
+    slab = read(rootdir / "slab.xyz")
+    slab_indices = [at_i for at_i, tag in enumerate(dfttraj[0].arrays["tags"]) if tag != 2]
+    reverse = None
+
+    fig, axs = plt.subplots(1, 2, figsize=(10, 5), dpi=300)
+    nebtrajs = {}
+
+    reverse = None
+
+    #    if reaction_type == "dissociation":
+    #        paths = Path(rootdir / f"2_NEB/00/0_geometry").glob("*")
+    #    elif reaction_type == "transfer":
+    paths = Path(rootdir / "2_NEB/00").glob("[!backup]*/")
+
+    for i, path in enumerate(paths):
+        alpha = 1
+        if Path(path / "NEB_climb.xyz").is_file():
+            nebtraj = read(path / "NEB_climb.xyz", ":")
+            reverse = determine_whether_reverse(dfttraj[0], nebtraj[0], nebtraj[-1], slab_indices)
+            if reverse is None:
+                reverse = determine_whether_reverse(
+                    dfttraj[0], nebtraj[0], nebtraj[-1], slab_indices
+                )
+            subdir = str(path).split("/")[-1]
+            nebtrajs[subdir] = nebtraj
+            energy_minhop = np.array([at.info["last_op__neb_energy"] for at in nebtraj])
+            if nebtraj[0].info["neb_config_type"] == "neb_last_converged":
+                converged = "T"
+            elif (
+                nebtraj[0].info["neb_config_type"] == "neb_last_unconverged"
+                and nebtraj[0].info["neb_n_steps"] < 500
+            ):
+                converged = "T"
+            else:
+                converged = "F"
+                alpha = 0.1
+
+            # if reverse:
+            #     print('reversed!')
+            #     axs[1].plot(get_neb_path(nebtraj), energy_minhop[::-1], marker="o", c=f"C{i}", alpha=alpha, label=f"permute#{i}-{converged}")
+            # else:
+            marker = markers[int(subdir.split("_")[0])]
+            axs[1].plot(
+                get_neb_path(nebtraj),
+                energy_minhop,
+                marker=marker,
+                c=f"C{i}",
+                alpha=alpha,
+                label=f"#{subdir}-{converged}",
+            )
+
+            if annotate and converged == "T":
+                for j, at in enumerate(nebtraj):
+                    axs[1].annotate(
+                        f"{j}", (get_neb_path(nebtraj)[j], energy_minhop[j]), fontsize=8
+                    )
+
+    #    reverse = determine_whether_reverse(dfttraj[0], nebtraj[0], nebtraj[-1], slab_indices)
+
+    # for i in range(3):
+    ocp_trajs = {}
+    for i, path in enumerate(Path(rootdir / "3_OCP_benchmark").glob("*")):
+        alpha = 1
+        if Path(rootdir / f"3_OCP_benchmark/{i}").is_dir():
+            nebtraj = read(rootdir / f"3_OCP_benchmark/{i}/NEB_climb.xyz", ":")
+            ocp_trajs[i] = nebtraj
+    same_with_dft = get_same_with_dft_traj(dfttraj, ocp_trajs, slab)
+    # print(f'{same_with_dft = }')
+
+    ocp_trajs = {}
+    for i, path in enumerate(Path(rootdir / "3_OCP_benchmark").glob("*")):
+        alpha = 1
+        if Path(rootdir / f"3_OCP_benchmark/{i}").is_dir():
+            nebtraj = read(rootdir / f"3_OCP_benchmark/{i}/NEB_climb.xyz", ":")
+            ocp_trajs[i] = nebtraj
+            if reverse:
+                print("reversed!")
+                nebtraj = nebtraj[::-1]
+            energy_eq2 = np.array([at.info["last_op__neb_energy"] for at in nebtraj])
+            if nebtraj[0].info["neb_config_type"] == "neb_last_converged":
+                converged = "T"
+            elif (
+                nebtraj[0].info["neb_config_type"] == "neb_last_unconverged"
+                and nebtraj[0].info["neb_n_steps"] < 500
+            ):
+                converged = "T"
+            else:
+                converged = "F"
+                alpha = 0.1
+
+            # if reverse:
+            #     print('reversed')
+            #     axs[0].plot(get_neb_path(nebtraj)[::-1], energy_eq2, alpha=alpha, c=f"C{i}", marker="o", label=f"permute#{i}-{converged}")
+            # else:
+            subdir = str(path).split("/")[-1]
+            if i == same_with_dft:
+                axs[0].plot(
+                    get_neb_path(nebtraj),
+                    energy_eq2,
+                    alpha=alpha,
+                    c=f"C{i}",
+                    marker="o",
+                    label=f"#{subdir}-{converged}**",
+                )
+            else:
+                axs[0].plot(
+                    get_neb_path(nebtraj),
+                    energy_eq2,
+                    alpha=alpha,
+                    c=f"C{i}",
+                    marker="o",
+                    label=f"#{subdir}-{converged}",
+                )
+
+    # for at in dfttraj[1:-1]:
+    #     at.info["dft_energy"] = at.get_potential_energy()
+    # for at in dfttraj:
+    #     at.calc = ocpcalc
+    #     at.info["ocp_energy"] = at.get_potential_energy()
+
+    # energy_dft = np.array([at.info["dft_energy"] for at in dfttraj[1:-1]])
+    # energy_sp = np.array([at.info["ocp_energy"] for at in dfttraj])
+    # axs[0].plot(get_neb_path(dfttraj), energy_sp, marker="s",c=f"C3", label="eq2-sp@DFT")
+    # axs[0].plot(get_neb_path(dfttraj)[1:-1], energy_dft-energy_dft[0]+energy_sp[1], c=f"C4", marker="x", label="DFT")
+    for i in range(2):
+        # axs[i].plot(get_neb_path(dfttraj), energy_sp, marker="s",c=f"C3", label="eq2-sp@DFT")
+        # axs[i].plot(get_neb_path(dfttraj)[1:-1], energy_dft-energy_dft[0]+energy_sp[1], c=f"C4", marker="x", label="DFT")
+        axs[i].set_ylabel("Energy [eV]", fontsize=12)
+        axs[i].set_xlabel("displacement [$\AA$]", fontsize=12)
+        axs[i].legend()
+
+    axs[0].set_title("OCP geometry")
+    axs[1].set_title("Minima hopping geometry")
+
+    if ylim is None:
+        axs[0].set_ylim(
+            (
+                min(min(axs[0].get_ylim()), min(axs[1].get_ylim())),
+                max(max(axs[0].get_ylim()), max(axs[1].get_ylim())),
+            )
+        )
+        ymin, ymax = (
+            min(min(axs[0].get_ylim()), min(axs[1].get_ylim())),
+            max(max(axs[0].get_ylim()), max(axs[1].get_ylim())),
+        )
+        axs[1].set_ylim(
+            (
+                min(min(axs[0].get_ylim()), min(axs[1].get_ylim())),
+                max(max(axs[0].get_ylim()), max(axs[1].get_ylim())),
+            )
+        )
+    else:
+        ymin, ymax = ylim
+        axs[0].set_ylim((ylim[0], ylim[1]))
+        axs[1].set_ylim((ylim[0], ylim[1]))
+    ### Partition neb profile
+    for j, profiles in enumerate([ocp_trajs, nebtrajs]):
+        highest_points = {}
+        for key, images in profiles.items():
+            if images[0].info["neb_config_type"] == "neb_last_converged":
+                highest_points[key] = np.max([at.info["last_op__neb_energy"] for at in images])
+            elif (
+                images[0].info["neb_config_type"] == "neb_last_unconverged"
+                and images[0].info["neb_n_steps"] < 500
+            ):
+                highest_points[key] = np.max([at.info["last_op__neb_energy"] for at in images])
+
+        if len(highest_points) != 0:
+            best_neb_profile = min(highest_points, key=highest_points.get)
+            groups, partitions = partition_neb_profile(profiles[best_neb_profile], slab_indices)
+
+            #            for indices in partitions:
+            #                axs[j].fill_between(x[indices], (ymax - ymin) / 40 + ymin, (ymax - ymin) / 20 + ymin, color=f"C{color_i}", alpha=0.5)
+            for i in range(len(partitions) - 1):
+                transition = (partitions[i][-1] + partitions[i + 1][0]) / 2
+                plt.axvline(x=transition, linestyle=":", c="k")
+    #            for indices in partitions[:-1]:
+    #                axs[j].axvline(x=x[indices][-1], c=f"C{color_i}", alpha=0.9, linestyle="--")
+
+    # axs[1].set_ylim([-1,2])
+    plt.suptitle(str(rootdir).split("/")[-1])
+    plt.savefig(rootdir / "neb.png")
+    plt.tight_layout()
+    plt.show()
+    return ocp_trajs, nebtrajs