def spline_redistribute(geoms):
szts = SimpleZTS.SimpleZTS(geoms)
pre_diffs = get_coords_diffs([image.coords for image in szts.images])
szts.reparametrize()
post_diffs = get_coords_diffs([image.coords for image in szts.images])
cds_str = lambda cds: " ".join([f"{cd:.2f}" for cd in cds])
print("Normalized path segments before splining:")
print(cds_str(pre_diffs))
print("Normalized path segments after redistribution along spline:")
print(cds_str(post_diffs))
return szts.images
def get_splined_hei(self):
self.log("Splining HEI")
# Interpolate energies
cart_coords = align_coords(
[image.cart_coords for image in self.images])
coord_diffs = get_coords_diffs(cart_coords)
self.log(f"\tCoordinate differences: {coord_diffs}")
energies = np.array(self.energy)
energies_spline = interp1d(coord_diffs, energies, kind="cubic")
x_fine = np.linspace(0, 1, 500)
energies_fine = energies_spline(x_fine)
# Determine index that yields the highest energy
hei_ind = energies_fine.argmax()
hei_x = x_fine[hei_ind]
self.log(f"Found splined HEI at x={hei_x:.4f}")
hei_frac_index = hei_x * (len(self.images) - 1)
hei_energy = energies_fine[hei_ind]
reshaped = cart_coords.reshape(-1, self.cart_coords_length)
# To use splprep we have to transpose the coords.
transp_coords = reshaped.transpose()
tcks, us = zip(*[
splprep(transp_coords[i:i + 9], s=0, k=3, u=coord_diffs)
for i in range(0, len(transp_coords), 9)
])
# Reparametrize mesh
hei_coords = np.vstack([splev(
[
hei_x,
],
tck,
) for tck in tcks])
hei_coords = hei_coords.flatten()
# Actually it looks like that splined tangents are really bad approximations
# to the actual imaginary mode. The Cartesian upwinding tangent is usually
# much much better. In 'run_tsopt_from_cos' we actually mix two "normal" tangents
# to obtain the HEI tangent.
hei_tangent = np.vstack(
[splev(
[
hei_x,
],
tck,
der=1,
) for tck in tcks]).T
hei_tangent = hei_tangent.flatten()
hei_tangent /= np.linalg.norm(hei_tangent)
return hei_coords, hei_energy, hei_tangent, hei_frac_index
def func(self, frame):
if self.title:
self.fig.suptitle("Cycle {}".format(frame))
images_x = self.coords[frame][:, 0]
images_y = self.coords[frame][:, 1]
self.images.set_xdata(images_x)
self.images.set_ydata(images_y)
if not self.growing:
# Update total forces quiver
forces_x = self.forces[frame][:, 0]
forces_y = self.forces[frame][:, 1]
offsets = np.stack((images_x, images_y), axis=-1).flatten()
# https://stackoverflow.com/questions/19329039
# https://stackoverflow.com/questions/17758942
self.total_forces_quiv.set_offsets(offsets)
self.total_forces_quiv.set_UVC(forces_x, forces_y)
# Update tangent quiver
tangents_x = self.tangents[frame][:, 0]
tangents_y = self.tangents[frame][:, 1]
self.tangent_quiv.set_offsets(offsets)
self.tangent_quiv.set_UVC(tangents_x, tangents_y)
if self.energy_profile:
coords_diffs = get_coords_diffs(self.coords[frame])
energies = self.energies[frame]
self.energies_plot.set_xdata(coords_diffs)
self.energies_plot.set_ydata(energies)
self.ax1.relim()
self.ax1.autoscale_view()
if self.tight_layout:
plt.tight_layout()
if self.save:
frame_fn = f"step{frame}.png"
if not os.path.exists(frame_fn):
self.fig.savefig(frame_fn)
def run(self):
if not self.restarted:
prep_start_time = time.time()
self.prepare_opt()
prep_end_time = time.time()
prep_time = prep_end_time - prep_start_time
print(f"Spent {prep_time:.1f} s preparing the first cycle.")
self.print_header()
self.stopped = False
# Actual optimization loop
for self.cur_cycle in range(self.last_cycle, self.max_cycles):
start_time = time.time()
self.log(highlight_text(f"Cycle {self.cur_cycle:03d}"))
if self.is_cos and self.check_coord_diffs:
image_coords = [
image.cart_coords for image in self.geometry.images
]
align = len(image_coords[0]) > 3
cds = get_coords_diffs(image_coords, align=align)
# Differences of coordinate differences ;)
cds_diffs = np.diff(cds)
min_ind = cds_diffs.argmin()
if cds_diffs[min_ind] < self.coord_diff_thresh:
similar_inds = min_ind, min_ind + 1
msg = (
f"Cartesian coordinates of images {similar_inds} are "
"too similar. Stopping optimization!")
# I should improve my logging :)
print(msg)
self.log(msg)
break
# Check if something considerably changed in the optimization,
# e.g. new images were added/interpolated. Then the optimizer
# should be reset.
reset_flag = False
if self.cur_cycle > 0 and self.is_cos:
reset_flag = self.geometry.prepare_opt_cycle(
self.coords[-1], self.energies[-1], self.forces[-1])
# Reset when number of coordinates changed
elif self.cur_cycle > 0:
reset_flag = reset_flag or (self.geometry.coords.size !=
self.coords[-1].size)
if reset_flag:
self.reset()
self.coords.append(self.geometry.coords.copy())
self.cart_coords.append(self.geometry.cart_coords.copy())
# Determine and store number of currenctly actively optimized images
try:
image_inds = self.geometry.image_inds
image_num = len(image_inds)
except AttributeError:
image_inds = [
0,
]
image_num = 1
self.image_inds.append(image_inds)
self.image_nums.append(image_num)
step = self.optimize()
if step is None:
# Remove the previously added coords
self.coords.pop(-1)
self.cart_coords.pop(-1)
continue
if self.is_cos:
self.tangents.append(self.geometry.get_tangents().flatten())
self.steps.append(step)
# Convergence check
self.is_converged = self.check_convergence()
end_time = time.time()
elapsed_seconds = end_time - start_time
self.cycle_times.append(elapsed_seconds)
if self.dump:
self.write_cycle_to_file()
with open(self.current_fn, "w") as handle:
handle.write(self.geometry.as_xyz())
if (self.dump and self.dump_restart
and (self.cur_cycle % self.dump_restart) == 0):
self.dump_restart_info()
self.print_opt_progress()
if self.is_converged:
print("Converged!")
print()
break
# Update coordinates
new_coords = self.geometry.coords.copy() + step
try:
self.geometry.coords = new_coords
# Use the actual step. It may differ from the proposed step
# when internal coordinates are used, as the internal-Cartesian
# transformation is done iteratively.
self.steps[-1] = self.geometry.coords - self.coords[-1]
except RebuiltInternalsException as exception:
print("Rebuilt internal coordinates")
with open("rebuilt_primitives.xyz", "w") as handle:
handle.write(self.geometry.as_xyz())
if self.is_cos:
for image in self.geometry.images:
image.reset_coords(exception.typed_prims)
self.reset()
if hasattr(self.geometry, "reparametrize"):
reparametrized = self.geometry.reparametrize()
cur_coords = self.geometry.coords
prev_coords = self.coords[-1]
if reparametrized and (cur_coords.size == prev_coords.size):
self.log("Did reparametrization")
rms = np.sqrt(np.mean((prev_coords - cur_coords)**2))
self.log(
f"rms of coordinates after reparametrization={rms:.6f}"
)
self.is_converged = rms < self.reparam_thresh
if self.is_converged:
print("Insignificant coordinate change after "
"reparametrization. Signalling convergence!")
print()
break
sys.stdout.flush()
sign = check_for_end_sign()
if sign == "stop":
self.stopped = True
break
elif sign == "converged":
self.converged = True
print("Operator indicated convergence!")
break
self.log("")
else:
print("Number of cycles exceeded!")
# Outside loop
if self.dump:
self.out_trj_handle.close()
if (not self.is_cos) and (not self.stopped):
print(self.final_summary())
# Remove 'current_geometry.xyz' file
try:
os.remove(self.current_fn)
except FileNotFoundError:
self.log(
f"Tried to delete '{self.current_fn}'. Couldn't find it.")
with open(self.final_fn, "w") as handle:
handle.write(self.geometry.as_xyz())
print(
f"Wrote final, hopefully optimized, geometry to '{self.final_fn.name}'"
)
sys.stdout.flush()
def __init__(self,
calculator,
optimizer,
xlim=None,
ylim=None,
levels=None,
num=100,
figsize=(8, 8),
interval=250,
energy_profile=True,
colorbar=True,
save=None,
title=True,
tight_layout=False):
self.calculator = calculator
self.optimizer = optimizer
self.interval = interval
if xlim is None:
try:
xlim = calculator.xlim
except AttributeError:
xlim = (-1, 1)
if ylim is None:
try:
ylim = calculator.ylim
except AttributeError:
ylim = (-1, 1)
if levels is None:
try:
lvls = calculator.levels
levels = (lvls.min(), lvls.max(), lvls.size)
except AttributeError:
levels = (-150, 5, 30)
self.energy_profile = energy_profile
self.colorbar = colorbar
self.save = save
self.title = title
self.tight_layout = tight_layout
self.growing = isinstance(self.optimizer.geometry,
GrowingChainOfStates)
self.coords = [c.reshape(-1, 3) for c in self.optimizer.coords]
self.forces = [f.reshape((-1, 3)) for f in self.optimizer.forces]
self.energies = self.optimizer.energies
self.tangents = [t.reshape((-1, 3)) for t in self.optimizer.tangents]
# ax: the contour plot
# ax1: energy along the path
if self.energy_profile:
self.fig, (self.ax, self.ax1) = plt.subplots(
2, figsize=figsize, gridspec_kw={'height_ratios': [3, 1]})
else:
self.fig, self.ax = plt.subplots(figsize=figsize)
self.pause = True
self.fig.canvas.mpl_connect('key_press_event', self.on_keypress)
# Calculate the potential
x = np.linspace(*xlim, 100)
y = np.linspace(*ylim, 100)
X, Y = np.meshgrid(x, y)
Z = np.full_like(X, 0)
fake_atoms = ("H", )
pot_coords = np.stack((X, Y, Z))
pot = self.calculator.get_energy(fake_atoms, pot_coords)["energy"]
# Draw the contourlines of the potential
levels = np.linspace(*levels)
contours = self.ax.contour(X, Y, pot, levels)
#self.ax.clabel(contours, inline=1, fontsize=5)
self.ax.set_xlabel("x")
self.ax.set_ylabel("y")
if self.colorbar:
# Create a colorbar
self.fig.subplots_adjust(right=0.8)
cbar_ax = self.fig.add_axes([0.85, 0.15, 0.05, 0.7])
self.fig.colorbar(contours, cax=cbar_ax)
images_x = self.coords[0][:, 0]
images_y = self.coords[0][:, 1]
forces_x = self.forces[0][:, 0]
forces_y = self.forces[0][:, 1]
tangents_x = self.tangents[0][:, 0]
tangents_y = self.tangents[0][:, 1]
energies = self.energies[0]
# Create artists, so we can update their data later
# Image positions
self.images, = self.ax.plot(images_x, images_y, "ro", ls="-")
if not self.growing:
# Total forces
self.total_forces_quiv = self.ax.quiver(images_x, images_y,
forces_x, forces_y)
# Tangents
self.tangent_quiv = self.ax.quiver(images_x,
images_y,
tangents_x,
tangents_y,
color="b")
# Energy along the path
if self.energy_profile:
self.energies_plot, = self.ax1.plot(get_coords_diffs(
self.coords[0]),
energies,
"ro",
ls="-")
self.ax1.set_xlabel("q(x, y)")
self.ax1.set_ylabel("f(x, y)")