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

pysisyphus/drivers/opt.py

@@ -0,0 +1,375 @@
+# [1] https://doi.org/10.1021/acs.jctc.0c01306
+#     Single-Point Hessian Calculations for Improved Vibrational Frequencies and
+#     Rigid-Rotor-Harmonic-Oscillator Thermodynamics
+#     Spicher, Grimme
+
+
+from dataclasses import dataclass
+from math import floor, ceil
+from pathlib import Path
+import shutil
+from typing import Callable, Dict, Optional, Tuple
+import sys
+
+import numpy as np
+
+from pysisyphus.calculators import Dimer, ExternalPotential
+from pysisyphus.cos.ChainOfStates import ChainOfStates
+from pysisyphus.config import T_DEFAULT, p_DEFAULT
+from pysisyphus.Geometry import Geometry
+from pysisyphus.helpers import do_final_hessian
+from pysisyphus.helpers_pure import highlight_text, report_frozen_atoms
+from pysisyphus.io import save_hessian
+from pysisyphus.modefollow import NormalMode, geom_davidson
+from pysisyphus.optimizers.cls_map import get_opt_cls, key_is_tsopt
+from pysisyphus.optimizers.Optimizer import Optimizer
+from pysisyphus.optimizers.HessianOptimizer import HessianOptimizer
+from pysisyphus.optimizers.hessian_updates import bfgs_update
+from pysisyphus.tsoptimizers.TSHessianOptimizer import TSHessianOptimizer
+
+
+def opt_davidson(opt, tsopt=True, res_rms_thresh=1e-4):
+    try:
+        H = opt.H
+    except AttributeError:
+        if tsopt:
+            raise Exception("Can't handle TS optimization without Hessian yet!")
+
+        # Create approximate updated Hessian
+        cart_coords = opt.cart_coords
+        cart_forces = opt.cart_forces
+        coord_diffs = np.diff(cart_coords, axis=0)
+        grad_diffs = -np.diff(cart_forces, axis=0)
+        H = np.eye(cart_coords[0].size)
+        for s, y in zip(coord_diffs, grad_diffs):
+            dH, _ = bfgs_update(H, s, y)
+            H += dH
+
+    geom = opt.geometry
+    if geom.coord_type != "cart":
+        H = geom.internal.backtransform_hessian(H)
+
+    masses_rep = geom.masses_rep
+    # Mass-weigh and project Hessian
+    H = geom.eckart_projection(geom.mass_weigh_hessian(H))
+    w, v = np.linalg.eigh(H)
+    inds = [0, 1] if tsopt else [6]
+    # Converge the lowest two modes for TS optimizations; for minimizations the lowest
+    # mode would is enough.
+    lowest = 2 if tsopt else 1
+    guess_modes = [NormalMode(l, masses_rep) for l in v[:, inds].T]
+
+    davidson_kwargs = {
+        "hessian_precon": H,
+        "guess_modes": guess_modes,
+        "lowest": lowest,
+        "res_rms_thresh": res_rms_thresh,
+        "remove_trans_rot": True,
+    }
+
+    result = geom_davidson(geom, **davidson_kwargs)
+    return result
+
+
+@dataclass
+class OptResult:
+    opt: Optimizer
+    geom: Geometry
+    fn: Path
+
+
+def run_opt(
+    geom,
+    calc_getter,
+    opt_key,
+    opt_kwargs=None,
+    iterative=False,
+    iterative_max_cycles=5,
+    iterative_thresh=-15,
+    iterative_scale=2.00,
+    cart_hessian=None,
+    print_thermo=False,
+    title="Optimization",
+    copy_final_geom=None,
+    level=0,
+    save_hessian=False,
+):
+    is_cos = issubclass(type(geom), ChainOfStates)
+    is_tsopt = key_is_tsopt(opt_key)
+    # Disallow iterative optimizations for COS objects
+    if opt_kwargs is None:
+        opt_kwargs = dict()
+    is_iterative = (not is_cos) and (iterative or opt_kwargs.pop("iterative", False))
+
+    if is_cos:
+        # Set calculators on all images
+        for image in geom.images:
+            image.set_calculator(calc_getter())
+            title = str(geom)
+    else:
+        geom.set_calculator(calc_getter())
+        geom.cart_hessian = cart_hessian
+
+    do_hess = opt_kwargs.pop("do_hess", False)
+    do_davidson = opt_kwargs.pop("do_davidson", False)
+    T = opt_kwargs.pop("T", T_DEFAULT)
+    p = opt_kwargs.pop("p", p_DEFAULT)
+    propagate = opt_kwargs.pop("propagate", False)
+
+    opt_cls = get_opt_cls(opt_key)
+    for i in range(iterative_max_cycles):
+        # Modify hessian_init in later cycles, te reuse the calculated Hessian
+        # from the final geometry of the previous optimization.
+        if (i > 0) and issubclass(opt_cls, HessianOptimizer):
+            opt_kwargs["hessian_init"] = "calc"
+        opt = opt_cls(geom, **opt_kwargs)
+        print(highlight_text(f"Running {title}", level=level) + "\n")
+        print(f"     Input geometry: {geom.describe()}")
+        print(f"  Coordinate system: {geom.coord_type}")
+        print(f"  Coordinate number: {len(geom.coords)}")
+        print(f"         Calculator: {geom.calculator}")
+        print(f"             Charge: {geom.calculator.charge}")
+        print(f"       Multiplicity: {geom.calculator.mult}")
+        print(f"          Optimizer: {opt_key}\n")
+        report_frozen_atoms(geom)
+        print()
+
+        # Try to propagate chkfiles along calculators in COS optimizations
+        if propagate and is_cos:
+            print("Propagating chkfiles along COS")
+            for j, image in enumerate(geom.images):
+                image.energy
+                cur_calc = image.calculator
+                try:
+                    next_calc = geom.images[j + 1].calculator
+                except IndexError:
+                    pass
+                try:
+                    next_calc.set_chkfiles(cur_calc.get_chkfiles())
+                except AttributeError:
+                    break
+
+        opt.run()
+
+        # Only do 1 cycle in non-iterative optimizations
+        if not is_iterative:
+            break
+
+        # Determine imaginary modes for subsequent displacements
+        nus, *_, cart_displs = geom.get_normal_modes()
+        below_thresh = nus < iterative_thresh
+        # Never displace along transition vector in ts-optimizations. Just skip it.
+        if is_tsopt:
+            below_thresh[0] = False
+        imag_nus = nus[below_thresh]
+        imag_displs = cart_displs[:, below_thresh].T
+
+        if len(imag_nus) == 0:
+            print(f"Iterative optimization converged in cycle {i}.")
+            break
+
+        # h5_fn = f"hess_calc_iter_{i:02d}.h5"
+        # save_hessian(h5_fn, geom)
+        # print(f"Saved HDF5-Hessian to {h5_fn}.")
+
+        print(f"\nImaginary modes below threshold of {iterative_thresh:.2f} cm⁻¹:")
+        for j, nu in enumerate(nus[below_thresh]):
+            print(f"\t{j:02d}: {nu:8.2f} cm⁻¹")
+        sys.stdout.flush()
+
+        print(f"\nGeometry after optimization cycle {i}:\n{geom.as_xyz()}")
+
+        # Displace along imaginary modes
+        for j, (nu, imag_displ) in enumerate(zip(imag_nus, imag_displs)):
+            step = iterative_scale * imag_displ
+            new_cart_coords = geom.cart_coords + step
+            geom.cart_coords = new_cart_coords
+
+        print(f"\nDisplaced geometry for optimization cycle {i+1}:\n{geom.as_xyz()}\n")
+
+    # # ChainOfStates specific
+    # if is_cos and (not opt.stopped):
+    #     hei_coords, hei_energy, hei_tangent, hei_frac_index = geom.get_splined_hei()
+    #     floor_ind = floor(hei_frac_index)
+    #     ceil_ind = ceil(hei_frac_index)
+    #     print(
+    #         f"Splined HEI is at {hei_frac_index:.2f}/{len(geom.images)-1:.2f}, "
+    #         f"between image {floor_ind} and {ceil_ind} (0-based indexing)."
+    #     )
+    #     hei_geom = Geometry(geom.images[0].atoms, hei_coords)
+    #     hei_geom.energy = hei_energy
+    #     hei_fn = "splined_hei.xyz"
+    #     with open(hei_fn, "w") as handle:
+    #         handle.write(hei_geom.as_xyz())
+    #     print(f"Wrote splined HEI to '{hei_fn}'")
+
+    if copy_final_geom and opt.is_converged:
+        copy_fn = copy_final_geom
+        shutil.copy(opt.final_fn, copy_fn)
+        print(f"Copied '{opt.final_fn}' to '{copy_fn}'.")
+
+    if do_davidson and (not opt.stopped):
+        tsopt = isinstance(opt, TSHessianOptimizer.TSHessianOptimizer) or isinstance(
+            geom.calculator, Dimer
+        )
+        type_ = "TS" if tsopt else "minimum"
+        print(highlight_text(f"Davidson after {type_} search", level=1))
+        opt_davidson(opt, tsopt=tsopt)
+    elif do_hess and (not opt.stopped):
+        print()
+        prefix = opt_kwargs.get("prefix", "")
+        out_dir = opt_kwargs.get("out_dir", None)
+        do_final_hessian(
+            geom,
+            save_hessian=save_hessian,
+            write_imag_modes=True,
+            prefix=prefix,
+            T=T,
+            p=p,
+            print_thermo=print_thermo,
+            is_ts=is_tsopt,
+            out_dir=out_dir,
+        )
+    print()
+
+    opt_result = OptResult(opt, opt.geometry, opt.final_fn)
+    return opt_result
+
+
+def get_optimal_bias(
+    ref_geom: Geometry,
+    calc_getter: Callable,
+    opt_key: str,
+    opt_kwargs: Dict,
+    k_max: float,
+    k_min: float = 0.0,
+    rmsd_target: float = 0.188973,
+    rmsd_thresh: Optional[float] = None,
+    rmsd_kwargs: Optional[Dict] = None,
+    k_thresh: float = 1e-3,
+    strict=True,
+) -> Tuple[OptResult, float, bool]:
+    """Driver to determine optimal bias value k for RMSD restraint.
+
+    Parameters
+    ----------
+    ref_geom
+        Reference geometry. Starting point of the optimizations and reference
+        for RMSD calculations.
+    calc_getter
+        Function that returns the actual calculator, providing the energy and
+        its derivatives.
+    opt_key
+        Determines optimizer type. See pysisyphus.optimizers.cls_map.
+    opt_kwargs
+        Optional dict of arguments passed to the optimizer.
+    k_max
+        Maximum absolute value of bias factor k. Must be a > k_min.
+    k_max
+        Minimum absolute value of bias factor k. Must be a positive number >= 0.0.
+        Defaults to 0.0.
+    rmsd_target
+        Target RMSD value in au. Defaults to 0.188973 a0 (approx. 0.1 Å).
+    rmsd_thresh
+        Allowed deviation from rmsd_target in au. If omitted, 5% of rmsd_target are used.
+    rmsd_kwargs
+        Additional keyword arguments that are passed to the RMSD class, e.g., atom_indices.
+    k_thresh:
+       When the absolute value of k_bias - k_min or k_max becomes smaller than
+       k_thresh, the bisection is aborted.
+    strict
+        If True, AssertionError is raised when an optimization did not converged.
+
+    Returns
+    -------
+    opt_result
+        OptimizationResult object containig the Optimizer object.
+    k_opt
+        Optimal value of k_bias.
+    valid_k
+        Whether an appropriate bias value k was found.
+    """
+    assert rmsd_target > 0.0
+    assert k_min >= 0.0
+    assert k_max > k_min
+    assert k_thresh >= 0.0
+    if rmsd_thresh is None:
+        rmsd_thresh = 0.05 * rmsd_target
+    if rmsd_kwargs is None:
+        rmsd_kwargs = dict()
+
+    opt_counter = 0
+
+    def run_biased_opt(k):
+        nonlocal opt_counter
+        geom = ref_geom.copy()
+
+        def biased_calc_getter():
+            act_calc = calc_getter()
+            _rmsd_kwargs = {
+                "type": "rmsd",
+                "k": k,
+            }
+            _rmsd_kwargs.update(rmsd_kwargs)
+            potentials = (_rmsd_kwargs,)
+            calc = ExternalPotential(
+                calculator=act_calc, potentials=potentials, geom=ref_geom
+            )
+            return calc
+
+        sys.stdout.flush()
+        prefix = f"biased_{opt_counter:02d}"
+        prefixed_opt_kwargs = opt_kwargs.copy()
+        prefixed_opt_kwargs.update(
+            {
+                "prefix": prefix,
+                "h5_group_name": f"{prefix}_opt",
+            }
+        )
+        print(f"@@@ Running optimization {opt_counter:02d} with {k=:.6f} au.")
+        opt_result = run_opt(geom, biased_calc_getter, opt_key, prefixed_opt_kwargs)
+        if strict:
+            assert opt_result.opt.is_converged
+
+        rmsd_current = ref_geom.rmsd(geom)
+        print(
+            f"@@@ Biased optimization {opt_counter:02d} with {k=:.6f} au yielded an "
+            f"RMSD of RMSD={rmsd_current:.4f} au."
+        )
+        opt_counter += 1
+        return opt_result, rmsd_current
+
+    k_bias = 0.0
+    opt_result, rmsd_current = run_biased_opt(
+        k=k_bias
+    )  # Start with unbiased optimization
+    unbiased_failed = rmsd_current > rmsd_target
+
+    k_min0 = k_min
+    k_max0 = k_max
+
+    def k_is_valid(k):
+        """Return whether k is too close to either k_min0 or k_max0."""
+        return all([abs(k_bias - k) >= k_thresh for k in (k_min0, k_max0)])
+
+    valid_k = True
+
+    # Only run the loop when the initial unbiased optimization failed. This way
+    # we don't have to check for an early return before the loop, but can keep
+    # one return statement at the end of the function.
+    while unbiased_failed and abs(rmsd_current - rmsd_target) > rmsd_thresh:
+        k_bias = -0.5 * (k_min + k_max)
+        if not (valid_k := k_is_valid(k_bias)):
+            break
+        # Biased geometry optimization
+        opt_result, rmsd_current = run_biased_opt(k_bias)
+        if rmsd_current < rmsd_target:
+            k_max = abs(k_bias)
+        elif rmsd_current > rmsd_target:
+            k_min = abs(k_bias)
+        print(
+            f"@@@ After biased optimization {opt_counter:02d}: {k_min=:.4f}, {k_bias=:.4f}, "
+            f"{k_max=:.4f}, {rmsd_current=:4f} au."
+        )
+    k_opt = k_bias
+    return opt_result, k_opt, valid_k

