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

pysisyphus/dynamics/mdp.py

@@ -0,0 +1,364 @@
+# [1] https://doi.org/10.1063/1.5082885
+#     Minimum dynamic path
+#     Unke, 2019
+
+from collections import namedtuple
+import operator
+import re
+
+import numpy as np
+
+from pysisyphus.constants import AU2KJPERMOL
+from pysisyphus.dynamics.driver import md, MDResult
+from pysisyphus.dynamics.helpers import (
+    dump_coords,
+    get_mb_velocities_for_geom,
+    temperature_for_kinetic_energy,
+)
+from pysisyphus.helpers_pure import highlight_text
+
+
+def parse_raw_term_func(raw_term_func):
+    funcs = {
+        "<": operator.lt,
+        ">": operator.gt,
+        "<=": operator.le,
+        ">=": operator.ge,
+        "==": operator.eq,
+    }
+
+    def comp_closure(indices, op, ref_value):
+        a_ind, b_ind = indices
+        func = funcs[op]
+
+        def comp_func(coords3d):
+            a = coords3d[a_ind]
+            b = coords3d[b_ind]
+            dist = np.linalg.norm(a - b)
+            return func(dist, ref_value)
+
+        return comp_func
+
+    operator_re = re.compile("([<>=]+)")
+    mobj = operator_re.split(raw_term_func)
+    if mobj is None:
+        print(f"Could not parse term_func '{raw_term_func}!'")
+        return None
+    indices, op, ref_value = mobj
+    ref_value = float(ref_value)
+    indices = [int(ind) for ind in indices.split(",")]
+    return comp_closure(indices, op, ref_value)
+
+
+def parse_raw_term_funcs(raw_term_funcs):
+    comp_funcs = {}
+    for k, v in raw_term_funcs.items():
+        comp_func = parse_raw_term_func(v)
+        if comp_func:
+            comp_funcs[k] = comp_func
+    return comp_funcs
+
+
+def run_md(geom, dt, steps, v0=None, term_funcs=None, external=False):
+    if external and hasattr(geom.calculator, "run_md"):
+        t = dt * steps
+        t_ps = t * 1e-3
+        md_kwargs = {
+            "atoms": geom.atoms,
+            "coords": geom.coords,
+            "t": t,
+            "dt": dt,
+            "velocities": v0,
+            "dump": dt,
+        }
+        print("Running MD with external calculator implementation.")
+        if term_funcs is not None:
+            print("Termination functions are not supported in external MD!")
+        geoms = geom.calculator.run_md(**md_kwargs)
+        md_result = MDResult(
+            coords=[geom.coords for geom in geoms],
+            t_ps=t_ps,
+            step=int(t / dt - 1),
+            terminated=None,
+            T=None,
+            E_tot=None,
+        )
+    else:
+        md_kwargs = {
+            "v0": v0,
+            "steps": steps,
+            "dt": dt,
+            "term_funcs": term_funcs,
+            "verbose": False,
+            "remove_com_v": False,
+        }
+        print("Running MD with internal implementation.")
+        md_result = md(geom, **md_kwargs)
+
+    return md_result
+
+
+MDPResult = namedtuple(
+    "MDResult",
+    "ascent_xs md_init_plus md_init_minus "
+    "md_fin_plus md_fin_minus "
+    "md_fin_plus_term md_fin_minus_term",
+)
+
+
+def mdp(
+    geom,
+    steps,
+    dt,
+    term_funcs=None,
+    steps_init=None,
+    E_excess=0.0,
+    displ_length=0.1,
+    epsilon=5e-4,
+    ascent_alpha=0.05,
+    max_ascent_steps=25,
+    max_init_trajs=10,
+    dump=True,
+    seed=None,
+    external_md=False,
+):
+    # Sanity checks and forcing some types
+    dt = float(dt)
+    assert dt > 0.0
+    steps = int(steps)
+    t = dt * steps
+    # assert t > dt
+    if steps_init is None:
+        steps_init = steps // 10
+        print(f"No 'steps_init' provided! Using {steps_init}")
+    E_excess = float(E_excess)
+    assert E_excess >= 0.0
+    displ_length = float(displ_length)
+    assert displ_length >= 0.0
+    if term_funcs is None:
+        term_funcs = {}
+    for k, v in term_funcs.items():
+        if callable(v):
+            continue
+        elif isinstance(v, str):
+            term_funcs[k] = parse_raw_term_func(v)
+        else:
+            raise Exception(f"Invalid term function '{k}: {v}' encountered!")
+
+    print(highlight_text("Minimum dynamic path calculation"))
+
+    if seed is None:
+        # 2**32 - 1
+        seed = np.random.randint(4294967295)
+    np.random.seed(seed)
+    print(f"Using seed {seed} to initialize the random number generator.\n")
+
+    E_TS = geom.energy
+    E_tot = E_TS + E_excess
+    # Distribute E_excess evenly on E_pot and E_kin
+    E_pot_diff = 0.5 * E_excess
+    E_pot_desired = E_TS + E_pot_diff
+
+    print(f"E_TS={E_TS:.6f} au")
+
+    # Determine transition vector
+    w, v = np.linalg.eigh(geom.hessian)
+    assert w[0] < -1e-8
+    trans_vec = v[:, 0]
+
+    # Disable removal of translation/rotation for analytical potentials
+    remove_com_v = remove_rot_v = geom.cart_coords.size > 3
+
+    if E_excess == 0.0:
+        print("MDP without excess energy.")
+        # Without excess energy we have to do an initial displacement along
+        # the transition vector to get a non-vanishing gradient.
+        initial_displacement = displ_length * trans_vec
+        x0_plus = geom.coords + initial_displacement
+        x0_minus = geom.coords - initial_displacement
+
+        v0_zero = np.zeros_like(geom.coords)
+        md_kwargs = {
+            "v0": v0_zero.copy(),
+            "t": t,
+            "dt": dt,
+            "term_funcs": term_funcs,
+            "external": external_md,
+        }
+
+        geom.coords = x0_plus
+        md_fin_plus = run_md(geom, **md_kwargs)
+
+        geom.coords = x0_minus
+        md_fin_minus = run_md(geom, **md_kwargs)
+
+        if dump:
+            dump_coords(geom.atoms, md_fin_plus.coords, "mdp_plus_trj.xyz")
+            dump_coords(geom.atoms, md_fin_minus.coords, "mdp_minus_trj.xyz")
+
+        mdp_result = MDPResult(
+            ascent_xs=None,
+            md_init_plus=None,
+            md_init_minus=None,
+            md_fin_plus=md_fin_plus,
+            md_fin_minus=md_fin_minus,
+        )
+        return mdp_result
+
+    print(f"E_excess={E_excess:.6f} au, ({E_excess*AU2KJPERMOL:.1f} kJ/mol)")
+    print(f"E_pot,desired=E_TS + {E_pot_diff*AU2KJPERMOL:.1f} kJ/mol")
+    print()
+
+    # Generate random vector perpendicular to transition vector
+    perp_vec = np.random.rand(*trans_vec.shape)
+    # Zero last element if we have an analytical surface
+    if perp_vec.size == 3:
+        perp_vec[2] = 0
+    # Orthogonalize vector
+    perp_vec = perp_vec - (perp_vec @ trans_vec) * trans_vec
+    perp_vec /= np.linalg.norm(perp_vec)
+
+    # Initial displacement from x_TS to x, generating a point with
+    # non-vanishing gradient.
+    x = geom.coords + epsilon * perp_vec
+    geom.coords = x
+
+    # Do steepest ascent until E_tot is reached
+    E_pot = geom.energy
+    ascent_xs = list()
+    for i in range(max_ascent_steps):
+        ascent_xs.append(geom.coords.copy())
+        ascent_converged = E_pot >= E_pot_desired
+        if ascent_converged:
+            break
+        gradient = geom.gradient
+        E_pot = geom.energy
+
+        direction = gradient / np.linalg.norm(gradient)
+        step = ascent_alpha * direction
+        new_coords = geom.coords + step
+        geom.coords = new_coords
+
+    # calc = geom.calculator
+    # class Opt:
+    # pass
+    # _opt = Opt()
+    # _opt.coords = np.array(ascent_xs)
+    # calc.plot_opt(_opt, show=True)
+
+    assert ascent_converged, "Steepest ascent didn't converge!"
+    assert (E_tot - E_pot) > 0.0, (
+        "Potential energy after steepst ascent is greater than the desired "
+        f"total energy ({E_pot:.6f} > {E_tot:.6f}). Maybe try a smaller epsilon? "
+        f"The current value Ɛ={epsilon:.6f} may be too big!"
+    )
+
+    ascent_xs = np.array(ascent_xs)
+    if dump:
+        dump_coords(geom.atoms, ascent_xs, "mdp_ee_ascent_trj.xyz")
+    x0 = geom.coords.copy()
+
+    print(highlight_text("Runninig initialization trajectories", level=1))
+    for i in range(max_init_trajs):
+        # Determine random momentum vector for the given kinetic energy
+        E_kin = E_tot - E_pot
+        T = temperature_for_kinetic_energy(len(geom.atoms), E_kin)
+        v0 = get_mb_velocities_for_geom(
+            geom, T, remove_com_v=remove_com_v, remove_rot_v=remove_rot_v
+        ).flatten()
+
+        # Zero last element if we have an analytical surface
+        if v0.size == 3:
+            v0[2] = 0
+
+        # Run initial MD to check if both trajectories run towards different
+        # basins of attraction.
+
+        # First MD with positive v0
+        md_init_kwargs = {
+            "v0": v0.copy(),
+            "steps": steps_init,
+            "dt": dt,
+            "external": external_md,
+        }
+        geom.coords = x0.copy()
+        md_init_plus = run_md(geom, **md_init_kwargs)
+        # Second MD with negative v0
+        geom.coords = x0.copy()
+        md_init_kwargs["v0"] = -v0.copy()
+        md_init_minus = run_md(geom, **md_init_kwargs)
+
+        dump_coords(geom.atoms, md_init_plus.coords, f"mdp_ee_init_plus_{i:02d}_trj.xyz")
+        dump_coords(geom.atoms, md_init_minus.coords, f"mdp_ee_init_minus_{i:02d}_trj.xyz")
+
+        # Check if both MDs run into different basins of attraction.
+        # We (try to) do this by calculating the overlap between the
+        # transition vector and the normalized vector defined by the
+        # difference between x0 and the endpoint of the respective
+        # test trajectory. Both overlaps should have different sings.
+        end_plus = md_init_plus.coords[-1]
+        pls = end_plus - x0
+        pls /= np.linalg.norm(pls)
+        end_minus = md_init_minus.coords[-1]
+        minus = end_minus - x0
+        minus /= np.linalg.norm(minus)
+        p = trans_vec @ pls
+        m = trans_vec @ minus
+        init_trajs_converged = np.sign(p) != np.sign(m)
+
+        if init_trajs_converged:
+            print("Trajectories ran into different basins. Breaking.")
+            break
+    if dump:
+        dump_coords(geom.atoms, md_init_plus.coords, "mdp_ee_init_plus_trj.xyz")
+        dump_coords(geom.atoms, md_init_minus.coords, "mdp_ee_init_minus_trj.xyz")
+    assert init_trajs_converged
+    print(f"Ran 2*{i+1} initialization trajectories.")
+    print()
+
+    # Run actual trajectories, using the supplied termination functions if possible.
+    print(highlight_text("Running actual full trajectories.", level=1))
+
+    def print_status(terminated, step):
+        if terminated:
+            msg = f"\tTerminated by '{terminated}' in step {step}."
+        else:
+            msg = "\tMax time steps reached!"
+        print(msg)
+
+    # "Production"/Final MDs
+    md_fin_kwargs = {
+        "v0": v0.copy(),
+        "steps": steps,
+        "dt": dt,
+        "term_funcs": term_funcs,
+        "external": external_md,
+    }
+    # MD with positive v0.
+    geom.coords = x0.copy()
+    md_fin_plus = run_md(geom, **md_fin_kwargs)
+    print_status(md_fin_plus.terminated, md_fin_plus.step)
+
+    # MD with negative v0.
+    geom.coords = x0.copy()
+    md_fin_kwargs["v0"] = -v0
+    md_fin_minus = run_md(geom, **md_fin_kwargs)
+    print_status(md_fin_minus.terminated, md_fin_minus.step)
+
+    md_fin_plus_term = md_fin_plus.terminated
+    md_fin_minus_term = md_fin_minus.terminated
+
+    if dump:
+        dump_coords(geom.atoms, md_fin_plus.coords, "mdp_ee_fin_plus_trj.xyz")
+        dump_coords(geom.atoms, md_fin_minus.coords, "mdp_ee_fin_minus_trj.xyz")
+
+    mdp_result = MDPResult(
+        ascent_xs=ascent_xs,
+        md_init_plus=md_init_plus,
+        md_init_minus=md_init_minus,
+        md_fin_plus=md_fin_plus,
+        md_fin_minus=md_fin_minus,
+        md_fin_plus_term=md_fin_plus_term,
+        md_fin_minus_term=md_fin_minus_term,
+    )
+    return mdp_result

