openmm-opes 0.1.0__py3-none-any.whl

openmm_opes/__init__.py

@@ -0,0 +1,19 @@
+"""OpenMM implementation of On-the-fly Probability Enhanced Sampling (OPES)."""
+
+__version__ = "0.1.0"
+
+__all__ = ["OPES", "__version__"]
+
+
+def __getattr__(name):
+    """Import OPES lazily (PEP 562).
+
+    An eager ``from .opes import OPES`` would also run on
+    ``import openmm_opes.kde``, dragging openmm in and breaking the guarantee
+    that kde and io stay OpenMM-free. Task 1's test enforces that.
+    """
+    if name == "OPES":
+        from .opes import OPES
+
+        return OPES
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

openmm_opes/io.py

@@ -0,0 +1,102 @@
+"""Multi-walker bias sharing through a shared directory.
+
+Pure NumPy. This module must not import openmm, and deliberately knows
+nothing about kernels: it moves flat dicts of arrays so it stays testable
+without a simulation.
+
+The format is ``np.savez``, not pickle. A walker reads files written by other
+processes out of a shared directory, and ``pickle.load`` on data from another
+process would be a remote-code-execution primitive. OpenMM's own Metadynamics
+writes .npy for the same reason.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+import warnings
+
+import numpy as np
+
+FILENAME_PATTERN = re.compile(r"kde_(\d+)_(\d+)\.npz")
+
+
+class _LoadedBias:
+    """A peer's most recently seen state. Mirrors Metadynamics' own namedtuple."""
+
+    __slots__ = ("index", "state", "walkerId")
+
+    def __init__(self, walkerId: int, index: int, state: dict):
+        self.walkerId = walkerId
+        self.index = index
+        self.state = state
+
+
+class BiasSharer:
+    """Reads and writes walker bias states in a shared directory.
+
+    Parameters
+    ----------
+    biasDir
+        Directory shared by every walker.
+    walkerId
+        This walker's identifier. Drawn at random when omitted, from a fresh
+        generator so the global NumPy random state is left undisturbed.
+    """
+
+    def __init__(self, biasDir: str, walkerId: int | None = None):
+        self.biasDir = biasDir
+        self.walkerId = (
+            int(np.random.default_rng().integers(0x7FFFFFFF))
+            if walkerId is None
+            else walkerId
+        )
+        self._saveIndex = 0
+        self._loaded: dict[int, _LoadedBias] = {}
+
+    def _path(self, prefix: str, index: int) -> str:
+        return os.path.join(self.biasDir, f"{prefix}_{self.walkerId}_{index}.npz")
+
+    def save(self, state: dict) -> None:
+        """Write this walker's state atomically and drop the previous index."""
+        oldName = self._path("kde", self._saveIndex)
+        self._saveIndex += 1
+        tempName = self._path("temp", self._saveIndex)
+        fileName = self._path("kde", self._saveIndex)
+        np.savez(tempName, **state)
+        os.rename(tempName, fileName)
+        if os.path.exists(oldName):
+            os.remove(oldName)
+
+    def load(self) -> dict[int, dict]:
+        """Return peer states that are new or have advanced since the last call."""
+        updated: dict[int, dict] = {}
+        for filename in os.listdir(self.biasDir):
+            match = FILENAME_PATTERN.match(filename)
+            if match is None:
+                continue
+            walkerId, index = int(match.group(1)), int(match.group(2))
+            if walkerId == self.walkerId:
+                continue
+            known = self._loaded.get(walkerId)
+            if known is not None and index <= known.index:
+                continue
+            try:
+                with np.load(
+                    os.path.join(self.biasDir, filename), allow_pickle=False
+                ) as data:
+                    state = {key: data[key] for key in data.files}
+            except (OSError, ValueError):
+                warnings.warn(
+                    f"The file {filename} seems to have been deleted. Using the "
+                    "latest loaded data from the same walker.",
+                    stacklevel=2,
+                )
+                continue
+            self._loaded[walkerId] = _LoadedBias(walkerId, index, state)
+            updated[walkerId] = state
+        return updated
+
+    def getLoadedStates(self) -> dict[int, dict]:
+        """Every peer state seen so far, whether or not it changed recently."""
+        return {walkerId: bias.state for walkerId, bias in self._loaded.items()}

openmm_opes/kde.py

