bfflearn 0.1.0.dev1__py3-none-any.whl

bff/io/utils.py

@@ -0,0 +1,208 @@
+import json
+import tarfile
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+import yaml
+
+PathLike = Union[str, Path]
+
+
+def path_to_str(path: PathLike) -> str:
+    """Convert a Path object to a string."""
+    return str(path)
+
+
+class NumpyYAMLEncoder(yaml.SafeDumper):
+    """YAML encoder for numpy arrays and numpy scalar types."""
+
+    def represent_numpy(self, obj):
+        if isinstance(obj, np.ndarray):
+            return self.represent_list(obj.tolist())
+        elif isinstance(obj, (np.generic, np.number)):
+            if isinstance(obj, np.floating):
+                return self.represent_float(obj.item())
+            else:
+                return self.represent_int(obj.item())
+        return super().represent_data(obj)
+
+
+# Add the custom representers for NumPy types
+NumpyYAMLEncoder.add_multi_representer(np.ndarray, NumpyYAMLEncoder.represent_numpy)
+NumpyYAMLEncoder.add_multi_representer(np.generic, NumpyYAMLEncoder.represent_numpy)
+
+
+def save_yaml(data: dict, fn: PathLike) -> None:
+    """Save a dictionary as YAML to a file."""
+    fn = str(fn) if isinstance(fn, Path) else fn
+    with open(fn, "w") as f:
+        yaml.dump(data, f, Dumper=NumpyYAMLEncoder, default_flow_style=False)
+
+
+def load_yaml(fn: PathLike) -> dict:
+    """Load .yaml file into a dictionary"""
+    fn = path_to_str(fn) if isinstance(fn, Path) else fn
+    with open(fn, "r") as f:
+        file = yaml.safe_load(f)
+    return file
+
+
+class NumpyArrayEncoder(json.JSONEncoder):
+    """JSON encoder for numpy arrays, scalars, and floats with 3 decimals."""
+
+    def default(self, obj):
+        if isinstance(obj, np.ndarray):
+            # Convert to nested lists, rounding elementwise
+            return self._round_nested(obj.tolist())
+        elif isinstance(obj, (np.generic, np.number)):
+            return round(obj.item(), 3)
+        elif isinstance(obj, float):
+            return round(obj, 3)
+        elif hasattr(obj, '__dict__'):
+            return obj.__dict__
+        return super().default(obj)
+
+    def _round_nested(self, obj):
+        """Recursively round floats in nested lists/structs."""
+        if isinstance(obj, list):
+            return [self._round_nested(x) for x in obj]
+        elif isinstance(obj, float):
+            return round(obj, 3)
+        return obj
+
+
+def save_json(data: dict, fn: PathLike) -> None:
+    """Save a dictionary as JSON with floats rounded to 3 decimals."""
+    fn = str(fn) if isinstance(fn, Path) else fn
+    with open(fn, "w") as f:
+        json.dump(data, f, cls=NumpyArrayEncoder)
+
+
+def load_json(fn: PathLike) -> dict:
+    """Save a dictionary as JSON to a file."""
+    fn = path_to_str(fn) if isinstance(fn, Path) else fn
+    with open(fn, "r") as f:
+        file = json.load(f)
+    return file
+
+
+def _to_json_compatible(data: object) -> object:
+    """Convert nested numpy objects into plain JSON-compatible values."""
+    if isinstance(data, dict):
+        return {str(key): _to_json_compatible(value) for key, value in data.items()}
+    if isinstance(data, (list, tuple)):
+        return [_to_json_compatible(value) for value in data]
+    if isinstance(data, np.ndarray):
+        return data.tolist()
+    if isinstance(data, (np.generic, np.number)):
+        return data.item()
+    return data
+
+
+def save_pt(data: object, fn: PathLike) -> None:
+    """Save a JSON-serializable object to a `.pt` file."""
+    path = Path(fn)
+    if path.suffix != ".pt":
+        raise ValueError(f"Expected a '.pt' file, got {path}.")
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(_to_json_compatible(data), f)
+
+
+def load_pt(fn: PathLike) -> object:
+    """Load an object from a `.pt` file.
+
+    The preferred format is plain JSON. Legacy Torch `.pt` archives are still
+    accepted when PyTorch is available.
+    """
+    path = Path(fn)
+    if path.suffix != ".pt":
+        raise ValueError(f"Expected a '.pt' file, got {path}.")
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except (UnicodeDecodeError, json.JSONDecodeError):
+        try:
+            import torch
+        except ModuleNotFoundError as exc:
+            raise RuntimeError(
+                f"Legacy binary .pt file detected at {path}, but PyTorch is not "
+                "installed. Install PyTorch or regenerate the file with the "
+                "current BFF version."
+            ) from exc
+        return torch.load(path, map_location="cpu", weights_only=False)
+
+
+def compress_results(source_dir: PathLike) -> None:
+    """Compress the results into a tarball."""
+
+    extensions = {'.tpr', '.xtc', '.yaml', '.top', '.gro'}
+    source_path = Path(source_dir).resolve()
+    tar_filename = source_path.with_suffix('.tar.gz')
+    with tarfile.open(tar_filename, "w:gz") as tar:
+        for file_path in source_path.rglob('*'):
+            if file_path.suffix in extensions:
+                arcname = file_path.relative_to(source_path.parent)
+                tar.add(file_path, arcname=str(arcname))
+
+
+def extract_tarball(fn: PathLike) -> None:
+    """Extract a tarball file into a directory."""
+    fn = path_to_str(fn) if isinstance(fn, Path) else fn
+    with tarfile.open(fn, "r:gz") as tar:
+        tar.extractall(path=fn.parent)
+
+
+def prepare_path(fn: PathLike) -> Path:
+    """Prepare a directory from a tarball or validate its existence."""
+    if fn.suffix == ".gz":
+        train_dir = fn.parent / fn.name.replace(".tar.gz", "")
+        if not train_dir.exists():
+            extract_tarball(fn)
+        return train_dir
+    elif fn.is_dir():
+        return fn
+    else:
+        raise ValueError(f"Unknown file type: {fn}")
+
+
+def extract_train_dir(
+    train_dir: PathLike
+) -> tuple[dict | None, list[dict], dict | None]:
+
+    """Extract training directory contents.
+
+    Parameters
+    ----------
+    train_dir : Path
+        Path to the training directory.
+
+    Returns
+    -------
+    tuple
+        A tuple containing:
+        - specs (dict | None): Specifications loaded from specs.yaml.
+        - systems (list[dict]): Staged system records from samples.yaml.
+        - samples (dict | None): Sample records loaded from samples.yaml.
+    """
+    if not train_dir.is_dir():
+        raise ValueError(f"Provided train_dir '{train_dir}' is not a valid directory.")
+
+    fn_specs = train_dir / "specs.yaml"
+    specs = load_yaml(fn_specs) if fn_specs.exists() else None
+
+    fn_samples = train_dir / "samples.yaml"
+    campaign = load_yaml(fn_samples) if fn_samples.exists() else None
+    systems = None if campaign is None else campaign.get("systems")
+    samples = None if campaign is None else campaign.get("samples")
+
+    if not systems:
+        raise ValueError(f"No staged simulation systems found in {train_dir}.")
+
+    fn_topol = [train_dir / system["topology"] for system in systems]
+    fn_coord = [train_dir / system["coordinates"] for system in systems]
+
+    if not all(path.exists() for path in fn_topol + fn_coord):
+        raise ValueError(f"Missing staged topology or coordinate files in {train_dir}.")
+
+    return specs, list(systems), samples

