import pathlib
from copy import copy
import ase.io
import ase.optimize
import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import zntrack
from ase.mep import NEB
from mlipx.abc import ComparisonResults, NodeWithCalculator, Optimizer
[docs]
class NEBinterpolate(zntrack.Node):
"""
Interpolates between two or three images to create a NEB path.
Parameters
----------
data : list[ase.Atoms]
List of atoms objects.
n_images : int
Number of images to interpolate.
mic : bool
Whether to use the minimum image convention.
add_constraints : bool
Whether to copy constraints from initial image to all the interpolated images
frames_path : pathlib.Path
Path to save the interpolated frames.
"""
data: list[ase.Atoms] = zntrack.deps()
n_images: int = zntrack.params(5)
mic: bool = zntrack.params(False)
add_constraints: bool = zntrack.params(True)
frames_path: pathlib.Path = zntrack.outs_path(zntrack.nwd / "initial_frames.xyz")
def run(self):
frames = []
if len(self.data) == 2:
atoms = self.data[0]
for i in range(self.n_images - 1):
atoms_copy = atoms.copy()
frames += [atoms_copy]
atoms_copy = self.data[1]
frames += [atoms_copy]
elif len(self.data) == 3:
atoms0 = self.data[0]
atoms1 = self.data[1]
atoms2 = self.data[2]
ts_index = self.n_images // 2
for i in range(ts_index - 1):
atoms_copy = atoms0.copy()
frames += [atoms_copy]
for i in range(ts_index, self.n_images):
atoms_copy = atoms1.copy()
frames += [atoms_copy]
atoms_copy = atoms2.copy()
frames += [atoms_copy]
if self.add_constraints is True:
_constraints = self.data[0].constraints
for image in frames:
image.set_constraint(_constraints)
neb = NEB(frames)
neb.interpolate(mic=self.mic, apply_constraint=self.add_constraints)
ase.io.write(self.frames_path, frames)
@property
def frames(self) -> list[ase.Atoms]:
with self.state.fs.open(self.frames_path, "r") as f:
return list(ase.io.iread(f, format="extxyz"))
[docs]
class NEBs(zntrack.Node):
"""
Runs NEB calculation on a list of images.
Parameters
----------
data : list[ase.Atoms]
List of atoms objects.
model : NodeWithCalculator
Node with a calculator.
relax : bool
Whether to relax the initial and final images.
optimizer : Optimizer
ASE optimizer to use.
fmax : float
Maximum force allowed.
n_steps : int
Maximum number of steps allowed.
frames_path : pathlib.Path
Path to save the final frames.
trajectory_path : pathlib.Path
Path to save the neb trajectory file.
Attributes
----------
results : pd.DataFrame
DataFrame with the data_id and potential energy of the NEB calculation
"""
data: list[ase.Atoms] = zntrack.deps()
model: NodeWithCalculator = zntrack.deps()
relax: bool = zntrack.params(True)
optimizer: Optimizer = zntrack.params(Optimizer.FIRE.value)
fmax: float = zntrack.params(0.09)
n_steps: int = zntrack.params(500)
frames_path: pathlib.Path = zntrack.outs_path(zntrack.nwd / "frames.xyz")
trajectory_path: pathlib.Path = zntrack.outs_path(
zntrack.nwd / "neb_trajectory.traj"
)
results: pd.DataFrame = zntrack.plots(y="potential_energy", x="data_id")
def run(self):
frames = []
# neb_trajectory = []
calc = self.model.get_calculator()
optimizer = getattr(ase.optimize, self.optimizer)
for atoms in self.data:
atoms_copy = atoms.copy()
atoms_copy.calc = copy(calc)
atoms_copy.get_potential_energy()
frames += [atoms_copy]
ase.io.write(self.frames_path, atoms_copy, format="extxyz", append=True)
if self.relax is True:
for i in [0, -1]:
dyn = optimizer(frames[0])
dyn.run(fmax=self.fmax)
neb = NEB(frames, allow_shared_calculator=False)
dyn = optimizer(neb, trajectory=self.trajectory_path.as_posix())
dyn.run(fmax=self.fmax, steps=self.n_steps)
row_dicts = []
for i, frame in enumerate(frames):
row_dicts.append(
{
"data_id": i,
"potential_energy": frame.get_potential_energy(),
"neb_adjusted_energy": frame.get_potential_energy()
- frames[0].get_potential_energy(),
},
)
self.results = pd.DataFrame(row_dicts)
@property
def trajectory_frames(self) -> list[ase.Atoms]:
with self.state.fs.open(self.trajectory_path, "rb") as f:
return list(ase.io.iread(f, format="traj"))
@property
def frames(self) -> list[ase.Atoms]:
with self.state.fs.open(self.frames_path, "r") as f:
return list(ase.io.iread(f, format="extxyz"))
@property
def figures(self) -> dict[str, go.Figure]:
fig = px.scatter(self.results, x="data_id", y="potential_energy")
fig.update_layout(title="NEB_path")
fig.update_traces(customdata=np.stack([np.arange(len(self.results))], axis=1))
return {"NEB_path": fig}
@property
def traj_plots(self) -> dict[str, go.Figure]:
trajectory_frames = self.trajectory_frames
total_iterations = len(trajectory_frames) // len(self.frames)
neb_length = len(self.frames)
figure = go.Figure()
for iteration in range(total_iterations):
images = trajectory_frames[
iteration * neb_length : (iteration + 1) * neb_length
]
energies = [image.get_potential_energy() for image in images]
offset = iteration * neb_length
figure.add_trace(
go.Scatter(
x=list(range(len(energies))),
y=energies,
mode="lines+markers",
name=f"{iteration}",
customdata=np.stack([np.arange(len(energies)) + offset], axis=1),
)
)
figure.update_layout(
title="Energy vs. NEB image",
xaxis_title="image number",
yaxis_title="Energy / eV",
)
return {"energy_vs_iteration": figure}
@staticmethod
def compare(*nodes: "NEBs") -> ComparisonResults:
frames = sum([node.frames for node in nodes], [])
offset = 0
fig = go.Figure()
for idx, node in enumerate(nodes):
energies = [atoms.get_potential_energy() for atoms in node.frames]
fig.add_trace(
go.Scatter(
x=list(range(len(energies))),
y=energies,
mode="lines+markers",
name=node.name.replace(f"_{node.__class__.__name__}", ""),
customdata=np.stack([np.arange(len(energies)) + offset], axis=1),
)
)
offset += len(energies)
fig.update_layout(
title="Energy vs. NEB image",
xaxis_title="image number",
yaxis_title="Energy / eV",
plot_bgcolor="rgba(0, 0, 0, 0)",
paper_bgcolor="rgba(0, 0, 0, 0)",
)
fig.update_xaxes(
showgrid=True,
gridwidth=1,
gridcolor="rgba(120, 120, 120, 0.3)",
zeroline=False,
)
fig.update_yaxes(
showgrid=True,
gridwidth=1,
gridcolor="rgba(120, 120, 120, 0.3)",
zeroline=False,
)
# Now adjusted
fig_adjusted = go.Figure()
offset = 0
for _, node in enumerate(nodes):
energies = np.array([atoms.get_potential_energy() for atoms in node.frames])
energies -= energies[0]
fig_adjusted.add_trace(
go.Scatter(
x=list(range(len(energies))),
y=energies,
mode="lines+markers",
name=node.name.replace(f"_{node.__class__.__name__}", ""),
customdata=np.stack([np.arange(len(energies)) + offset], axis=1),
)
)
offset += len(energies)
fig_adjusted.update_layout(
title="Adjusted energy vs. NEB image",
xaxis_title="Image number",
yaxis_title="Adjusted energy / eV",
plot_bgcolor="rgba(0, 0, 0, 0)",
paper_bgcolor="rgba(0, 0, 0, 0)",
)
fig_adjusted.update_xaxes(
showgrid=True,
gridwidth=1,
gridcolor="rgba(120, 120, 120, 0.3)",
zeroline=False,
)
fig_adjusted.update_yaxes(
showgrid=True,
gridwidth=1,
gridcolor="rgba(120, 120, 120, 0.3)",
zeroline=False,
)
return ComparisonResults(
frames=frames,
figures={
"energy_vs_neb_image": fig,
"adjusted_energy_vs_neb_image": fig_adjusted,
},
)