boltzmann-generators 0.2.1__py3-none-any.whl

boltzmann_generators/__init__.py

@@ -0,0 +1,18 @@
+"""Boltzmann Generators — from-scratch PyTorch implementation."""
+
+from . import analysis, base, energies, flows, io, losses, mcmc, sampling, services, train, training
+
+__version__ = "0.2.1"
+__all__ = [
+    "analysis",
+    "base",
+    "energies",
+    "flows",
+    "io",
+    "losses",
+    "mcmc",
+    "sampling",
+    "services",
+    "train",
+    "training",
+]

boltzmann_generators/analysis.py

@@ -0,0 +1,36 @@
+"""Analysis helpers for weighted and unweighted population estimates."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from torch import Tensor
+
+from .services.analysis import AnalysisSuite
+
+RegionFn = Callable[[Tensor], Tensor]
+
+__all__ = ["AnalysisSuite", "RegionFn", "basin_populations", "rectangular_region"]
+
+_suite = AnalysisSuite()
+
+
+def basin_populations(
+    x: Tensor,
+    region_fns: dict[str, RegionFn],
+    *,
+    log_w: Tensor | None = None,
+) -> dict[str, float]:
+    """Compute basin populations from point assignments or importance weights."""
+    return _suite.basin_populations(x, region_fns, log_w=log_w)
+
+
+def rectangular_region(
+    *,
+    x_min: float,
+    x_max: float,
+    y_min: float,
+    y_max: float,
+) -> RegionFn:
+    """Create a rectangular region predicate over 2D coordinates."""
+    return AnalysisSuite.rectangular_region(x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max)

boltzmann_generators/base/__init__.py

@@ -0,0 +1,6 @@
+"""Core abstract base classes for energies and density models."""
+
+from .density import BaseDensityModel, DensityModel
+from .energy import EnergyModel
+
+__all__ = ["BaseDensityModel", "DensityModel", "EnergyModel"]

boltzmann_generators/base/density.py

@@ -0,0 +1,36 @@
+"""Abstract base class for Boltzmann generator density models."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Protocol, runtime_checkable
+
+import torch
+from torch import Tensor, nn
+
+
+@runtime_checkable
+class DensityModel(Protocol):
+    """Structural typing contract for density models used in losses and sampling."""
+
+    def sample(self, n: int, device: torch.device | str = "cpu") -> tuple[Tensor, Tensor]:
+        """Draw ``n`` samples; return ``(x, log_q(x))``."""
+
+    def log_prob(self, x: Tensor) -> Tensor:
+        """Log-density ``log q(x)`` for batch ``x``."""
+
+
+class BaseDensityModel(nn.Module, ABC):
+    """PyTorch module implementing a tractable approximate Boltzmann density."""
+
+    @abstractmethod
+    def sample(self, n: int, device: torch.device | str = "cpu") -> tuple[Tensor, Tensor]:
+        """Draw ``n`` samples in data space; return ``(x, log_q(x))``."""
+
+    @abstractmethod
+    def log_prob(self, x: Tensor) -> Tensor:
+        """Log-density ``log q(x)`` for batch ``x``."""
+
+    def nll(self, x: Tensor) -> Tensor:
+        """Negative log-likelihood (forward KL up to constant). Mean over batch."""
+        return -self.log_prob(x).mean()

boltzmann_generators/base/energy.py

