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

pysisyphus/dynamics/driver.py

@@ -0,0 +1,297 @@
+from collections import namedtuple
+import logging
+
+import numpy as np
+
+from pysisyphus.config import T_DEFAULT
+from pysisyphus.constants import FORCE2ACC
+from pysisyphus.dynamics.helpers import (
+    kinetic_energy_from_velocities,
+    temperature_for_kinetic_energy,
+    energy_forces_getter_closure,
+    kinetic_energy_for_temperature,
+    remove_com_velocity,
+)
+from pysisyphus.dynamics.thermostats import csvr_closure, csvr_closure_2, berendsen_closure
+from pysisyphus.dynamics.rattle import rattle_closure
+from pysisyphus.helpers import check_for_end_sign
+from pysisyphus.helpers_pure import log
+from pysisyphus.io.hdf5 import get_h5_group
+
+logger = logging.getLogger("dynamics")
+
+MDResult = namedtuple(
+    "MDResult",
+    "coords t_ps step terminated T E_tot",
+)
+
+THERMOSTATS = {
+    "csvr": csvr_closure,
+    "csvr_2": csvr_closure_2,
+    "berendsen": berendsen_closure,
+}
+
+
+def get_data_model(atoms, dump_steps):
+    coord_size = 3 * len(atoms)
+    _1d = (dump_steps,)
+    _2d = (dump_steps, coord_size)
+
+    data_model = {
+        "cart_coords": _2d,
+        "step": _1d,
+        "energy_tot": _1d,
+        "energy_pot" : _1d,
+        "energy_kin": _1d,
+        "energy_conserved": _1d,
+        "T": _1d,
+        "T_avg": _1d,
+        "velocity": _2d,
+    }
+
+    return data_model
+
+
+def md(
+    geom,
+    v0,
+    steps,
+    dt,
+    remove_com_v=True,
+    thermostat=None,
+    T=T_DEFAULT,
+    timecon=100,
+    term_funcs=None,
+    constraints=None,
+    constraint_kwargs=None,
+    gaussians=None,
+    verbose=True,
+    print_stride=50,
+    dump_stride=None,
+    h5_group_name="run",
+):
+    """Velocity verlet integrator.
+
+    Parameters
+    ----------
+    geom : Geometry
+        The system for which the dynamics are to be run.
+    v0 : np.array, floats
+        Initial velocities in Bohr/fs.
+    steps : float
+        Number of simulation steps.
+    dt : float
+        Timestep in fs.
+    remove_com_v : bool, default=True
+        Remove center-of-mass velocity.
+    thermostat : str, optional, default None
+        Which and whether to use a thermostat.
+    T : float, optional, default=None
+        Desired temperature in thermostated runs.
+    timecon : float
+        Timeconsanst of the thermostat in fs.
+    term_funcs : dict, optional
+        Iterable of functions that are called with the atomic
+        coordinates in every MD cycle and result in termination
+    constraints : 2d iterable, optional
+        2D iterable containing atom indices describing constrained
+        bond lengths.
+        of the MD integration when they evaluate to true.
+    constraint_kwargs : dict, optional
+        Keyword arguments for the constraint algorithm.
+    gaussians: list, optional, default=None
+        List of Gaussians to be used in a metadynamics run.
+    verbose : bool, default=True
+        Do additional printing when True.
+    print_stride : int, default=50
+        Report every n-th step.
+    dump_stride : int, default=None
+        If given, MD data will be dumped to a HDF5 file every n-th step.
+    h5_group_name = str, optional
+        Name of the HDF5 group, used for dumping.
+    """
+
+    assert geom.coord_type == "cart"
+
+    if term_funcs is None:
+        term_funcs = dict()
+
+    # if dump_stride:
+    #     dump_steps = steps // dump_stride
+    #     data_model = get_data_model(geom.atoms, dump_steps)
+    #     h5_group = get_h5_group("md.h5", h5_group_name, data_model, reset=True)
+    #     h5_group.attrs["atoms"] = geom.atoms
+    #     h5_group.attrs["dt"] = dt
+    #     h5_group.attrs["steps"] = steps
+    #     h5_group.attrs["masses"] = geom.masses
+    #     h5_group.attrs["T_target"] = T
+
+    if verbose:
+        t_ps = steps * dt * 1e-3  # Total simulation time
+        print(f"Doing {steps} steps of {dt:.4f} fs for a total of {t_ps:.2f} ps.\n")
+
+    energy_forces_getter = energy_forces_getter_closure(geom)
+
+    if gaussians is None:
+        gaussians = tuple()
+    gau_centers = list()
+    gau_center_num = np.zeros(len(gaussians), dtype=int)
+    for _, gaussian, gau_stride in gaussians:
+        num_centers = steps // gau_stride
+        gau_centers.append(np.zeros(num_centers))
+
+    def update_gaussians(step, coords):
+        for i, (gau_key, gaussian, gau_stride) in enumerate(gaussians):
+            if (step + 1) % gau_stride == 0:
+                new_center = gaussian.colvar.value(coords.reshape(-1, 3))
+                gau_centers[i][gau_center_num[i]] = new_center
+                gau_center_num[i] += 1
+                log(
+                    logger,
+                    f"Added {gau_center_num[i]: >6d}. '{gau_key}' Gaussian in step {step}.",
+                )
+                gaussian.dump(step, gaussian.s, gaussian.w, new_center)
+
+    def gaussian_wrapper(coords):
+        E_pot, forces = energy_forces_getter(geom.coords)
+        for i, (_, gaussian, _) in enumerate(gaussians):
+            center_num = gau_center_num[i]
+            gau_pot, gau_grad = gaussian.eval(coords, gau_centers[i][:center_num])
+            # print("\tgau_pot", gau_pot)
+            gau_forces = -gau_grad.flatten()
+            E_pot += gau_pot
+            forces += gau_forces
+        return E_pot, forces
+
+    if constraint_kwargs is None:
+        constraint_kwargs = dict()
+
+    # Fixed degrees of freedom
+    fixed_dof = 0
+    if remove_com_v:
+        print("Removing center-of-mass velocity.")
+        fixed_dof += 3
+    constrained_md = constraints is not None
+    # Get RATTLE function from closure for constrained MD
+    if constrained_md:
+        fixed_dof += len(constraints)
+        rattle = rattle_closure(
+            geom,
+            constraints,
+            dt,
+            energy_forces_getter=energy_forces_getter,
+            **constraint_kwargs,
+        )
+    print(f"Fixed degrees of freedom: {fixed_dof}")
+    dof = len(geom.coords) - fixed_dof
+
+    if thermostat is not None:
+        sigma = kinetic_energy_for_temperature(len(geom.atoms), T, fixed_dof=fixed_dof)
+        thermo_func = THERMOSTATS[thermostat](sigma, dof, dt=dt, tau=timecon)
+    # In amu
+    masses = geom.masses
+    masses_rep = geom.masses_rep
+
+    x = geom.cart_coords
+    # v is given in Bohr/fs
+    v = v0
+    a_prev = np.zeros_like(x)
+    xs = list()
+    Ts = list()
+    E_tots = list()
+
+    E_pot, forces = energy_forces_getter(geom.coords)
+
+    t_cur = 0
+    terminate = False
+    terminate_key = None
+    T_avg = 0
+    log(logger, f"Running MD with Δt={dt:.2f} fs for {steps} steps.")
+    print()
+    thermo_corr = 0.0
+    for step in range(steps):
+        xs.append(x.copy())
+
+        E_kin = kinetic_energy_from_velocities(masses, v.reshape(-1, 3))
+        T = temperature_for_kinetic_energy(len(masses), E_kin, fixed_dof=fixed_dof)
+        T_avg += T
+        Ts.append(T)
+        E_tot = E_pot + E_kin
+        E_tots.append(E_tot)
+        E_conserved = E_pot + E_kin - thermo_corr
+
+        status_msg = (
+            f"Step {step:06d}  {t_cur*1e-3: >6.2f} ps  E={E_tot: >8.6f} E_h  "
+            f"T={T: >8.2f} K <T>={T_avg/(step+1): >8.2f} K"
+        )
+        if step % print_stride == 0:
+            log(logger, status_msg)
+            if verbose:
+                print(status_msg)
+
+        # if dump_stride and (step % dump_stride == 0):
+        #     ind = step // dump_stride
+        #     h5_group["step"][ind] = step
+        #     h5_group["cart_coords"][ind] = x
+        #     h5_group["velocity"][ind] = v
+        #     h5_group["energy_tot"][ind] = E_tot
+        #     h5_group["energy_pot"][ind] = E_pot
+        #     h5_group["energy_kin"][ind] = E_kin
+        #     h5_group["energy_conserved"][ind] = E_conserved
+        #     h5_group["T"][ind] = T
+        #     h5_group["T_avg"][ind] = T_avg / (step + 1)
+
+        if thermostat:
+            alpha = thermo_func(E_kin)
+            thermo_corr += (alpha**2 - 1) * E_kin
+            v *= alpha
+
+        update_gaussians(step, x)
+
+        # RATTLE algorithm
+        if constrained_md:
+            x, v, E_pot, forces = rattle(x, v, forces)
+        # Simple Velocity-Verlet integration
+        else:
+            # E_pot, forces = energy_forces_getter(geom.coords)
+            E_pot, forces = gaussian_wrapper(geom.coords)
+            # Acceleration, convert from Hartree / (Bohr * amu) to Bohr/fs²
+            a = forces / masses_rep * FORCE2ACC
+            v += 0.5 * (a + a_prev) * dt
+
+            if remove_com_v:
+                v = remove_com_velocity(v.reshape(-1, 3), masses).flatten()
+            # v*dt = Bohr/fs * fs -> Bohr
+            # a*dt**2 = Bohr/fs² * fs² -> Bohr
+            x += v * dt + 0.5 * a * dt ** 2
+            a_prev = a
+
+        # Update coordinates
+        geom.coords = x
+
+        for name, func in term_funcs.items():
+            if func(x.reshape(-1, 3)):
+                terminate = True
+                terminate_key = name
+                break
+        if terminate:
+            log(logger, f"Termination function '{name}' evaluted to True. Breaking.")
+            break
+
+        if check_for_end_sign():
+            break
+
+        # Advance time
+        t_cur += dt
+    log(logger, "")
+
+    md_result = MDResult(
+        coords=np.array(xs),
+        t_ps=t_cur * 1e-3,
+        step=step,
+        terminated=terminate_key,
+        T=np.array(Ts),
+        E_tot=np.array(E_tots),
+    )
+
+    return md_result