@@ -0,0 +1,450 @@
+"""Online kernel density estimation for OPES.
+
+Pure NumPy/SciPy. This module must not import openmm.
+"""
+
+from __future__ import annotations
+
+import functools
+from collections import namedtuple
+from copy import copy
+
+import numpy as np
+
+
+class CVSpace:
+    """The domain of the collective variables: grid, periodicity, boundaries.
+
+    Parameters
+    ----------
+    variables
+        Sequence of collective variable descriptions. Each must expose
+        ``minValue``, ``maxValue``, ``gridWidth`` and ``periodic``; an
+        ``openmm.app.BiasVariable`` qualifies.
+    bounded
+        Whether non-periodic variables have reflective boundaries. When true the
+        internal grid is tripled with mirrored copies on each side, and
+        :meth:`foldedGrid` folds them back.
+    """
+
+    CV = namedtuple("CV", ["minValue", "maxValue", "gridWidth", "periodic"])
+
+    def __init__(self, variables, bounded: bool = False):
+        self.variables = [
+            self.CV(cv.minValue, cv.maxValue, cv.gridWidth, cv.periodic)
+            for cv in variables
+        ]
+        self.bounded = bounded
+        self._periodic = any(cv.periodic for cv in self.variables)
+        self._grid = []
+        for cv in self.variables:
+            a, b, n = cv.minValue, cv.maxValue, cv.gridWidth
+            points = np.linspace(a, b, n)
+            if bounded and not cv.periodic:
+                left = np.linspace(2 * a - b, a, n)
+                right = np.linspace(b, 2 * b - a, n)
+                left[-1] = right[0] = np.inf
+                points = np.concatenate((points, np.flip(left), np.flip(right)))
+            self._grid.append(points)
+        self._widths = np.array([cv.gridWidth for cv in self.variables])
+        self._lbounds = np.array([cv.minValue for cv in self.variables])
+        self._lengths = np.array([cv.maxValue for cv in self.variables]) - self._lbounds
+        if self._periodic:
+            self._pdims = tuple(i for i, cv in enumerate(self.variables) if cv.periodic)
+            self._plbounds = self._lbounds[list(self._pdims)]
+            self._plengths = self._lengths[list(self._pdims)]
+
+    @property
+    def gridShape(self) -> tuple[int, ...]:
+        """Shape of the CV-space grid, slowest axis first."""
+        return tuple(int(width) for width in reversed(self._widths))
+
+    @property
+    def numDimensions(self) -> int:
+        """Number of collective variables."""
+        return len(self.variables)
+
+    def displacement(self, position, endpoint):
+        """Displacement from ``position`` to ``endpoint``, minimum-image if periodic."""
+        disp = endpoint - position
+        if self._periodic:
+            disp[..., self._pdims] -= self._plengths * np.rint(
+                disp[..., self._pdims] / self._plengths
+            )
+        return disp
+
+    def endpoint(self, position, displacement):
+        """Endpoint reached from ``position``, wrapped into the periodic domain."""
+        end = position + displacement
+        if self._periodic:
+            end[..., self._pdims] = (
+                self._plbounds
+                + (end[..., self._pdims] - self._plbounds) % self._plengths
+            )
+        return end
+
+    def gridDistances(self, position):
+        """Per-axis distances from ``position`` to every grid node."""
+        distances = [
+            centers - x for centers, x in zip(self._grid, position, strict=True)
+        ]
+        if self._periodic:
+            for dim, length in zip(self._pdims, self._plengths, strict=True):
+                distances[dim] -= length * np.rint(distances[dim] / length)
+                # a periodic grid's last node is the same point as its first
+                distances[dim][-1] = distances[dim][0]
+        return distances
+
+    def closestNode(self, position) -> tuple[int, ...]:
+        """Index of the grid node nearest to ``position``, in gridShape order."""
+        indices = np.rint(
+            (self._widths - 1) * (position - self._lbounds) / self._lengths
+        ).astype(int)
+        if self._periodic:
+            indices[list(self._pdims)] %= self._widths[list(self._pdims)]
+        indices = np.clip(indices, 0, self._widths - 1)
+        return tuple(reversed(indices))
+
+    def foldedGrid(self, values):
+        """Fold a tripled bounded grid back onto the physical domain, in log space."""
+        if self.bounded:
+            for i, cv in enumerate(reversed(self.variables)):
+                if not cv.periodic:
+                    values, left, right = np.array_split(values, 3, axis=i)
+                    values = np.logaddexp(left, np.logaddexp(values, right))
+        return values
+
+
+class KernelShape:
+    """A kernel profile: its per-dimension log normalization and its exponent."""
+
+    def __init__(self, name: str, logNorm: float, exponents):
+        self.name = name
+        self.logNorm = logNorm
+        self._exponents = exponents
+
+    def exponents(self, x):
+        """Log of the kernel profile at scaled distances ``x``."""
+        return self._exponents(x)
+
+
+def _gaussianExponents(x):
+    return -0.5 * x**2
+
+
+def _compactExponents(x):
+    values = 9 - x**2
+    mask = values > 0
+    result = np.empty_like(values)
+    result[mask] = 4 * np.log(values[mask])
+    result[~mask] = -np.inf
+    return result
+
+
+#: Unbounded Gaussian profile; the default and the one both papers use.
+GAUSSIAN = KernelShape("gaussian", np.log(2 * np.pi) / 2, _gaussianExponents)
+
+#: Compact quartic profile with support of +/- 3 bandwidths. Its normalization
+#: constant is exactly the integral of (9 - x**2)**4 over [-3, 3].
+COMPACT = KernelShape("compact", np.log(559872 / 35), _compactExponents)
+
+KERNEL_SHAPES = {shape.name: shape for shape in (GAUSSIAN, COMPACT)}
+
+
+class Kernel:
+    """A multivariate kernel with diagonal covariance."""
+
+    def __init__(
+        self, cvSpace, position, bandwidth, logWeight, numSamples=1, shape=GAUSSIAN
+    ):
+        self.cvSpace = cvSpace
+        self.position = np.array(position, dtype=float)
+        self.bandwidth = np.array(bandwidth, dtype=float)
+        self.logWeight = logWeight
+        self.numSamples = numSamples
+        self.shape = shape
+        self.logHeight = self._computeLogHeight()
+
+    def __copy__(self):
+        return Kernel(
+            self.cvSpace,
+            self.position,
+            self.bandwidth,
+            self.logWeight,
+            self.numSamples,
+            self.shape,
+        )
+
+    def _computeLogHeight(self):
+        if np.any(self.bandwidth == 0):
+            return -np.inf
+        d = self.cvSpace.numDimensions
+        return self.logWeight - d * self.shape.logNorm - np.sum(np.log(self.bandwidth))
+
+    def _scaledDistances(self, points, bandwidths):
+        return self.cvSpace.displacement(self.position, points) / bandwidths
+
+    def findNearest(self, centers, bandwidths, ignore=()):
+        """Index of and squared Mahalanobis distance to the nearest center."""
+        if centers.size == 0:
+            return -1, np.inf
+        sqDistances = np.sum(self._scaledDistances(centers, bandwidths) ** 2, axis=-1)
+        if len(ignore):
+            sqDistances[list(ignore)] = np.inf
+        index = int(np.argmin(sqDistances))
+        return index, sqDistances[index]
+
+    def merge(self, other) -> None:
+        """Absorb ``other``, preserving total weight, mean and second moment."""
+        logSumWeights = np.logaddexp(self.logWeight, other.logWeight)
+        w1 = np.exp(self.logWeight - logSumWeights)
+        w2 = np.exp(other.logWeight - logSumWeights)
+        disp = self.cvSpace.displacement(self.position, other.position)
+        self.position = self.cvSpace.endpoint(self.position, w2 * disp)
+        self.bandwidth = np.sqrt(
+            w1 * self.bandwidth**2 + w2 * other.bandwidth**2 + w1 * w2 * disp**2
+        )
+        self.logWeight = logSumWeights
+        self.numSamples += other.numSamples
+        self.logHeight = self._computeLogHeight()
+
+    def evaluate(self, points):
+        """Log of the kernel at the given point or points."""
+        return self.logHeight + np.sum(
+            self.shape.exponents(self._scaledDistances(points, self.bandwidth)), axis=-1
+        )
+
+    def evaluateOnGrid(self):
+        """Log of the kernel on the CV-space grid, in gridShape order."""
+        distances = self.cvSpace.gridDistances(self.position)
+        exponents = [
+            self.shape.exponents(dist / sigma)
+            for dist, sigma in zip(distances, self.bandwidth, strict=True)
+        ]
+        return self.cvSpace.foldedGrid(
+            self.logHeight + functools.reduce(np.add.outer, reversed(exponents))
+        )
+
+
+class OnlineKDE:
+    """Online kernel density estimate with on-the-fly kernel compression.
+
+    Parameters
+    ----------
+    cvSpace
+        The :class:`CVSpace` the estimate lives on.
+    compressionThreshold
+        Merge a new kernel into any existing kernel closer than this, measured
+        as a Mahalanobis distance. Zero disables merging.
+    useExistingBandwidths
+        Whether the nearest-neighbour search scales distances by the existing
+        kernels' bandwidths rather than the incoming kernel's.
+    kernelShape
+        Either ``"gaussian"`` or ``"compact"``.
+    """
+
+    def __init__(
+        self,
+        cvSpace,
+        compressionThreshold: float = 1.0,
+        useExistingBandwidths: bool = True,
+        kernelShape: str = "gaussian",
+    ):
+        if kernelShape not in KERNEL_SHAPES:
+            raise ValueError(
+                f"Unknown kernelShape {kernelShape!r}; "
+                f"expected one of {sorted(KERNEL_SHAPES)}"
+            )
+        self._cvSpace = cvSpace
+        self._compressionThreshold = compressionThreshold
+        self._useExistingBandwidths = useExistingBandwidths
+        self._shape = KERNEL_SHAPES[kernelShape]
+        self._kernels: list[Kernel] = []
+        self._logSumW = -np.inf
+        self._logSumWSq = -np.inf
+        self._logPK = np.empty(0)
+        self._logPG = np.full(cvSpace.gridShape, -np.inf)
+        self._d = cvSpace.numDimensions
+
+    def __bool__(self) -> bool:
+        return bool(self._kernels)
+
+    def __copy__(self):
+        new = self.__class__(
+            self._cvSpace,
+            self._compressionThreshold,
+            self._useExistingBandwidths,
+            self._shape.name,
+        )
+        new._kernels = list(map(copy, self._kernels))
+        new._logSumW = self._logSumW
+        new._logSumWSq = self._logSumWSq
+        new._logPK = self._logPK.copy()
+        new._logPG = self._logPG.copy()
+        return new
+
+    def __iadd__(self, other):
+        for kernel in other._kernels:
+            self._addKernel(
+                kernel.position,
+                kernel.bandwidth,
+                kernel.logWeight,
+                kernel.numSamples,
+                adjustBandwidth=False,
+            )
+        return self
+
+    @staticmethod
+    def _logsubexp(x, y):
+        """log(exp(x) - exp(y)), elementwise and numerically stable."""
+        result = np.full_like(x, -np.inf)
+        valid = y < x
+        inner = -np.exp(y[valid] - x[valid])
+        representable = inner > -1.0
+        inner[representable] = np.log1p(inner[representable])
+        inner[~representable] = -np.inf
+        result[valid] = inner + x[valid]
+        return result
+
+    def _removeKernels(self, centers, toRemove):
+        """Remove kernels by index, subtracting their contributions in log space.
+
+        Callers MUST add any replacement kernel's contribution to _logPK and
+        _logPG *before* calling this. _logsubexp loses precision as its two
+        arguments converge, and the replacement is what keeps the removed
+        kernel from dominating the density at its own center. See spec 12.4;
+        tests/test_kde.py pins the resulting margin.
+        """
+        toRemove = sorted(toRemove, reverse=True)
+        removed = []
+        for index in toRemove:
+            kernel = self._kernels.pop(index)
+            self._logPK = self._logsubexp(self._logPK, kernel.evaluate(centers))
+            self._logPG = self._logsubexp(self._logPG, kernel.evaluateOnGrid())
+            removed.append(kernel)
+        self._logPK = np.delete(self._logPK, toRemove)
+        return removed
+
+    def _pushKernel(self, newKernel):
+        centers = np.stack([k.position for k in self._kernels])
+        bandwidths = (
+            np.stack([k.bandwidth for k in self._kernels])
+            if self._useExistingBandwidths
+            else newKernel.bandwidth
+        )
+        threshold = self._compressionThreshold
+        index, minSqDist = newKernel.findNearest(centers, bandwidths)
+        toRemove = []
+        # threshold > 0 guard: without it, coincident kernels satisfy 0 <= 0
+        # and merge even when compression is meant to be disabled.
+        while threshold > 0 and index >= 0 and minSqDist <= threshold**2:
+            toRemove.append(index)
+            newKernel.merge(self._kernels[index])
+            index, minSqDist = newKernel.findNearest(centers, bandwidths, toRemove)
+        self._logPK = np.logaddexp(self._logPK, newKernel.evaluate(centers))
+        self._logPG = np.logaddexp(self._logPG, newKernel.evaluateOnGrid())
+        if toRemove:
+            self._removeKernels(centers, toRemove)
+        self._kernels.append(newKernel)
+        self._logPK = np.append(
+            self._logPK,
+            np.logaddexp.reduce(
+                [k.evaluate(newKernel.position) for k in self._kernels]
+            ),
+        )
+
+    def _addKernel(
+        self, position, bandwidth, logWeight, numSamples=1, adjustBandwidth=True
+    ):
+        self._logSumW = np.logaddexp(self._logSumW, logWeight)
+        self._logSumWSq = np.logaddexp(self._logSumWSq, 2 * logWeight)
+        if adjustBandwidth:
+            neff = np.exp(2 * self._logSumW - self._logSumWSq)
+            silverman = (neff * (self._d + 2) / 4) ** (-1 / (self._d + 4))
+            bandwidth = bandwidth * silverman
+        newKernel = Kernel(
+            self._cvSpace, position, bandwidth, logWeight, numSamples, self._shape
+        )
+        if self._kernels:
+            self._pushKernel(newKernel)
+        else:
+            self._kernels = [newKernel]
+            self._logPG = newKernel.evaluateOnGrid()
+            self._logPK = np.array([newKernel.logHeight])
+
+    def update(self, position, logWeight, variance) -> None:
+        """Deposit a kernel of the given log weight and per-CV variance."""
+        self._addKernel(position, np.sqrt(variance), logWeight)
+
+    def getNumKernels(self) -> int:
+        """Number of compressed kernels currently stored."""
+        return len(self._kernels)
+
+    def getLogPDF(self):
+        """Log of the normalized probability density on the grid."""
+        return self._logPG - self._logSumW
+
+    def getLogMeanDensity(self) -> float:
+        """Log of Z_n, the mean density over the compressed kernel centers."""
+        return (
+            np.logaddexp.reduce(self._logPK)
+            - np.log(len(self._kernels))
+            - self._logSumW
+        )
+
+    def evaluate(self, point) -> float:
+        """Log of the normalized density at a single point."""
+        return (
+            np.logaddexp.reduce([k.evaluate(point) for k in self._kernels])
+            - self._logSumW
+        )
+
+    def getState(self) -> dict:
+        """Flat, npz-writable snapshot of this estimate.
+
+        The CVSpace is deliberately excluded: all walkers share identical CV
+        definitions, so a reader constructs its own and calls setState.
+        """
+        d = self._cvSpace.numDimensions
+        if self._kernels:
+            positions = np.stack([k.position for k in self._kernels])
+            bandwidths = np.stack([k.bandwidth for k in self._kernels])
+            logWeights = np.array([k.logWeight for k in self._kernels])
+            numSamples = np.array([k.numSamples for k in self._kernels])
+        else:
+            positions = np.empty((0, d))
+            bandwidths = np.empty((0, d))
+            logWeights = np.empty(0)
+            numSamples = np.empty(0, dtype=int)
+        return {
+            "positions": positions,
+            "bandwidths": bandwidths,
+            "logWeights": logWeights,
+            "numSamples": numSamples,
+            "logSumW": float(self._logSumW),
+            "logSumWSq": float(self._logSumWSq),
+            "logPK": self._logPK.copy(),
+        }
+
+    def setState(self, state) -> None:
+        """Restore from a :meth:`getState` snapshot, discarding current contents."""
+        self._kernels = [
+            Kernel(self._cvSpace, position, bandwidth, logWeight, int(n), self._shape)
+            for position, bandwidth, logWeight, n in zip(
+                state["positions"],
+                state["bandwidths"],
+                state["logWeights"],
+                state["numSamples"],
+                strict=True,
+            )
+        ]
+        self._logSumW = float(state["logSumW"])
+        self._logSumWSq = float(state["logSumWSq"])
+        self._logPK = np.asarray(state["logPK"]).copy()
+        # Seeded with -inf so an empty kernel list restores cleanly; the source
+        # used a bare reduce here and raised TypeError on an empty sequence.
+        self._logPG = functools.reduce(
+            np.logaddexp,
+            (k.evaluateOnGrid() for k in self._kernels),
+            np.full(self._cvSpace.gridShape, -np.inf),
+        )

