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

pysisyphus/optimizers/poly_fit.py

@@ -0,0 +1,285 @@
+from collections import namedtuple
+import logging
+from math import sqrt
+from pprint import pprint
+
+import numpy as np
+import sympy as sym
+
+from pysisyphus.helpers_pure import log
+
+import torch
+
+logger = logging.getLogger("optimizer")
+
+
+def gen_solutions():
+    """
+    Given two energies (e0, e1) and corresponding gradients (g0, g1) we
+    can (try to) fit a quartic polynomial
+        f(x) = a0 + a1*x + a2*x**2 + a3*x**3 + a4*x**4
+    s.t. the constraint f''(x) >= 0, with the equality being fullfilled
+    at only one point.
+    There are five unknowns (a0 - a4) to be determined. Four equations can
+    be derived from f(x) and its first derivative
+        f'(x) = a1 + 2*a2*x + 3*a3*x**2 + 4*a4*x**3 .
+    With (e0, g0) being given at x=0 and (e1, g1) being given at x=1 we can
+    setup the following equations:
+        f (0) = a0                          (1)
+        f'(0) = a1                          (2)
+    using e0 and g0 at x=0, and
+        f (1) = a0 + a1 + a2 + a3 + a4      (3)
+        f'(1) = a1 + 2*a2 + 3*a3 + 4*a4 .   (4)
+    The missing last equation can be derived from the constraint. The second
+    derivative of f(x) is
+        f''(x) = 2*a2 + 6*a3*x + 12*a4*x**2
+    and shall be positive except at one point where it is allowed to be 0, that
+    its two roots (f''(x) = 0) must be degenerate. This is fullfilled when the
+    discriminant D of the quadratic polynomial a*x**2 + b*x + c is zero.
+        D = b**2 – 4*a*c = 0
+    With
+        a = 12*a4
+        b =  6*a3
+        c =  2*a2
+    we get
+        0 = (6*a3)**2 - 4*12*a4*2*a2
+        0 = 36*a3**2 - 96*a4*a2
+        0 = 3*a3**2 - 8*a4*a2               (5)
+        or
+        a4 = 3/8 * a3**2 / a2
+    Using (1) - (5) we can solve the set of equations for a0 - a4.
+    """
+
+    e0, e1, g0, g1, a0, a1, a2, a3 = sym.symbols("e0 e1 g0 g1 a:4")
+
+    a4 = sym.Rational(3, 8) * a3 ** 2 / a2
+    s0, s1 = sym.solve(
+        (
+            e0 - a0,
+            g0 - a1,
+            e1 - a0 - a1 - a2 - a3 - a4,
+            g1 - a1 - 2 * a2 - 3 * a3 - 4 * a4,
+            3 * a3 ** 2 - 8 * a2 * a4,
+        ),
+        (a0, a1, a2, a3),
+    )
+    print("Solution 0")
+    print("\t", s0)
+    print()
+    print("Solution 1")
+    print("\t", s1)
+    print()
+    # There will be two solutions (s0, s1), both containing a big sqrt(...) term
+    # that can be computed once and reused.
+    s0_cse = sym.cse(s0)
+    s1_cse = sym.cse(s1)
+    print("Solution 0 after CSE")
+    pprint(s0_cse)
+    print("Solution 1 after CSE")
+    pprint(s1_cse)
+    print()
+    # The terms in the sqrt-term correspond to binomial expansions and can be further
+    # simplified.
+    ref_term = (
+        -12 * e0 ** 2
+        + 24 * e0 * e1
+        - 12 * e0 * g0
+        - 12 * e0 * g1
+        - 12 * e1 ** 2
+        + 12 * e1 * g0
+        + 12 * e1 * g1
+        - 2 * g0 ** 2
+        - 8 * g0 * g1
+        - 2 * g1 ** 2
+    )
+    sqrt_term = -2 * (
+        6 * (e0 - e1) ** 2 + 6 * (e0 - e1) * (g0 + g1) + (g0 + g1) ** 2 + 2 * g0 * g1
+    )
+    assert sym.simplify(sym.expand(sqrt_term) - ref_term) == 0
+
+
+def get_minimum(poly):
+    roots = np.roots(np.polyder(poly))
+    real_roots = np.real(roots[np.isreal(roots)])
+    vals = poly(real_roots)
+    min_ind = vals.argmin()
+    min_root = real_roots[vals.argmin()]
+    min_val = vals[min_ind]
+    return min_root, min_val
+
+
+def get_maximum(poly):
+    roots = np.roots(np.polyder(poly))
+    real_roots = np.real(roots[np.isreal(roots)])
+    vals = poly(real_roots)
+    max_ind = vals.argmax()
+    max_root = real_roots[max_ind]
+    max_val = vals[max_ind]
+    return max_root, max_val
+
+
+FitResult = namedtuple("FitResult", "x y polys")
+
+
+def quintic_fit(e0, e1, g0, g1, H0, H1):
+    a = -H0 / 2 + H1 / 2 - 6 * e0 + 6 * e1 - 3 * g0 - 3 * g1
+    b = 3 * H0 / 2 - H1 + 15 * e0 - 15 * e1 + 8 * g0 + 7 * g1
+    c = -3 * H0 / 2 + H1 / 2 - 10 * e0 + 10 * e1 - 6 * g0 - 4 * g1
+    d = H0 / 2
+    e = g0
+    f = e0
+
+    poly = np.poly1d((a, b, c, d, e, f))
+    try:
+        mr, mv = get_minimum(poly)
+    except ValueError:
+        return None
+
+    fit_result = FitResult(mr, mv, (poly,))
+    return fit_result
+
+
+def quartic_fit(e0, e1, g0, g1, maximize=False):
+    """See gen_solutions() for derivation."""
+    a0 = e0
+    a1 = g0
+    try:
+        sqrt_term = sqrt(
+            -2
+            * (
+                6 * (e0 - e1) ** 2
+                + 6 * (e0 - e1) * (g0 + g1)
+                + (g0 + g1) ** 2
+                + 2 * g0 * g1
+            )
+        )
+    except ValueError:
+        # In these cases there is no intermediate minimum between 0 and 1 and the term
+        # under the square root becomes negative.
+        return None
+
+    a2_pre = -3 * (e0 - e1) - 5 * g0 / 2 - g1 / 2
+    a3_pre = 2 * e0 - 2 * e1 + 2 * g0
+
+    def get_poly(a3, a2, a1, a0):
+        a4 = 3 / 8 * a3 ** 2 / a2
+        return np.poly1d((a4, a3, a2, a1, a0))
+
+    a2 = a2_pre - sqrt_term / 2
+    a3 = a3_pre + sqrt_term
+    poly0 = get_poly(a3, a2, a1, a0)
+
+    a2 = a2_pre + sqrt_term / 2
+    a3 = a3_pre - sqrt_term
+    poly1 = get_poly(a3, a2, a1, a0)
+
+    get_func = get_maximum if maximize else get_minimum
+    mr0, mv0 = get_func(poly0)
+    mr1, mv1 = get_func(poly1)
+
+    if maximize:
+        mr, mv = (mr0, mv0) if mv0 > mv1 else (mr1, mv1)
+    else:
+        mr, mv = (mr0, mv0) if mv0 < mv1 else (mr1, mv1)
+
+    # Shorter sympy implementation. Probably slower? But shouldn't matter...
+    # ... of course it does ;)
+    # a0, a1, a2, a3 = sym.symbols("a:4")
+    # a4 = sym.Rational(3, 8) * a3**2 / a2
+    # s0, s1 = sym.solve((e0-a0,
+    # g0-a1,
+    # e1-a0-a1-a2-a3-a4,
+    # g1-a1-2*a2-3*a3-4*a4,
+    # 3*a3**2 - 8*a2*a4),
+    # (a0, a1, a2, a3)
+    # )
+    # N = lambda exprs: [sym.N(expr) for expr in exprs]
+    # sym_poly0 = get_poly(*N(s0[::-1]))
+    # sym_poly1 = get_poly(*N(s1[::-1]))
+
+    fit_result = FitResult(mr, mv, (poly0, poly1))
+    return fit_result
+
+
+def cubic_fit(e0, e1, g0, g1):
+    # # Shorter sympy implementation. Probably slower? But shouldn't matter...
+    # # Ok it is really slow ... and it's gone.
+    # a0, a1, a2, a3 = sym.symbols("a:4")
+    # s = sym.solve((e0-a0,
+    # g0-a1,
+    # e1-a0-a1-a2-a3,
+    # g1-a1-2*a2-3*a3),
+    # (a0, a1, a2, a3),
+    # )
+    # coeffs = [float(sym.N(expr)) for expr in (s[a3], s[a2], s[a1], s[a0])]
+    d = e0
+    c = g0
+    b = -(g1 + 2 * g0 + 3 * e0 - 3 * e1)
+    a = 2 * (e0 - e1) + g0 + g1
+    # np.testing.assert_allclose([a, b, c, d], coeffs, atol=1e-10)
+    poly = np.poly1d((a, b, c, d))
+    try:
+        mr, mv = get_minimum(poly)
+    except ValueError:
+        return None
+
+    fit_result = FitResult(mr, mv, (poly,))
+    return fit_result
+
+
+def poly_line_search(
+    cur_energy,
+    prev_energy,
+    cur_grad,
+    prev_grad,
+    prev_step,
+    cubic_max_x=2.0,
+    quartic_max_x=4.0,
+    logger=None,
+):
+    """Generate directional gradients by projecting them on the previous step."""
+    return_torch = False
+    if isinstance(cur_grad, torch.Tensor):
+        return_torch = True
+        device, dtype = cur_grad.device, cur_grad.dtype
+        cur_grad = cur_grad.cpu().numpy()
+    prev_grad_proj = prev_step @ prev_grad
+    cur_grad_proj = prev_step @ cur_grad
+    cubic_result = cubic_fit(prev_energy, cur_energy, prev_grad_proj, cur_grad_proj)
+    quartic_result = quartic_fit(prev_energy, cur_energy, prev_grad_proj, cur_grad_proj)
+    accept = {
+        "cubic": lambda x: (x > 0.0) and (x < cubic_max_x),
+        "quartic": lambda x: (x > 0.0) and (x <= quartic_max_x),
+    }
+
+    fit_result = None
+    if quartic_result and accept["quartic"](quartic_result.x):
+        fit_result = quartic_result
+        deg = "quartic"
+    elif cubic_result and accept["cubic"](cubic_result.x):
+        fit_result = cubic_result
+        deg = "cubic"
+
+    fit_energy = None
+    fit_grad = None
+    fit_step = None
+    if fit_result and fit_result.y < prev_energy:
+        x = fit_result.x
+        fit_energy = fit_result.y
+        log(logger, f"Did {deg} interpolation with x={x:.6f}.")
+
+        # Interpolate coordinates and gradient. 'fit_step' applied to the current
+        # coordinates yields interpolated coordinates.
+        #
+        # x == 0 would take us to the previous coordinates:
+        #  (1-0) * -prev_step = -prev_step (we revert the last step)
+        # x == 1 would preserve the current coordinates:
+        #  (1-1) * -prev_step = 0 (we stay at the current coordinates)
+        # x > 1 extrapolate along previous step direction:
+        #  with x=2, (1-2) * -prev_step = -1*-prev_step = prev_step
+        fit_step = (1 - x) * -prev_step
+        fit_grad = (1 - x) * prev_grad + x * cur_grad
+        if return_torch:
+            fit_grad = torch.tensor(fit_grad, dtype=dtype, device=device)
+            fit_step = torch.tensor(fit_step, dtype=dtype, device=device)
+    return fit_energy, fit_grad, fit_step