pysisyphus/dynamics/helpers.py

@@ -0,0 +1,256 @@
+import numpy as np
+
+from pysisyphus.xyzloader import make_trj_str
+from pysisyphus.Geometry import get_trans_rot_projector
+from pysisyphus.constants import BOHR2ANG, KBAU, VELO2E
+
+
+def dump_coords(atoms, coords, trj_fn):
+    coords = np.array(coords)
+    coords = coords.reshape(-1, len(atoms), 3) * BOHR2ANG
+    trj_str = make_trj_str(atoms, coords)
+
+    with open(trj_fn, "w") as handle:
+        handle.write(trj_str)
+
+
+def kinetic_energy_from_velocities(masses, velocities):
+    """Kinetic energy for given velocities and masses.
+
+    Parameters
+    ----------
+    masses : 1d array, shape (number of atoms, )
+        Atomic masses in amu.
+    velocities : 2d array, (number of atoms, 3)
+        Atomic velocities in Bohr/fs.
+
+    Returns
+    -------
+    E_kin : float
+        Kinetic energy in Hartree.
+    """
+    return np.sum(masses[:,None] / 2 * velocities**2) * VELO2E
+
+
+def kinetic_energy_for_temperature(atom_number, T, fixed_dof=0):
+    """Kinetic energy for given temperature and number of atoms.
+
+    Each atom has three degrees of freedom (1/2 * 3 == 3/2).
+
+    Parameters
+    ----------
+    atom_number : int
+        Number of atoms. Each atom has three degrees of freedom.
+    T : float
+        Temperature in Kelvin.
+    fixed_dof : int, optional, default=0
+        Number of fixed degrees of freedom, e.g. 3 when the center-of-mass
+        velocity is removed.
+
+    Returns
+    -------
+    E_kin : float
+        Kinetic energy in Hartree.
+    """
+    return (3*atom_number - fixed_dof) / 2 * T * KBAU
+
+
+def temperature_for_kinetic_energy(atom_number, E_kin, fixed_dof=0):
+    """Temperature for given kinetic energy and atom number.
+
+    Each atom has three degrees of freedom (1/2 * 3 == 3/2).
+
+    Parameters
+    ----------
+    atom_number : int
+        Number of atoms. Each atom has three degrees of freedom.
+    E_kin : float
+        Kinetic energy in Hartree.
+    fixed_dof : int, optional, default=0
+        Number of fixed degrees of freedom, e.g. 3 when the center-of-mass
+        velocity is removed.
+
+    Returns
+    -------
+    Temperature : float
+        Temperature in Kelvin.
+    """
+    return 2 * E_kin / ((3*atom_number - fixed_dof) * KBAU)
+
+
+def remove_com_velocity(v, masses, keep_norm=True):
+    """Remove center-of-mass velocity.
+
+    Returned units vary with the given input units.
+
+    Parameters
+    ----------
+    v : np.array, 2d, shape (number of atoms, 3)
+        Velocities.
+    masses : np.array, 1d, shape (number of atoms, )
+        Atomic masses.
+    keep_norm : bool, default=True
+        Whether to rescale v to its original norm, after removal of v_com.
+
+    Returns
+    -------
+    v : np.array
+        Velocities without center-of-mass velocity.
+    """
+    org_norm = np.linalg.norm(v)
+    v_com = (v * masses[:,None] / np.sum(masses)).sum(axis=0)
+    v -= v_com
+    if keep_norm:
+        v /= np.linalg.norm(v) / org_norm
+    return v
+
+
+def scale_velocities_to_temperatue(masses, v, T_desired, fixed_dof=0):
+    """Scale velocities to a given temperature.
+
+    Parameters
+    ----------
+    masses : np.array, 1d, shape (number of atoms, )
+        Atomic masses in amu.
+    v : np.array, 2d, shape (number of atoms, 3)
+        (Unscaled) velocities in Bohr/fs.
+    T_desired : float
+        Desired temperature in Kelvin.
+    fixed_dof : int, optional, default=0
+        Number of fixed degrees of freedom, e.g. 3 when the center-of-mass
+        velocity is removed.
+
+    Returns
+    -------
+    v : np.array, 2d, shape (number of atoms, 3)
+        Scaled velocities in Bohr/fs.
+    """
+
+    E_ref = kinetic_energy_for_temperature(len(masses), T_desired,
+                                           fixed_dof=fixed_dof)  # in Hartree
+    E_cur = kinetic_energy_from_velocities(masses, v)  # in Hartree
+    scale = (E_ref / E_cur)**0.5
+    v *= scale
+    # E_now = kinetic_energy_from_velocities(masses, v)
+    # np.testing.assert_allclose(E_now, E_ref)
+    return v
+
+
+def scale_velocities_to_energy(masses, v, E_desired):
+    """Scale velocities to a given temperature.
+
+    Parameters
+    ----------
+    masses : np.array, 1d, shape (number of atoms, )
+        Atomic masses in amu.
+    v : np.array, 2d, shape (number of atoms, 3)
+        (Unscaled) velocities in Bohr/fs.
+    E_desired : float
+        Desired kinetic energy in Hartree.
+
+    Returns
+    -------
+    v : np.array, 2d, shape (number of atoms, 3)
+        Scaled velocities in Bohr/fs.
+    """
+
+    E_cur = kinetic_energy_from_velocities(masses, v)  # in Hartree
+    scale = (E_desired / E_cur)**0.5
+    v *= scale
+    return v
+
+
+def unscaled_velocity_distribution(masses, T, seed=None):
+    """
+    ρ ∝ exp(- v² * m / (2 * kT))
+    ln(ρ) ∝ - 1/2 * v² * m / kT
+    v² ∝ -2 * ln(ρ) * kT / m
+    v ∝ ±(-2 * ln(ρ) * kT / m)**0.5
+    v ∝ ±(-2 * k)**0.5 * (ln(ρ) * T / m)**0.5
+
+    The first term of the RHS is constant. As we later scale the
+    velocities we neglect it. Don't use these velocities unscaled!
+
+    v ∝ ±(ln(ρ) * T / m)**0.5
+    """
+
+    if seed is not None:
+        np.random.seed(seed)
+
+    ps = np.random.standard_normal((len(masses), 3))
+    v = ps * np.sqrt(T / masses[:,None])
+    return v
+
+
+def get_mb_velocities(masses, cart_coords, T, remove_com_v=True, remove_rot_v=True,
+                      seed=None):
+    """Initial velocities from Maxwell-Boltzmann distribution.
+
+    Parameters
+    ----------
+    masses : np.array, 1d, shape (number of atoms, )
+        Atomic masses in amu.
+    cart_coords : iterable, 1d, shape (3 * number of atoms, )
+        Atomic cartesian coordinates. Needed for removal of rotation.
+    T : float
+        Temperature in Kelvin.
+    remove_com_v : bool, default=True, optional
+        Whether to remove center-of-mass velocity.
+    remove_rot_v : bool, default=True, optional
+        Whether to remove rotational velocity.
+    seed : int, default=None, optional
+        Seed for the random-number-generator.
+
+    Returns
+    -------
+    v : np.array, 2d, shape (number of atoms, 3)
+        Initial velocities in Bohr/fs.
+    """
+
+    masses = np.array(masses)
+
+    # Initial velocities
+    v = unscaled_velocity_distribution(masses, T, seed=seed)
+
+    if (len(masses) == 1) and remove_com_v:
+        raise Exception("Removing COM velocity with only 1 atom is a bad idea!")
+
+    fixed_dof = 0
+    # Remove center-of-mass velocity
+    if remove_com_v:
+        v = remove_com_velocity(v, masses)
+        fixed_dof += 3
+
+    if remove_rot_v:
+        # Right now this also removes the translational components
+        P = get_trans_rot_projector(cart_coords, masses, full=True)
+        v = P.dot(v.flatten()).reshape(-1, 3)
+        fixed_dof = 6
+
+    # In Bohr/fs
+    v = scale_velocities_to_temperatue(masses, v, T, fixed_dof=fixed_dof)
+
+    return v
+
+
+def get_mb_velocities_for_geom(geom, T, remove_com_v=True, remove_rot_v=True,
+                               seed=None):
+    """Initial velocities from Maxwell-Boltzmann distribution.
+
+    See 'get_mb_velocities' for explanation.
+    """
+
+    return get_mb_velocities(geom.masses, geom.cart_coords, T,
+                             remove_com_v=remove_com_v,
+                             remove_rot_v=remove_rot_v,
+                             seed=seed,
+    )
+
+
+def energy_forces_getter_closure(geom):
+    def energy_forces_getter(coords):
+        results = geom.get_energy_and_forces_at(coords)
+        energy = results["energy"]
+        forces = results["forces"]
+        return energy, forces
+    return energy_forces_getter