bff/mcmc/__init__.py

@@ -0,0 +1,3 @@
+from . import convergence, proposal, sampler
+
+__all__ = ["convergence", "proposal", "sampler"]

bff/mcmc/convergence.py

@@ -0,0 +1,168 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+
+@dataclass
+class ConvergenceInfo:
+    rhat: Optional[torch.Tensor] = None
+    tau: Optional[torch.Tensor] = None   # shape (n_walkers, n_dim)
+    ess: Optional[torch.Tensor] = None   # shape (n_dim,)
+
+    @property
+    def max_rhat(self) -> Optional[float]:
+        return torch.max(self.rhat).item() if self.rhat is not None else None
+
+    @property
+    def min_ess(self) -> Optional[float]:
+        return torch.min(self.ess).item() if self.ess is not None else None
+
+
+def split_rhat(chain: torch.Tensor) -> torch.Tensor:
+    """
+    Compute split R-hat for each parameter.
+
+    Parameters
+    ----------
+    chain : torch.Tensor
+        Shape (n_samples, n_walkers, n_dim).
+
+    Returns
+    -------
+    torch.Tensor
+        Shape (n_dim,).
+    """
+    n_samples, n_walkers, n_dim = chain.shape
+    half = n_samples // 2
+    if half < 2:
+        raise ValueError("Need at least 4 samples per walker for split R-hat.")
+
+    chains = chain.permute(1, 0, 2)  # (n_walkers, n_samples, n_dim)
+    chains = torch.cat([chains[:, :half], chains[:, -half:]], dim=0)
+
+    _, n, _ = chains.shape
+    chain_means = chains.mean(dim=1)
+    chain_vars = chains.var(dim=1, unbiased=True)
+
+    W = chain_vars.mean(dim=0)
+    B = n * chain_means.var(dim=0, unbiased=True)
+    var_hat = ((n - 1) / n) * W + B / n
+
+    return torch.sqrt(var_hat / W)
+
+
+def rank_normalize(values: torch.Tensor) -> torch.Tensor:
+    """Rank-normalize pooled samples along the first dimension."""
+    flat = values.reshape(values.shape[0], -1)
+    normalized = torch.empty_like(flat)
+    n = flat.shape[0]
+    denom = float(n) + 0.25
+    sqrt_two = torch.sqrt(torch.tensor(2.0, device=values.device))
+
+    for i in range(flat.shape[1]):
+        order = torch.argsort(flat[:, i], stable=True)
+        ranks = torch.empty(n, device=values.device, dtype=values.dtype)
+        ranks[order] = torch.arange(
+            1,
+            n + 1,
+            device=values.device,
+            dtype=values.dtype,
+        )
+        u = (ranks - 0.375) / denom
+        u = torch.clamp(u, 1e-12, 1 - 1e-12)
+        normalized[:, i] = sqrt_two * torch.erfinv(2 * u - 1)
+
+    return normalized.reshape(values.shape)
+
+
+def rank_normalized_split_rhat(chain: torch.Tensor) -> torch.Tensor:
+    """Compute rank-normalized split R-hat with folded refinement."""
+    pooled = chain.reshape(-1, chain.shape[-1])
+    ranked = rank_normalize(pooled).reshape(chain.shape)
+    rhat_rank = split_rhat(ranked)
+
+    pooled_median = pooled.median(dim=0).values
+    folded = torch.abs(chain - pooled_median.view(1, 1, -1))
+    ranked_folded = rank_normalize(
+        folded.reshape(-1, folded.shape[-1])
+    ).reshape(chain.shape)
+    rhat_folded = split_rhat(ranked_folded)
+    return torch.maximum(rhat_rank, rhat_folded)
+
+
+def autocorrelation(x: torch.Tensor) -> torch.Tensor:
+    """
+    Compute autocorrelation along axis 0 using FFT.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Shape (n_samples, ...).
+
+    Returns
+    -------
+    torch.Tensor
+        Shape (n_samples, ...).
+    """
+    if x.ndim == 0:
+        raise ValueError("x must have at least one dimension")
+    if x.shape[0] < 2:
+        return torch.ones_like(x)
+
+    n = x.shape[0]
+    x = x - x.mean(dim=0, keepdim=True)
+
+    nfft = 1 << (2 * n - 1).bit_length()
+    fx = torch.fft.rfft(x, n=nfft, dim=0)
+    acov = torch.fft.irfft(fx * torch.conj(fx), n=nfft, dim=0)[:n]
+
+    norm = torch.arange(n, 0, -1, device=x.device, dtype=x.dtype)
+    norm = norm.view((n,) + (1,) * (x.ndim - 1))
+    acov = acov / norm
+
+    acf = torch.zeros_like(acov)
+    var0 = acov[0]
+    mask = var0 > 0
+    acf[:, mask] = acov[:, mask] / var0[mask]
+    return acf
+
+
+def integrated_autocorr_time(
+    x: torch.Tensor,
+    max_lag: Optional[int] = None
+) -> torch.Tensor:
+    """
+    Estimate integrated autocorrelation time along axis 0.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Shape (n_samples, ...).
+    max_lag : int, optional
+        Maximum lag included.
+
+    Returns
+    -------
+    torch.Tensor
+        Shape x.shape[1:].
+    """
+    rho = autocorrelation(x)
+    n = rho.shape[0]
+
+    if max_lag is None:
+        max_lag = n - 1
+    max_lag = min(max_lag, n - 1)
+
+    tau = torch.ones(rho.shape[1:], device=x.device, dtype=x.dtype)
+    active = torch.ones_like(tau, dtype=torch.bool)
+
+    for k in range(1, max_lag + 1):
+        positive = rho[k] > 0
+        mask = active & positive
+        tau[mask] += 2.0 * rho[k][mask]
+        active &= positive
+        if not active.any():
+            break
+
+    return torch.clamp(tau, min=1.0)

