mlmm-toolkit 0.2.2.dev0__py3-none-any.whl

pysisyphus/drivers/scan.py

@@ -0,0 +1,212 @@
+import h5py
+import numpy as np
+
+from pysisyphus.drivers.opt import run_opt
+from pysisyphus.helpers_pure import highlight_text
+from pysisyphus.optimizers.RFOptimizer import RFOptimizer
+from pysisyphus.TablePrinter import TablePrinter
+
+
+def relaxed_scan(
+    geom,
+    calc_getter,
+    constrain_prims,
+    target_values,
+    title,
+    max_cycles=25,
+    trust_radius=0.5,
+    thresh=1e-2,
+    dump=True,
+):
+    """Relaxed scan, allowing fixing of multiple primitive internals."""
+
+    # Constrain desired primitives
+    copy_kwargs = {
+        "coord_type": "redund",
+        "coord_kwargs": {"constrain_prims": constrain_prims},
+    }
+    scan_geom = geom.copy(**copy_kwargs)
+    calc = calc_getter()
+    scan_geom.set_calculator(calc)
+
+    lengths_ = [len(inds) for prim_type, *inds in constrain_prims]
+    if all([True for l in lengths_ if l == 2]):
+        unit = "au"
+    elif all([True for l in lengths_ if l in (3, 4)]):
+        unit = "rad"
+    else:
+        unit = "au (rad)"
+
+    typed_prims = scan_geom.internal.typed_prims
+    constr_inds = [typed_prims.index(cp) for cp in constrain_prims]
+    constr_fmt = lambda arr: f"({np.array2string(arr, precision=4)}) {unit}"
+
+    def print_constraints():
+        print(f"Desired constraints: {constr_fmt(target_values)}")
+        print(f" Actual constraints: {constr_fmt(scan_geom.coords[constr_inds])}")
+
+    trj = ""
+    scan_cart_coords = [scan_geom.cart_coords.copy()]
+    scan_energies = [scan_geom.energy]
+    actual_values = [scan_geom.coords[constr_inds]]
+
+    for i in range(max_cycles):
+        print(highlight_text(f"Step {i}", level=1))
+        cur_diff = target_values - scan_geom.coords[constr_inds]
+        step = cur_diff
+        step_norm = np.linalg.norm(step)
+        print_constraints()
+        print(f"         Difference: {step_norm:.6f} {unit}\n")
+        if step_norm <= thresh:
+            print(
+                f"Relaxed scan converged! Norm of proposed step <= {thresh:.4f} {unit}"
+            )
+            break
+        if step_norm > trust_radius:
+            step = trust_radius * step / step_norm
+        new_coords = scan_geom.coords.copy()
+        new_coords[constr_inds] += step
+        scan_geom.set_coords(new_coords, update_constraints=True)
+
+        prefix = f"{title}_step_{i}"
+        opt_kwargs = {
+            "prefix": prefix,
+            "h5_group_name": prefix,
+            "dump": dump,
+        }
+        opt = RFOptimizer(scan_geom, **opt_kwargs)
+        opt.run()
+
+        scan_cart_coords.append(scan_geom.cart_coords.copy())
+        scan_energies.append(scan_geom.energy)
+        actual_values.append(scan_geom.coords[constr_inds])
+        trj += scan_geom.as_xyz() + "\n"
+
+    scan_cart_coords = np.array(scan_cart_coords)
+    scan_energies = np.array(scan_energies)
+    actual_values = np.array(actual_values)
+
+    if dump:
+        trj_fn = f"{title}_scan_trj.xyz"
+        with open(trj_fn, "w") as handle:
+            handle.write(trj)
+        print(f"Dumped optimized geometries to '{trj_fn}'.")
+
+        group_name = f"scan_{title}"
+        # with h5py.File("scan.h5", "a") as handle:
+        #     try:
+        #         del handle[group_name]
+        #     except KeyError:
+        #         pass
+        #     group = handle.create_group(group_name)
+        #     group.create_dataset("energies", data=scan_energies)
+        #     group.create_dataset("cart_coords", data=scan_cart_coords)
+        #     group.create_dataset("target_values", data=target_values)
+        #     group.create_dataset("actual_values", data=actual_values)
+        #     dt = h5py.vlen_dtype(np.dtype("int32"))
+        #     cp = np.array([(pt.value, *ind) for pt, *ind in constrain_prims])
+        #     cp_dataset = group.create_dataset(
+        #         "constrain_prims", (len(target_values),), dtype=dt
+        #     )
+        #     cp_dataset[:] = cp
+    return scan_geom, scan_cart_coords, scan_energies
+
+
+def relaxed_1d_scan(
+    geom,
+    calc_getter,
+    constrain_prims,
+    start,
+    step_size,
+    steps,
+    opt_key,
+    opt_kwargs,
+    pref=None,
+    callback=None,
+):
+    assert geom.coord_type == "redund", "'coord_type: redund' is required!"
+    if pref is None:
+        pref = ""
+    else:
+        pref = f"{pref}_"
+
+    constr_prim = constrain_prims[0]
+    copy_kwargs = {
+        "coord_type": "redund",
+        "coord_kwargs": {
+            "constrain_prims": constrain_prims,
+            "recalc_B": True,
+        },
+    }
+    constr_geom = geom.copy(**copy_kwargs)
+    constr_ind = constr_geom.internal.get_index_of_typed_prim(constrain_prims[0])
+    calc = calc_getter()
+    # Always return same calculator, so orbitals can be reused etc.
+    def _calc_getter():
+        return calc
+
+    # Drop PrimType at index 0
+    unit = "au" if len(constr_prim[1:]) == 2 else "rad"
+
+    org_val = constr_geom.coords[constr_ind]
+    cur_val = start
+    end_val = start + step_size * steps
+    target_scan_vals = np.linspace(cur_val, end_val, steps + 1)
+    print(
+        f"    Coordinate: {constr_prim}\n"
+        f"Original value: {org_val:.4f} {unit}\n"
+        f" Initial value: {cur_val:.4f} {unit}\n"
+        f"   Final value: {end_val:.4f} {unit}\n"
+        f"         Steps: {steps}+1\n"
+        f"     Step size: {step_size:.4f} {unit}\n"
+    )
+    scan_vals = list()
+    scan_geoms = list()
+    scan_energies = list()
+    scan_xyzs = list()
+
+    for cycle, cur_val in enumerate(target_scan_vals):
+        opt_kwargs_ = opt_kwargs.copy()
+        name = f"{pref}relaxed_scan_{cycle:04d}"
+        opt_kwargs_["prefix"] = name
+        opt_kwargs_["h5_group_name"] = name
+        # Keep a copy
+        scan_geoms.append(constr_geom.copy())
+
+        # Update constrained coordinate
+        new_coords = constr_geom.coords
+        new_coords[constr_ind] = cur_val
+        constr_geom.set_coords(new_coords, update_constraints=True)
+
+        title = f"{pref}Step {cycle:02d}, coord={cur_val:.4f} {unit}"
+        opt_result = run_opt(
+            constr_geom, _calc_getter, opt_key, opt_kwargs_, title=title, level=1
+        )
+        if callback is not None:
+            callback(opt_result)
+        scan_xyzs.append(constr_geom.as_xyz())
+        scan_vals.append(constr_geom.coords[constr_ind])
+        scan_energies.append(constr_geom.energy)
+
+        if not opt_result.opt.is_converged:
+            print(f"Step {cycle} did not converge. Breaking!")
+            break
+
+    with open(f"{pref}relaxed_scan_trj.xyz", "w") as handle:
+        handle.write("\n".join(scan_xyzs))
+
+    scan_data = np.stack((scan_vals, scan_energies), axis=1)
+    np.savetxt(f"{pref}relaxed_scan.dat", scan_data)
+    scan_vals, scan_energies = scan_data.T
+
+    print(highlight_text(f"Scan summary", level=1))
+    col_fmts = "int float float float float".split()
+    header = ("Step", f"Target / {unit}", f"Actual / {unit}", f"Δ / {unit}", "E / au")
+    table = TablePrinter(header, col_fmts, width=14)
+    table.print_header()
+    for step, (target, act, en) in enumerate(
+        zip(target_scan_vals, scan_vals, scan_energies)
+    ):
+        delta = target - act
+        table.print_row((step, target, act, delta, en))
+    return scan_geoms, scan_vals, scan_energies