pysisyphus/dynamics/lincs.py

@@ -0,0 +1,105 @@
+import numpy as np
+
+
+def lincs_closure(geom, constraints, order=4):
+    """
+    Drop conn_nums and keep conns_ in a list of list, instead of an array.
+    We could do everything with
+        for i, conn_constr in enumerate(conns)
+            pass
+    instead of the pseudo-code way as given in the paper.
+    """
+    # Number of constraints
+    K = len(constraints)
+    # Inverse atom masses
+    inv_masses = 1 / geom.masses
+
+    # List of constraint sets
+    constraint_sets = [set(c) for c in constraints]
+    constraints = np.array(constraints, dtype=int)
+    # Indices of the atoms that make up the bond constraints
+    atom_1, atom_2 = constraints.T
+
+    # Original bond lengths
+    lengths = np.linalg.norm(geom.coords3d[atom_1] - geom.coords3d[atom_2], axis=1)
+
+    # Determine number of constraints that are connected to every constraint
+    conn_nums = np.zeros(K, dtype=int)
+    conns_ = [list() for _ in range(K)]
+    for i, con_1 in enumerate(constraint_sets):
+        for j, con_2 in enumerate(constraint_sets):
+            if i == j:
+                continue
+
+            if con_1 & con_2:
+                conn_nums[i] += 1
+                conns_[i].append(j)
+    # Maximum number of connected constraints
+    c_max = conn_nums.max()
+    S_diag = 1 / np.sqrt(inv_masses[atom_1] + inv_masses[atom_2])
+
+    # Array containing the indices of the connected constraints for every constraint
+    conns = np.zeros((K, c_max), dtype=int)
+    for i, c in enumerate(conns_):
+        conns[i, :len(c)] = c
+
+    coefs = np.zeros((K, c_max))
+    signs = {
+        True: -1,
+        False: 1,
+    }
+    for i, con_1 in enumerate(constraint_sets):
+        for j in range(c_max):
+            conns_ij = conns[i, j]
+            sign = signs[(atom_1[i] == atom_1[conns_ij]) or (atom_2[i] == atom_2[conns_ij])]
+            # Connected atom
+            c = tuple(con_1 & constraint_sets[conns_ij])
+            assert len(c) == 1
+            c = int(c[0])
+            coefs[i, j] = sign * inv_masses[c] * S_diag[i] * S_diag[conns_ij]
+
+    def solve(rhs, sol, new_coords3d, A, B):
+        w = 1
+        for _ in range(order):
+            for i, _ in enumerate(constraints):
+                rhs[w, i] = 0.
+                for j in range(conn_nums[i]):
+                    # rhs[w, i] += A[i, j] * rhs[2-w, conns_[i][j]]
+                    rhs[w, i] += A[i, j] * rhs[2-w, conns[i, j]]
+                sol[i] += rhs[w, i]
+            w = 2 - w
+
+        for i, (a1, a2) in enumerate(constraints):
+            new_coords3d[a1] -= inv_masses[a1] * S_diag[i] * sol[i] * B[i].sum()
+            new_coords3d[a2] += inv_masses[a2] * S_diag[i] * sol[i] * B[i].sum()
+
+    def lincs(prev_coords3d, new_coords3d):
+        # Calculate constraint directions
+        B = prev_coords3d[atom_1] - prev_coords3d[atom_2]
+        B /= np.linalg.norm(B, axis=1)[:, None]
+        print("B", B)
+
+        A = np.zeros((K, c_max))
+        rhs = np.zeros((2, K))
+        sol = np.zeros(K)
+
+        for i, (a1, a2) in enumerate(constraints):
+            print(i, a1, a2)
+            for j, k in enumerate(conns[i]):
+                A[i, j] = coefs[i, j] * (B[i] * B[k]).sum()
+            
+            rhs[0, i] = S_diag[i] * ((B[i] * (new_coords3d[a1] - new_coords3d[a2])).sum() - lengths[i])
+            sol[i] = rhs[0, i]
+        solve(rhs, sol, new_coords3d, A, B)
+
+        # Correction for rotational lenghthening
+        for i, (a1, a2) in enumerate(constraints):
+            length_i = lengths[i]
+            p = (2 * length_i**2 - (-(new_coords3d[a1] - new_coords3d[a2])**2).sum())**0.5
+            rhs[0, i] = S_diag[i] * (length_i - p)
+            sol[i] = rhs[0, i]
+
+        solve(rhs, sol, new_coords3d, A, B)
+        # import pdb; pdb.set_trace()
+        return new_coords3d
+    return lincs