Source code for ccu.workflows.scan

"""This module defines the bond scan workflow.

Example:
    Relaxation using an ASE :class:`~ase.optimize.optimize.Optimizer`

    .. code-block:: python

        import logging

        from ase.build import molecule
        from ase.calculators.emt import EMT
        from ccu.workflows.calculation import run_molecular_bond_scan

        logging.basicConfig(level=logging.DEBUG)

        atoms = molecule("CO2")
        atoms.calc = EMT()
        run_calculation(atoms, opt=opt, fmax=0.01, steps=10)
"""

import logging
from pathlib import Path
from typing import NamedTuple
import warnings

from ase.atoms import Atoms
from ase.calculators.calculator import Calculator
from ase.calculators.calculator import all_changes
import ase.io
from ase.io.trajectory import Trajectory
from ase.neighborlist import NewPrimitiveNeighborList
from ase.neighborlist import build_neighbor_list

from ccu.structure.geometry import fragment_atoms

logger = logging.getLogger(__name__)


[docs] class BondScanParams(NamedTuple): """Parameters for modifying bond scans. ``a0``, ``a1``, ``mask``, ``indices``, and ``fix`` map to the parameters in :meth:`~ase.Atoms.set_distance``. ``bond_lims`` specifies the minimum and maximum bond lengths for the bond scan. The endpoint may be excluded if the scan step does not evenly divide the difference between the minimum and maximum bond lengths. """ a0: int a1: int mask: list[bool] | None = None indices: list[int] | None = None fix: float = 0.0 bond_lims: tuple[float, float] = (0.7, 4.0)
[docs] def get_bond_scan_parameters( atoms: Atoms, *, cutoffs: list[float] | None = None ) -> list[BondScanParams]: """Determine all bonded pairs and corresponding bond scan parameters. Args: atoms: An :class:`~ase.Atoms` object. cutoffs: A list of cutoff radii used to determine which atoms are bonded. Returns: A list of bond modification parameters suitable for the ``bonds`` parameter of :func:`run_molecular_bond_scans`. Warns: UserWarning: Bonded atoms are not bond-separable. Notes: Two atoms in a molecule are bond-separable if the distance between them can be increased without affecting the bond lengths of any other bonds in the molecule. In graph theory terms, this means that both atoms are articulation points in the graph representation of the molecule. Bond scan parameters are defined to preferentially fix the more bonded atom in the bond. Example: Stretch an intramolecular bond while preserving other bond lengths >>> from ccu.adsorption.adsorbates import get_adsorbate >>> from ccu.workflows.scan import get_bond_scan_parameters >>> hcooh = get_adsorbate("HCOOH") >>> bp = get_bond_scan_parameters(hcooh)[0] # This atom should move together with bp.a1 >>> moving = next(a for a in bp.indices if a not in (bp.a0, bp.a1)) >>> d1 = hcooh.get_distance(bp.a0, bp.a1) >>> d2 = hcooh.get_distance(bp.a1, moving) >>> hcooh.set_distance(bp.a0, bp.a1, d1 * 2, indices=bp.indices) >>> hcooh.get_distance(bp.a0, bp.a1) == d1 * 2 True >>> d2 == hcooh.get_distance(bp.a1, moving) True """ params: list[BondScanParams] = [] nl = build_neighbor_list( atoms, cutoffs=cutoffs, primitive=NewPrimitiveNeighborList, self_interaction=False, bothways=True, ) for atom in atoms: a0 = atom.index a0_neighbors, _ = nl.get_neighbors(a0) if len(a0_neighbors) == 1: # neighborlists return np.int64's a1 = int(a0_neighbors[0]) # Skip duplicate bond pairs if not any(sorted(p[:2]) == sorted([a0, a1]) for p in params): # Move singly bonded atom by default params.append(BondScanParams(a0=a1, a1=a0, indices=[a0])) else: for a1 in a0_neighbors: # Skip duplicate bond pairs if any(sorted(p[:2]) == sorted([a0, a1]) for p in params): continue # Move fragment of less bonded atom frag1 = fragment_atoms(atoms, a0, a1, nl=nl) frag2 = fragment_atoms(atoms, a1, a0, nl=nl) if len(frag1) > len(frag2): pair = to_fix, to_move = a0, a1 frag = frag2 else: pair = to_fix, to_move = a1, a0 frag = frag1 for a in frag: frag_nl, _ = nl.get_neighbors(a) # check if moving fragment is doubly bonded to fixed # fragment if any(n not in frag and n not in pair for n in frag_nl): msg = ( f"Bonded atoms {atoms.symbols[a0]}{a0} and " f"{atoms.symbols[a1]}{a1} are not bond-separable. " "Extending this bond will affect other bonds." ) warnings.warn(msg, UserWarning, stacklevel=2) params.append( BondScanParams( a0=to_fix, a1=to_move, indices=frag, ) ) return params
[docs] def _remove_existing_bond_lengths( filename: Path, a0: int, a1: int, bond_lengths: list[float], scan_step: float, ) -> tuple[list[float], list[Atoms]]: if filename.exists(): old_images: list[Atoms] = ase.io.read(filename=filename, index=":") # type: ignore[assignment] completed_distances = [a.get_distance(a0, a1) for a in old_images] incomplete_bond_lengths = [] for bl in bond_lengths: if not any( abs(bl - d) < (scan_step / 10) for d in completed_distances ): incomplete_bond_lengths.append(bl) else: old_images = [] incomplete_bond_lengths = bond_lengths return incomplete_bond_lengths, old_images
[docs] def run_molecular_bond_scans( atoms: Atoms, calc: Calculator, *, bond_scan_params: list[tuple] | list[BondScanParams] | None = None, scan_step: float = 0.1, write_traj: bool = True, traj_template: str = "scan_{}_{}.traj", ) -> dict[tuple[int, int], list[Atoms]]: """Run linear scan calculations for molecular bonds. Args: atoms: An :class:`~ase.Atoms` object for which a hookean constraint is to be parametrized. calc: A preconfigured ASE calculator to use to run the bond scans. bond_scan_params: A list of bond scan parameters specifying which and how bonds will be linearly scanned. Bond scan parameters can be specified as a plain tuple or a :class:`BondScanParams` instance. If the former, they will be converted into the latter. Defaults to None in which bond scan parameters will be automatically generated using :func:`get_bond_scan_parameters`. scan_step: The step size (in Angstroms) to use for the linear scan calculations. Defaults to 0.1. write_traj: Whether or not to write trajectory files for each bond scan containing the images for which energies are calculated. Defaults to True. traj_template: A format string to be used to name the trajectory files of each linear scan. The format string must accept two fields, which correspond to the atomic indices of the atoms in the bond for which the linear bond scan is being performed. Defaults to ``"scan_{}_{}.traj"``. Note: This function can resume from a previous call if ``write_traj=True`` and the same ``traj_template`` string is used. In that case, calculations for bond separations that differ by less than one-tenth of the step size of an existing image will be skipped. Returns: A dictionary mapping atomic indices to lists, each atom in the list being the image of a completed step in the linear bond scan specified by the atomic indices. """ logger.debug("Running molecular bond scan calculation") directory = getattr(calc, "directory", Path().cwd()) if bond_scan_params: bond_scan_params = [BondScanParams(*bp) for bp in bond_scan_params] else: bond_scan_params = get_bond_scan_parameters(atoms) # Construct images images: dict[tuple[int, int], list[Atoms]] = {} for bond in bond_scan_params: filename = Path(directory, traj_template.format(bond.a0, bond.a1)) min_bond, max_bond = bond.bond_lims d = min_bond num_steps = int(-(-(max_bond - min_bond) // scan_step)) bond_lengths = [min_bond + i * scan_step for i in range(num_steps)] incomplete_bond_lengths, old_images = _remove_existing_bond_lengths( filename, bond.a0, bond.a1, bond_lengths, scan_step ) images[(bond.a0, bond.a1)] = old_images for d in incomplete_bond_lengths: logger.debug( "Scanning %s%s-%s%s bond with distance %.2f", atoms.symbols[bond.a0], bond.a0, atoms.symbols[bond.a1], bond.a1, d, ) image = atoms.copy() image.set_distance( bond.a0, bond.a1, distance=d, fix=bond.fix, mask=bond.mask, indices=bond.indices, ) image.calc = calc image.calc.calculate( image, properties=["energy"], system_changes=all_changes, ) images[(bond.a0, bond.a1)].append(image) if write_traj: with Trajectory(filename=filename, mode="a") as traj: traj.write(image) logger.info("Successfully ran molecular bond scan calculation") return images