diff --git a/docs/source/examples.rst b/docs/source/examples.rst index 0bbdbcd..2d80645 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -103,7 +103,7 @@ the gradient provider. Below is a minimal example using the UMA Small 1.2 model. .. code-block:: python - from pathlib import Path + from pathlib import Path from fairchem.core import FAIRChemCalculator, pretrained_mlip from mlfsm.utils import load_xyz from mlfsm.cos import FreezingString diff --git a/docs/source/installation.rst b/docs/source/installation.rst index da614da..409d29b 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -31,7 +31,7 @@ To access the most recent features and bug fixes, install from source: Optional: Calculator Dependencies -------------------------------------- -ML-FSM works with any ASE-compatible calculator including traditional QM and semi-empirical potential energy surfaces. +ML-FSM works with any ASE-compatible calculator including traditional QM and semi-empirical potential energy surfaces. The following potentials have been tested and are supported out of the box via ``examples/fsm_example.py``: .. list-table:: diff --git a/examples/fsm_example.py b/examples/fsm_example.py index 71b966e..fb01ce3 100644 --- a/examples/fsm_example.py +++ b/examples/fsm_example.py @@ -146,7 +146,10 @@ def parse_indices(text): calc = XTB(method="GFN2-xTB") elif calculator == "uma": import torch - from fairchem.core import FAIRChemCalculator, pretrained_mlip # type: ignore [import-not-found] + from fairchem.core import ( # type: ignore [import-not-found] + FAIRChemCalculator, + pretrained_mlip, + ) dev = "cuda" if torch.cuda.is_available() else "cpu" predictor = pretrained_mlip.get_predict_unit("uma-s-1", device=dev) @@ -214,10 +217,18 @@ def parse_indices(text): parser = argparse.ArgumentParser() parser.add_argument("reaction_dir", type=Path, help="absolute path to reaction") parser.add_argument( - "--optcoords", type=str, default="cart", choices=["cart", "ric"], help="Coordinates for optimization" + "--optcoords", + type=str, + default="cart", + choices=["cart", "ric"], + help="Coordinates for optimization", ) parser.add_argument( - "--interp", type=str, default="ric", choices=["cart", "lst", "ric"], help="Interpolation method" + "--interp", + type=str, + default="ric", + choices=["cart", "lst", "ric"], + help="Interpolation method", ) parser.add_argument("--nnodes_min", type=int, default=18, help="Minimum number of nodes in the FSM string") parser.add_argument( @@ -229,7 +240,11 @@ def parse_indices(text): parser.add_argument("--ninterp", type=int, default=50, help="Number of interpolation points between nodes") parser.add_argument("--suffix", type=str, default=None, help="Suffix for output directory") parser.add_argument( - "--method", type=str, default="L-BFGS-B", choices=["L-BFGS-B", "CG"], help="Optimization method" + "--method", + type=str, + default="L-BFGS-B", + choices=["L-BFGS-B", "CG"], + help="Optimization method", ) parser.add_argument("--maxls", type=int, default=3, help="Maximum number of line search iterations") parser.add_argument("--maxiter", type=int, default=2, help="Maximum number of optimization iterations") @@ -248,7 +263,10 @@ def parse_indices(text): help="Checkpoint for calculator", ) parser.add_argument( - "--fixed", type=str, default="", help="Fix atoms, 1-indexed. usage: 1-12 fixes the first 12 atoms" + "--fixed", + type=str, + default="", + help="Fix atoms, 1-indexed. usage: 1-12 fixes the first 12 atoms", ) parser.add_argument("--chg", type=int, default=0, help="Charge of the system") parser.add_argument("--mult", type=int, default=1, help="Multiplicity of the system") diff --git a/src/mlfsm/coords.py b/src/mlfsm/coords.py index 10b5ab2..c9a9161 100644 --- a/src/mlfsm/coords.py +++ b/src/mlfsm/coords.py @@ -1,7 +1,7 @@ """Coordinate generation and transformation tools for FSM optimization.""" import itertools -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict import networkx as nx import numpy as np @@ -31,9 +31,16 @@ class Coordinates: """Base class for internal coordinate systems used in FSM.""" - def __init__(self, atoms1: Atoms, atoms2: Optional[Atoms] = None, verbose: bool = False) -> None: + def __init__( + self, + atoms1: Atoms, + atoms2: Atoms | None = None, + verbose: bool = False, + raise_on_backtransf_fail: bool = True, + ) -> None: self.atoms1 = atoms1 self.atoms2 = atoms2 + self.raise_on_backtransf_fail = raise_on_backtransf_fail c = atoms1.constraints # constraint indicies must be identical between R&P therefor only one is needed if len(c) > 0: self.fixed_atoms = c[0].get_indices() @@ -63,7 +70,10 @@ def q(self, xyz: NDArray[np.float64]) -> NDArray[np.float64]: """Return internal coordinate values from Cartesian positions.""" xyzb = xyz * angs_to_bohr # return np.array([coord.value(xyzb) for coord in self.coords.values()], dtype=np.float64) - return np.fromiter((coord.value(xyzb) for coord in self.coords.values()), dtype=np.float64) + return np.fromiter( + (coord.value(xyzb) for coord in self.coords.values()), + dtype=np.float64, + ) def dqprint(self, atoms1: Atoms, atoms2: Atoms) -> None: """Print differences in internal coordinates between two structures.""" @@ -93,7 +103,11 @@ def u_matrix(self, Bprim: NDArray[np.float64]) -> NDArray[np.float64]: # noqa: evals, evecs = np.linalg.eigh(Bprim @ Bprim.T) return evecs[:, evals > EIGENVAL_CUTOFF] - def x(self, xyz: NDArray[np.float64], qtarget: NDArray[np.float64]) -> NDArray[np.float64]: + def x( + self, + xyz: NDArray[np.float64], + qtarget: NDArray[np.float64], + ) -> NDArray[np.float64]: """Back-transform internal coordinate displacements to Cartesian updates.""" xyz1 = xyz.copy() @@ -150,6 +164,10 @@ def x(self, xyz: NDArray[np.float64], qtarget: NDArray[np.float64]) -> NDArray[n print(f"\tRMS(dx) = {rms_dx:10.5e}") if self.verbose: print(f"\tRMS(dq) = {rms_dq:10.5e}") + if self.raise_on_backtransf_fail: + raise RuntimeError( + f"Back transformation did not converge after {MAX_ITERATIONS} iterations.", + ) return np.array(xyz_backup, dtype=np.float64) @@ -177,13 +195,23 @@ def construct(self) -> Dict[str, Any]: class Redundant(Coordinates): """Redundant internal coordinate system including bond, angle, torsion, etc.""" - def checkstre(self, A: NDArray[np.float64], B: NDArray[np.float64], eps: float = 1e-08) -> bool: # noqa: N803 + def checkstre( + self, + A: NDArray[np.float64], # noqa: N803 + B: NDArray[np.float64], # noqa: N803 + eps: float = 1e-08, + ) -> bool: """Check if distance between two atoms is significant (non-zero within tolerance).""" v0 = A - B n = np.maximum(1e-12, v0.dot(v0)) return n >= eps - def checkangle(self, A: NDArray[np.float64], B: NDArray[np.float64], C: NDArray[np.float64]) -> bool: # noqa: N803 + def checkangle( + self, + A: NDArray[np.float64], # noqa: N803 + B: NDArray[np.float64], # noqa: N803 + C: NDArray[np.float64], # noqa: N803 + ) -> bool: """Check if angle defined by three atoms is physically valid.""" return self.checkstre(A, B) and self.checkstre(B, C) @@ -197,14 +225,24 @@ def checktors( """Check if torsion angle defined by four atoms is physically valid.""" return self.checkstre(A, B) and self.checkstre(B, C) and self.checkstre(C, D) - def get_fragments(self, A: NDArray[np.int_]) -> List[NDArray[np.int_]]: # noqa: N803 + def get_fragments( + self, + A: NDArray[np.int_], # noqa: N803 + ) -> list[NDArray[np.int_]]: """Return list of fragments as connected components in adjacency matrix.""" G: nx.Graph = nx.to_networkx_graph(A) return [np.array(list(d)) for d in nx.connected_components(G)] def connectivity( - self, atoms: Atoms - ) -> Tuple[List[NDArray[np.int64]], NDArray[np.int64], NDArray[np.int64], NDArray[np.int64], NDArray[np.int64]]: + self, + atoms: Atoms, + ) -> tuple[ + list[NDArray[np.int64]], + NDArray[np.int64], + NDArray[np.int64], + NDArray[np.int64], + NDArray[np.int64], + ]: """Compute connectivity matrices from atomic positions.""" # this is done in Angstrom z = atoms.get_atomic_numbers() @@ -429,12 +467,12 @@ def construct(self) -> Dict[str, Any]: to_delete.append(name) if ("tors" in name) and (np.cos(q1[i]) < -tors_thresh or np.cos(q2[i]) < -tors_thresh): to_delete.append(name) - to_add["stre_{}_{}".format(coord.a, coord.d)] = Distance(coord.a, coord.d) + to_add[f"stre_{coord.a}_{coord.d}"] = Distance(coord.a, coord.d) if ("linearbnd" in name) and ((np.cos(q1[i]) < lb_thresh) or (np.cos(q2[i]) < lb_thresh)): basecoord = name[:-2] to_delete.append(basecoord + "_0") to_delete.append(basecoord + "_1") - to_add["bend_{}_{}_{}".format(coord.a, coord.b, coord.c)] = Angle(coord.a, coord.b, coord.c) + to_add[f"bend_{coord.a}_{coord.b}_{coord.c}"] = Angle(coord.a, coord.b, coord.c) if "linearbnd" in name: a, b, c = coord.a, coord.b, coord.c ang = Angle(a, b, c) @@ -444,7 +482,7 @@ def construct(self) -> Dict[str, Any]: basecoord = name[:-2] to_delete.append(basecoord + "_0") to_delete.append(basecoord + "_1") - to_add["bend_{}_{}_{}".format(coord.a, coord.b, coord.c)] = Angle(coord.a, coord.b, coord.c) + to_add[f"bend_{coord.a}_{coord.b}_{coord.c}"] = Angle(coord.a, coord.b, coord.c) for k in set(to_delete): del coords[k] diff --git a/src/mlfsm/cos.py b/src/mlfsm/cos.py index 7031f56..8f80685 100644 --- a/src/mlfsm/cos.py +++ b/src/mlfsm/cos.py @@ -61,6 +61,9 @@ class FreezingString: ``nnodes_min``. The step size is measured along the Cartesian arc length of the initial linear interpolation. Default is 0.0 (derive step size from ``nnodes_min``). + raise_on_backtransf_fail : bool, optional + If ``True``, raise a RuntimeError if the RIC back transformation fails + to converge during interpolation or node growth. Default is ``True``. Attributes ---------- @@ -82,6 +85,7 @@ def __init__( interp_method: str = "ric", ninterp: int = 100, stepsize: float = 0.0, + raise_on_backtransf_fail: bool = True, output: Optional["FSMOutput"] = None, ) -> None: self.output = output @@ -90,6 +94,7 @@ def __init__( self.nnodes_min = int(nnodes_min) self.ninterp = int(ninterp) self.use_cartesian_distance = True if stepsize > 0 else False + self.raise_on_backtransf_fail = raise_on_backtransf_fail if interp_method == "cart": self.interp = Linear @@ -176,7 +181,11 @@ def interpolate(self, outdir: Path | str) -> None: for i, atoms in enumerate(path): f.write(f"{self.natoms}\n") f.write(f"{s[i]:.5f}\n") - for atom, xyz in zip(atoms.get_chemical_symbols(), atoms.get_positions(), strict=True): + for atom, xyz in zip( + atoms.get_chemical_symbols(), + atoms.get_positions(), + strict=True, + ): x, y, z = map(float, xyz) f.write(f"{atom} {x:.8f} {y:.8f} {z:.8f}\n") @@ -203,7 +212,11 @@ def grow(self) -> None: try: self.coordsobj = interp.coords except Exception: - self.coordsobj = Cartesian(r_atoms, p_atoms) + self.coordsobj = Cartesian( + r_atoms, + p_atoms, + raise_on_backtransf_fail=self.raise_on_backtransf_fail, + ) if self.use_cartesian_distance and self.interp_method == "ric": string = interp() diff --git a/src/mlfsm/interp.py b/src/mlfsm/interp.py index f501d15..e6a5d3a 100644 --- a/src/mlfsm/interp.py +++ b/src/mlfsm/interp.py @@ -1,7 +1,7 @@ """Interpolation methods for constructing paths between endpoint geometries.""" +from collections.abc import Callable from dataclasses import dataclass -from typing import Callable import numpy as np from ase import Atoms @@ -35,6 +35,10 @@ class Interpolate: If ``True``, return interpolated internal coordinate vectors instead of Cartesian positions. Only meaningful for :class:`RIC`. Default is ``False``. + raise_on_backtransf_fail : bool, optional + If ``True``, raise a RuntimeError if the RIC back transformation fails + to converge during interpolation or node growth. Only meaningful for + :class:`RIC`. Default is ``True``. """ atoms1: Atoms @@ -42,6 +46,7 @@ class Interpolate: ninterp: int gtol: float = 1e-4 return_q: bool = False + raise_on_backtransf_fail: bool = True def interpolate(self) -> NDArray[np.float32]: """Abstract interpolation routine — must be overridden by subclasses.""" @@ -184,7 +189,9 @@ class RIC(Interpolate): def __post_init__(self) -> None: """Build the shared redundant internal coordinate system.""" - self.coords = Redundant(self.atoms1, self.atoms2, verbose=False) + self.coords = Redundant( + self.atoms1, self.atoms2, verbose=False, raise_on_backtransf_fail=self.raise_on_backtransf_fail + ) def interpolate(self) -> NDArray[np.float32]: """Return RIC-interpolated path in Cartesian (or internal) coordinates.