@@ -0,0 +1,20 @@
+"""Abstract base class for benchmark energy functions."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from torch import Tensor
+
+
+class EnergyModel(ABC):
+    """Reduced energy u(x) = U(x) / (kT) in dimension ``dim``."""
+
+    dim: int
+
+    @abstractmethod
+    def energy(self, x: Tensor) -> Tensor:
+        """Evaluate reduced energy on batch ``x`` of shape ``(..., dim)``."""
+
+    def __call__(self, x: Tensor) -> Tensor:
+        return self.energy(x)

boltzmann_generators/energies/__init__.py

@@ -0,0 +1,5 @@
+from .dipeptide import RamachandranDipeptide
+from .double_well import DoubleWell1D, DoubleWell2D
+from .muller import MullerBrown
+
+__all__ = ["DoubleWell1D", "DoubleWell2D", "MullerBrown", "RamachandranDipeptide"]

boltzmann_generators/energies/dipeptide.py

@@ -0,0 +1,81 @@
+"""Synthetic Ramachandran-like potential — a stand-in for alanine dipeptide.
+
+Real alanine dipeptide has a ~50-atom Cartesian description and a CHARMM/AMBER
+force field. Implementing that here would pull in OpenMM or a full from-scratch
+MM energy, both of which are out of scope.
+
+Instead we model the *free-energy surface* in dihedral space (phi, psi) directly,
+using a sum of Gaussian wells at the canonical alanine dipeptide minima. This
+gives a 2D periodic potential with the right qualitative structure (alpha_R,
+alpha_L, beta/PPII, C5/C7eq) so we can demonstrate BG/CFM training on a
+molecular-flavored target without molecular machinery.
+
+Coordinates: (phi, psi) in degrees, periodic in [-180, 180].
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+
+from ..base.energy import EnergyModel
+
+# Approximate alanine dipeptide minima (degrees), depths in kT
+_MINIMA = torch.tensor(
+    [
+        [-65.0, -40.0],  # alpha_R    (deepest)
+        [-150.0, 155.0],  # C5 / beta
+        [-80.0, 80.0],  # PPII / C7eq
+        [65.0, 40.0],  # alpha_L    (shallow)
+    ]
+)
+_DEPTHS = torch.tensor([6.0, 5.0, 4.5, 2.5])  # in kT
+_WIDTHS = torch.tensor([22.0, 28.0, 30.0, 30.0])  # degrees
+
+
+def _wrap_deg(d: Tensor) -> Tensor:
+    """Wrap angle differences to (-180, 180]."""
+    return (d + 180.0) % 360.0 - 180.0
+
+
+class RamachandranDipeptide(EnergyModel):
+    """Synthetic 2D dipeptide free-energy surface in (phi, psi) degrees.
+
+    u(phi, psi) = -log sum_k exp(-d_k(phi, psi) / w_k^2 + log depth_k) + const
+
+    Each well is a Gaussian in periodic-angle distance. Total potential is a
+    smooth log-sum-exp combination so derivatives are well-defined.
+    """
+
+    def __init__(self) -> None:
+        self.dim = 2
+
+    def __call__(self, x: Tensor) -> Tensor:
+        return self.energy(x)
+
+    def energy(self, x: Tensor) -> Tensor:
+        """x: (..., 2) in degrees. Returns reduced energy (kT units)."""
+        device = x.device
+        minima = _MINIMA.to(device)
+        depths = _DEPTHS.to(device)
+        widths = _WIDTHS.to(device)
+        # Periodic squared distance to each minimum
+        dphi = _wrap_deg(x[..., 0:1] - minima[:, 0])  # (..., K)
+        dpsi = _wrap_deg(x[..., 1:2] - minima[:, 1])
+        d2 = (dphi.pow(2) + dpsi.pow(2)) / widths.pow(2)
+        # log-sum-exp combination: u = -log sum exp(depth - d2)
+        logits = depths - d2
+        u = -torch.logsumexp(logits, dim=-1)
+        return u
+
+    def grid(self, n: int = 200) -> tuple[Tensor, Tensor, Tensor]:
+        xs = torch.linspace(-180, 180, n)
+        ys = torch.linspace(-180, 180, n)
+        gx, gy = torch.meshgrid(xs, ys, indexing="xy")
+        grid = torch.stack([gx.flatten(), gy.flatten()], dim=-1)
+        u = self.energy(grid).reshape(n, n)
+        return gx, gy, u
+
+    @property
+    def minima(self) -> Tensor:
+        return _MINIMA.clone()

boltzmann_generators/energies/double_well.py

