bfflearn 0.1.0.dev1__py3-none-any.whl

bff/io/__init__.py

@@ -0,0 +1,12 @@
+from . import colvars, cp2k, logs, mdp, plumed, progress, schedulers, utils
+
+__all__ = [
+    "colvars",
+    "cp2k",
+    "logs",
+    "mdp",
+    "plumed",
+    "progress",
+    "schedulers",
+    "utils",
+]

bff/io/colvars.py

@@ -0,0 +1,186 @@
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+import MDAnalysis as mda
+
+from .mdp import patch_mdp
+
+PathLike = str | Path
+
+
+@dataclass(frozen=True, slots=True)
+class ColvarsDistanceBias:
+    """Parsed distance bias from a COLVARS file.
+
+    Parameters
+    ----------
+    name
+        Colvar name.
+    atom_indices
+        One-based atom indices of the restrained pair.
+    center
+        Optional harmonic center in nm.
+    """
+
+    name: str
+    atom_indices: tuple[int, int]
+    center: float | None = None
+
+
+def _extract_named_blocks(text: str, name: str) -> list[str]:
+    """Extract balanced top-level blocks of one COLVARS section type."""
+    pattern = re.compile(rf"\b{name}\s*\{{", re.MULTILINE)
+    blocks: list[str] = []
+    start = 0
+    while True:
+        match = pattern.search(text, start)
+        if match is None:
+            return blocks
+
+        i = match.end() - 1
+        depth = 0
+        end = i
+        while end < len(text):
+            if text[end] == "{":
+                depth += 1
+            elif text[end] == "}":
+                depth -= 1
+                if depth == 0:
+                    blocks.append(text[i + 1:end])
+                    start = end + 1
+                    break
+            end += 1
+        else:
+            raise ValueError(f"Unbalanced '{name}' block in COLVARS file.")
+
+
+def _extract_atom_numbers(block: str, group_name: str) -> tuple[int, ...]:
+    """Extract atom numbers from one COLVARS group block."""
+    groups = _extract_named_blocks(block, group_name)
+    if not groups:
+        return ()
+    match = re.search(r"\batomNumbers\s+([0-9\s]+)", groups[0])
+    if match is None:
+        return ()
+    return tuple(int(value) for value in match.group(1).split())
+
+
+def parse_distance_biases(fn_colvars: PathLike) -> list[ColvarsDistanceBias]:
+    """Parse simple distance/harmonic restraints from a COLVARS file.
+
+    Parameters
+    ----------
+    fn_colvars
+        Path to the COLVARS input file.
+
+    Returns
+    -------
+    list of ColvarsDistanceBias
+        Parsed single-pair distance colvars. Only colvars with one atom in
+        ``group1`` and one atom in ``group2`` are returned.
+    """
+    text = Path(fn_colvars).read_text()
+
+    colvars: dict[str, tuple[int, int]] = {}
+    for block in _extract_named_blocks(text, "colvar"):
+        name_match = re.search(r"^\s*name\s+(\S+)", block, re.MULTILINE)
+        if name_match is None:
+            continue
+        distance_blocks = _extract_named_blocks(block, "distance")
+        if not distance_blocks:
+            continue
+        group_1 = _extract_atom_numbers(distance_blocks[0], "group1")
+        group_2 = _extract_atom_numbers(distance_blocks[0], "group2")
+        if len(group_1) != 1 or len(group_2) != 1:
+            continue
+        colvars[name_match.group(1)] = (group_1[0], group_2[0])
+
+    centers: dict[str, float] = {}
+    for block in _extract_named_blocks(text, "harmonic"):
+        name_match = re.search(r"^\s*colvars\s+(\S+)", block, re.MULTILINE)
+        center_match = re.search(r"^\s*centers\s+([-+0-9.eE]+)", block, re.MULTILINE)
+        if name_match is None or center_match is None:
+            continue
+        centers[name_match.group(1)] = float(center_match.group(1))
+
+    return [
+        ColvarsDistanceBias(
+            name=name,
+            atom_indices=indices,
+            center=centers.get(name),
+        )
+        for name, indices in colvars.items()
+    ]
+
+
+def resolve_distance_bias_metadata(
+    fn_colvars: PathLike,
+    fn_system: PathLike,
+) -> tuple[tuple[str, ...], tuple[float, ...]]:
+    """Resolve atom-name pairs and centers from a COLVARS file.
+
+    Parameters
+    ----------
+    fn_colvars
+        Path to the COLVARS input file.
+    fn_system
+        Structure file used to map atom indices to atom names.
+
+    Returns
+    -------
+    tuple
+        Tuple of ``(pair_labels, centers)`` where labels are atom-name pairs
+        such as ``"C2 CAL"`` and centers are harmonic centers in nm.
+    """
+    universe = mda.Universe(str(fn_system))
+    distances = parse_distance_biases(fn_colvars)
+
+    labels: list[str] = []
+    centers: list[float] = []
+    for distance in distances:
+        index_1, index_2 = distance.atom_indices
+        atom_1 = universe.atoms[index_1 - 1].name
+        atom_2 = universe.atoms[index_2 - 1].name
+        labels.append(f"{atom_1} {atom_2}")
+        if distance.center is not None:
+            centers.append(float(distance.center))
+
+    return tuple(labels), tuple(centers)
+
+
+def write_mdp_with_colvars(
+    fn_mdp: PathLike,
+    fn_colvars: PathLike,
+    fn_out: PathLike,
+    *,
+    seed: int | None = None,
+) -> None:
+    """Write an MDP file with GROMACS Colvars enabled.
+
+    Parameters
+    ----------
+    fn_mdp
+        Input MDP file.
+    fn_colvars
+        Colvars configuration file referenced from the generated MDP.
+    fn_out
+        Output MDP file with Colvars options injected.
+    seed
+        Optional seed passed to the Colvars module.
+    """
+    updates = {
+        "colvars-active": "yes",
+        "colvars-configfile": f"./{Path(fn_colvars).name}",
+    }
+    if seed is not None:
+        updates["colvars-seed"] = str(int(seed))
+    patch_mdp(fn_mdp, updates, fn_out)
+    with open(fn_out, "r+") as handle:
+        content = handle.read()
+        block = "\n; colvars\n"
+        if "colvars-active" in content:
+            content = content.replace("colvars-active", block + "colvars-active", 1)
+        handle.seek(0)
+        handle.write(content)
+        handle.truncate()

