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

mlmm/harmonic_constraints.py

@@ -0,0 +1,140 @@
+import warnings
+
+import numpy as np
+from ase.calculators.calculator import Calculator
+from ase.constraints import FixAtoms, FixBondLengths
+
+
+class HarmonicFixAtoms(Calculator):
+    """
+    Harmonic position restraint on a subset of atoms.
+
+    E = 1/2 * k_fix * sum_i |r_i - r_i^ref|^2
+    """
+
+    implemented_properties = ["energy", "forces"]
+
+    def __init__(self, indices, ref_positions, k_fix=300.0):
+        super().__init__()
+        idx = np.asarray(indices, dtype=int).ravel()
+        if idx.size == 0:
+            raise ValueError("HarmonicFixAtoms requires at least one index.")
+        ref_pos = np.asarray(ref_positions, dtype=float)
+        if ref_pos.shape != (idx.size, 3):
+            raise ValueError(
+                f"ref_positions must have shape ({idx.size}, 3), got {ref_pos.shape}"
+            )
+        self.indices = idx
+        self.ref_positions = ref_pos
+        self.k_fix = float(k_fix)
+
+    def calculate(self, atoms, properties, system_changes):
+        super().calculate(atoms, properties, system_changes)
+        pos = atoms.get_positions().astype(float)
+        disp = pos[self.indices] - self.ref_positions
+        energy = 0.5 * self.k_fix * np.sum(disp ** 2)
+        forces = np.zeros_like(pos, dtype=float)
+        forces[self.indices] = -self.k_fix * disp
+        self.results = {
+            "energy": float(energy),
+            "forces": forces,
+        }
+
+
+class HarmonicFixBondLengths(Calculator):
+    """
+    Harmonic bond-length restraint for selected atom pairs.
+
+    E = 1/2 * k_bond * sum_ij (d_ij - d_ij^ref)^2
+    """
+
+    implemented_properties = ["energy", "forces"]
+
+    def __init__(self, pairs, ref_distances, k_bond=300.0):
+        super().__init__()
+        pairs = np.asarray(pairs, dtype=int)
+        if pairs.ndim != 2 or pairs.shape[1] != 2:
+            raise ValueError("pairs must have shape (n_pairs, 2)")
+        ref_distances = np.asarray(ref_distances, dtype=float).ravel()
+        if ref_distances.shape[0] != pairs.shape[0]:
+            raise ValueError(
+                "ref_distances must have same length as pairs"
+            )
+        self.pairs = pairs
+        self.ref_distances = ref_distances
+        self.k_bond = float(k_bond)
+
+    def calculate(self, atoms, properties, system_changes):
+        super().calculate(atoms, properties, system_changes)
+        pos = atoms.get_positions().astype(float)
+        i_idx = self.pairs[:, 0]
+        j_idx = self.pairs[:, 1]
+        rij = pos[i_idx] - pos[j_idx]
+        dist = np.linalg.norm(rij, axis=1)
+        diff = dist - self.ref_distances
+        energy = 0.5 * self.k_bond * np.sum(diff ** 2)
+        inv_dist = np.zeros_like(dist)
+        nonzero = dist > 0.0
+        inv_dist[nonzero] = 1.0 / dist[nonzero]
+        force_scale = -self.k_bond * diff * inv_dist
+        fij = (force_scale[:, None] * rij).astype(float)
+        forces = np.zeros_like(pos, dtype=float)
+        np.add.at(forces, i_idx, fij)
+        np.add.at(forces, j_idx, -fij)
+        self.results = {
+            "energy": float(energy),
+            "forces": forces,
+        }
+
+
+def _iter_constraints(atoms):
+    constraints = getattr(atoms, "constraints", None)
+    if constraints is None:
+        return []
+    if isinstance(constraints, (list, tuple)):
+        return list(constraints)
+    return [constraints]
+
+
+def build_harmonic_constraint_calculators(atoms, k_fix=300.0, k_bond=300.0):
+    """
+    Build harmonic-restraint calculators from ASE constraints.
+
+    Supported constraints:
+    - FixAtoms
+    - FixBondLength / FixBondLengths
+    """
+
+    calculators = []
+    constraints = _iter_constraints(atoms)
+    remaining_constraints = []
+    for constraint in constraints:
+        if isinstance(constraint, FixAtoms):
+            indices = np.asarray(constraint.index, dtype=int)
+            ref_positions = atoms.get_positions()[indices]
+            calculators.append(
+                HarmonicFixAtoms(indices, ref_positions, k_fix=k_fix)
+            )
+        elif isinstance(constraint, FixBondLengths):
+            pairs = np.asarray(constraint.pairs, dtype=int)
+            if constraint.bondlengths is None:
+                pos = atoms.get_positions()
+                rij = pos[pairs[:, 0]] - pos[pairs[:, 1]]
+                ref_distances = np.linalg.norm(rij, axis=1)
+            else:
+                ref_distances = np.asarray(constraint.bondlengths, dtype=float)
+            calculators.append(
+                HarmonicFixBondLengths(pairs, ref_distances, k_bond=k_bond)
+            )
+            remaining_constraints.append(constraint)
+        else:
+            warnings.warn(
+                f"Unsupported ASE constraint {type(constraint).__name__}; "
+                "ignored for harmonic restraints."
+            )
+            remaining_constraints.append(constraint)
+    if remaining_constraints != constraints:
+        atoms.set_constraint(remaining_constraints or None)
+    return calculators
+
+