pysisyphus/dynamics/rattle.py

@@ -0,0 +1,121 @@
+import numpy as np
+
+from pysisyphus.constants import FORCE2ACC
+from pysisyphus.dynamics.helpers import energy_forces_getter_closure, remove_com_velocity
+
+
+def rattle_closure(geom, constraints, dt, tol=1e-3, max_cycles=25,
+                   energy_forces_getter=None, remove_com_v=True):
+    # Inverse atom masses
+    masses = geom.masses
+    inv_masses = 1 / masses
+    inv_masses_rep = np.repeat(inv_masses, 3)
+
+    dt2 = 0.5 * dt
+    dt2_sq = dt2 * dt
+    tol_sq = tol**2
+
+    constraints = np.array(constraints, dtype=int)
+    inv_masses_sums = inv_masses[constraints[:, 0]] + inv_masses[constraints[:, 1]]
+    coords3d = geom.coords3d
+
+    def get_bond_vecs(coords3d):
+        return coords3d[constraints[:, 0]] - coords3d[constraints[:, 1]]
+
+    # Original bond lengths
+    lengths = np.linalg.norm(get_bond_vecs(coords3d), axis=1)
+    lengths_sq = lengths**2
+
+    if energy_forces_getter is None:
+        energy_forces_getter = energy_forces_getter_closure(geom)
+        # def forces_getter(coords):
+            # return geom.get_energy_and_forces_at(coords)["forces"]
+
+    def rattle(coords, velocities, forces):
+        acceleration = forces * inv_masses_rep * FORCE2ACC
+        # Bond vectors at time t
+        bond_vecs = get_bond_vecs(coords.reshape(-1, 3))
+
+        # Unconstrained position update
+        coords3d_updated = (coords + dt*velocities + dt2_sq * acceleration).reshape(-1, 3)
+        # Unconstrained velocity update
+        velocities3d_updated = (velocities + dt2 * acceleration).reshape(-1, 3)
+        if remove_com_v:
+            velocities3d_updated = remove_com_velocity(velocities3d_updated, masses)
+
+        # First part of RATTLE algorithm.
+        #   Yields updated positions for t+dt and half-updated velocities for t+dt/2
+        for i in range(max_cycles):
+            corrected = False
+            for j, (atom_1, atom_2) in enumerate(constraints):
+                bond_vec_updated = coords3d_updated[atom_1] - coords3d_updated[atom_2]
+                bond_length_sq = bond_vec_updated.dot(bond_vec_updated)
+                ref_length_sq = lengths_sq[j]
+                diff_sq = ref_length_sq - bond_length_sq
+
+                # We have to correct coordinates and velocities if the deviations is
+                # above the tolerance.
+                if abs(diff_sq) <= (ref_length_sq * tol_sq):
+                    continue
+
+                corrected = True
+                # Otherwise try to satisfy the constraint by calculating the
+                # approximate lagrange multiplier 'g'.
+                dot = bond_vec_updated.dot(bond_vecs[j])
+                # TODO: test for constraint failure
+                g = diff_sq / (2 * inv_masses_sums[j] * dot)
+                
+                # Update positions to satify constraint
+                atom_1_factor = g * bond_vecs[j] * inv_masses[atom_1]
+                atom_2_factor = g * bond_vecs[j] * inv_masses[atom_2]
+                coords3d_updated[atom_1] += atom_1_factor
+                coords3d_updated[atom_2] -= atom_2_factor
+                # Update velocities to satify constraint
+                velocities3d_updated[atom_1] += atom_1_factor / dt
+                velocities3d_updated[atom_2] -= atom_2_factor / dt
+
+            # Stop macro-iterations when no correction was done
+            if not corrected:
+                # print(f"RATTLE_1 finished after {i+1} macro cycle(s)!")
+                break
+        else:
+            raise Exception("First part of RATTLE did not converge!")
+
+        # Calculate energy (E_pot) and forces at updated coordinates
+        energy_new, forces_new = energy_forces_getter(coords3d_updated.flatten())
+        velocities3d_updated += dt2 * (forces_new * inv_masses_rep * FORCE2ACC).reshape(-1, 3)
+        if remove_com_v:
+            velocities3d_updated = remove_com_velocity(velocities3d_updated, masses)
+        # The coordinates aren't updated anymore in the second part, so we can
+        # calculate the bond vectors here once and store them.
+        bond_vecs = get_bond_vecs(coords3d_updated)
+
+        # Second part of RATTLE algorithm
+        #   Update velocities to full time step t+dt
+        for i in range(max_cycles):
+            corrected = False
+            for j, (atom_1, atom_2) in enumerate(constraints):
+                velocity_diff = (velocities3d_updated[atom_1]
+                                 - velocities3d_updated[atom_2])
+                dot = bond_vecs[j].dot(velocity_diff)
+                # Approximate Lagrange multiplier
+                k = dot / (inv_masses_sums[j] * lengths_sq[j])
+
+                if abs(k) <= tol_sq:
+                    continue
+
+                corrected = True
+                # Update velocities to satify constraint
+                velocities3d_updated[atom_1] -= k * bond_vecs[j] * inv_masses[atom_1]
+                velocities3d_updated[atom_2] += k * bond_vecs[j] * inv_masses[atom_2]
+
+            # Stop macro-iterations when no correction was done
+            if not corrected:
+                # print(f"RATTLE_2 finished after {i+1} macro cycle(s)!")
+                break
+        else:
+            raise Exception("Second part of RATTLE did not converge!")
+
+        return coords3d_updated.flatten(), velocities3d_updated.flatten(), energy_new, forces_new
+
+    return rattle

