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

mlmm/xtb_embedcharge_correction.py

@@ -0,0 +1,475 @@
+"""xTB point-charge embedding correction for ONIOM ML/MM.
+
+The correction term accounts for the electrostatic interaction between
+the ML (QM) region and the MM point charges:
+
+    dE = E_xTB(QM + MM_charges) - E_xTB(QM_only)
+    dF_Q = F_Q(embed) - F_Q(no-embed)
+    dF_M = F_M(embed) - 0
+
+where Q are QM/ML atoms and M are external MM point charges.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+import logging
+import os
+import re
+import shlex
+import shutil
+import subprocess
+import tempfile
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+# ---- Constants ----
+BOHR_PER_ANG = 1.8897261254578281
+EV_PER_HARTREE = 27.211386245988
+ANG_PER_BOHR = 0.529177210903
+FORCE_EV_ANG_PER_HA_BOHR = EV_PER_HARTREE / ANG_PER_BOHR
+
+_TOTAL_ENERGY_RE = re.compile(
+    r"TOTAL ENERGY\s+([-+]?(?:\d+\.\d*|\d*\.\d+|\d+)(?:[DdEe][-+]?\d+)?)"
+)
+
+
+class XTBEmbedError(RuntimeError):
+    """Raised for xTB embedding correction failures."""
+
+
+def _parse_float_token(text):
+    return float(str(text).replace("D", "E").replace("d", "e"))
+
+
+def _resolve_ncores(explicit_ncores=None, fallback=4):
+    candidates = []
+    if explicit_ncores is not None:
+        candidates.append(explicit_ncores)
+    candidates.extend([
+        os.getenv("OMP_NUM_THREADS"),
+        os.getenv("SLURM_CPUS_PER_TASK"),
+    ])
+    for item in candidates:
+        try:
+            val = int(item)
+            if val > 0:
+                return val
+        except Exception:
+            continue
+    return max(1, int(fallback))
+
+
+def _make_run_dir(xtb_workdir):
+    base = None
+    mode = str(xtb_workdir or "tmp").strip()
+    if mode and mode.lower() != "tmp":
+        base = os.path.abspath(mode)
+        os.makedirs(base, exist_ok=True)
+    return tempfile.mkdtemp(prefix="xtb-embed-", dir=base)
+
+
+def _cleanup_run_dir(path, keep_files):
+    if keep_files:
+        return
+    try:
+        shutil.rmtree(path)
+    except Exception:
+        logger.debug("Failed to clean up xTB run directory %s", path, exc_info=True)
+
+
+def _write_xyz(path, symbols, coords_ang):
+    coords = np.asarray(coords_ang, dtype=np.float64).reshape(-1, 3)
+    if len(symbols) != coords.shape[0]:
+        raise XTBEmbedError(
+            f"XYZ write failed: symbol count {len(symbols)} != coord count {coords.shape[0]}"
+        )
+    with open(path, "w") as fh:
+        fh.write(f"{coords.shape[0]}\n")
+        fh.write("generated by mlmm embedcharge correction\n")
+        for sym, xyz in zip(symbols, coords):
+            fh.write(f"{sym:<2s} {xyz[0]: .12f} {xyz[1]: .12f} {xyz[2]: .12f}\n")
+
+
+def _write_pcharge(path, mm_coords_ang, mm_charges):
+    coords = np.asarray(mm_coords_ang, dtype=np.float64).reshape(-1, 3)
+    charges = np.asarray(mm_charges, dtype=np.float64).reshape(-1)
+    if coords.shape[0] != charges.shape[0]:
+        raise XTBEmbedError(
+            f"Point-charge write failed: coord count {coords.shape[0]} != charge count {charges.shape[0]}"
+        )
+    coords_bohr = coords * BOHR_PER_ANG
+    with open(path, "w") as fh:
+        fh.write(f"{coords_bohr.shape[0]}\n")
+        for q, xyz in zip(charges, coords_bohr):
+            fh.write(f"{q: .12f} {xyz[0]: .12f} {xyz[1]: .12f} {xyz[2]: .12f}\n")
+
+
+def _build_xtb_cmd(xtb_cmd, xyz_filename, charge, multiplicity, xtb_acc, mode):
+    tokens = shlex.split(str(xtb_cmd or "xtb"))
+    if not tokens:
+        raise XTBEmbedError("Empty xtb_cmd value.")
+    tokens.extend([str(xyz_filename), "--acc", str(float(xtb_acc)), "--norestart"])
+    m = str(mode).strip().lower()
+    if m == "grad":
+        tokens.append("--grad")
+    elif m == "sp":
+        pass
+    else:
+        raise XTBEmbedError(f"Unknown xTB mode '{mode}'")
+    if charge is not None:
+        tokens.extend(["--chrg", str(int(charge))])
+    if multiplicity is not None:
+        tokens.extend(["--uhf", str(max(0, int(multiplicity) - 1))])
+    return tokens
+
+
+def _run_xtb(run_dir, xyz_filename, charge, multiplicity, xtb_cmd, xtb_acc, mode, ncores):
+    cmd = _build_xtb_cmd(xtb_cmd, xyz_filename, charge, multiplicity, xtb_acc, mode)
+    env = os.environ.copy()
+    env["OMP_NUM_THREADS"] = str(_resolve_ncores(ncores))
+
+    try:
+        proc = subprocess.run(
+            cmd, cwd=run_dir, env=env,
+            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+            text=True, check=False,
+        )
+    except FileNotFoundError as exc:
+        raise XTBEmbedError(f"xTB command not found: '{xtb_cmd}'.") from exc
+    except Exception as exc:
+        raise XTBEmbedError(f"Failed to run xTB: {exc}")
+
+    if proc.returncode != 0:
+        out_tail = "\n".join((proc.stdout or "").splitlines()[-20:])
+        err_tail = "\n".join((proc.stderr or "").splitlines()[-20:])
+        raise XTBEmbedError(
+            f"xTB failed (embedcharge mode={mode}, rc={proc.returncode}).\n"
+            f"CMD: {' '.join(cmd)}\nSTDOUT:\n{out_tail}\nSTDERR:\n{err_tail}"
+        )
+    return proc.stdout or ""
+
+
+def _parse_energy_from_stdout(stdout_text):
+    matches = _TOTAL_ENERGY_RE.findall(str(stdout_text or ""))
+    if not matches:
+        return None
+    return _parse_float_token(matches[-1])
+
+
+def _parse_engrad(path, natoms):
+    with open(path) as fh:
+        lines = fh.readlines()
+
+    n = len(lines)
+    energy_ha = None
+    grad_vals = []
+
+    i = 0
+    while i < n:
+        if "the current total energy in eh" in lines[i].strip().lower():
+            i += 1
+            while i < n:
+                text = lines[i].strip()
+                if text and not text.startswith("#"):
+                    energy_ha = _parse_float_token(text.split()[0])
+                    break
+                i += 1
+            break
+        i += 1
+
+    grad_start = None
+    for idx, line in enumerate(lines):
+        if "the current gradient in eh/bohr" in line.strip().lower():
+            grad_start = idx + 1
+            break
+
+    if grad_start is not None:
+        j = grad_start
+        while j < n and len(grad_vals) < 3 * int(natoms):
+            text = lines[j].strip()
+            if text and not text.startswith("#"):
+                grad_vals.append(_parse_float_token(text.split()[0]))
+            j += 1
+
+    if energy_ha is None:
+        raise XTBEmbedError(f"Could not parse energy from engrad file: {path}")
+    if len(grad_vals) != 3 * int(natoms):
+        raise XTBEmbedError(
+            f"Gradient parse failed in {path}: expected {3 * int(natoms)}, got {len(grad_vals)}"
+        )
+
+    grad = np.asarray(grad_vals, dtype=np.float64).reshape(int(natoms), 3)
+    return float(energy_ha), grad
+
+
+def _parse_pcgrad(path, ncharges):
+    if int(ncharges) <= 0:
+        return np.zeros((0, 3), dtype=np.float64)
+    if not os.path.isfile(path):
+        raise XTBEmbedError(f"xTB point-charge gradient file not found: {path}")
+
+    rows = []
+    with open(path) as fh:
+        for line in fh:
+            text = line.strip()
+            if not text:
+                continue
+            parts = text.split()
+            if len(parts) < 3:
+                raise XTBEmbedError(f"Malformed pcgrad row in {path}: '{text}'")
+            rows.append([_parse_float_token(p) for p in parts[:3]])
+
+    if len(rows) != int(ncharges):
+        raise XTBEmbedError(
+            f"pcgrad row count mismatch: expected {int(ncharges)}, got {len(rows)} in {path}"
+        )
+    return np.asarray(rows, dtype=np.float64).reshape(int(ncharges), 3)
+
+
+def _convert_units(energy_ha=None, gradient_ha_bohr=None):
+    """Convert xTB atomic units to eV / eV/Å."""
+    energy_ev = None
+    forces_ev_ang = None
+    if energy_ha is not None:
+        energy_ev = float(energy_ha) * EV_PER_HARTREE
+    if gradient_ha_bohr is not None:
+        grad = np.asarray(gradient_ha_bohr, dtype=np.float64).reshape(-1, 3)
+        forces_ev_ang = -grad * FORCE_EV_ANG_PER_HA_BOHR
+    return energy_ev, forces_ev_ang
+
+
+def _run_xtb_state(
+    symbols, coords_q_ang, charge, multiplicity, mode,
+    with_embedding, mm_coords_ang, mm_charges,
+    xtb_cmd, xtb_acc, xtb_workdir, xtb_keep_files, ncores,
+):
+    """Run xTB with or without point-charge embedding."""
+    run_dir = _make_run_dir(xtb_workdir)
+    xyz_name = "xtb_input.xyz"
+    nm = int(np.asarray(mm_charges, dtype=np.float64).reshape(-1).size)
+    try:
+        _write_xyz(os.path.join(run_dir, xyz_name), symbols, coords_q_ang)
+        if with_embedding and nm > 0:
+            _write_pcharge(
+                os.path.join(run_dir, "pcharge"),
+                mm_coords_ang=mm_coords_ang,
+                mm_charges=mm_charges,
+            )
+
+        stdout = _run_xtb(
+            run_dir=run_dir, xyz_filename=xyz_name,
+            charge=charge, multiplicity=multiplicity,
+            xtb_cmd=xtb_cmd, xtb_acc=xtb_acc,
+            mode=mode, ncores=ncores,
+        )
+
+        if str(mode).strip().lower() == "sp":
+            energy_ha = _parse_energy_from_stdout(stdout)
+            if energy_ha is None:
+                raise XTBEmbedError("Could not parse xTB energy from stdout.")
+            energy_ev, _ = _convert_units(energy_ha=energy_ha)
+            return float(energy_ev), None, None
+
+        engrad_path = os.path.join(
+            run_dir, f"{os.path.splitext(xyz_name)[0]}.engrad"
+        )
+        if not os.path.isfile(engrad_path):
+            raise XTBEmbedError(f"xTB engrad file not found: {engrad_path}")
+        energy_ha, grad_q_ha_bohr = _parse_engrad(engrad_path, len(symbols))
+        energy_ev, forces_q_ev_ang = _convert_units(
+            energy_ha=energy_ha, gradient_ha_bohr=grad_q_ha_bohr,
+        )
+
+        forces_m_ev_ang = np.zeros((nm, 3), dtype=np.float64)
+        if with_embedding and nm > 0:
+            pcgrad_ha_bohr = _parse_pcgrad(os.path.join(run_dir, "pcgrad"), nm)
+            _, forces_m_ev_ang = _convert_units(gradient_ha_bohr=pcgrad_ha_bohr)
+
+        return (
+            float(energy_ev),
+            np.asarray(forces_q_ev_ang, dtype=np.float64).reshape(len(symbols), 3),
+            np.asarray(forces_m_ev_ang, dtype=np.float64).reshape(nm, 3),
+        )
+    finally:
+        _cleanup_run_dir(run_dir, xtb_keep_files)
+
+
+def _delta_embed_minus_vac(
+    symbols, coords_q_ang, mm_coords_ang, mm_charges,
+    charge, multiplicity, need_forces,
+    xtb_cmd, xtb_acc, xtb_workdir, xtb_keep_files, ncores,
+):
+    """Compute dE = E(embed) - E(vac), dF = F(embed) - F(vac)."""
+    nm = int(np.asarray(mm_charges, dtype=np.float64).reshape(-1).size)
+    if nm <= 0:
+        zq = np.zeros((len(symbols), 3), dtype=np.float64) if need_forces else None
+        zm = np.zeros((0, 3), dtype=np.float64) if need_forces else None
+        return 0.0, zq, zm
+
+    mode = "grad" if need_forces else "sp"
+    common_kw = dict(
+        symbols=symbols, coords_q_ang=coords_q_ang,
+        charge=charge, multiplicity=multiplicity, mode=mode,
+        xtb_cmd=xtb_cmd, xtb_acc=xtb_acc,
+        xtb_workdir=xtb_workdir, xtb_keep_files=xtb_keep_files, ncores=ncores,
+    )
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
+        fut_embed = pool.submit(
+            _run_xtb_state, with_embedding=True,
+            mm_coords_ang=mm_coords_ang, mm_charges=mm_charges, **common_kw,
+        )
+        fut_vac = pool.submit(
+            _run_xtb_state, with_embedding=False,
+            mm_coords_ang=np.zeros((0, 3), dtype=np.float64),
+            mm_charges=np.zeros((0,), dtype=np.float64), **common_kw,
+        )
+        embed_e, embed_fq, embed_fm = fut_embed.result()
+        vac_e, vac_fq, _ = fut_vac.result()
+
+    de_ev = float(embed_e - vac_e)
+    if not need_forces:
+        return de_ev, None, None
+
+    df_q = (
+        np.asarray(embed_fq, dtype=np.float64).reshape(len(symbols), 3)
+        - np.asarray(vac_fq, dtype=np.float64).reshape(len(symbols), 3)
+    )
+    df_m = np.asarray(embed_fm, dtype=np.float64).reshape(nm, 3)
+    return de_ev, df_q, df_m
+
+
+def _assemble_full_force(df_q, df_m):
+    fq = np.asarray(df_q, dtype=np.float64).reshape(-1, 3)
+    fm = np.asarray(df_m, dtype=np.float64).reshape(-1, 3)
+    if fm.shape[0] == 0:
+        return fq
+    return np.vstack([fq, fm])
+
+
+def _numerical_hessian_from_forces(
+    eval_full_forces, coords_q_ang, mm_coords_ang, step_ang, force0_ev_ang=None,
+):
+    """Central-difference numerical Hessian from embedding force deltas."""
+    q0 = np.asarray(coords_q_ang, dtype=np.float64).reshape(-1, 3)
+    m0 = np.asarray(mm_coords_ang, dtype=np.float64).reshape(-1, 3)
+    nq = q0.shape[0]
+    nm = m0.shape[0]
+    nall = nq + nm
+    dof = nall * 3
+    step = float(step_ang)
+
+    if force0_ev_ang is None:
+        f0 = np.asarray(eval_full_forces(q0, m0), dtype=np.float64).reshape(nall, 3)
+    else:
+        f0 = np.asarray(force0_ev_ang, dtype=np.float64).reshape(nall, 3)
+
+    hess = np.zeros((dof, dof), dtype=np.float64)
+    nqdof = 3 * nq
+    for k in range(dof):
+        qp = q0.copy()
+        qm = q0.copy()
+        mp = m0.copy()
+        mm_m = m0.copy()
+
+        if k < nqdof:
+            a, c = k // 3, k % 3
+            qp[a, c] += step
+            qm[a, c] -= step
+        else:
+            kk = k - nqdof
+            a, c = kk // 3, kk % 3
+            mp[a, c] += step
+            mm_m[a, c] -= step
+
+        f_plus = np.asarray(eval_full_forces(qp, mp), dtype=np.float64).reshape(-1)
+        f_minus = np.asarray(eval_full_forces(qm, mm_m), dtype=np.float64).reshape(-1)
+        hess[:, k] = (-f_plus + f_minus) / (2.0 * step)
+
+    hess = 0.5 * (hess + hess.T)
+    return f0, hess
+
+
+def delta_embedcharge_minus_noembed(
+    symbols,
+    coords_q_ang,
+    mm_coords_ang,
+    mm_charges,
+    charge,
+    multiplicity,
+    need_forces,
+    need_hessian,
+    xtb_cmd="xtb",
+    xtb_acc=0.2,
+    xtb_workdir="tmp",
+    xtb_keep_files=False,
+    ncores=4,
+    hessian_step=1.0e-3,
+):
+    """Return xTB point-charge embedding correction in eV / eV/Š/ eV/Ų units.
+
+    Returns
+    -------
+    tuple
+        ``(dE_ev, dF_full_ev_ang, dH_full_ev_ang2)``, where full index order is
+        concatenated as ``[QM atoms, MM point charges]``.
+    """
+    q_symbols = list(symbols or [])
+    q_coords = np.asarray(coords_q_ang, dtype=np.float64).reshape(-1, 3)
+    m_coords = np.asarray(mm_coords_ang, dtype=np.float64).reshape(-1, 3)
+    m_charges = np.asarray(mm_charges, dtype=np.float64).reshape(-1)
+
+    nq = q_coords.shape[0]
+    nm = m_coords.shape[0]
+    nall = nq + nm
+    if nm <= 0:
+        zero_force = np.zeros((nall, 3), dtype=np.float64) if (need_forces or need_hessian) else None
+        zero_hess = np.zeros((3 * nall, 3 * nall), dtype=np.float64) if need_hessian else None
+        return 0.0, zero_force, zero_hess
+
+    run_ncores = _resolve_ncores(ncores)
+    if need_hessian:
+        run_ncores = 1
+
+    req_force = bool(need_forces) or bool(need_hessian)
+    de_ev, df_q, df_m = _delta_embed_minus_vac(
+        symbols=q_symbols, coords_q_ang=q_coords,
+        mm_coords_ang=m_coords, mm_charges=m_charges,
+        charge=charge, multiplicity=multiplicity,
+        need_forces=req_force,
+        xtb_cmd=xtb_cmd, xtb_acc=xtb_acc,
+        xtb_workdir=xtb_workdir, xtb_keep_files=xtb_keep_files,
+        ncores=run_ncores,
+    )
+
+    df_full = None
+    if req_force:
+        df_full = _assemble_full_force(df_q, df_m)
+
+    dh_full = None
+    if need_hessian:
+        step = float(hessian_step)
+
+        def _eval_full_forces(cur_q, cur_m):
+            _de, _fq, _fm = _delta_embed_minus_vac(
+                symbols=q_symbols, coords_q_ang=cur_q,
+                mm_coords_ang=cur_m, mm_charges=m_charges,
+                charge=charge, multiplicity=multiplicity,
+                need_forces=True,
+                xtb_cmd=xtb_cmd, xtb_acc=xtb_acc,
+                xtb_workdir=xtb_workdir, xtb_keep_files=xtb_keep_files,
+                ncores=run_ncores,
+            )
+            return _assemble_full_force(_fq, _fm)
+
+        _f0, dh_full = _numerical_hessian_from_forces(
+            eval_full_forces=_eval_full_forces,
+            coords_q_ang=q_coords, mm_coords_ang=m_coords,
+            step_ang=step, force0_ev_ang=df_full,
+        )
+
+    return float(de_ev), df_full, dh_full