pysisyphus/drivers/spectrum.py

@@ -0,0 +1,166 @@
+# [1] https://gaussian.com/uvvisplot/
+
+import dataclasses
+from pathlib import Path
+from typing import Tuple
+
+import matplotlib.pyplot as plt
+import numpy as np
+from numpy.typing import NDArray
+
+from pysisyphus.constants import AU2J, C, M_E, NA, PLANCK, AU2EV
+
+
+# Computation of prefactor from Gaussian whitepaper
+Q_E_ESU = 4.803204e-10  # Crazy unit, cm**1.5 g**0.5 s⁻¹
+C_CM = C * 100  # Speed of light in cm/s
+M_E_G = M_E * 1000  # Electron mass in gram
+NM2CM = 1e7  # Nanometer to centimeter
+PREFACTOR = (  # Prefactor in eq. (5) of [1]
+    np.sqrt(np.pi) * NA * Q_E_ESU**2 / (1e3 * np.log(10) * C_CM**2 * M_E_G)
+) / NM2CM
+_04EV = 0.4 / AU2EV  # in Hartree
+# Factor used in converting energy in Hartree to wavelength
+_AU2NM = PLANCK * C * 1e9 / AU2J
+
+
+@dataclasses.dataclass
+class Spectrum:
+    exc_ens: NDArray[float]
+    exc_ens_nm: NDArray[float]
+    fosc: NDArray[float]
+    nm: NDArray[float]
+    epsilon: NDArray[float]
+
+    def plot(self, **kwargs):
+        plot_spectrum(self.nm, self.epsilon, self.exc_ens_nm, self.fosc, **kwargs)
+
+    def save(self, fn):
+        act_fn = Path(fn).with_suffix(".npz")
+        data = dataclasses.asdict(self)
+        np.savez(act_fn, **data)
+        return act_fn
+
+    @staticmethod
+    def load(fn):
+        handle = np.load(fn)
+        kwargs = {key: handle[key] for key in handle.files}
+        return Spectrum(**kwargs)
+
+    @staticmethod
+    def from_orca(wf_fn, cis_fn, log_fn, **kwargs):
+        from pysisyphus.calculators.ORCA import get_exc_ens_fosc
+
+        exc_ens, fosc = get_exc_ens_fosc(wf_fn, cis_fn, log_fn)
+        return spectrum_from_ens_fosc(exc_ens, fosc, **kwargs)
+
+
+def au2nm(au):
+    return _AU2NM / au
+
+
+def get_grid(resolution, exc_ens, padding, min_, max_):
+    from_ = min(min_, int(exc_ens.min() - padding))
+    to_ = max(max_, int(exc_ens.max() + padding))
+    return np.arange(from_, to_ + resolution, step=resolution)
+
+
+def homogeneous_broadening(
+    exc_ens: NDArray[float],
+    osc_strengths: NDArray[float],
+    resolution: float = 0.5,
+    stddev: float = _04EV,
+    from_to=None,
+) -> Tuple[NDArray[float], NDArray[float]]:
+    """Homogeneous broadening of stick spectra as outlined in Gaussian
+    whitepaper [1].
+
+    σ = 0.0147 au corresponds to about 0.4 eV. The function yields molar
+    extinction coefficients in l mol cm⁻¹."""
+    exc_ens_nm = au2nm(exc_ens)
+    if from_to is None:
+        from_to = np.array((exc_ens_nm[0], exc_ens_nm[-1]))
+    nm = get_grid(resolution, from_to, padding=100, min_=300, max_=900)
+    stddev_nm = au2nm(stddev)
+
+    quot = stddev_nm * (1 / nm[None, :] - (1 / exc_ens_nm[:, None]))
+    exp_ = np.exp(-(quot**2))
+    epsilon = PREFACTOR * osc_strengths[:, None] * stddev_nm * exp_
+    epsilon = epsilon.sum(axis=0)
+    return nm, epsilon
+
+
+def spectrum_from_ens_fosc(exc_ens, fosc, **kwargs):
+    exc_ens_nm = au2nm(exc_ens)
+    nm, epsilon = homogeneous_broadening(exc_ens, fosc, **kwargs)
+    spectrum = Spectrum(
+        exc_ens=exc_ens,
+        exc_ens_nm=exc_ens_nm,
+        fosc=fosc,
+        nm=nm,
+        epsilon=epsilon,
+    )
+    return spectrum
+
+
+def homogeneous_broadening_eV(
+    exc_ens: NDArray[float],
+    osc_strengths: NDArray[float],
+    resolution: float = 0.05,
+    stddev: float = 0.4,
+    from_to=None,
+) -> NDArray[float]:
+    """
+    Homogeneous broadening of stick spectra as outlined in eq. (5) and (6)
+    of https://doi.org/10.1063/1.4948471.
+
+    Parameters
+    ----------
+    exc_ens
+        Excitation energies in Hartree.
+    osc_strengths
+        Unitless oscillator strengths.
+    resolution
+        Resolution of broadened spectra in eV.
+    stddev
+        Standard deviation, sigma in eq. (6).
+    from_to
+        Limits for abscissa of spectrum in eV.
+
+    Returns
+    -------
+    eV
+        Spectral grid in eV.
+    epsilon
+        Molar extinction coefficient.
+    """
+    exc_ens_eV = exc_ens * AU2EV
+    if from_to is None:
+        from_to = np.array((exc_ens_eV[0], exc_ens_eV[-1]))
+    eV = get_grid(resolution, from_to, padding=2, min_=0, max_=5)
+    epsilon = (osc_strengths / (3.7922e33 * 4 * 2.926e-39 * np.sqrt(np.pi) * stddev))[
+        None, :
+    ] * np.exp(-(((eV[None, :] - exc_ens_eV) / stddev) ** 2))
+    epsilon = epsilon.sum(axis=0)
+    return eV, epsilon
+
+
+def plot_spectrum(nm, epsilon, exc_ens_nm=None, fosc=None, show=False):
+    fig, ax = plt.subplots()
+    ax.plot(nm, epsilon)
+    ax.set_ylabel(r"$\epsilon$")
+    ax.set_xlabel(r"$\lambda$ / nm")
+    axs = [
+        ax,
+    ]
+
+    if (exc_ens_nm is not None) and (fosc is not None):
+        ax2 = ax.twinx()
+        ax2.stem(exc_ens_nm, fosc, "r", markerfmt=" ", basefmt=" ")
+        ax2.set_ylabel("fosc")
+        axs.append(ax2)
+    fig.tight_layout()
+    if show:
+        plt.show()
+
+    return fig, axs