bff/io/commands.py

@@ -0,0 +1,20 @@
+import shlex
+from collections.abc import Sequence
+from pathlib import Path
+
+
+def split_command(command: str | Sequence[str]) -> list[str]:
+    """Normalize a shell-style command into a subprocess argument vector."""
+    if isinstance(command, str):
+        parts = shlex.split(command)
+    else:
+        parts = [str(part) for part in command]
+
+    if not parts:
+        raise ValueError("Command must not be empty.")
+    return parts
+
+
+def build_command(command: str | Sequence[str], *args: str | Path) -> list[str]:
+    """Append stringified arguments to a normalized command vector."""
+    return [*split_command(command), *(str(arg) for arg in args)]

bff/io/cp2k.py

@@ -0,0 +1,330 @@
+from copy import deepcopy
+from pathlib import Path
+from typing import Any
+
+import MDAnalysis as mda
+import numpy as np
+
+from ..data import CATIONS, CP2K_INPUT_TEMPLATE, CP2K_KIND_DEFAULTS
+
+
+def format_cp2k_input(data: dict[str, Any], indent: int = 0) -> list[str]:
+    """Format a CP2K input tree as text.
+
+    Parameters
+    ----------
+    data
+        Nested mapping representing the CP2K input tree.
+    indent
+        Current indentation level.
+
+    Returns
+    -------
+    list of str
+        Formatted CP2K input lines.
+    """
+    lines: list[str] = []
+    prefix = " " * indent
+
+    for key, value in data.items():
+        section_name = key.upper()
+        section_end = key.split()[0].upper()
+
+        if isinstance(value, dict):
+            lines.append(f"{prefix}&{section_name}")
+            lines.extend(format_cp2k_input(value, indent + 2))
+            lines.append(f"{prefix}&END {section_end}")
+            continue
+
+        if isinstance(value, list):
+            for item in value:
+                if isinstance(item, dict):
+                    lines.extend(format_cp2k_input({key: item}, indent))
+                elif isinstance(item, list):
+                    lines.append(f"{prefix}{section_name} {' '.join(map(str, item))}")
+                else:
+                    lines.append(f"{prefix}{section_name} {item}")
+            continue
+
+        if isinstance(value, bool):
+            value = ".TRUE." if value else ".FALSE."
+
+        lines.append(f"{prefix}{section_name} {value}")
+
+    return lines
+
+
+def get_cp2k_elements(fn_pos: str | Path) -> np.ndarray:
+    """Return element symbols in coordinate-file order without duplicates."""
+    universe = mda.Universe(fn_pos)
+    elements, index = np.unique(universe.atoms.elements, return_index=True)
+    return elements[np.argsort(index)]
+
+
+def build_cp2k_kinds(elements: np.ndarray) -> dict[str, dict[str, str]]:
+    """Build CP2K KIND sections for the elements present in a system."""
+    return {
+        f"kind {element.capitalize()}": CP2K_KIND_DEFAULTS[element.lower()]
+        for element in elements
+    }
+
+
+def build_cp2k_d3_exclusions(elements: np.ndarray) -> list[list[int]]:
+    """Build D3 kind-pair exclusions for cation-containing systems."""
+    return [
+        [i + 1, j + 1]
+        for i, element in enumerate(elements)
+        if element.lower() in CATIONS
+        for j in range(len(elements))
+    ]
+
+
+def configure_cp2k_restart(
+    tree: dict[str, Any],
+    *,
+    project: str,
+    restart: bool,
+) -> None:
+    """Configure CP2K SCF and restart sections."""
+    if restart:
+        tree["ext_restart"] = {"restart_file_name": f"{project}-1.restart"}
+        tree["force_eval"]["dft"]["scf"]["scf_guess"] = "RESTART"
+        tree["force_eval"]["dft"]["wfn_restart_file_name"] = (
+            f"{project}-RESTART.wfn"
+        )
+        return
+
+    tree["force_eval"]["dft"]["scf"]["scf_guess"] = "ATOMIC"
+
+
+def configure_cp2k_ensemble(
+    tree: dict[str, Any],
+    *,
+    equilibration: bool,
+) -> None:
+    """Configure the MD ensemble for equilibration or production."""
+    tree["motion"]["md"]["ensemble"] = "LANGEVIN" if equilibration else "NVT"
+    if equilibration:
+        tree["motion"]["md"]["langevin"] = {"gamma": 0.02}
+        tree["motion"]["md"].pop("thermostat", None)
+        return
+
+    tree["motion"]["md"]["thermostat"] = {
+        "region": "GLOBAL",
+        "type": "CSVR",
+        "csvr": {"timecon": "[fs] 100"},
+    }
+    tree["motion"]["md"].pop("langevin", None)
+
+
+def configure_cp2k_plumed(
+    tree: dict[str, Any],
+    *,
+    plumed_input_file: str | Path,
+) -> None:
+    """Configure CP2K to use a PLUMED input file for biased MD.
+
+    Parameters
+    ----------
+    tree
+        CP2K input tree to update.
+    plumed_input_file
+        PLUMED input filename referenced from the CP2K input.
+    """
+    tree["motion"]["free_energy"] = {
+        "metadyn": {
+            "use_plumed": True,
+            "plumed_input_file": Path(plumed_input_file).name,
+        }
+    }
+
+
+def configure_cp2k_single_point(tree: dict[str, Any]) -> None:
+    """Configure a CP2K input tree for single-point energy/forces."""
+    tree["global"]["run_type"] = "ENERGY_FORCE"
+    tree.pop("motion", None)
+
+
+def build_cp2k_tree(
+    *,
+    project: str,
+    charge: int,
+    multiplicity: int,
+    unitcell: list[float],
+    fn_pos: str | Path,
+) -> dict[str, Any]:
+    """Build the common CP2K input tree for one structure."""
+    fn_pos = Path(fn_pos)
+    elements = get_cp2k_elements(fn_pos)
+    tree = deepcopy(CP2K_INPUT_TEMPLATE)
+
+    tree["global"]["project_name"] = project
+    tree["force_eval"]["subsys"]["topology"] = {
+        "coord_file_name": fn_pos.name,
+        "coord_file_format": "XYZ",
+    }
+    tree["force_eval"]["subsys"]["cell"] = {
+        "ABC": f"[angstrom] {unitcell[0]} {unitcell[1]} {unitcell[2]}",
+        "periodic": "XYZ",
+    }
+    tree["force_eval"]["subsys"].update(build_cp2k_kinds(elements))
+
+    tree["force_eval"]["dft"]["charge"] = charge
+    tree["force_eval"]["dft"]["multiplicity"] = multiplicity
+    tree["force_eval"]["dft"]["uks"] = multiplicity != 1
+
+    if any(element.lower() in CATIONS for element in elements):
+        tree["force_eval"]["dft"]["xc"]["vdw_potential"]["pair_potential"][
+            "d3_exclude_kind_pair"
+        ] = build_cp2k_d3_exclusions(elements)
+
+    return tree
+
+
+def make_cp2k_input(
+    project: str,
+    charge: int,
+    multiplicity: int,
+    unitcell: list[float],
+    fn_pos: str | Path,
+    equilibration: bool,
+    restart: bool,
+    fn_out: str | Path,
+    plumed_input_file: str | Path | None = None,
+) -> None:
+    """Generate a CP2K input file for one simulation window.
+
+    Parameters
+    ----------
+    project
+        Base project name.
+    charge
+        Total charge of the system.
+    multiplicity
+        Spin multiplicity.
+    unitcell
+        Unit-cell dimensions in Angstrom.
+    fn_pos
+        Path to the XYZ coordinate file.
+    equilibration
+        Whether to generate an equilibration input.
+    restart
+        Whether the CP2K run restarts from a previous step.
+    fn_out
+        Output CP2K input file.
+    plumed_input_file
+        Optional PLUMED input file used for biased CP2K MD.
+    """
+    tree = build_cp2k_tree(
+        project=project,
+        charge=charge,
+        multiplicity=multiplicity,
+        unitcell=unitcell,
+        fn_pos=fn_pos,
+    )
+    configure_cp2k_restart(tree, project=project, restart=restart)
+    configure_cp2k_ensemble(tree, equilibration=equilibration)
+    if plumed_input_file is not None:
+        configure_cp2k_plumed(tree, plumed_input_file=plumed_input_file)
+
+    with open(fn_out, "w") as handle:
+        handle.write("\n".join(format_cp2k_input(tree)) + "\n")
+
+
+def make_cp2k_single_point_input(
+    project: str,
+    charge: int,
+    multiplicity: int,
+    unitcell: list[float],
+    fn_pos: str | Path,
+    fn_out: str | Path,
+) -> None:
+    """Generate a CP2K single-point input file."""
+    tree = build_cp2k_tree(
+        project=project,
+        charge=charge,
+        multiplicity=multiplicity,
+        unitcell=unitcell,
+        fn_pos=fn_pos,
+    )
+    configure_cp2k_single_point(tree)
+    tree["global"]["print_level"] = "MEDIUM"
+
+    with open(fn_out, "w") as handle:
+        handle.write("\n".join(format_cp2k_input(tree)) + "\n")
+
+
+def write_cp2k_md_slurm_script(
+    fn_out: str | Path,
+    *,
+    uses_plumed: bool = False,
+) -> None:
+    """Write a restart-aware Slurm script for CP2K MD."""
+    plumed_block = ""
+    if uses_plumed:
+        plumed_block = """
+if [ -f plumed.dat ]; then
+  echo "Using PLUMED input: plumed.dat"
+fi
+"""
+
+    script = f"""#!/bin/bash
+#SBATCH --job-name=cp2k-md
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --time=24:00:00
+#SBATCH --output=slurm-%j.out
+
+set -euo pipefail
+
+CP2K_CMD="${{CP2K_CMD:-cp2k.psmp}}"
+{plumed_block}
+
+if [ ! -f .md-eq-start.done ]; then
+  ${{CP2K_CMD}} -i md-eq-start.inp -o md-eq-start.out
+  touch .md-eq-start.done
+fi
+
+if [ ! -f .md-eq-restart.done ]; then
+  ${{CP2K_CMD}} -i md-eq-restart.inp -o md-eq-restart.out
+  touch .md-eq-restart.done
+fi
+
+${{CP2K_CMD}} -i md-prod.inp -o md-prod.out
+"""
+    Path(fn_out).write_text(script)
+
+
+def write_cp2k_single_point_slurm_script(
+    fn_out: str | Path,
+    *,
+    snapshots_dirname: str = "snapshots",
+) -> None:
+    """Write a Slurm script for batched CP2K single-point evaluations."""
+    script = f"""#!/bin/bash
+#SBATCH --job-name=cp2k-sp
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --time=24:00:00
+#SBATCH --output=slurm-%j.out
+
+set -euo pipefail
+
+CP2K_CMD="${{CP2K_CMD:-cp2k.psmp}}"
+SNAPSHOT_DIR="{snapshots_dirname}"
+INPUT_TEMPLATE="single-point.inp"
+
+for xyz in "${{SNAPSHOT_DIR}}"/snapshot-*.xyz; do
+  [ -e "${{xyz}}" ] || continue
+  stem=$(basename "${{xyz}}" .xyz)
+  run_dir="runs/${{stem}}"
+  mkdir -p "${{run_dir}}"
+  cp "${{INPUT_TEMPLATE}}" "${{run_dir}}/single-point.inp"
+  cp "${{xyz}}" "${{run_dir}}/snapshot.xyz"
+  (
+    cd "${{run_dir}}"
+    ${{CP2K_CMD}} -i single-point.inp -o single-point.out
+  )
+done
+"""
+    Path(fn_out).write_text(script)