pysisyphus/optimizers/precon.py

@@ -0,0 +1,153 @@
+import itertools as it
+from typing import Literal, get_args
+
+import numpy as np
+from scipy.spatial.distance import pdist, squareform
+from scipy.sparse import dok_matrix
+
+from pysisyphus.helpers_pure import log
+from pysisyphus.intcoords.setup import get_pair_covalent_radii
+from pysisyphus.intcoords.setup_fast import (
+    find_bonds_for_geom,
+    find_bonds_bends,
+    find_bonds_bends_dihedrals,
+)
+from pysisyphus.intcoords.derivatives import dq_b, dq_a, dq_d
+from pysisyphus.intcoords import RedundantCoords
+from pysisyphus.optimizers.guess_hessians import get_lindh_alpha
+
+
+# [1] https://www.nature.com/articles/s41598-018-32105-x
+#     Mones, 2018
+
+
+def get_lindh_k(atoms, coords3d, bonds=None, angles=None, torsions=None):
+    if bonds is None:
+        bonds = list()
+    if angles is None:
+        angles = list()
+    if torsions is None:
+        torsions = list()
+
+    atoms = [a.lower() for a in atoms]
+
+    alphas = [get_lindh_alpha(a1, a2) for a1, a2 in it.combinations(atoms, 2)]
+    pair_cov_radii = get_pair_covalent_radii(atoms)
+    cdm = pdist(coords3d)
+    rhos = squareform(np.exp(alphas * (pair_cov_radii ** 2 - cdm ** 2)))
+
+    k_dict = {
+        2: 0.45,  # Stretches/bonds
+        3: 0.15,  # Bends/angles
+        4: 0.005,  # Torsions/dihedrals
+    }
+    ks = list()
+    for inds in it.chain(bonds, angles, torsions):
+        rho_product = 1
+        for i in range(len(inds) - 1):
+            i1, i2 = inds[i : i + 2]
+            rho_product *= rhos[i1, i2]
+        ks.append(k_dict[len(inds)] * rho_product)
+    return ks
+
+
+def get_lindh_precon(
+    atoms,
+    coords,
+    bonds=None,
+    bends=None,
+    dihedrals=None,
+    c_stab=0.0103,
+    logger=None,
+):
+    """c_stab = 0.00103 hartree/bohr² corresponds to 0.1 eV/Ų as
+    given in the paper."""
+
+    if bonds is None:
+        bonds = list()
+    if bends is None:
+        bends = list()
+    if dihedrals is None:
+        dihedrals = list()
+
+    dim = coords.size
+    c3d = coords.reshape(-1, 3)
+
+    # Calculate Lindh force constants
+    ks = get_lindh_k(atoms, c3d, bonds, bends)
+
+    grad_funcs = {
+        2: dq_b,  # Bond
+        3: dq_a,  # Bend
+        4: dq_d,  # Dihedral
+    }
+
+    P = dok_matrix((dim, dim))
+    for inds, k in zip(it.chain(bonds, bends, dihedrals), ks):
+        # First derivatives of internal coordinates w.r.t cartesian coordinates
+        int_grad = grad_funcs[len(inds)](*c3d[inds].flatten())
+        # Assign to the correct cartesian indices
+        cart_inds = np.array(list(it.chain(*[range(3 * i, 3 * i + 3) for i in inds])))
+        P[cart_inds[:, None], cart_inds[None, :]] += abs(k) * np.outer(
+            int_grad, int_grad
+        )
+
+    # Add stabilization to diagonal
+    P[np.diag_indices(dim)] += c_stab
+    P = P.tocsc()
+    filled = P.size / dim ** 2
+    log(logger, f"Preconditioner P has {P.size} entries ({filled:.2%} filled)")
+
+    return P
+
+
+PreconKind = Literal["full", "full_fast", "bonds", "bonds_bends"]
+
+
+def precon_getter(geom, c_stab=0.0103, kind="full", logger=None):
+    valid_kinds = get_args(PreconKind)
+    assert kind in valid_kinds, f"Invalid kind='{kind}'! Valid kinds are: {valid_kinds}"
+
+    atoms = geom.moving_atoms
+    if len(geom.freeze_atoms) > 0:
+        assert (
+            kind == "full"
+        ), "Preconditioning with frozen atoms is only supported for kind='full'"
+    # Default empty lists for coordinates that may be skipped
+    # for kind != "full".
+    bends = list()
+    dihedrals = list()
+    if kind == "full":
+        internal = RedundantCoords(
+            atoms,
+            geom.coords,
+        )
+        bonds = internal.bond_atom_indices
+        bends = internal.bend_atom_indices
+        dihedrals = internal.dihedral_atom_indices
+    elif kind == "full_fast":
+        bonds, bends, dihedrals = find_bonds_bends_dihedrals(geom)
+    elif kind == "bonds_bends":
+        bonds, bends = find_bonds_bends(geom)
+    elif kind == "bonds":
+        bonds = find_bonds_for_geom(geom)
+
+    msg = (
+        f"Constructing preconditioner from {len(bonds)} bonds, {len(bends)} bends "
+        f"and {len(dihedrals)} dihedrals (kind='{kind}')."
+    )
+    log(logger, msg)
+
+    def wrapper(coords):
+        P = get_lindh_precon(
+            atoms,
+            coords,
+            bonds,
+            bends,
+            dihedrals,
+            c_stab=c_stab,
+            logger=logger,
+        )
+        return P
+
+    return wrapper

