bfflearn 0.1.0.dev1__py3-none-any.whl

bff/io/logs.py

@@ -0,0 +1,213 @@
+import logging
+from typing import Optional
+
+import numpy as np
+
+from ..mcmc.sampler import Sampler
+
+
+class Logger:
+    """
+    A logging utility that supports logging to a file
+    or console without duplicate output.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        fn_log: Optional[str] = None,
+        width: Optional[int] = 100,
+        verbose: bool = True,
+        mode: str = "a",
+    ) -> None:
+
+        """
+        Parameters
+        ----------
+        name : str
+            Name of the logger.
+        fn_log : Optional[str], optional
+            Filename to log messages to. If None, logs to console.
+        width : Optional[int], default=100
+            Width for message formatting.
+        verbose : bool, default=True
+            If False, suppresses logging output.
+        mode : str, default="a"
+            File open mode used when ``fn_log`` is provided.
+        """
+
+        self.fn_log = fn_log
+        self.width = width
+        self.verbose = verbose
+
+        self.logger = logging.getLogger(name)
+        self.logger.setLevel(logging.INFO)
+
+        # Remove old handlers if re-running in Jupyter
+        if self.logger.hasHandlers():
+            self.logger.handlers.clear()
+
+        formatter = logging.Formatter("%(message)s")
+
+        if fn_log:
+            handler = logging.FileHandler(fn_log, mode=mode)
+        else:
+            handler = logging.StreamHandler()
+            handler.terminator = ""
+        handler.setFormatter(formatter)
+        self.logger.addHandler(handler)
+
+        self.logger.propagate = False  # Prevent double logging
+
+    def info(self, message: str, level: int, overwrite: bool = False) -> None:
+        """
+        Logs a message to a file or console, optionally overwriting the console output.
+
+        Parameters
+        ----------
+        message : str
+            The message to log.
+        overwrite : bool, default=False
+            If True, the message overwrites the previous console output (stdout only).
+        """
+
+        # determine width of the line automatically if not set
+        width = len(message) + 10 if self.width is None else self.width
+
+        if self.verbose:
+            if level < 1:
+                start = ''
+            elif level == 1:
+                start = '> '
+            else:
+                start = ' ' * 2 * (level - 1) + '- '
+            message = start + f"{message}".ljust(width)
+            if self.fn_log:
+                self.logger.info(message)  # Always log to file if specified
+
+            else:
+                terminator = '\r' if overwrite else '\n'
+                self.logger.info(message + terminator)  # Log to logger (console)
+
+
+def print_progress_mcmc(
+    sampler: Sampler,
+    p0: np.ndarray,
+    *,
+    total_steps: int,
+    restart: bool = True,
+    fn_checkpoint: Optional[str] = None,
+    logger: Optional[Logger] = None,
+    **kwargs
+) -> None:
+    """
+    Monitors and logs the convergence progress of an MCMC sampler.
+
+    Parameters
+    ----------
+    sampler : Sampler
+        The MCMC sampler with methods
+        `sample`, `get_autocorr_time`, and attribute `iteration`.
+    p0 : array-like
+        Initial parameter vector for the MCMC posterior sampling.
+    total_steps : int
+        Total number of MCMC iterations including warmup.
+    stride : int, default=100
+        Frequency of autocorrelation time checks and printing.
+    min_chain_length : int, default=100
+        Factor to multiply the autocorrelation time with in order to
+        get the minimum chain length to have proper sampling.
+    rtol : float, default=0.01
+        Relative tolerance of the autocorrelation time fluctuations
+        for the MCMC sampling to be considered converged.
+    logger : Logger, optional
+        Logger instance for logging progress messages.
+    **kwargs
+        Additional arguments passed to the sampler's `sample` method.
+    """
+
+    logger = logger or Logger("mcmc-progress")
+
+    rhat_tol = kwargs.get("rhat_tol", 1.01)
+    ess_target = kwargs.get("ess_min", 100)
+    total_digits = len(str(total_steps))
+    warmup_digits = len(str(kwargs.get("warmup", total_steps)))
+    sampling_digits = len(str(max(total_steps - kwargs.get("warmup", 0), 0)))
+    ess_digits = max(len(str(int(ess_target))), 1)
+
+    def _fmt_progress(current: int, total: int, width: int, label: str) -> str:
+        return f"{label}: {current:>{width}d}/{total:<{width}d}"
+
+    if fn_checkpoint is not None:
+        logger.info(f"Checkpoint: {fn_checkpoint}", level=2)
+
+    logger.info(
+        (
+            f"Posterior sampling: it. {0:>{total_digits}d}/"
+            f"{total_steps:<{total_digits}d}"
+        ),
+        level=1,
+        overwrite=True,
+    )
+
+    line = ""
+    for state in sampler.run(
+        p0,
+        total_steps=total_steps,
+        restart=restart,
+        fn_checkpoint=fn_checkpoint,
+        **kwargs,
+    ):
+        total_steps = state.total_steps
+        sampling_steps = total_steps - state.warmup
+        if state.phase == "warmup":
+            phase_progress = _fmt_progress(
+                state.step,
+                state.warmup,
+                warmup_digits,
+                "warmup",
+            )
+        else:
+            sampling_step = max(state.step - state.warmup, 0)
+            phase_progress = _fmt_progress(
+                sampling_step,
+                sampling_steps,
+                sampling_digits,
+                "sampling",
+            )
+
+        if state.phase == "sampling" and state.convergence is not None:
+            line = (
+                f"Posterior sampling: it. "
+                f"{state.step:>{total_digits}d}/{total_steps:<{total_digits}d} | "
+                f"{phase_progress} | "
+                f"R-hat max: {state.convergence.max_rhat:.4f}/{rhat_tol:.4f} | "
+                f"ESS min: {state.convergence.min_ess:>{ess_digits}.0f}/"
+                f"{ess_target:<{ess_digits}.0f}"
+            )
+            if state.acceptance_rate is not None:
+                line += f" | acc: {state.acceptance_rate:.3f}"
+            if state.it_per_sec is not None:
+                line += f" | {state.it_per_sec:>3.0f} it/s"
+        else:
+            line = (
+                f"Posterior sampling: it. "
+                f"{state.step:>{total_digits}d}/{total_steps:<{total_digits}d} | "
+                f"{phase_progress}"
+            )
+            if state.acceptance_rate is not None:
+                line += f" | acc: {state.acceptance_rate:.3f}"
+            if state.it_per_sec is not None:
+                line += f" | {state.it_per_sec:>3.0f} it/s"
+        logger.info(line, level=1, overwrite=True)
+
+    if line and logger.fn_log is None:
+        logger.info(line, level=1)
+    logger.info("", level=0)
+    if sampler.converged:
+        logger.info("Posterior sampling: Done.", level=1)
+    else:
+        logger.info(
+            "Posterior sampling: Failed to converge within the maximum iterations.",
+            level=1,
+        )

bff/io/mdp.py

@@ -0,0 +1,98 @@
+from collections import OrderedDict
+from pathlib import Path
+from typing import Mapping, Union
+
+PathLike = Union[str, Path]
+
+
+def read_mdp(fn_mdp: PathLike) -> OrderedDict[str, str]:
+    """Read a GROMACS MDP file while preserving comments and blank lines.
+
+    Parameters
+    ----------
+    fn_mdp
+        Path to the input MDP file.
+
+    Returns
+    -------
+    collections.OrderedDict
+        Ordered mapping of MDP keys to values. Comment and blank lines are
+        stored under synthetic keys starting with ``C`` and ``B``.
+    """
+    content: OrderedDict[str, str] = OrderedDict()
+    i_comment = 0
+    i_blank = 0
+    with open(fn_mdp, "r") as handle:
+        for line in handle:
+            if line.startswith(";"):
+                content[f"C{i_comment:03d}"] = line
+                i_comment += 1
+            elif line.strip() == "":
+                content[f"B{i_blank:03d}"] = line
+                i_blank += 1
+            else:
+                key, value = line.split("=", 1)
+                content[key.strip()] = value.strip()
+    return content
+
+
+def write_mdp(content: Mapping[str, str], fn_out: PathLike) -> None:
+    """Write an MDP mapping back to disk.
+
+    Parameters
+    ----------
+    content
+        Ordered MDP mapping produced by :func:`read_mdp`.
+    fn_out
+        Output MDP path.
+    """
+    with open(fn_out, "w") as handle:
+        for key, value in content.items():
+            if key.startswith("C") or key.startswith("B"):
+                handle.write(value)
+            else:
+                handle.write(f"{key:<25} = {value}\n")
+
+
+def patch_mdp(
+    fn_mdp: PathLike,
+    updates: Mapping[str, str],
+    fn_out: PathLike,
+) -> None:
+    """Write a modified MDP file with a small set of updated parameters.
+
+    Parameters
+    ----------
+    fn_mdp
+        Input MDP file.
+    updates
+        Mapping of MDP keys to replacement values.
+    fn_out
+        Output MDP file.
+    """
+    content = read_mdp(fn_mdp)
+    for key, value in updates.items():
+        content[str(key)] = str(value)
+    write_mdp(content, fn_out)
+
+
+def get_n_frames_target(fn_mdp: PathLike) -> tuple[int | None, int | None]:
+    """Extract the expected number of saved trajectory frames.
+
+    Parameters
+    ----------
+    fn_mdp
+        Input MDP file.
+
+    Returns
+    -------
+    tuple
+        Number of saved frames and trajectory stride. If ``nsteps`` is zero or
+        missing, both entries are returned as ``None``.
+    """
+    mdp_data = read_mdp(fn_mdp)
+    n_steps = int(mdp_data.get("nsteps", 0))
+    if n_steps <= 0:
+        return None, None
+    stride = int(mdp_data["nstxout-compressed"])
+    return int(n_steps / stride), stride

bff/io/plumed.py

@@ -0,0 +1,184 @@
+import os
+import re
+import shutil
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import MDAnalysis as mda
+
+PathLike = str | Path
+
+
+@dataclass(frozen=True, slots=True)
+class PlumedDistanceBias:
+    """Parsed distance bias from a PLUMED file.
+
+    Parameters
+    ----------
+    name
+        Name of the distance collective variable.
+    atom_indices
+        One-based atom indices defining the distance.
+    center
+        Optional restraint center in nm.
+    """
+
+    name: str
+    atom_indices: tuple[int, int]
+    center: float | None = None
+
+
+def _parse_arguments(text: str) -> dict[str, str]:
+    """Parse key-value arguments from one PLUMED line."""
+    pairs = re.findall(r"([A-Za-z_]+)=([^\s]+)", text)
+    return {key.upper(): value for key, value in pairs}
+
+
+def find_plumed_kernel() -> Path | None:
+    """Locate a PLUMED kernel library if one is available.
+
+    Returns
+    -------
+    pathlib.Path or None
+        Path to a loadable PLUMED kernel library, or ``None`` if one could
+        not be located from the current environment.
+    """
+    env_value = os.environ.get("PLUMED_KERNEL")
+    if env_value:
+        env_path = Path(env_value).expanduser().resolve()
+        if env_path.exists():
+            return env_path
+
+    candidate_dirs = [
+        Path(sys.prefix) / "lib",
+        Path(sys.prefix).resolve().parent / "lib",
+    ]
+
+    plumed_executable = shutil.which("plumed")
+    if plumed_executable is not None:
+        candidate_dirs.append(Path(plumed_executable).resolve().parents[1] / "lib")
+
+    seen: set[Path] = set()
+    for lib_dir in candidate_dirs:
+        lib_dir = lib_dir.resolve()
+        if lib_dir in seen or not lib_dir.exists():
+            continue
+        seen.add(lib_dir)
+        candidates = sorted(lib_dir.glob("libplumedKernel*"))
+        if candidates:
+            return candidates[0]
+
+    return None
+
+
+def ensure_plumed_kernel() -> Path:
+    """Ensure that a PLUMED kernel is available for biased MD runs.
+
+    Returns
+    -------
+    pathlib.Path
+        Path to the PLUMED kernel library.
+
+    Raises
+    ------
+    RuntimeError
+        If no PLUMED kernel can be located.
+    """
+    kernel = find_plumed_kernel()
+    if kernel is None:
+        raise RuntimeError(
+            "PLUMED biasing requires a loadable PLUMED kernel. "
+            "Set the PLUMED_KERNEL environment variable or install PLUMED "
+            "so that libplumedKernel is discoverable."
+        )
+    return kernel
+
+
+def parse_distance_biases(fn_plumed: PathLike) -> list[PlumedDistanceBias]:
+    """Parse simple distance restraints from a PLUMED input file.
+
+    Parameters
+    ----------
+    fn_plumed
+        Path to the PLUMED input file.
+
+    Returns
+    -------
+    list of PlumedDistanceBias
+        Parsed distance collective variables with optional restraint centers.
+    """
+    lines = Path(fn_plumed).read_text().splitlines()
+    distances: dict[str, tuple[int, int]] = {}
+    centers: dict[str, float] = {}
+
+    for raw_line in lines:
+        line = raw_line.split("#", maxsplit=1)[0].strip()
+        if not line or ":" not in line:
+            continue
+
+        name, expression = [part.strip() for part in line.split(":", maxsplit=1)]
+        keyword = expression.split(maxsplit=1)[0].upper()
+        args = _parse_arguments(expression)
+
+        if keyword == "DISTANCE" and "ATOMS" in args:
+            atom_indices = tuple(int(value) for value in args["ATOMS"].split(","))
+            if len(atom_indices) == 2:
+                distances[name] = atom_indices  # type: ignore[assignment]
+            continue
+
+        if keyword in {"RESTRAINT", "UPPER_WALLS", "LOWER_WALLS"}:
+            arg_name = args.get("ARG")
+            if arg_name is None:
+                continue
+            center_value = args.get("AT")
+            if center_value is None:
+                continue
+            try:
+                centers[arg_name] = float(center_value.split(",")[0])
+            except ValueError:
+                continue
+
+    return [
+        PlumedDistanceBias(
+            name=name,
+            atom_indices=atom_indices,
+            center=centers.get(name),
+        )
+        for name, atom_indices in distances.items()
+    ]
+
+
+def resolve_distance_bias_metadata(
+    fn_plumed: PathLike,
+    fn_system: PathLike,
+) -> tuple[tuple[str, ...], tuple[float, ...]]:
+    """Resolve atom-name pairs and centers from a PLUMED file.
+
+    Parameters
+    ----------
+    fn_plumed
+        Path to the PLUMED 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
+        and centers are restraint centers in nm.
+    """
+    universe = mda.Universe(str(fn_system))
+    distances = parse_distance_biases(fn_plumed)
+
+    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(distance.center)
+
+    return tuple(labels), tuple(centers)

bff/io/progress.py

@@ -0,0 +1,74 @@
+import time
+from typing import Any, Iterable, Iterator, TypeVar
+
+T = TypeVar("T")
+
+
+def format_time(seconds: float) -> str:
+    """Format seconds as hours, minutes, and seconds."""
+    hours, remainder = divmod(seconds, 3600)
+    minutes, seconds = divmod(remainder, 60)
+
+    if hours > 0:
+        return f"{int(hours)}h {int(minutes)}m {int(seconds)}s"
+    if minutes > 0:
+        return f"{int(minutes)}m {int(seconds)}s"
+    return f"{int(seconds)}s"
+
+
+def iter_progress(
+    iterable: Iterable[T],
+    *,
+    total: int,
+    logger: Any,
+    label: str,
+    stride: int = 1,
+) -> Iterator[T]:
+    """Yield items while logging coarse progress updates."""
+    if stride < 1:
+        raise ValueError("'stride' must be a positive integer.")
+    if total < 0:
+        raise ValueError("'total' must be non-negative.")
+    if total == 0:
+        return
+
+    start_time = time.time()
+    pad = len(str(total))
+    template = (
+        f"{label}: " + "{i:>{pad}}/{total} "
+        "({percent:3.0f}%) | {elapsed} < {eta}"
+    )
+
+    logger.info(
+        template.format(
+            i=0,
+            pad=pad,
+            total=total,
+            percent=0,
+            elapsed="0s",
+            eta="0s",
+        ),
+        level=1,
+        overwrite=True,
+    )
+
+    for i, item in enumerate(iterable, start=1):
+        yield item
+
+        if (i % stride != 0) and (i != total):
+            continue
+
+        elapsed_time = time.time() - start_time
+        eta = 0.0 if i == 0 else (elapsed_time / i) * max(total - i, 0)
+        logger.info(
+            template.format(
+                i=i,
+                pad=pad,
+                total=total,
+                percent=100 * i / total if total else 100,
+                elapsed=format_time(elapsed_time),
+                eta=format_time(eta),
+            ),
+            level=1,
+            overwrite=True,
+        )

bff/io/schedulers.py

@@ -0,0 +1,118 @@
+import subprocess
+from pathlib import Path
+
+
+class Slurm:
+    """
+    A class to generate and save SLURM job submission scripts.
+
+    Parameters
+    ----------
+    job_name : str
+        Name of the SLURM job.
+    nodes : int
+        Number of nodes to allocate.
+    ntasks_per_node : int
+        Number of tasks per node.
+    cpus_per_task : int
+        Number of CPUs per task.
+    time : str
+        Maximum wall time in the format 'HH:MM:SS'.
+    partition : str
+        Partition or queue name.
+    mem : str
+        Memory allocation (e.g., '4G', '16G').
+    output : str, optional
+        Path to the output file, default is './slurm-files/slurm-%x.%j.out'.
+
+    Attributes
+    ----------
+    commands : list
+        List of shell commands to execute in the SLURM script.
+    """
+
+    def __init__(self, **sbatch):
+        """
+        Initialize the class with SBATCH parameters.
+
+        Parameters
+        ----------
+        **sbatch : dict
+            Arbitrary keyword arguments representing SBATCH parameters.
+            Each key-value pair will be formatted as '#SBATCH --key=value'.
+        """
+
+        self.sbatch = [
+            f"#SBATCH --{key.replace('_', '-')}={value}"
+            for key, value in sbatch.items()
+        ]
+        self.commands = []
+        self.fname = None
+
+    def add_command(self, command: str) -> None:
+        """
+        Add a shell command to the script.
+
+        Parameters
+        ----------
+        command : str
+            The shell command to add.
+        """
+        self.commands.append(command)
+
+    def generate(self) -> str:
+        """
+        Generate the SLURM script content as a string.
+
+        Returns
+        -------
+        str
+            The complete SLURM script.
+        """
+        script = "#!/bin/bash\n"
+        if self.sbatch:
+            script += "\n".join(self.sbatch) + "\n"
+        if self.commands:
+            script += "\n" + "\n".join(self.commands) + "\n"
+        return script
+
+    def save(self, filename: str | Path) -> None:
+        """
+        Save the SLURM script to a file.
+
+        Parameters
+        ----------
+        filename : str
+            Path to save the SLURM script.
+        """
+        self.fname = Path(filename)
+        with open(self.fname, "w") as file:
+            file.write(self.generate())
+
+    def submit(self, filename: str | Path) -> int:
+        """
+        Submit the SLURM script to the queue.
+        """
+        self.save(filename)
+        out = subprocess.run(
+            ["sbatch", str(self.fname)],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        if out.returncode != 0:
+            message = out.stderr.strip() or out.stdout.strip() or "unknown sbatch error"
+            raise RuntimeError(f"SLURM submission failed: {message}")
+
+        tokens = out.stdout.strip().split()
+        if not tokens:
+            raise RuntimeError(
+                "SLURM submission failed: could not parse job ID from sbatch output."
+            )
+        try:
+            return int(tokens[-1])
+        except ValueError as exc:
+            raise RuntimeError(
+                "SLURM submission failed: could not parse job ID from "
+                f"sbatch output {out.stdout!r}."
+            ) from exc