pysisyphus/drivers/thermo.py

@@ -0,0 +1,31 @@
+import argparse
+import sys
+
+from pysisyphus.config import p_DEFAULT, T_DEFAULT
+from pysisyphus.thermo import get_thermoanalysis_from_hess_h5, print_thermoanalysis
+
+
+def parse_args(args):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("hess_h5", help="HDF5 Hessian file from pysisyphus.")
+    parser.add_argument("-T", default=T_DEFAULT, type=float, help="Temperature")
+    parser.add_argument("-p", default=p_DEFAULT, type=float, help="Pressure")
+    parser.add_argument("--pg", default="c1", help="Point group.")
+    parser.add_argument(
+        "--calorie",
+        action="store_true",
+        help="Output in kcal mol⁻¹ instead of kJ mol⁻¹.",
+    )
+
+    return parser.parse_args(args)
+
+
+def run_thermo():
+    args = parse_args(sys.argv[1:])
+    hess_h5 = args.hess_h5
+    T = args.T
+    p = args.p
+    point_group = args.pg
+    unit = "calorie" if args.calorie else "joule"
+    thermo = get_thermoanalysis_from_hess_h5(hess_h5, T=T, p=p, point_group=point_group)
+    print_thermoanalysis(thermo, unit=unit)

pysisyphus/dynamics/Gaussian.py