pysisyphus/optimizers/restrict_step.py

@@ -0,0 +1,24 @@
+import numpy as np
+
+
+def scale_by_max_step(steps, max_step):
+    steps_max = np.abs(steps).max()
+    if steps_max > max_step:
+        steps *= max_step / steps_max
+    return steps
+
+
+def get_scale_max(max_element):
+    def scale_max(step):
+        step_max = np.abs(step).max()
+        if step_max > max_element:
+            step *= max_element / step_max
+        return step
+    return scale_max
+
+
+def restrict_step(steps, max_step):
+    too_big = np.abs(steps) > max_step
+    signs = np.sign(steps[too_big])
+    steps[too_big] = signs * max_step
+    return steps

pysisyphus/pack.py

@@ -0,0 +1,172 @@
+import argparse
+from math import pi as PI, ceil
+from pathlib import Path
+import sys
+
+from pysisyphus.constants import AMU2KG
+from pysisyphus.helpers import geom_loader
+from pysisyphus.helpers_pure import get_input, highlight_text
+from pysisyphus.io.pdb import geom_to_pdb_str
+from pysisyphus.wrapper.packmol import make_input, call_packmol
+
+from pysisyphus.db import LEVELS, MOLECULES
+from pysisyphus.db.helpers import get_path as db_get_path
+
+
+AMU2G = AMU2KG * 1e3
+CM2ANG = 1e-8
+
+
+def parse_args(args):
+    parser = argparse.ArgumentParser()
+
+    # Solvent
+    solvent_group = parser.add_mutually_exclusive_group(required=True)
+    solvent_group.add_argument("--solv", help="Filename of solvent geometry.")
+    solvent_group.add_argument(
+        "--db", action="store_true", help="Choose from internal database."
+    )
+
+    parser.add_argument(
+        "--solv_num", type=int, help="Number of solvent molecules to pack."
+    )
+    parser.add_argument("--solv_dens", type=float, help="Solvent density in g/cm³.")
+
+    parser.add_argument(
+        "--output", default="output.pdb", help="Filename of packed molecules."
+    )
+
+    # Solute
+    parser.add_argument("--solute", default=None, help="Filename of solute geometry.")
+    parser.add_argument(
+        "--solute_num", type=int, default=1, help="Number of solute molecules to pack."
+    )
+
+    return parser.parse_args(args)
+
+
+def as_pdb(fn):
+    if not str(fn).endswith(".pdb"):
+        geom = geom_loader(fn)
+        pdb_str = geom_to_pdb_str(geom)
+        cwd = Path(".")
+        pdb_fn = cwd / Path(fn).with_suffix(".pdb").name
+        with open(pdb_fn, "w") as handle:
+            handle.write(pdb_str)
+        print(f"Converted '{fn}' to PDB format ('{pdb_fn}')")
+    return pdb_fn
+
+
+def sphere_radius_from_volume(volume):
+    radius = (3 / 4 * volume / PI) ** (1 / 3)
+    return radius
+
+
+def volume_for_density(molecule_num, mol_mass, density):
+    # Convert density from g/cm³ to amu/ų
+    density_au = density / AMU2G * CM2ANG ** 3
+    # The molar mass in g/mol is numerically equal to the value in AMU (dalton)
+    # so we can use it as it is.
+    total_mass = mol_mass * molecule_num
+    # Volume in Å
+    volume = total_mass / density_au
+
+    return volume
+
+
+def print_info(title, geom):
+    print(title)
+    print(f"\t{geom}")
+    print(f"\tMolar mass: {geom.total_mass:.2f} g mol⁻¹")
+    print()
+
+
+def run():
+    args = parse_args(sys.argv[1:])
+
+    solute_fn = args.solute
+    if solute_fn:
+        solute = geom_loader(solute_fn)
+        solute_num = args.solute_num
+        solute_mass = solute.total_mass
+        print_info("Solute", solute)
+    else:
+        solute = None
+        solute_mass = 0.0
+
+    solv_fn = args.solv
+    if solv_fn:
+        solv_dens = args.solv_dens
+    # Load from internal db
+    else:
+        print(highlight_text("Interactive solvent selection"))
+        level = get_input(LEVELS, "Level of theory", lbl_func=lambda lvl: lvl[0])
+        print()
+        molecule = get_input(MOLECULES, "Molecule", lbl_func=lambda mol: mol.name)
+        print()
+
+        solv_fn = db_get_path(molecule.name, level[0])
+        solv_dens = molecule.density
+
+    solv = geom_loader(solv_fn)
+    solv_num = args.solv_num
+    solv_mass = solv.total_mass
+    print_info("Solvent", solv)
+
+    solute_solv_mass = solute_mass + solv_num * solv_mass
+    print(f"Total mass of solute(s) and solvent(s): {solute_solv_mass:.2f} amu")
+    print()
+
+    # Solvent volume
+    solv_vol = volume_for_density(solv_num, solv_mass, solv_dens)
+    print(f"Solvent volume: {solv_vol:>10.2f} ų")
+
+    # Solute volume; Use the solvent density for this calculation
+    if solute:
+        solute_vol = volume_for_density(solute_num, solute_mass, solv_dens)
+        print(f" Solute volume: {solute_vol:>10.2f} ų")
+    else:
+        solute_vol = 0.0
+
+    total_vol = solv_vol + solute_vol
+    print(f"  Total volume: {total_vol:>10.2f} ų")
+    print()
+
+    radius = sphere_radius_from_volume(total_vol)
+    print(f"     Sphere radius: {radius:>8.2f} Å")
+    cradius = ceil(radius)
+    print(f"Using ceil(radius): {cradius:>8.2f} Å")
+    print()
+
+    # Create solute/solvent PDBs if needed
+
+    inp_kwargs = {
+        "output_fn": args.output,
+        "solvent_fn": as_pdb(solv_fn),
+        "solvent_num": solv_num,
+        "sphere_radius": cradius,
+    }
+    if solute:
+        inp_kwargs.update({"solute_fn": as_pdb(solute_fn), "solute_num": solute_num})
+
+    inp = make_input(**inp_kwargs)
+    inp_fn = "packmol.inp"
+    with open(inp_fn, "w") as handle:
+        handle.write(inp)
+    print(f"Wrote packmol input to '{inp_fn}'")
+
+    proc = call_packmol(inp)
+
+    log_fn = "packmol.log"
+    with open(log_fn, "w") as handle:
+        handle.write(proc.stdout)
+    print(f"Wrote packmol ouput to '{log_fn}'")
+    print()
+
+    return_ = proc.returncode
+    if return_ != 0:
+        print(proc.stdout)
+    else:
+        print("packmol returned successfully!")
+
+    return return_