mlmm/hessian_cache.py

@@ -0,0 +1,44 @@
+"""In-process Hessian cache for the ``all`` workflow.
+
+Stores Hessian matrices as CPU tensors (or numpy arrays) in a module-level
+dict so that subsequent CLI stages executed via ``_run_cli_main()`` can
+reuse them without recomputation.
+
+Each entry may carry ``active_dofs`` — a list of DOF indices that the stored
+(partial) Hessian spans.  Consumers set ``geometry.within_partial_hessian``
+before assigning partial Hessians to ``geometry.cart_hessian``.
+"""
+
+import numpy as np
+import torch
+from typing import Any, Dict, Optional, Sequence
+
+_cache: Dict[str, Any] = {}
+
+
+def store(
+    key: str,
+    H,
+    active_dofs: Optional[Sequence[int]] = None,
+    meta: Optional[dict] = None,
+) -> None:
+    """Cache a (possibly partial) Hessian on CPU."""
+    if isinstance(H, torch.Tensor):
+        h_cpu = H.detach().cpu().clone()
+    else:
+        h_cpu = np.array(H, copy=True)
+    _cache[key] = {
+        "hessian": h_cpu,
+        "active_dofs": list(active_dofs) if active_dofs is not None else None,
+        "meta": meta or {},
+    }
+
+
+def load(key: str) -> Optional[Dict[str, Any]]:
+    """Return cached entry or *None*."""
+    return _cache.get(key)
+
+
+def clear() -> None:
+    """Drop all cached Hessians."""
+    _cache.clear()

mlmm/hessian_calc.py