bff/mcmc/proposal.py

@@ -0,0 +1,254 @@
+from typing import Optional
+
+import torch
+
+
+class Proposal:
+    """Base class for proposal mechanisms."""
+
+    def __init__(
+        self,
+        device: str | torch.device = "cpu",
+        dtype: torch.dtype = torch.float64,
+    ) -> None:
+        self.device = torch.device(device)
+        self.dtype = dtype
+
+    def initialize(self, n_dim: int) -> None:
+        raise NotImplementedError
+
+    def propose(self, x: torch.Tensor, rng: torch.Generator) -> torch.Tensor:
+        raise NotImplementedError
+
+    def adapt(self, step: int, x: torch.Tensor, accepted: torch.Tensor) -> None:
+        pass
+
+    def info(self) -> dict:
+        return {}
+
+    def state_dict(self) -> dict:
+        return {}
+
+    def load_state_dict(self, state: dict) -> None:
+        pass
+
+
+class AdaptiveGaussianProposal(Proposal):
+    """
+    Multivariate Gaussian random-walk proposal with optional adaptive covariance.
+
+    The proposal has the form
+
+    ``x_new = x + z @ L.T``,
+
+    where ``z ~ N(0, I)`` and ``L`` is the Cholesky factor of the proposal
+    covariance matrix.
+
+    During warmup, the empirical covariance of visited states can be used to
+    adapt the proposal shape and scale.
+
+    Parameters
+    ----------
+    proposal_cov : torch.Tensor, optional
+        Initial proposal covariance matrix of shape ``(n_dim, n_dim)``.
+        If ``None``, the identity matrix is used.
+    adapt : bool, default=True
+        Whether to adapt the proposal during warmup.
+    adapt_start : int, default=100
+        First warmup step at which adaptation is allowed.
+    adapt_interval : int, default=100
+        Adaptation frequency in number of warmup steps.
+    target_acceptance : float, default=0.234
+        Target acceptance rate used for multiplicative scale adaptation.
+    device : str or torch.device, default="cpu"
+        Device used internally by the proposal.
+    dtype : torch.dtype, default=torch.float64
+        Tensor dtype used internally.
+
+    Notes
+    -----
+    The covariance adaptation uses the empirical covariance of all walker
+    positions seen during warmup, together with the standard scaling factor
+    ``2.38**2 / n_dim``.
+    """
+
+    def __init__(
+        self,
+        proposal_cov: Optional[torch.Tensor] = None,
+        adapt: bool = True,
+        adapt_start: int = 100,
+        adapt_interval: int = 10,
+        target_acceptance: float = 0.234,
+        device: str | torch.device = "cpu",
+        dtype: torch.dtype = torch.float32,
+    ) -> None:
+        super().__init__(device=device, dtype=dtype)
+
+        self.proposal_cov = proposal_cov
+        self.do_adapt = adapt
+        self.adapt_start = adapt_start
+        self.adapt_interval = adapt_interval
+        self.target_acceptance = target_acceptance
+
+        self.eps = 1e-6
+        self.n_dim: Optional[int] = None
+        self.scale = 1.0
+        self.L: Optional[torch.Tensor] = None
+        self.n = 0
+        self.sum_x: Optional[torch.Tensor] = None
+        self.sum_xx: Optional[torch.Tensor] = None
+
+    def _to_tensor(self, x) -> torch.Tensor:
+        """Convert input to the proposal device and dtype."""
+        if isinstance(x, torch.Tensor):
+            return x.to(device=self.device, dtype=self.dtype)
+        return torch.as_tensor(x, device=self.device, dtype=self.dtype)
+
+    def _update_cov(self, x: torch.Tensor) -> None:
+        """
+        Update running first and second moments from a batch of walker positions.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Walker positions with shape ``(n_walkers, n_dim)``.
+        """
+        x = self._to_tensor(x)
+        self.n += x.shape[0]
+        self.sum_x += x.sum(dim=0)
+        self.sum_xx += x.T @ x
+
+    def _randn(self, shape, rng: torch.Generator) -> torch.Tensor:
+        return torch.randn(shape, generator=rng, device=self.device, dtype=self.dtype)
+
+    @property
+    def cov(self) -> torch.Tensor:
+        """
+        Current empirical covariance estimate.
+
+        Returns
+        -------
+        torch.Tensor
+            Empirical covariance matrix of shape ``(n_dim, n_dim)``.
+            Returns the identity matrix if fewer than two samples have been seen.
+        """
+        if self.n < 2:
+            return torch.eye(self.n_dim, device=self.device, dtype=self.dtype)
+        mean = self.sum_x / self.n
+        return (self.sum_xx - self.n * torch.outer(mean, mean)) / (self.n - 1)
+
+    def initialize(self, n_dim: int) -> None:
+        """
+        Initialize proposal state for a given dimensionality.
+
+        Parameters
+        ----------
+        n_dim : int
+            Dimensionality of the sampled parameter vector.
+        """
+        self.n_dim = n_dim
+
+        if self.proposal_cov is None:
+            cov = torch.eye(n_dim, device=self.device, dtype=self.dtype)
+        else:
+            cov = self._to_tensor(self.proposal_cov)
+
+        unit_matrix = torch.eye(n_dim, device=self.device, dtype=self.dtype)
+        cov = 0.5 * (cov + cov.T) + self.eps * unit_matrix
+
+        self.scale = 1.0
+        self.L = torch.linalg.cholesky(cov)
+
+        self.n = 0
+        self.sum_x = torch.zeros(n_dim, device=self.device, dtype=self.dtype)
+        self.sum_xx = torch.zeros((n_dim, n_dim), device=self.device, dtype=self.dtype)
+
+    def propose(self, x: torch.Tensor, rng: torch.Generator) -> torch.Tensor:
+        """
+        Generate Gaussian random-walk proposals.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Current walker positions with shape ``(n_walkers, n_dim)``.
+        rng : torch.Generator
+            Random number generator.
+
+        Returns
+        -------
+        torch.Tensor
+            Proposed walker positions with shape ``(n_walkers, n_dim)``.
+
+        Raises
+        ------
+        ValueError
+            If the supplied positions have the wrong dimensionality.
+        """
+        x = self._to_tensor(x)
+        if x.ndim != 2 or x.shape[1] != self.n_dim:
+            raise ValueError(
+                f"Expected shape (n_walkers, {self.n_dim}), "
+                f"got {tuple(x.shape)}"
+            )
+        z = self._randn(x.shape, rng)
+        return x + z @ self.L.T
+
+    def adapt(self, step: int, x: torch.Tensor, accepted: torch.Tensor) -> None:
+        """
+        Adapt the proposal covariance and scale during warmup.
+
+        Parameters
+        ----------
+        step : int
+            Current warmup step number, starting from 1.
+        x : torch.Tensor
+            Current walker positions with shape ``(n_walkers, n_dim)``.
+        accepted : torch.Tensor
+            Boolean mask of shape ``(n_walkers,)`` indicating accepted walkers.
+        """
+        if not self.do_adapt:
+            return
+
+        self._update_cov(x)
+
+        if step < self.adapt_start or step % self.adapt_interval != 0:
+            return
+
+        accept_mean = accepted.to(self.dtype).mean().item()
+        gamma = 1.0 / step**0.5
+        self.scale *= torch.exp(
+            torch.tensor(
+                gamma * (accept_mean - self.target_acceptance),
+                device=self.device,
+                dtype=self.dtype,
+            )
+        ).item()
+
+        cov = (2.38**2 / self.n_dim) * self.cov
+        unit_matrix = torch.eye(self.n_dim, device=self.device, dtype=self.dtype)
+        cov = cov + self.eps * unit_matrix
+        self.L = torch.linalg.cholesky((self.scale**2) * cov)
+
+    def state_dict(self) -> dict:
+        return {
+            "n_dim": self.n_dim,
+            "scale": self.scale,
+            "L": None if self.L is None else self.L.detach().cpu(),
+            "n": self.n,
+            "sum_x": None if self.sum_x is None else self.sum_x.detach().cpu(),
+            "sum_xx": None if self.sum_xx is None else self.sum_xx.detach().cpu(),
+        }
+
+    def load_state_dict(self, state: dict) -> None:
+        self.n_dim = state["n_dim"]
+        self.scale = state["scale"]
+        self.L = None if state["L"] is None else state["L"].to(
+            self.device, self.dtype
+        )
+        self.n = state["n"]
+        self.sum_x = None if state["sum_x"] is None else state["sum_x"].to(
+            self.device, self.dtype
+        )
+        self.sum_xx = None if state["sum_xx"] is None else (
+            state["sum_xx"].to(self.device, self.dtype)
+        )