openmm_opes/opes.py

@@ -0,0 +1,465 @@
+"""The OPES bias, integrated with OpenMM as a CustomCVForce.
+
+Closely mirrors ``openmm.app.Metadynamics``: same constructor shape, same
+unit-coercion idiom, same force-group selection, same ``step`` contract.
+"""
+
+from __future__ import annotations
+
+import warnings
+from copy import copy
+
+import numpy as np
+import openmm as mm
+from openmm import unit
+
+from .io import BiasSharer
+from .kde import CVSpace, OnlineKDE
+
+
+class RunningAverage:
+    """Accumulates a running mean of per-CV squared deviations."""
+
+    def __init__(self, numDimensions: int = 1):
+        self._num = 0
+        self._total = np.zeros(numDimensions)
+
+    def __iadd__(self, other):
+        self._total = self._total + other._total
+        self._num += other._num
+        return self
+
+    def copy(self):
+        """An independent copy. The array is copied, not aliased."""
+        new = self.__class__(self._total.shape[0])
+        new._num = self._num
+        # The source aliased this array, so merging peers into a copy silently
+        # corrupted the original. See spec 7.4.1.
+        new._total = self._total.copy()
+        return new
+
+    def update(self, sample) -> None:
+        """Fold one sample into the average."""
+        self._num += 1
+        self._total = self._total + sample
+
+    def get(self):
+        """The running average."""
+        return self._total / max(1, self._num)
+
+    def getState(self) -> dict:
+        """Flat, npz-writable snapshot."""
+        return {"num": float(self._num), "total": self._total.copy()}
+
+    def setState(self, state) -> None:
+        """Restore from a :meth:`getState` snapshot."""
+        self._num = int(state["num"])
+        self._total = np.asarray(state["total"]).copy()
+
+
+class OPES:
+    """On-the-fly Probability Enhanced Sampling.
+
+    Implements OPES and its exploratory variant, following Invernizzi and
+    Parrinello, https://doi.org/10.1021/acs.jpclett.0c00497 and
+    https://doi.org/10.1021/acs.jctc.2c00152.
+
+    Parameters
+    ----------
+    system: System
+        The System to simulate. A CustomCVForce implementing the bias is
+        created and added to it.
+    variables: list of BiasVariable
+        The collective variables whose sampling should be enhanced.
+    temperature: temperature
+        The temperature at which the simulation is run.
+    barrier: energy
+        The free energy barrier the simulation should overcome.
+    frequency: int
+        Interval in time steps at which to deposit a kernel.
+    varianceFrequency: int or None
+        Interval in time steps at which to update the CV variance. When None,
+        the bandwidth is fixed and taken from each variable's ``biasWidth``,
+        which is interpreted as the standard deviation of the UNBIASED
+        distribution, exactly as both papers define sigma^(0).
+    biasFactor: float, optional
+        Defaults to ``barrier / kT``.
+    exploreMode: bool
+        Whether to apply the OPES-explore variant.
+    bounded: bool
+        Whether non-periodic CVs have reflective boundaries.
+    saveFrequency: int, optional
+        Interval in time steps at which to share the bias on disk. Must be a
+        multiple of ``frequency``.
+    biasDir: str, optional
+        Directory shared with other walkers.
+    warmupSteps: int, optional
+        When given, run this many steps without depositing kernels while
+        measuring the CV variance, then freeze it. See the spec, section 7.7.
+    compressionThreshold: float
+        Mahalanobis distance below which kernels are merged.
+    useExistingBandwidths: bool
+        Whether the merge search scales by existing kernels' bandwidths.
+    kernelShape: str
+        ``"gaussian"`` or ``"compact"``.
+    statsWindowSize: int
+        Window, in deposition strides, of the running CV-mean estimate.
+    """
+
+    def __init__(
+        self,
+        system,
+        variables,
+        temperature,
+        barrier,
+        frequency,
+        varianceFrequency,
+        biasFactor=None,
+        exploreMode: bool = False,
+        bounded: bool = False,
+        saveFrequency=None,
+        biasDir=None,
+        warmupSteps=None,
+        compressionThreshold: float = 1.0,
+        useExistingBandwidths: bool = True,
+        kernelShape: str = "gaussian",
+        statsWindowSize: int = 10,
+    ):
+        if not unit.is_quantity(temperature):
+            temperature = temperature * unit.kelvin
+        if not unit.is_quantity(barrier):
+            barrier = barrier * unit.kilojoules_per_mole
+
+        self.variables = variables
+        self.temperature = temperature
+        self.barrier = barrier
+        self.frequency = frequency
+        self.varianceFrequency = varianceFrequency
+        self.exploreMode = exploreMode
+        self.bounded = bounded
+        self.saveFrequency = saveFrequency
+        self.biasDir = biasDir
+        self.warmupSteps = warmupSteps
+        self.statsWindowSize = statsWindowSize
+
+        d = len(variables)
+        kbt = unit.MOLAR_GAS_CONSTANT_R * temperature
+        userSuppliedBiasFactor = biasFactor is not None
+        if not userSuppliedBiasFactor:
+            biasFactor = barrier / kbt
+        self.biasFactor = biasFactor
+        numPeriodics = sum(v.periodic for v in variables)
+        freeGroups = set(range(32)) - {f.getForceGroup() for f in system.getForces()}
+        self._validate(d, biasFactor, numPeriodics, freeGroups, userSuppliedBiasFactor)
+
+        prefactor = (1 - 1 / biasFactor) * kbt
+        if exploreMode:
+            prefactor *= biasFactor
+        self._kbt = kbt.in_units_of(unit.kilojoules_per_mole)
+        self._biasFactor = biasFactor
+        self._prefactor = prefactor
+        self._logEpsilon = -barrier / prefactor
+
+        self._cvSpace = CVSpace(variables, bounded)
+        self._cases = ("total",) + ("self",) * bool(saveFrequency)
+        kdeOptions = {
+            "compressionThreshold": compressionThreshold,
+            "useExistingBandwidths": useExistingBandwidths,
+            "kernelShape": kernelShape,
+        }
+        self._kde = {}
+        for case in self._cases:
+            self._kde[case] = OnlineKDE(self._cvSpace, **kdeOptions)
+            self._kde[f"{case}.rw"] = OnlineKDE(self._cvSpace, **kdeOptions)
+
+        self._adaptiveVariance = varianceFrequency is not None
+        self._interval = varianceFrequency or frequency
+        self._variance = {case: RunningAverage(d) for case in self._cases}
+        self._warmupComplete = warmupSteps is None
+        if self._adaptiveVariance:
+            self._tau = statsWindowSize * frequency // varianceFrequency
+            self._counter = 0
+            self._sampleMean = np.zeros(d)
+        else:
+            # Spec 7.5: biasWidth is the unbiased sigma^(0), while _variance
+            # holds a sampled-distribution variance, which is gamma times wider.
+            self._setFixedVariance(
+                biasFactor * np.array([v.biasWidth**2 for v in variables])
+            )
+
+        self._sharer = BiasSharer(biasDir) if biasDir is not None else None
+
+        gridWidths = [v.gridWidth for v in variables]
+        self._widths = [] if d == 1 else gridWidths
+        self._limits = [limit for v in variables for limit in (v.minValue, v.maxValue)]
+        periodic = numPeriodics == d
+        initial = np.full(int(np.prod(gridWidths)), -barrier / unit.kilojoules_per_mole)
+
+        energyFunction = "table(" + ",".join(f"cv{i}" for i in range(d)) + ")"
+        self._force = mm.CustomCVForce(energyFunction)
+        for i, variable in enumerate(variables):
+            self._force.addCollectiveVariable(f"cv{i}", variable.force)
+        # Explicit branch rather than a dynamic getattr, so ty can check it.
+        if d == 1:
+            table = mm.Continuous1DFunction(initial, *self._limits, periodic)
+        elif d == 2:
+            table = mm.Continuous2DFunction(
+                *self._widths, initial, *self._limits, periodic
+            )
+        else:
+            table = mm.Continuous3DFunction(
+                *self._widths, initial, *self._limits, periodic
+            )
+        self._force.addTabulatedFunction("table", table)
+        self._force.setForceGroup(max(freeGroups))
+        system.addForce(self._force)
+
+    def _validate(
+        self, d, biasFactor, numPeriodics, freeGroups, userSuppliedBiasFactor
+    ):
+        if self.varianceFrequency is not None and self.varianceFrequency <= 0:
+            raise ValueError("varianceFrequency must be positive or None")
+        if self.varianceFrequency and self.frequency % self.varianceFrequency != 0:
+            raise ValueError("varianceFrequency must be a divisor of frequency")
+        if (self.saveFrequency is None) != (self.biasDir is None):
+            raise ValueError("Must specify both saveFrequency and biasDir")
+        if self.saveFrequency is not None:
+            if self.saveFrequency <= 0:
+                raise ValueError("saveFrequency must be positive")
+            if self.saveFrequency % self.frequency != 0:
+                raise ValueError("saveFrequency must be a multiple of frequency")
+        if self.warmupSteps is not None:
+            if self.varianceFrequency is None:
+                raise ValueError("warmupSteps requires varianceFrequency")
+            if self.warmupSteps <= 0:
+                raise ValueError("warmupSteps must be positive")
+            if self.warmupSteps % self.varianceFrequency != 0:
+                raise ValueError("warmupSteps must be a multiple of varianceFrequency")
+            if self.warmupSteps // self.varianceFrequency < 2:
+                raise ValueError(
+                    "warmupSteps must span at least two varianceFrequency intervals"
+                )
+        if biasFactor <= 1.0:
+            raise ValueError(
+                "biasFactor must be greater than 1"
+                if userSuppliedBiasFactor
+                else "barrier must be greater than 1 kT"
+            )
+        if numPeriodics not in [0, d]:
+            raise ValueError("OPES cannot handle mixed periodic/non-periodic variables")
+        if not 1 <= d <= 3:
+            raise ValueError("OPES requires 1, 2, or 3 collective variables")
+        if not freeGroups:
+            raise RuntimeError("OPES requires a free force group, but all are in use.")
+
+    def _setFixedVariance(self, value) -> None:
+        """Replace every accumulator with a frozen value (spec 7.5 and 7.7)."""
+        for case in self._cases:
+            average = RunningAverage(len(self.variables))
+            average.update(np.asarray(value, dtype=float))
+            self._variance[case] = average
+        self._adaptiveVariance = False
+
+    def getVariance(self):
+        """Variance currently used to size newly deposited kernels."""
+        return self._variance["total"].get()
+
+    def getNumKernels(self) -> int:
+        """Number of kernels in the estimate that defines the bias."""
+        return self._kde["total" if self.exploreMode else "total.rw"].getNumKernels()
+
+    def getBias(self):
+        """The OPES bias potential on the grid."""
+        kde = self._kde["total" if self.exploreMode else "total.rw"]
+        return self._prefactor * np.logaddexp(
+            kde.getLogPDF() - kde.getLogMeanDensity(), self._logEpsilon
+        )
+
+    def getFreeEnergy(self):
+        """Free energy as a function of the collective variables.
+
+        Returned as an N-dimensional array in kJ/mole. The i'th position along
+        an axis corresponds to ``minValue + i*(maxValue-minValue)/gridWidth``.
+        Always the importance-sampling estimate, which converges better than
+        the direct one in explore mode.
+        """
+        return -self._kbt * self._kde["total.rw"].getLogPDF()
+
+    def getAverageDensity(self):
+        """Z_n, the mean density over the explored CV space."""
+        kde = self._kde["total" if self.exploreMode else "total.rw"]
+        return np.exp(kde.getLogMeanDensity())
+
+    def getCollectiveVariables(self, simulation):
+        """Current values of all collective variables in a Simulation."""
+        return self._force.getCollectiveVariableValues(simulation.context)
+
+    def updateContext(self, context) -> None:
+        """Push the current bias into a Context."""
+        bias = self.getBias().value_in_unit(unit.kilojoules_per_mole)
+        self._force.getTabulatedFunction(0).setFunctionParameters(
+            *self._widths, bias.ravel(), *self._limits
+        )
+        self._force.updateParametersInContext(context)
+
+    def addKernel(self, values, biasEnergy, variance=None) -> None:
+        """Deposit a kernel into the probability estimates.
+
+        This does not refresh any Context; call :meth:`updateContext`
+        afterwards if a simulation is running.
+        """
+        if not unit.is_quantity(biasEnergy):
+            biasEnergy = biasEnergy * unit.kilojoules_per_mole
+        if variance is None:
+            variance = self._variance["total"].get()
+        if np.any(np.asarray(variance) <= 0):
+            # Spec 7.4.4: a zero bandwidth poisons the estimate with NaN.
+            warnings.warn(
+                "Skipping kernel deposition: the CV variance estimate is not "
+                "yet positive. Use a varianceFrequency smaller than frequency, "
+                "or set warmupSteps.",
+                stacklevel=2,
+            )
+            return
+        for case in self._cases:
+            self._kde[case].update(values, 0.0, variance)
+            self._kde[f"{case}.rw"].update(
+                values, biasEnergy / self._kbt, variance / self._biasFactor
+            )
+
+    def _updateSampleStats(self, values) -> None:
+        self._counter += 1
+        delta = self._cvSpace.displacement(self._sampleMean, values)
+        x = 1 / min(self._tau, self._counter)
+        self._sampleMean = self._cvSpace.endpoint(self._sampleMean, x * delta)
+        sqdev = delta * self._cvSpace.displacement(self._sampleMean, values)
+        for case in self._cases:
+            self._variance[case].update(sqdev)
+
+    def _syncWithDisk(self) -> None:
+        # Only ever called when saveFrequency is set, which _validate ties to
+        # biasDir being set, which is what constructs self._sharer.
+        assert self._sharer is not None
+        self._sharer.save(self._getSharedState())
+        if not self._sharer.load():
+            return
+        self._kde["total"] = copy(self._kde["self"])
+        self._kde["total.rw"] = copy(self._kde["self.rw"])
+        self._variance["total"] = self._variance["self"].copy()
+        for state in self._sharer.getLoadedStates().values():
+            self._kde["total"] += self._kdeFromState(state, "kde")
+            self._kde["total.rw"] += self._kdeFromState(state, "kdeRW")
+            peer = RunningAverage(len(self.variables))
+            peer.setState({"num": state["var_num"], "total": state["var_total"]})
+            self._variance["total"] += peer
+
+    def _getSharedState(self) -> dict:
+        state = {}
+        for prefix, key in (("kde", "self"), ("kdeRW", "self.rw")):
+            for name, value in self._kde[key].getState().items():
+                state[f"{prefix}_{name}"] = value
+        variance = self._variance["self"].getState()
+        state["var_num"] = variance["num"]
+        state["var_total"] = variance["total"]
+        return state
+
+    def _kdeFromState(self, state, prefix):
+        kde = OnlineKDE(self._cvSpace)
+        strip = len(prefix) + 1
+        kde.setState(
+            {
+                key[strip:]: value
+                for key, value in state.items()
+                if key.startswith(f"{prefix}_")
+            }
+        )
+        return kde
+
+    def step(self, simulation, steps) -> None:
+        """Advance the simulation by a number of time steps.
+
+        Parameters
+        ----------
+        simulation: Simulation
+            The Simulation to advance.
+        steps: int
+            The number of time steps to integrate.
+        """
+        stepsToGo = steps
+        while stepsToGo > 0:
+            nextSteps = min(
+                stepsToGo,
+                self._interval - simulation.currentStep % self._interval,
+            )
+            simulation.step(nextSteps)
+            if simulation.currentStep % self._interval == 0:
+                self._onInterval(simulation)
+            stepsToGo -= nextSteps
+
+    def _onInterval(self, simulation) -> None:
+        position = self.getCollectiveVariables(simulation)
+        if not self._warmupComplete:
+            self._updateSampleStats(position)
+            if simulation.currentStep >= self.warmupSteps:
+                self._finishWarmup()
+            return
+        if self._adaptiveVariance:
+            self._updateSampleStats(position)
+        if simulation.currentStep % self.frequency == 0:
+            groups = {self._force.getForceGroup()}
+            energy = simulation.context.getState(
+                getEnergy=True, groups=groups
+            ).getPotentialEnergy()
+            self.addKernel(position, energy)
+            self.updateContext(simulation.context)
+            if (
+                self.saveFrequency is not None
+                and simulation.currentStep % self.saveFrequency == 0
+            ):
+                self._syncWithDisk()
+
+    def _finishWarmup(self) -> None:
+        """Freeze sigma^(0) measured from the unbiased warm-up segment.
+
+        Nothing was deposited during warm-up, so the trajectory was unbiased
+        and this measurement is the unbiased variance. _variance holds a
+        sampled-distribution variance by convention, which is gamma times
+        wider, so that is what gets stored. This is the same transformation
+        the fixed-bandwidth path applies to biasWidth. See spec 7.5 and 7.7.
+        """
+        measured = self._variance["total"].get()
+        if np.any(measured <= 0):
+            raise ValueError(
+                "Warm-up produced a non-positive CV variance. Increase "
+                "warmupSteps or decrease varianceFrequency."
+            )
+        self._setFixedVariance(self._biasFactor * measured)
+        self._warmupComplete = True
+
+    def getState(self) -> dict:
+        """Flat, npz-writable snapshot of this sampler's accumulated state."""
+        state = self._getSharedState() if "self" in self._cases else {}
+        for prefix, key in (("total", "total"), ("totalRW", "total.rw")):
+            for name, value in self._kde[key].getState().items():
+                state[f"{prefix}_{name}"] = value
+        variance = self._variance["total"].getState()
+        state["totalVar_num"] = variance["num"]
+        state["totalVar_total"] = variance["total"]
+        # Persisted explicitly rather than re-derived from the step counter:
+        # re-deriving would risk rescaling an already-frozen variance by gamma
+        # a second time on reload. See spec 7.7.
+        state["warmupComplete"] = float(self._warmupComplete)
+        return state
+
+    def setState(self, state) -> None:
+        """Restore from a :meth:`getState` snapshot."""
+        for prefix, key in (("total", "total"), ("totalRW", "total.rw")):
+            self._kde[key] = self._kdeFromState(state, prefix)
+        average = RunningAverage(len(self.variables))
+        average.setState(
+            {"num": state["totalVar_num"], "total": state["totalVar_total"]}
+        )
+        self._variance["total"] = average
+        self._warmupComplete = bool(float(state["warmupComplete"]))
+        if self._warmupComplete and self.warmupSteps is not None:
+            self._adaptiveVariance = False