@@ -0,0 +1,174 @@
+"""Hessian utilities and vibrational analysis helpers for the ML/MM workflow.
+
+These routines construct Hessians, compute frequencies and handle I/O.
+"""
+
+import os
+from contextlib import nullcontext
+import numpy as np
+import time
+import torch
+from typing import Sequence
+from ase.data import atomic_masses
+from ase.constraints import FixAtoms
+from ase.atoms import Atoms
+from pysisyphus.constants import AMU2AU
+
+
+def clone_atoms(atoms):
+    """Return a shallow copy of *atoms* keeping constraints (no PBC)."""
+    new = Atoms(
+        symbols=atoms.get_chemical_symbols(),
+        positions=atoms.get_positions(),
+        pbc=False,
+    )
+
+    if atoms.constraints:
+        fix_idx = [
+            i for c in atoms.constraints if isinstance(c, FixAtoms)
+            for i in c.get_indices()
+        ]
+        if fix_idx:
+            new.set_constraint(FixAtoms(indices=fix_idx))
+
+    return new
+
+
+def hessian_calc(
+    atoms,
+    calc,
+    delta: float = 0.01,
+    info_path: str | None = None,
+    *,
+    dtype: np.dtype = np.float64,
+):
+    """
+    Numerically compute the full Cartesian Hessian (second derivatives).
+
+    Finite differences are taken by displacing each *unfixed* atom by
+    ``±delta`` Å along x, y, z.  Central differences give the sub-Hessian,
+    which is **always expanded to the full (3N, 3N) matrix** by padding
+    zeros for fixed atoms before returning.
+
+    Parameters
+    ----------
+    atoms : ase.Atoms
+        Structure to differentiate.
+    calc : ase.Calculator
+        Calculator providing forces.
+    delta : float, default 0.01 Å
+        Displacement size.
+    info_path : str | None
+        If given, progress info is appended to this file.
+    dtype : numpy dtype, default float64
+        Data type of the returned Hessian.
+
+    Returns
+    -------
+    ndarray
+        Cartesian Hessian of shape ``(3N, 3N)`` where
+        ``N == len(atoms)`` (fixed-atom rows/cols are zero).
+    """
+    fixed = {
+        i for c in atoms.constraints if isinstance(c, FixAtoms)
+        for i in c.get_indices()
+    }
+    movable = np.asarray([i for i in range(len(atoms)) if i not in fixed])
+    m = len(movable)
+    n_dof = 3 * m
+
+    # Short-circuit: all atoms frozen → return zeros
+    if m == 0:
+        N = len(atoms)
+        return np.zeros((3 * N, 3 * N))
+
+    H_sub = np.empty((n_dof, n_dof), dtype=dtype)
+    row = 0
+
+    log_cm = nullcontext(None)
+    if info_path is not None:
+        os.makedirs(os.path.dirname(info_path), exist_ok=True)
+        log_cm = open(info_path, "w", encoding="utf-8")
+
+    with log_cm as log:
+        if info_path is not None:
+            log.write("Hessian calculation by numerical differentiation\n")
+            log.write("------------------------------------------------\n")
+            t0 = time.time()
+            checkpoints = [int(m * k / 10) for k in range(1, 11)]  # every 10 %
+
+        for count, a in enumerate(movable, start=1):
+            for coord in range(3):
+                plus = clone_atoms(atoms)
+                plus.calc = calc
+                plus.positions[a, coord] += delta
+                Fp = plus.calc.get_forces(plus)
+
+                minus = clone_atoms(atoms)
+                minus.calc = calc
+                minus.positions[a, coord] -= delta
+                Fm = minus.calc.get_forces(minus)
+
+                # Central difference:  −∂F/∂x = ∂²E/∂x²
+                H_sub[row] = (Fm - Fp)[movable].ravel() / (2 * delta)
+                row += 1
+
+            if info_path is not None and count in checkpoints:
+                pct = checkpoints.index(count) * 10 + 10
+                elapsed = time.time() - t0
+                total_est = elapsed * 100.0 / pct
+                remaining = total_est - elapsed
+                me, se = divmod(int(elapsed), 60)
+                mr, sr = divmod(int(remaining), 60)
+                log.write(
+                    f"{pct:3d}% ({count}/{m})   Elapsed {me}m{se}s   ETA {mr}m{sr}s\n"
+                )
+                log.flush()
+
+        if info_path is not None:
+            elapsed = time.time() - t0
+            me, se = divmod(int(elapsed), 60)
+            log.write(f"Done in {me}m{se}s\n")
+
+    N = len(atoms)
+    H = np.zeros((3 * N, 3 * N), dtype=dtype)
+    for i_local, i_atom in enumerate(movable):
+        for j_local, j_atom in enumerate(movable):
+            H[
+                3 * i_atom : 3 * i_atom + 3,
+                3 * j_atom : 3 * j_atom + 3,
+            ] = H_sub[
+                3 * i_local : 3 * i_local + 3,
+                3 * j_local : 3 * j_local + 3,
+            ]
+    return H
+
+
+def _get_masses(Z):
+    """Return atomic masses (amu) for a list/array of atomic numbers."""
+    return np.array([atomic_masses[z] for z in Z])
+
+
+def _build_tr_basis(coords_bohr: torch.Tensor, masses_amu: torch.Tensor) -> torch.Tensor:
+    """Return (3N,6) mass-weighted basis for Tx,Ty,Tz,Rx,Ry,Rz."""
+    device, dtype = coords_bohr.device, coords_bohr.dtype
+    N = coords_bohr.shape[0]
+
+    m_au = masses_amu.to(dtype=dtype, device=device) * AMU2AU
+    m_sqrt = torch.sqrt(m_au).reshape(-1, 1)
+
+    com = (m_au.reshape(-1, 1) * coords_bohr).sum(0) / m_au.sum()
+    x = coords_bohr - com
+
+    eye3 = torch.eye(3, dtype=dtype, device=device)
+    cols = []
+
+    for i in range(3):                                   # Tx, Ty, Tz
+        cols.append((eye3[i].repeat(N, 1) * m_sqrt).reshape(-1, 1))
+    for i in range(3):                                   # Rx, Ry, Rz
+        rot = torch.cross(x, eye3[i].expand_as(x), dim=1) * m_sqrt
+        cols.append(rot.reshape(-1, 1))
+
+    return torch.cat(cols, dim=1)
+
+