@@ -0,0 +1,56 @@
+"""Double-well potentials in 1D and 2D.
+
+All energies are returned as reduced (unitless) energies u(x) = U(x)/(kT).
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+
+from ..base.energy import EnergyModel
+
+
+class DoubleWell2D(EnergyModel):
+    """u(x, y) = a*(x^2 - 1)^2 + 0.5/sigma^2 * y^2.
+
+    Two minima at (±1, 0). The y direction is a harmonic well of width sigma.
+    Parameter `a` controls the barrier height: barrier ≈ a (in kT).
+    """
+
+    def __init__(self, a: float = 4.0, sigma_y: float = 0.5) -> None:
+        self.a = a
+        self.sigma_y = sigma_y
+        self.dim = 2
+
+    def __call__(self, x: Tensor) -> Tensor:
+        return self.energy(x)
+
+    def energy(self, x: Tensor) -> Tensor:
+        assert x.shape[-1] == 2
+        xx = x[..., 0]
+        yy = x[..., 1]
+        return self.a * (xx.pow(2) - 1.0).pow(2) + 0.5 * (yy / self.sigma_y).pow(2)
+
+    def grid(self, n: int = 200, span: float = 2.5) -> tuple[Tensor, Tensor, Tensor]:
+        xs = torch.linspace(-span, span, n)
+        ys = torch.linspace(-span, span, n)
+        gx, gy = torch.meshgrid(xs, ys, indexing="xy")
+        grid = torch.stack([gx.flatten(), gy.flatten()], dim=-1)
+        u = self.energy(grid).reshape(n, n)
+        return gx, gy, u
+
+
+class DoubleWell1D(EnergyModel):
+    """u(x) = a*(x^2 - 1)^2. Minima at ±1, barrier height = a (in kT)."""
+
+    def __init__(self, a: float = 4.0) -> None:
+        self.a = a
+        self.dim = 1
+
+    def __call__(self, x: Tensor) -> Tensor:
+        return self.energy(x)
+
+    def energy(self, x: Tensor) -> Tensor:
+        xx = x[..., 0] if x.ndim > 1 else x
+        return self.a * (xx.pow(2) - 1.0).pow(2)

boltzmann_generators/energies/muller.py

@@ -0,0 +1,65 @@
+"""Müller-Brown potential.
+
+Standard 2D benchmark for rare-event sampling. Three minima with two saddle
+points connecting them. Parameters from Müller & Brown (1979).
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+
+from ..base.energy import EnergyModel
+
+_A = torch.tensor([-200.0, -100.0, -170.0, 15.0])
+_a = torch.tensor([-1.0, -1.0, -6.5, 0.7])
+_b = torch.tensor([0.0, 0.0, 11.0, 0.6])
+_c = torch.tensor([-10.0, -10.0, -6.5, 0.7])
+_x0 = torch.tensor([1.0, 0.0, -0.5, -1.0])
+_y0 = torch.tensor([0.0, 0.5, 1.5, 1.0])
+
+
+class MullerBrown(EnergyModel):
+    """Reduced energy u(x) = U(x) / scale, with scale tuning barrier heights.
+
+    Native U has barriers ~100-200 (arbitrary units). For BG training we need
+    barriers of moderate height (a few kT), so we rescale by `scale`. Default
+    scale=20 yields barriers ~5-10 kT, reasonable for training.
+    """
+
+    def __init__(self, scale: float = 20.0) -> None:
+        self.scale = scale
+        self.dim = 2
+
+    def __call__(self, x: Tensor) -> Tensor:
+        return self.energy(x)
+
+    def energy(self, x: Tensor) -> Tensor:
+        assert x.shape[-1] == 2
+        xx = x[..., 0:1]  # (..., 1)
+        yy = x[..., 1:2]
+        device = x.device
+        A = _A.to(device)
+        a = _a.to(device)
+        b = _b.to(device)
+        c = _c.to(device)
+        x0 = _x0.to(device)
+        y0 = _y0.to(device)
+        dx = xx - x0
+        dy = yy - y0
+        terms = A * torch.exp(a * dx.pow(2) + b * dx * dy + c * dy.pow(2))
+        U = terms.sum(dim=-1)
+        return U / self.scale
+
+    def grid(
+        self,
+        n: int = 200,
+        x_span: tuple[float, float] = (-1.7, 1.2),
+        y_span: tuple[float, float] = (-0.4, 2.1),
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        xs = torch.linspace(*x_span, n)
+        ys = torch.linspace(*y_span, n)
+        gx, gy = torch.meshgrid(xs, ys, indexing="xy")
+        grid = torch.stack([gx.flatten(), gy.flatten()], dim=-1)
+        u = self.energy(grid).reshape(n, n)
+        return gx, gy, u

boltzmann_generators/extensions/__init__.py

@@ -0,0 +1 @@
+"""Extension points for future transferable and molecular BG capabilities."""

boltzmann_generators/extensions/equivariant.py

@@ -0,0 +1,7 @@
+"""Roadmap stub: E(3)-equivariant flow architectures.
+
+Planned direction:
+- EGNN-based velocity fields.
+- SO(3)/E(3)-aware coupling transformations.
+- Support for coordinates + atom/token conditioning.
+"""

boltzmann_generators/extensions/molecular.py

@@ -0,0 +1,7 @@
+"""Roadmap stub: molecular-system interfaces.
+
+Planned direction:
+- OpenMM-backed energy wrappers.
+- Internal-coordinate preprocessing (bond/angle/dihedral).
+- Dataset adapters for trajectory sources (e.g., mdshare).
+"""

boltzmann_generators/extensions/transferable.py

@@ -0,0 +1,7 @@
+"""Roadmap stub: transferable Boltzmann Generator components.
+
+Planned direction:
+- Tokenized chemistry conditioning.
+- Shared backbone across molecules.
+- Joint CFM + reweighting pipelines for zero-shot transfer.
+"""

boltzmann_generators/flows/__init__.py

@@ -0,0 +1,20 @@
+from .base import Flow, FlowModel, GaussianPrior
+from .cnf import CNFFlowModel, CNFModel, VelocityField
+from .coupling import AffineCoupling
+from .periodic import PeriodicEmbedding, periodic_inverse
+from .realnvp import RealNVP, alternating_mask, halves_mask
+
+__all__ = [
+    "Flow",
+    "FlowModel",
+    "GaussianPrior",
+    "AffineCoupling",
+    "RealNVP",
+    "alternating_mask",
+    "halves_mask",
+    "CNFModel",
+    "CNFFlowModel",
+    "VelocityField",
+    "PeriodicEmbedding",
+    "periodic_inverse",
+]

boltzmann_generators/flows/base.py

@@ -0,0 +1,68 @@
+"""Base classes for normalizing flows.
+
+Convention used everywhere in this package:
+
+- `forward(z)` maps prior space → data space (sampling direction).
+  Returns `(x, log_det)` where `log_det = log|det df/dz|`.
+- `inverse(x)` maps data space → prior space (density direction).
+  Returns `(z, log_det)` where `log_det = log|det df^-1/dx| = -log|det df/dz|`.
+
+With this convention:
+    log p_X(x) = log p_Z(z) + log|det df^-1/dx|   # the inverse log-det
+"""
+
+from __future__ import annotations
+
+import math
+from abc import ABC, abstractmethod
+
+import torch
+from torch import Tensor, nn
+
+from ..base.density import BaseDensityModel
+
+
+class Flow(nn.Module, ABC):
+    """Invertible transformation z <-> x with tractable log-determinant."""
+
+    @abstractmethod
+    def forward(self, z: Tensor) -> tuple[Tensor, Tensor]: ...
+
+    @abstractmethod
+    def inverse(self, x: Tensor) -> tuple[Tensor, Tensor]: ...
+
+
+class GaussianPrior(nn.Module):
+    """Standard normal prior N(0, I) of given dimension."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.dim = dim
+        self.register_buffer("_log_norm", torch.tensor(0.5 * dim * math.log(2 * math.pi)))
+
+    def sample(self, n: int, device: torch.device | str = "cpu") -> Tensor:
+        return torch.randn(n, self.dim, device=device)
+
+    def log_prob(self, z: Tensor) -> Tensor:
+        return -0.5 * z.pow(2).sum(dim=-1) - self._log_norm
+
+
+class FlowModel(BaseDensityModel):
+    """Flow stack + prior. Provides sample, log_prob, and forward KL loss."""
+
+    def __init__(self, prior: GaussianPrior, flow: Flow) -> None:
+        super().__init__()
+        self.prior = prior
+        self.flow = flow
+
+    def sample(self, n: int, device: torch.device | str = "cpu") -> tuple[Tensor, Tensor]:
+        """Draw n samples in data space. Returns (x, log_prob_x)."""
+        z = self.prior.sample(n, device=device)
+        log_pz = self.prior.log_prob(z)
+        x, log_det_fwd = self.flow.forward(z)
+        log_px = log_pz - log_det_fwd
+        return x, log_px
+
+    def log_prob(self, x: Tensor) -> Tensor:
+        z, log_det_inv = self.flow.inverse(x)
+        return self.prior.log_prob(z) + log_det_inv