pysisyphus/drivers/perf.py

@@ -0,0 +1,91 @@
+from datetime import timedelta
+import sys
+import time
+
+from pysisyphus.helpers_pure import highlight_text
+from pysisyphus.TablePrinter import TablePrinter
+
+
+def run_perf(
+    geom,
+    calc_getter,
+    pals=None,
+    mems=None,
+    pal_range=None,
+    mem_range=None,
+    repeat=1,
+    kind="forces",
+):
+    assert repeat > 0
+    assert pals or pal_range
+    assert mems or mem_range
+
+    func_names = {
+        "energy": "get_energy",
+        "forces": "get_forces",
+        "hessian": "get_hessian",
+    }
+
+    if pal_range is not None:
+        pals = range(*pal_range)
+    try:
+        pal_iter = tuple(pals)
+    except TypeError:  # When 'pals' is a single integer
+        pal_iter = tuple((pals,))
+
+    if mem_range is not None:
+        mems = range(*mem_range)
+    try:
+        mem_iter = tuple(mems)
+    except TypeError:  # When 'mems' is a single integer
+        mem_iter = tuple((mems,))
+
+    results = {}
+    for pal in pal_iter:
+        print(f"{pal=} and ", end="")
+        for mem in mem_iter:
+            print(f"{mem=}:")
+            for i in range(repeat):
+                print(f"\tRunning {kind} cycle {i:02d} ... ", end="")
+                calc = calc_getter()
+                calc.pal = pal
+                calc.mem = mem
+                start = time.time()
+                func = getattr(calc, func_names[kind])
+                _ = func(geom.atoms, geom.cart_coords)
+                dur = time.time() - start
+                td = timedelta(seconds=dur)
+                print(f"finished in {td} h")
+                sys.stdout.flush()
+                key = (pal, mem)
+                results.setdefault(key, list()).append(td)
+            print()
+    return results
+
+
+def print_perf_results(results):
+    pal_avgs = dict()
+    pal_min = 1_000_000_000  # Dummy value
+    for i, (key, tds) in enumerate(sorted(results.items())):
+        pal, mem = key
+        pal_min = min(pal, pal_min)
+        print(highlight_text(f"{i:03d}: pal={pal}, mem={mem} MB"))
+        for j, td in enumerate(tds):
+            print(f"Run {j:03d}: {td} h")
+        avg = sum(tds, timedelta(0)) / len(tds)
+        # pal_avgs.setdefault(pal, list()).append(avg)
+        pal_avgs[pal] = avg
+        print(f"    Avg: {avg} h")
+        print()
+
+    avg_min = pal_avgs[pal_min]
+    print(f"   pal_min: {pal_min} core(s)")
+    print("efficiency: speedup / (pal / pal_min)\n")
+    header = ("pal", "avg. / h", "speedup", "efficieny")
+    col_fmts = ("int", "str", "float", "float")
+    table = TablePrinter(header, col_fmts, width=20)
+    table.print_header()
+    for pal, avg in pal_avgs.items():
+        speedup = avg_min / avg
+        speedup_per_pal = speedup / (pal / pal_min)
+        table.print_row((pal, str(avg), speedup, speedup_per_pal))