@@ -0,0 +1,103 @@
+import numpy as np
+
+
+class DummyColvar:
+    def eval(self, x):
+        assert np.isscalar(
+            x
+        ), f"DummyColvar can only be used for scalar inputs, but got {x}!"
+        return x, 1
+
+
+class Gaussian:
+    def __init__(self, w=1, s=1, x0=None, colvar=None, dump_name=None):
+        """
+        See:
+            https://doi.org/10.1016/j.cpc.2018.02.017
+
+        V(f(x)) = w * exp(-(f(x) - f0)**2 / (2*s**2))
+
+        F(x) = -dV/dx = -dV/df * df/dx
+
+        See also:
+            https://doi.org/10.1103/PhysRevLett.90.238302
+        """
+        # Gaussian height w and standard deviation s are read-only properties.
+        # This way they can't be accidentally altered, which would invalidate the
+        # values precomputed below.
+        self._w = w
+        self._s = s
+        if x0 is None:
+            x0 = np.array(())
+        self.x0 = np.ravel(x0)
+
+        # Collective variable
+        if colvar is None:
+            colvar = DummyColvar()
+        self.colvar = colvar
+        self.dump_name = dump_name
+        if self.dump_name:
+            self.dump_fn = (
+                dump_name
+                if self.dump_name.endswith(".gau")
+                else f"{self.dump_name}.gau"
+            )
+        else:
+            self.dump_fn = None
+
+        # Reset file and write header
+        if self.dump_fn:
+            with open(self.dump_fn, "w") as handle:
+                handle.write(f"# {dump_name}\n# step s w center\n")
+
+        # Store some values, to avoid recalculation
+        self.s2 = self.s ** 2
+        self.one_over_2s2 = 1 / (2 * self.s2)
+        self.minus_w_over_s2 = -self.w / self.s2
+
+    @property
+    def w(self):
+        return self._w
+
+    @property
+    def s(self):
+        return self._s
+
+    def calculate(self, coords, x0=None, gradient=False):
+        """Return potential and gradient for Gaussian(s) centered at x0."""
+
+        if x0 is None:
+            x0 = self.x0
+        x, cr_grad = self.colvar.eval(coords)
+        # TODO: allow colvar to supply a callback, that overrides the difference
+        # below. This will be needed to support periodic colvars, like torsion.
+        # There we can't just calculate the naive difference.
+        diff = x - np.atleast_1d(x0)
+        exp_ = np.exp(-(diff ** 2) * self.one_over_2s2)
+        to_return = self.w * exp_.sum()
+
+        if gradient:
+            try:
+                grad = np.einsum("i,i,ijk->jk", diff, exp_, cr_grad[None, :, :])
+            except TypeError:
+                grad = diff * exp_ * cr_grad
+
+            # Finalize & append gradient
+            grad *= self.minus_w_over_s2
+            to_return = to_return, grad
+        return to_return
+
+    def value(self, coords, x0=None):
+        return self.calculate(coords, x0)
+
+    def gradient(self, coords, x0=None):
+        _, grad = self.calculate(coords, x0, gradient=True)
+        return grad
+
+    def eval(self, coords, x0=None):
+        return self.calculate(coords, x0, gradient=True)
+
+    def dump(self, step, s, w, center):
+        line = f"{step:08d} {s:.12f} {w:.12f} {center:.12f}\n"
+        with open(self.dump_fn, "a") as handle:
+            handle.write(line)