pysisyphus/dynamics/thermostats.py

@@ -0,0 +1,128 @@
+from math import exp, sqrt
+
+import numpy as np
+from numpy.random import default_rng
+
+from pysisyphus.constants import KBAU
+
+
+"""
+[1] https://aip.scitation.org/doi/10.1063/1.2408420
+[2] https://dx.doi.org/10.1016/j.cpc.2008.01.006
+    Reformulation fo the algorithm. This is implemented e.g., in YAFF.
+    https://github.com/molmod/yaff/blob/master/yaff/sampling/nvt.py
+
+csvr_closure() is based on the implementation provided on Bussis homepage:
+    https://sites.google.com/site/giovannibussi/downloads/resamplekin.tgz
+    (At least in my implementation) there seems to be a problem with
+    the conserved quantity, which is not conserved at all ...
+csvr_closure_2() is based on [2]
+"""
+
+
+RNG = default_rng()
+
+
+def sum_noises(num, rng=None):
+    """
+    Parameters
+    ----------
+    num : int
+        Number of independent Gaussian noises to be squared.
+    rng : numpy.random.Generator, optional
+        Instances of a random number generator (RNG). If it is not provided the module-level
+        RNG will be used.
+    """
+
+    if rng is None:
+        rng = RNG
+
+    if num == 0:
+        sum_ = 0.0
+    elif num == 1:
+        sum_ = rng.normal()**2
+    # nn even, dof - 1 odd
+    elif (num % 2) == 0:
+        sum_ = 2.0 * rng.gamma(shape=num/2)
+    # nn odd, dof - 1 even
+    else:
+        sum_ = 2.0 * rng.gamma(shape=(num-1)/2) + rng.normal()**2
+
+    return sum_
+
+
+def csvr_closure(sigma, dof, dt, tau=100, rng=None):
+    """
+    Parameters
+    ----------
+    sigma : float
+        Target average value of the kinetic energy (1/2 dof k_b T) in the same units as
+        cur_kinetic_energy.
+    dof : int
+        Degrees of freedom.
+    tau : float
+        Timeconstant of the thermostat.    tau : float
+        Timeconstant of the thermostat.
+    rng : numpy.random.Generator, optional
+        Instances of a random number generator (RNG). If it is not provided the module-level
+        RNG will be used.
+    """
+    # Relaxation time of the thermostat in units of "how often this routine
+    # is called" (dt / timeconstant).
+    tau_t = dt / tau
+
+    if tau_t > 0.1:
+        factor = exp(-1.0 / tau_t)
+    else:
+        factor = 0.0
+
+    if rng is None:
+        rng = RNG
+
+    def resample_kin(cur_kinetic_energy):
+        """
+        Parameters
+        ----------
+        cur_kinetic_energy : float
+            Present value of the kinetic energy of the atoms to be thermalized
+            in arbitrary units.
+        """
+        rr = rng.normal()
+        new_kinetic_energy = (
+            cur_kinetic_energy
+            + (1.0 - factor)
+            * (sigma * (sum_noises(dof-1) + rr**2) / dof - cur_kinetic_energy)
+            + 2.0 * rr * sqrt(cur_kinetic_energy * sigma / dof * (1.0 - factor) * factor)
+        )
+        alpha = sqrt(new_kinetic_energy / cur_kinetic_energy)
+        return alpha
+    return resample_kin
+
+
+def csvr_closure_2(sigma, dof, dt, tau=100, rng=None):
+    if rng is None:
+        rng = RNG
+    c = exp(-dt / tau)
+
+    def resample_kin(cur_kinetic_energy):
+        """Canonical stocastical velocity rescaling.
+
+        See dx.doi.org/10.1016/j.cpc.2008.01.006
+        """
+        R = rng.normal()
+        S = np.sum(rng.normal(size=dof-1)**2)
+        quot = (1 - c) * sigma / (dof * cur_kinetic_energy)
+        alpha =  sqrt(c + quot * (S + R**2) + 2 * R * sqrt(c*quot))
+        sign = np.sign(R + sqrt(c / quot))
+        return sign * alpha
+    return resample_kin
+
+
+def berendsen_closure(sigma, dof, dt, tau=100, rng=None):
+    """ https://doi.org/10.1063/1.448118"""
+    tau_t = dt / tau
+
+    def resample_kin(cur_kinetic_energy):
+      alpha = sqrt(1 + tau_t * (sigma / cur_kinetic_energy - 1))
+      return alpha
+    return resample_kin