Source code for pysisyphus.drivers.opt

# [1]
#     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 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

[docs] 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
[docs] @dataclass class OptResult: opt: Optimizer geom: Geometry fn: Path
[docs] 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, ): 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): 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 # 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_energy hei_fn = "" 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, 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
[docs] 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