pysisyphus/dynamics/__init__.py

@@ -0,0 +1,20 @@
+import logging
+
+from pysisyphus.dynamics.colvars import get_colvar
+from pysisyphus.dynamics.Gaussian import Gaussian
+from pysisyphus.dynamics.helpers import get_mb_velocities_for_geom
+from pysisyphus.dynamics.mdp import mdp
+from pysisyphus.dynamics.rattle import rattle_closure
+from pysisyphus.dynamics.driver import md
+from pysisyphus.dynamics.wigner import get_wigner_sampler
+
+
+logger = logging.getLogger("dynamics")
+logger.setLevel(logging.DEBUG)
+# delay = True prevents creation of empty logfiles
+handler = logging.FileHandler("dynamics.log", mode="w", delay=True)
+# fmt_str = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+fmt_str = "%(asctime)s - %(message)s"
+formatter = logging.Formatter(fmt_str)
+handler.setFormatter(formatter)
+logger.addHandler(handler)

pysisyphus/dynamics/colvars.py

@@ -0,0 +1,136 @@
+import abc
+import types
+
+import numpy as np
+import autograd
+import autograd.numpy as anp
+
+from pysisyphus.intcoords.derivatives import dq_b, dq_a
+
+
+class Colvar(metaclass=abc.ABCMeta):
+    def __init__(self, force_agrad=False):
+        try:
+            getattr(self, "_gradient")
+        except AttributeError:
+            force_agrad = True
+
+        # Set autograd gradient method, if no _gradient is not implemented
+        # a autograd is forced.
+        if force_agrad:
+            grad_func = autograd.grad(self.value)
+
+            def wrapped(self, c3d):
+                return grad_func(c3d)
+
+            self._gradient = types.MethodType(wrapped, self)
+        # Store a flag to indicate use of autograd
+        self.agrad = force_agrad
+
+    @abc.abstractmethod
+    def value(self, c3d):
+        pass
+
+    def gradient(self, c3d):
+        return self._gradient(c3d)
+
+    def eval(self, coords):
+        c3d = coords.reshape(-1, 3)
+        return self.value(c3d), self.gradient(c3d)
+
+    def _wilson_gradient(self, func, c3d):
+        """Gradient of primitive internal w.r.t. Cartesians."""
+        grad = np.zeros_like(c3d)
+        grad[self.indices] = func(*c3d[self.indices].flatten()).reshape(-1, 3)
+        return grad
+
+    def __str__(self):
+        if hasattr(self, "indices"):
+            str_ = f"{self.__class__.__name__}({self.indices})"
+        else:
+            str_ = super().__str__()
+        return str_
+
+
+class CVDistance(Colvar):
+    def __init__(self, indices, **kwargs):
+        self.indices = list(indices)
+        self.i, self.j = self.indices
+        super().__init__(**kwargs)
+
+    def value(self, c3d):
+        return anp.linalg.norm(c3d[self.i] - c3d[self.j])
+
+    def _gradient(self, c3d):  # lgtm [py/attribute-shadows-method]
+        return self._wilson_gradient(dq_b, c3d)
+
+
+class CVBend(Colvar):
+    def __init__(self, indices, **kwargs):
+        self.indices = list(indices)
+        # Bonded like
+        # i---j <- central atom
+        #     |
+        #     k
+        self.i, self.j, self.k = self.indices
+        super().__init__(**kwargs)
+
+    def value(self, c3d):
+        u_dash = c3d[self.i] - c3d[self.j]
+        v_dash = c3d[self.k] - c3d[self.j]
+        u_norm = anp.linalg.norm(u_dash)
+        v_norm = anp.linalg.norm(v_dash)
+        u = u_dash / u_norm
+        v = v_dash / v_norm
+        rad = anp.arccos(anp.dot(u, v))
+        return rad
+
+    def _gradient(self, c3d):  # lgtm [py/attribute-shadows-method]
+        return self._wilson_gradient(dq_a, c3d)
+
+
+class CVTorsion(Colvar):
+    def __init__(self, indices, **kwargs):
+        self.indices = list(indices)
+        # Bonded like
+        # i--u--j
+        #       |
+        #       w
+        #       |
+        #       k--v--l
+        self.i, self.j, self.k, self.l = self.indices
+        super().__init__(**kwargs)
+
+    def value(self, c3d):
+        u_dash = c3d[self.i] - c3d[self.j]
+        v_dash = c3d[self.l] - c3d[self.k]
+        w_dash = c3d[self.k] - c3d[self.j]
+        u_norm = anp.linalg.norm(u_dash)
+        v_norm = anp.linalg.norm(v_dash)
+        w_norm = anp.linalg.norm(w_dash)
+        u = u_dash / u_norm
+        v = v_dash / v_norm
+        w = w_dash / w_norm
+        phi_u = anp.arccos(anp.dot(u, w))
+        phi_v = anp.arccos(-anp.dot(w, v))
+        uxw = anp.cross(u, w)
+        vxw = anp.cross(v, w)
+        cos_dihed = anp.dot(uxw, vxw)/(anp.sin(phi_u)*anp.sin(phi_v))
+        # Restrict cos_dihed to the allowed interval for arccos [-1, 1]
+        cos_dihed = min(1, max(cos_dihed, -1))
+
+        rad = anp.arccos(cos_dihed)
+        if (rad != np.pi) and (anp.dot(vxw, u) < 0):
+            rad *= -1
+        return rad
+
+
+COLVARS = {
+    "distance": CVDistance,
+    "bend": CVBend,
+    "torsion": CVTorsion,
+}
+
+
+def get_colvar(key, kwargs):
+    return COLVARS[key](**kwargs)