import numpy as np
from pysisyphus.Geometry import Geometry
from pysisyphus.helpers import align_geoms
from pysisyphus.xyzloader import write_geoms_to_trj
[docs]
class Interpolator:
def __init__(
self,
geoms,
between,
extrapolate=0,
extrapolate_before=0,
extrapolate_after=0,
extrapolate_damp=1.0,
align=False,
):
self.geoms = geoms
self.between = between
self.extrapolate = extrapolate
self.extrapolate_before = (
extrapolate_before if extrapolate_before else self.extrapolate
)
self.extrapolate_after = (
extrapolate_after if extrapolate_after else self.extrapolate
)
self.extrapolate_damp = extrapolate_damp
one_atom_geom = any([len(geom.atoms) == 1 for geom in geoms])
# Don't try to align one atom species
self.align = align and not one_atom_geom
assert len(geoms) >= 2, "Need at least two geometries to interpolate!"
# When the geometries have different coord_types we can't be sure which one
# to pick, when we create intermediate geometries.
self.coord_type_ref = geoms[0].coord_type
coord_types = [geom.coord_type for geom in geoms]
assert all(
[coord_type == self.coord_type_ref for coord_type in coord_types]
), "Geometries must have the same coord_type, but they have '{}"
# Check for consistent atom ordering
for i, geom in enumerate(self.geoms[:-1]):
next_geom = self.geoms[i + 1]
atoms = geom.atoms
next_atoms = next_geom.atoms
assert len(atoms) == len(
next_atoms
), f"Geometries {i} and {i+1} have a different number of atoms!"
assert (
atoms == next_atoms
), f"Different atom ordering in geometries {i} and {i+1}!"
self.atoms = self.geoms[0].atoms
if self.align:
align_geoms(geoms)
self.all_geoms = None
[docs]
def interpolate_all(self):
all_geoms = list()
if self.extrapolate_before:
print("Extrapolating before")
geoms_extrapol_before = self.interpolate(
self.geoms[0],
self.geoms[1],
interpolate_only=self.extrapolate_before,
extrapolate=True,
)
all_geoms += geoms_extrapol_before[::-1]
print("Interpolating between")
# Interpolate between all pairs of geometries
for i, initial_geom in enumerate(self.geoms[:-1]):
final_geom = self.geoms[i + 1]
interpolated_geoms = self.interpolate(initial_geom, final_geom)
assert len(interpolated_geoms) == self.between, (
"Something is wrong with the number of interpolated "
"geometries. Maybe you accidentally also return the "
"initial and final geometry?"
)
all_geoms.append(self.geoms[i])
all_geoms.extend(interpolated_geoms)
# As we only added the i-th geometry and the new interpolated
# geometries of every geometry pair we also have to add the last
# ((i+1)-th) geometry at the end.
all_geoms.append(self.geoms[-1])
print("Extrapolate between")
if self.extrapolate_after:
geoms_extrapol_after = self.interpolate(
self.geoms[-1],
self.geoms[-2],
interpolate_only=self.extrapolate_after,
extrapolate=True,
)
all_geoms += geoms_extrapol_after
self.all_geoms = all_geoms
if self.align:
align_geoms(self.all_geoms)
return all_geoms
[docs]
def interpolate(
self, initial_geom, final_geom, interpolate_only=0, extrapolate=False
):
initial_coords = initial_geom.coords
coord_type = self.coord_type_ref
# Linear interpolation
step = (final_geom.coords - initial_geom.coords) / (self.between + 1)
if extrapolate:
step *= -1
step *= self.extrapolate_damp
# When we extrapolate we probably want a number of geometries that is
# different from self.between
interpolations = interpolate_only if interpolate_only else self.between
# initial + i*step
i_array = np.arange(1, interpolations + 1)
new_coords = initial_coords + i_array[:, None] * step
return [Geometry(self.atoms, nc, coord_type=coord_type) for nc in new_coords]
[docs]
def all_geoms_to_trj(self, trj_fn):
write_geoms_to_trj(self.all_geoms, trj_fn)