# [1] https://pubs.acs.org/doi/pdf/10.1021/ct200290m?rand=dcfwsf09
# [2] https://onlinelibrary.wiley.com/doi/epdf/10.1002/jcc.23481
# [3] https://onlinelibrary.wiley.com/doi/epdf/10.1002/tcr.201600043
import itertools as it
from math import isclose
from typing import List
import autograd
import autograd.numpy as anp
import numpy as np
import numpy.typing as npt
from pysisyphus.calculators.Calculator import Calculator
from pysisyphus.elem_data import COVALENT_RADII
from pysisyphus.finite_diffs import finite_difference_hessian
from pysisyphus.Geometry import Geometry
from pysisyphus.helpers import complete_fragments
from pysisyphus.helpers_pure import log
from pysisyphus.io.hdf5 import get_h5_group, resize_h5_group
[docs]
def get_data_model(atoms, max_cycles):
coord_size = 3 * len(atoms)
_1d = (max_cycles,)
_2d = (max_cycles, coord_size)
_3d = (max_cycles, coord_size, coord_size)
data_model = {
"cart_coords": _2d,
"energy": _1d,
"forces": _2d,
"hessian": _3d,
"true_energy": _1d,
"true_forces": _2d,
"true_hessian": _3d,
}
return data_model
[docs]
class CovRadiiSumZero(Exception):
pass
[docs]
def afir_closure(
fragment_indices, cov_radii, gamma, rho=1, p=6, prefactor=1.0, logger=None
):
"""rho=1 pushes fragments together, rho=-1 pulls fragments apart."""
# See https://onlinelibrary.wiley.com/doi/full/10.1002/qua.24757
# Eq. (9) for extension to 3 fragments.
assert len(fragment_indices) == 2
inds = np.array(list(it.product(*fragment_indices)))
cov_rad_sums = cov_radii[inds].sum(axis=1)
if isclose(cov_rad_sums.sum(), 0.0, abs_tol=1e-10):
raise CovRadiiSumZero("Sum of covalent radii is 0.0!")
# 3.8164 Angstrom in Bohr
R0 = 7.21195
# 1.0061 kJ/mol to Hartree
epsilon = 0.000383203368
# Avoid division by zero for gamma = 0.
if gamma == 0.0:
alpha = 0.0
else:
alpha = gamma / (
(2 ** (-1 / 6) - (1 + (1 + gamma / epsilon) ** 0.5) ** (-1 / 6)) * R0
)
rho_verbose = {1: ("pushing", "together"), -1: ("pulling", "apart")}
w1, w2 = rho_verbose[rho]
log(
logger,
f"Creating AFIR closure with α={alpha:.6f}, prefactor {prefactor:.6f}, "
f"rho={rho}, {w1} framgents {w2}",
)
def afir_func(coords3d):
diffs = anp.diff(coords3d[inds], axis=1).reshape(-1, 3)
rs = anp.linalg.norm(diffs, axis=1)
omegas = (cov_rad_sums / rs) ** p
f = prefactor * alpha * rho * (omegas * rs).sum() / omegas.sum()
return f
return afir_func
[docs]
class AFIR(Calculator):
[docs]
def __init__(
self,
calculator: Calculator,
fragment_indices: List[List[int]],
gamma: npt.ArrayLike,
rho: npt.ArrayLike = 1,
p: int = 6,
ignore_hydrogen: bool = False,
zero_hydrogen: bool = True,
complete_fragments: bool = True,
dump: bool = True,
h5_fn: str = "afir.h5",
h5_group_name: str = "afir",
**kwargs,
):
"""
Artifical Force Induced Reaction calculator.
Currently, there are no automated drivers to run large-scale AFIR calculations
with many different initial orientations and/or increasing collision energy
parameter γ. Nontheless, selected AFIR calculations can be carried out by hand.
After convergence, artificial potential & forces, as well as real energies
and forces can be plotted with 'pysisplot --afir'. The highest energy point
along the AFIR path can then be selected for a subsequent TS-optimization,
e.g. via 'pysistrj --get [index] optimzation.trj'.
Future versions of pysisyphus may provide drivers for more automatted
AFIR calculations.
Parameters
----------
calculator
Actual QC calculator that provides energies and its derivatives,
that are modified by the AFIR calculator, e.g., ORCA or Psi4.
fragment_indices
List of lists of integers, specifying the separate fragments. If
the indices in theses lists don't comprise all atoms in the molecule,
the reamining indices will be added as a separate fragment. If a AFIR
calculation is carried out with 2 fragments and 'complete_fragments' is
True (see below) it is enough to specify only
the indices of one fragment, e.g., for a system of 10 atoms
'fragment_indices=[[0,1,2,3]]' is enough. The second system will be set
up automatically with indices [4,5,6,7,8,9].
gamma
Collision energy parameter γ in au. For 2 fragments it can be a single
integer, while for > 2 fragments a list of gammas must be given, specifying
the pair-wise collision energy parameters. For 3 fragments 3 gammas must be
given [γ_01, γ_02, γ_12], for 4 fragments 6 gammas would be required
[γ_01, γ_02, γ_03, γ_12, γ_13, γ_23] and so on.
rho
Direction of the artificial force, either 1 or -1. The same comments
as for gamma apply. For 2 fragments a single integer is enough, for
> 2 fragments a list of rhos must be given (see above). For rho=1
fragments are pushed together, for rho=-1 fragments are pulled apart.
p
Exponent p used in the calculation of the weight function ω. Defaults
to 6 and probably does not have to be changed.
ignore_hydrogen
Whether hydrogens are ignored in the calculation of the artificial force.
All weights between atom pairs containing hydrogen will be set to 0.
zero_hydrogen
Whether to use 0.0 as covalent radius for hydrogen in the weight function.
Compared to 'ignore_hydrogen', which results in zero weights for all atom
pairs involving hydrogen, 'zero_hydrogen' may be non-zero, depending on the
covalent radius of the second atom in the pair.
complete_fragments
Whether an incomplete specification in 'fragment_indices' is automatically
completed.
dump
Whether an HDF5 file is created.
h5_fn
Filename of the HDF5 file used for dumping.
h5_group_name
HDF5 group name used for dumping.
Other Parameters
----------------
**kwargs
Keyword arguments passed to the Calculator baseclass.
"""
super().__init__(**kwargs)
self.calculator = calculator
# Update own charge and multiplicity with values from the wrapped
# calculator.
self.charge = calculator.charge
self.mult = calculator.mult
self.fragment_indices = fragment_indices
assert len(self.fragment_indices) > 0
self.gamma = np.atleast_1d(gamma).astype(float)
self.rho = np.atleast_1d(rho).astype(int)
np.testing.assert_allclose(np.abs(self.rho), np.ones_like(self.rho))
self.p = p
self.ignore_hydrogen = ignore_hydrogen
self.zero_hydrogen = zero_hydrogen
self.complete_fragments = complete_fragments
self.dump = dump
self.h5_fn = self.out_dir / h5_fn
self.h5_group_name = h5_group_name
# We can't initialize the HDF5 group as we don't know the shape of
# atoms/coords yet. So we wait until after the first calculation.
self.h5_cycles = 50
if self.ignore_hydrogen:
self.log("No artificial force contribution from hydrogens!")
self.atoms = None
self.calc_counter = 0
self.h5_group_initialized = False
"""We try to keep charge and multiplicity consistent between the AFIR
calculator and the actual wrapped calculator. But we will always return
the charge and multiplicity of the underlying calculator."""
@property
def charge(self):
return self.calculator.charge
@charge.setter
def charge(self, charge):
try:
self.calculator.charge = charge
except AttributeError:
pass
self._charge = charge
@property
def mult(self):
return self.calculator.mult
@mult.setter
def mult(self, mult):
try:
self.calculator.mult = mult
except AttributeError:
pass
self._mult = mult
[docs]
def init_h5_group(self, atoms, max_cycles=None):
if max_cycles is None:
max_cycles = self.h5_cycles
self.data_model = get_data_model(atoms, max_cycles)
h5_group = get_h5_group(
self.h5_fn, self.h5_group_name, self.data_model, reset=True
)
self.h5_group_initialized = True
return h5_group
[docs]
def dump_h5(self, atoms, coords, results):
# Initialize if not done yet
if not self.h5_group_initialized:
h5_group = self.init_h5_group(atoms)
# Write atoms only once
h5_group.attrs["atoms"] = atoms
h5_group = get_h5_group(self.h5_fn, self.h5_group_name)
# Check if HDF5 datasets have to be resized
cur_max_cycles = h5_group["cart_coords"].shape[0]
need_resize = self.calc_counter > cur_max_cycles - 5
if need_resize:
new_max_cycles = cur_max_cycles + self.h5_cycles
resize_h5_group(h5_group, max_cycles=new_max_cycles)
for k, v in results.items():
h5_group[k][self.calc_counter] = v
h5_group["cart_coords"][self.calc_counter] = coords
h5_group.attrs["cur_cycle"] = self.calc_counter
h5_group.file.close()
[docs]
def log_fragments(self):
self.log(f"Using {len(self.fragment_indices)} fragments")
for i, frag in enumerate(self.fragment_indices):
self.log(f"Fragment {i:02d}, {len(frag)} atoms:")
self.log(f"\t{frag}")
[docs]
def write_fragment_geoms(self, atoms, coords):
geom = Geometry(atoms, coords)
for i, frag in enumerate(self.fragment_indices):
frag_geom = geom.get_subgeom(frag)
fn = self.make_fn(f"frag_geom_{i:02d}.xyz")
with open(fn, "w") as handle:
handle.write(frag_geom.as_xyz())
self.log(f"Wrote geometry of fragment {i:02d} to {fn}.")
[docs]
def set_atoms_and_funcs(self, atoms, coords):
"""Initially atoms was also an argument to the constructor of AFIR.
I removed it so creation becomes easier.
The first time a calculation is requested with a proper atom set
everything is set up (cov. radii, afir function and corresponding
gradient). Afterwards there is only a check if atoms != None and it
is expected that all functions are properly set.
fragment_indices can also be incomplete w.r.t. to the number of
atoms. If the sum of the specified fragment atoms is less than the
number of atoms present then all remaining unspecified atoms will
be gathered in one fragment.
"""
if self.atoms is not None:
assert self.atoms == atoms
return
self.log("Setting atoms on AFIR calculator")
self.atoms = atoms
if self.complete_fragments:
self.fragment_indices = complete_fragments(
self.atoms, self.fragment_indices
)
self.log_fragments()
self.write_fragment_geoms(atoms, coords)
hydrogen_inds = [i for i, atom in enumerate(atoms) if atom.lower() == "h"]
if self.ignore_hydrogen:
fragment_indices = list()
for i, frag_inds in enumerate(self.fragment_indices):
frag_inds_no_h = [j for j in frag_inds if j not in hydrogen_inds]
fragment_indices.append(frag_inds_no_h)
dropped_hydrogens = len(frag_inds) - len(frag_inds_no_h)
self.log(f"Ignoring {dropped_hydrogens} hydrogen(s) from fragment {i}.")
self.fragment_indices = fragment_indices
self.cov_radii_org = np.array([COVALENT_RADII[atom.lower()] for atom in atoms])
self.cov_radii = self.cov_radii_org.copy()
if self.zero_hydrogen:
self.cov_radii[hydrogen_inds] = 0.0
self.log("Set covalent radii")
self.afir_funcs = list()
self.afir_grad_funcs = list()
pairs = list(it.combinations(self.fragment_indices, 2))
prefactor = 1 / len(pairs)
self.log(
f"Doing AFIR with {len(pairs)} fragment pairs and prefactor={prefactor:.4f}"
)
try:
self.gamma = [float(self.gamma)] * len(pairs)
except TypeError:
assert len(self.gamma) == len(pairs)
try:
self.rho = [float(self.rho)] * len(pairs)
except TypeError:
assert len(self.rho) == len(pairs)
self.log(f"Using gamma(s): {self.gamma}")
self.log(f" Using rho(s): {self.rho}")
for (frag1, frag2), gamma, rho in zip(pairs, self.gamma, self.rho):
afir_func = afir_closure(
(frag1, frag2),
self.cov_radii,
gamma=gamma,
rho=rho,
p=self.p,
prefactor=prefactor,
logger=self.logger,
)
afir_grad_func = autograd.grad(afir_func)
self.afir_funcs.append(afir_func)
self.afir_grad_funcs.append(afir_grad_func)
self.log("Created and set AFIR function & gradient function.")
[docs]
def get_energy(self, atoms, coords, **prepare_kwargs):
self.set_atoms_and_funcs(atoms, coords)
true_results = self.calculator.get_energy(atoms, coords, **prepare_kwargs)
true_energy = true_results["energy"]
coords3d = coords.reshape(-1, 3)
# Iterate over all fragment pairs
afir_energy = sum([afir_func(coords3d) for afir_func in self.afir_funcs])
self.log()
results = {
"energy": true_energy + afir_energy,
"true_energy": true_energy,
}
if self.dump:
self.dump_h5(atoms, coords, results)
self.calc_counter += 1
return results
[docs]
def get_forces(self, atoms, coords, **prepare_kwargs):
self.set_atoms_and_funcs(atoms, coords)
true_results = self.calculator.get_forces(atoms, coords, **prepare_kwargs)
true_energy = true_results["energy"]
true_forces = true_results["forces"]
coords3d = coords.reshape(-1, 3)
afir_energy = 0.0
afir_forces = np.zeros_like(coords)
# Iterate over all fragment pairs
for afir_func, afir_grad_func in zip(self.afir_funcs, self.afir_grad_funcs):
afir_energy += afir_func(coords3d)
# Add negative of the gradient (the force)
afir_forces += -afir_grad_func(coords3d).flatten()
true_norm = np.linalg.norm(true_forces)
afir_norm = np.linalg.norm(afir_forces)
self.log(
f"\ntrue_energy={true_energy:.6f} au\n"
f"afir_energy={afir_energy:.6f} au\n"
f" sum_energy={true_energy+afir_energy:.6f} au\n"
f"norm(true_forces)={true_norm:.6f} au/bohr\n"
f"norm(afir_forces)={afir_norm:.6f} au/bohr\n"
)
results = {
"energy": true_energy + afir_energy,
"forces": true_forces + afir_forces,
"true_forces": true_forces,
"true_energy": true_energy,
}
if self.dump:
self.dump_h5(atoms, coords, results)
self.calc_counter += 1
return results
[docs]
def afir_fd_hessian_wrapper(self, coords3d, afir_grad_func):
coords = coords3d.flatten()
def grad_func(coords):
afir_grad = afir_grad_func(coords.reshape(-1, 3))
return afir_grad.flatten()
fd_hessian = finite_difference_hessian(coords, grad_func, acc=4)
return fd_hessian
[docs]
def get_hessian(self, atoms, coords, **prepare_kwargs):
self.set_atoms_and_funcs(atoms, coords)
true_results = self.calculator.get_hessian(atoms, coords, **prepare_kwargs)
true_energy = true_results["energy"]
true_hessian = true_results["hessian"]
coords3d = coords.reshape(-1, 3)
afir_energy = 0.0
afir_hessian = np.zeros_like(true_hessian)
# Iterate over all fragment pairs
for afir_func, afir_grad_func in zip(self.afir_funcs, self.afir_grad_funcs):
afir_energy += afir_func(coords3d)
# AFIR Hessian from finite differences
afir_hessian += self.afir_fd_hessian_wrapper(coords3d, afir_grad_func)
results = {
"energy": true_energy + afir_energy,
"hessian": true_hessian + afir_hessian,
"true_hessian": true_hessian,
"true_energy": true_energy,
}
if self.dump:
self.dump_h5(atoms, coords, results)
self.calc_counter += 1
return results
def __repr__(self):
return self.__str__()
def __str__(self):
return f"AFIR({self.calculator})"