openmm_opes-0.1.0.dist-info/METADATA

@@ -0,0 +1,64 @@
+Metadata-Version: 2.4
+Name: openmm-opes
+Version: 0.1.0
+Summary: OpenMM implementation of On-the-fly Probability Enhanced Sampling (OPES)
+Author: Charlles Abreu
+License-Expression: MIT
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.24
+Requires-Dist: scipy>=1.10
+Requires-Dist: openmm>=8.1
+Provides-Extra: dev
+Requires-Dist: pytest>=7.4; extra == "dev"
+Requires-Dist: ruff>=0.16; extra == "dev"
+Requires-Dist: ty==0.0.72; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: mkdocs>=1.6; extra == "docs"
+Requires-Dist: mkdocs-material>=9.5; extra == "docs"
+Requires-Dist: mkdocstrings[python]>=0.24; extra == "docs"
+Dynamic: license-file
+
+# openmm-opes
+
+[![CI](https://github.com/craabreu/openmm-opes/actions/workflows/ci.yml/badge.svg)](https://github.com/craabreu/openmm-opes/actions/workflows/ci.yml)
+
+An [OpenMM](https://openmm.org) implementation of On-the-fly Probability
+Enhanced Sampling (OPES) and OPES-explore.
+
+## Installation
+
+```bash
+pip install openmm-opes
+```
+
+## Usage
+
+```python
+from openmm_opes import OPES
+
+sampler = OPES(system, variables, temperature, barrier, frequency=500,
+               varianceFrequency=50)
+sampler.step(simulation, 1_000_000)
+fes = sampler.getFreeEnergy()
+```
+
+Full documentation: <https://craabreu.github.io/openmm-opes>
+
+## Origin
+
+This library packages and hardens the prototype implementation developed in
+[craabreu/opes-simulations](https://github.com/craabreu/opes-simulations),
+porting its `OPES`/`OnlineKDE` classes into a tested, documented package.
+
+## References
+
+- Invernizzi & Parrinello, J. Phys. Chem. Lett. 2020,
+  [doi:10.1021/acs.jpclett.0c00497](https://doi.org/10.1021/acs.jpclett.0c00497)
+- Invernizzi & Parrinello, J. Chem. Theory Comput. 2022,
+  [doi:10.1021/acs.jctc.2c00152](https://doi.org/10.1021/acs.jctc.2c00152)
+
+## License
+
+MIT

openmm_opes-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+openmm_opes/__init__.py,sha256=lod_MXaEyOiV-UQE9vDxeTJK5nJlbdTfmedodfmVLbQ,562
+openmm_opes/io.py,sha256=-K1grAYSua2afV1eehrFwLvooDBX5eGTOREmqXpF3AQ,3664
+openmm_opes/kde.py,sha256=85XUrSL5db9e2iRkO7u4P2YaLCdCQC0Gu7qf2wV7UrY,17265
+openmm_opes/opes.py,sha256=-Rge8wDelDI2p8WXFuLwpcGW3CXWkNBVoyaUsDDOD4s,19391
+openmm_opes-0.1.0.dist-info/licenses/LICENSE,sha256=TpNLslqYwa8hLxWToDJdeBd-Ld3Zz0qWo9pi3wCsgDo,1071
+openmm_opes-0.1.0.dist-info/METADATA,sha256=soT_eT2qffOKq7wpcLQDJpnp3p9_JEJKvisW4CSkfQI,1862
+openmm_opes-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+openmm_opes-0.1.0.dist-info/top_level.txt,sha256=Ql0Vg4GnHDMON9G0HFjkvl2B-EtUaX0WXujAMXX2hK0,12
+openmm_opes-0.1.0.dist-info/RECORD,,

openmm_opes-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

openmm_opes-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Charlles Abreu
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

openmm_opes-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+openmm_opes