pysisyphus/drivers/pka.py

@@ -0,0 +1,52 @@
+from math import log
+
+from pysisyphus.config import p_DEFAULT, T_DEFAULT
+from pysisyphus.constants import AU2KCALPERMOL, KBAU
+from pysisyphus.helpers_pure import standard_state_corr
+from pysisyphus import logger
+from pysisyphus.io import geom_from_hessian
+
+
+def G_aq_from_h5_hessian(h5_hessian, solv_en, T=T_DEFAULT, p=p_DEFAULT):
+    geom = geom_from_hessian(h5_hessian)
+    thermo = geom.get_thermoanalysis(T=T, p=p)
+    dG_solv = solv_en - thermo.U_el
+    G_aq = float(thermo.G + dG_solv)
+    return G_aq
+
+
+def direct_cycle(
+    acid_h5,
+    base_h5,
+    acid_solv_en,
+    base_solv_en,
+    G_aq_H=None,
+    G_gas_H=-6.28,
+    dG_solv_H=-265.9,
+    T=T_DEFAULT,
+    p=p_DEFAULT,
+):
+
+    G_aq_acid = G_aq_from_h5_hessian(acid_h5, acid_solv_en, T=T, p=p)
+    G_aq_base = G_aq_from_h5_hessian(base_h5, base_solv_en, T=T, p=p)
+
+    if G_aq_H is None:
+        G_aq_H = G_gas_H + dG_solv_H
+    G_aq_H /= AU2KCALPERMOL
+
+    dG_aq = G_aq_H + G_aq_base - G_aq_acid
+    dG_aq_corr = dG_aq + standard_state_corr(T=T, p=p)
+
+    pKa = dG_aq_corr / (KBAU * T * log(10))
+
+    logger.debug(
+        f"pKa calculation using direct cycle:\n"
+        f"\tAcid HDF5: {acid_h5}\n\tBase HDF5: {base_h5}\n"
+        f"\tT={T:.4f}, p={p:.1f} Pa\n"
+        f"\tG_aq_acid={G_aq_acid:.6f} au\n"
+        f"\tG_aq_base={G_aq_base:.6f} au\n"
+        f"\tG_aq_H={G_aq_H:.6f} au\n"
+        f"\tpKa={pKa:.6f}"
+    )
+
+    return pKa