Source code for pysisyphus.drivers.scan

import h5py
import numpy as np

from pysisyphus.drivers.opt import run_opt
from pysisyphus.helpers_pure import highlight_text
from pysisyphus.optimizers.RFOptimizer import RFOptimizer
from pysisyphus.TablePrinter import TablePrinter

[docs] def relaxed_scan( geom, calc_getter, constrain_prims, target_values, title, max_cycles=25, trust_radius=0.5, thresh=1e-2, dump=True, ): """Relaxed scan, allowing fixing of multiple primitive internals.""" # Constrain desired primitives copy_kwargs = { "coord_type": "redund", "coord_kwargs": {"constrain_prims": constrain_prims}, } scan_geom = geom.copy(**copy_kwargs) calc = calc_getter() scan_geom.set_calculator(calc) lengths_ = [len(inds) for prim_type, *inds in constrain_prims] if all([True for l in lengths_ if l == 2]): unit = "au" elif all([True for l in lengths_ if l in (3, 4)]): unit = "rad" else: unit = "au (rad)" typed_prims = scan_geom.internal.typed_prims constr_inds = [typed_prims.index(cp) for cp in constrain_prims] constr_fmt = lambda arr: f"({np.array2string(arr, precision=4)}) {unit}" def print_constraints(): print(f"Desired constraints: {constr_fmt(target_values)}") print(f" Actual constraints: {constr_fmt(scan_geom.coords[constr_inds])}") trj = "" scan_cart_coords = [scan_geom.cart_coords.copy()] scan_energies = [] actual_values = [scan_geom.coords[constr_inds]] for i in range(max_cycles): print(highlight_text(f"Step {i}", level=1)) cur_diff = target_values - scan_geom.coords[constr_inds] step = cur_diff step_norm = np.linalg.norm(step) print_constraints() print(f" Difference: {step_norm:.6f} {unit}\n") if step_norm <= thresh: print( f"Relaxed scan converged! Norm of proposed step <= {thresh:.4f} {unit}" ) break if step_norm > trust_radius: step = trust_radius * step / step_norm new_coords = scan_geom.coords.copy() new_coords[constr_inds] += step scan_geom.set_coords(new_coords, update_constraints=True) prefix = f"{title}_step_{i}" opt_kwargs = { "prefix": prefix, "h5_group_name": prefix, "dump": dump, } opt = RFOptimizer(scan_geom, **opt_kwargs) scan_cart_coords.append(scan_geom.cart_coords.copy()) scan_energies.append( actual_values.append(scan_geom.coords[constr_inds]) trj += scan_geom.as_xyz() + "\n" scan_cart_coords = np.array(scan_cart_coords) scan_energies = np.array(scan_energies) actual_values = np.array(actual_values) if dump: trj_fn = f"{title}_scan.trj" with open(trj_fn, "w") as handle: handle.write(trj) print(f"Dumped optimized geometries to '{trj_fn}'.") group_name = f"scan_{title}" with h5py.File("scan.h5", "a") as handle: try: del handle[group_name] except KeyError: pass group = handle.create_group(group_name) group.create_dataset("energies", data=scan_energies) group.create_dataset("cart_coords", data=scan_cart_coords) group.create_dataset("target_values", data=target_values) group.create_dataset("actual_values", data=actual_values) dt = h5py.vlen_dtype(np.dtype("int32")) cp = np.array([(pt.value, *ind) for pt, *ind in constrain_prims]) cp_dataset = group.create_dataset( "constrain_prims", (len(target_values),), dtype=dt ) cp_dataset[:] = cp return scan_geom, scan_cart_coords, scan_energies
[docs] def relaxed_1d_scan( geom, calc_getter, constrain_prims, start, step_size, steps, opt_key, opt_kwargs, pref=None, callback=None, ): assert geom.coord_type == "redund", "'coord_type: redund' is required!" if pref is None: pref = "" else: pref = f"{pref}_" constr_prim = constrain_prims[0] copy_kwargs = { "coord_type": "redund", "coord_kwargs": { "constrain_prims": constrain_prims, "recalc_B": True, }, } constr_geom = geom.copy(**copy_kwargs) constr_ind = constr_geom.internal.get_index_of_typed_prim(constrain_prims[0]) calc = calc_getter() # Always return same calculator, so orbitals can be reused etc. def _calc_getter(): return calc # Drop PrimType at index 0 unit = "au" if len(constr_prim[1:]) == 2 else "rad" org_val = constr_geom.coords[constr_ind] cur_val = start end_val = start + step_size * steps target_scan_vals = np.linspace(cur_val, end_val, steps + 1) print( f" Coordinate: {constr_prim}\n" f"Original value: {org_val:.4f} {unit}\n" f" Initial value: {cur_val:.4f} {unit}\n" f" Final value: {end_val:.4f} {unit}\n" f" Steps: {steps}+1\n" f" Step size: {step_size:.4f} {unit}\n" ) scan_vals = list() scan_geoms = list() scan_energies = list() scan_xyzs = list() for cycle, cur_val in enumerate(target_scan_vals): opt_kwargs_ = opt_kwargs.copy() name = f"{pref}relaxed_scan_{cycle:04d}" opt_kwargs_["prefix"] = name opt_kwargs_["h5_group_name"] = name # Keep a copy scan_geoms.append(constr_geom.copy()) # Update constrained coordinate new_coords = constr_geom.coords new_coords[constr_ind] = cur_val constr_geom.set_coords(new_coords, update_constraints=True) title = f"{pref}Step {cycle:02d}, coord={cur_val:.4f} {unit}" opt_result = run_opt( constr_geom, _calc_getter, opt_key, opt_kwargs_, title=title, level=1 ) if callback is not None: callback(opt_result) scan_xyzs.append(constr_geom.as_xyz()) scan_vals.append(constr_geom.coords[constr_ind]) scan_energies.append( if not opt_result.opt.is_converged: print(f"Step {cycle} did not converge. Breaking!") break with open(f"{pref}relaxed_scan.trj", "w") as handle: handle.write("\n".join(scan_xyzs)) scan_data = np.stack((scan_vals, scan_energies), axis=1) np.savetxt(f"{pref}relaxed_scan.dat", scan_data) scan_vals, scan_energies = scan_data.T print(highlight_text(f"Scan summary", level=1)) col_fmts = "int float float float float".split() header = ("Step", f"Target / {unit}", f"Actual / {unit}", f"Δ / {unit}", "E / au") table = TablePrinter(header, col_fmts, width=14) table.print_header() for step, (target, act, en) in enumerate( zip(target_scan_vals, scan_vals, scan_energies) ): delta = target - act table.print_row((step, target, act, delta, en)) return scan_geoms, scan_vals, scan_energies