stochkin.potentials¶

Analytic model potentials and basin-detection utilities.

stochkin.potentials¶

Analytic model potentials and basin-partitioning utilities.

This module provides:

Analytic potentials – standard 1D and 2D test surfaces commonly used in molecular-dynamics and stochastic-kinetics studies. Every potential is a callable potential(x) -> (U, F) returning the energy U and the force vector F = −∇U.
- double_well_2d() – quartic double-well in 2D.
- simple_double_well_2d() – symmetric two-state double-well with harmonic y-confinement.
- double_well_1d() / make_double_well_1d() – 1D double-well.
- mexican_hat_potential() – radially symmetric sombrero potential.
- central_well_barrier_ring_potential() – ring minimum with central Gaussian well and barrier.
- muller_potential() – Müller–Brown four-Gaussian surface.
- StringPotential / make_potential_from_string() – user-defined potential from a string expression.
Basin detection – automatic identification of metastable basins on a gridded free-energy surface and construction of BasinNetwork / BasinNetwork1D data structures used by downstream MFPT, CTMC, and committor analyses.
- sample_potential_grid() – evaluate a potential on a 2D grid.
- build_basin_network_from_potential() / detect_basins_for_mfpt() – detect 2D basins.
- build_basin_network_from_fes_1d() / detect_basins_for_mfpt_1d() – detect 1D basins.
- build_core_labels_from_full_labels() – shrink basins to low-energy cores for CTMC boundary-value problems.

Notes

All potentials are picklable so that they can be used with multiprocessing.Pool for parallel trajectory calculations.

class stochkin.potentials.StringPotential(expr, eps=1e-06)[source]¶

Bases: object

Picklable potential built from a string expression.

The expression must use x[i] to refer to the i-th component of the position vector and may call any function available in NumPy via the np namespace.

Parameters:

expr (str) – A valid Python expression, e.g. "0.5*(x[0]**2 + 2*x[1]**2)".
eps (float, optional) – Step size for the central-difference gradient (default 1e-6).

Examples

>>> pot = StringPotential("0.5*(x[0]**2 + 2*x[1]**2)")
>>> U, F = pot(np.array([1.0, 0.5]))

stochkin.potentials.make_potential_from_string(expr, eps=1e-06)[source]¶

Create a picklable potential from a Python math expression.

Parameters:

expr (str) – Expression using x[i] for coordinates and np for NumPy, e.g. "0.5*(x[0]**2 + 2*x[1]**2)".
eps (float, optional) – Central-difference step for force evaluation (default 1e-6).

Returns:

Callable pot(x) -> (U, F) where F = −∇U.

Return type:

StringPotential

stochkin.potentials.double_well_2d(x, a=1.0, b=1.3)[source]¶

Quartic double-well in 2D.

\[U(x_1, x_2) = a(x_1^4 + x_2^4) - b(x_1^2 + x_2^2)\]

Parameters:

x (array_like, shape (2,)) – Position [x1, x2].
a (float) – Quartic coefficient (default 1.0).
b (float) – Quadratic coefficient (default 1.3).

Returns:

U (float) – Potential energy.
F (ndarray, shape (2,)) – Force vector F = −∇U.

stochkin.potentials.mexican_hat_potential(x, a=-1.0, b=1.0)[source]¶

Radially symmetric 2D Mexican-hat (sombrero) potential.

\[U(r) = a \, r^2 + b \, r^4, \quad r = |x|\]

With a < 0 and b > 0 the potential has a ring of minima at \(r_{\min} = \sqrt{-a/(2b)}\).

Parameters:

x (array_like, shape (2,)) – Position [x1, x2].
a (float) – Quadratic coefficient (negative for sombrero shape, default −1).
b (float) – Quartic coefficient (positive for stability, default 1).

Returns:

U (float) – Potential energy.
F (ndarray, shape (2,)) – Force vector F = −∇U.

stochkin.potentials.central_well_barrier_ring_potential(x, a=1.0, b=1.0, A=0.25, sigma=0.2)[source]¶

2D ring potential with a central Gaussian well and an annular barrier.

\[U(r) = b\,r^4 - a\,r^2 - A\,\exp(-r^2/\sigma^2), \quad r = |x|\]

The potential supports three distinct regions: a central well (depth controlled by A and σ), an annular barrier, and a ring of minima at roughly \(r_{\min} \approx \sqrt{a/(2b)}\).

Parameters:

x (array_like, shape (2,)) – Position [x1, x2].
a (float) – Quadratic coefficient (> 0, controls ring radius).
b (float) – Quartic coefficient (> 0, controls ring steepness).
A (float) – Depth of the central Gaussian well (> 0).
sigma (float) – Width of the central Gaussian well (> 0).

Returns:

U (float) – Potential energy.
F (ndarray, shape (2,)) – Force vector F = −∇U.

stochkin.potentials.muller_potential(x, E_clip=50.0)[source]¶

Müller–Brown potential (four-Gaussian surface).

The Müller–Brown potential is a classic benchmark for transition-path and rare-event methods. It features three local minima connected by two saddle points.

Parameters:

x (array_like, shape (2,)) – Position [x, y].
E_clip (float, optional) – Exponent clipping threshold to prevent overflow (default 50).

Returns:

U (float) – Potential energy.
F (ndarray, shape (2,)) – Force vector F = −∇U.

References

Müller and L. D. Brown, Theor. Chim. Acta 53, 75 (1979).

stochkin.potentials.simple_double_well_2d(x, a=10.0, x0=1.0, k_y=10)[source]¶

Symmetric two-state double well in 2D with harmonic y-confinement.

\[U(x_1, x_2) = a\,(x_1^2 - x_0^2)^2 + \tfrac{1}{2} k_y\, x_2^2\]

Two minima at \((\pm x_0, 0)\) separated by a barrier of height \(a\,x_0^4\) at the origin.

Parameters:

x (array_like, shape (2,)) – Position [x1, x2].
a (float) – Double-well strength along x (default 10).
x0 (float) – Location of the two minima along x (default 1).
k_y (float) – Harmonic stiffness along y (default 10).

Returns:

U (float) – Potential energy.
F (ndarray, shape (2,)) – Force vector F = −∇U.

stochkin.potentials.double_well_1d(x, a=1.0, x0=1.0)[source]¶

Simple symmetric 1D double-well potential.

U(x) = a (x^2 - x0^2)^2

Two minima at x = ±x0
Barrier at x = 0 of height a * x0^4

Parameters:

x (scalar or array-like, shape (1,)) – Position along the 1D coordinate.
a (float) – Double-well strength.
x0 (float) – Position of the minima.

Returns:

U (float) – Potential energy at x.
F (ndarray, shape (1,)) – Force = -dU/dx as a 1D vector.

stochkin.potentials.make_double_well_1d(a=1.0, x0=1.0)[source]¶

Convenience wrapper that returns a picklable 1D double-well potential.

We use functools.partial on the top-level function double_well_1d so that the result works with multiprocessing.Pool.

class stochkin.potentials.Basin(id, minimum, f_min, radius, bounds)[source]¶

Bases: object

Description of a 2D metastable basin on a gridded free-energy surface.

Instances are created automatically by build_basin_network_from_potential() and stored inside BasinNetwork.

Parameters:

id (int)
minimum (ndarray)
f_min (float)
radius (float)
bounds (ndarray)

id¶

Basin index (0 … N−1).

Type:: int

minimum¶

Coordinates (x, y) of the local free-energy minimum.

Type:: ndarray, shape (2,)

f_min¶

Free-energy value at the minimum.

Type:: float

radius¶

Approximate spatial extent (max distance from minimum to any assigned grid point).

Type:: float

bounds¶

Rectangular bounding box [[xmin, xmax], [ymin, ymax]].

Type:: ndarray, shape (2, 2)

id: int¶

minimum: ndarray¶

f_min: float¶

radius: float¶

bounds: ndarray¶

class stochkin.potentials.BasinNetwork(basins, xs, ys, U, labels, labels_full=None, core_labels=None)[source]¶

Bases: object

Collection of metastable basins on a regular 2D grid.

This is the central data structure for multi-basin analyses (MFPT, CTMC generator estimation, committor maps). It holds the basin definitions together with the underlying grid and label arrays.

Parameters:

basins (list)
xs (ndarray)
ys (ndarray)
U (ndarray)
labels (ndarray)
labels_full (ndarray | None)
core_labels (ndarray | None)

basins¶

List of detected Basin objects.

Type:: list[Basin]

xs, ys

1D coordinate arrays defining the rectangular grid.

Type:: ndarray

U¶

Free-energy / potential values on the grid.

Type:: ndarray, shape (nx, ny)

labels¶

Basin assignment for each grid point. labels[i, j] gives the basin id of grid point (xs[i], ys[j]), or −1 if the point is unassigned (e.g. NaN/inf regions).

Type:: ndarray of int, shape (nx, ny)

labels_full¶

Copy of the full-partition labels before any core-shrinkage.

Type:: ndarray or None

core_labels¶

Labels restricted to low-energy basin cores (set by build_core_labels_from_full_labels()).

Type:: ndarray or None

basins: list¶

xs: ndarray¶

ys: ndarray¶

U: ndarray¶

labels: ndarray¶

labels_full: ndarray | None = None¶

core_labels: ndarray | None = None¶

property n_basins: int¶

which_basin(x)[source]¶: Return basin id (int) for a position x = [x, y], or None if x is outside the grid or in an unlabeled region.

sample_point_in_basin(basin_id, rng=None)[source]¶: Sample a random (x, y) from grid points assigned to basin_id. Returns np.array([x, y]) or None if basin is empty.

class stochkin.potentials.Basin1D(id, minimum, f_min, radius, bounds)[source]¶

Bases: object

Description of a 1D metastable basin on a gridded free-energy profile.

Parameters:

id (int)
minimum (float)
f_min (float)
radius (float)
bounds (ndarray)

id¶

Basin index (0 … N−1).

Type:: int

minimum¶

Coordinate of the local free-energy minimum.

Type:: float

f_min¶

Free-energy value at the minimum.

Type:: float

radius¶

Half-width: max distance from the minimum among assigned grid points.

Type:: float

bounds¶

[xmin, xmax] of assigned grid points.

Type:: ndarray, shape (2,)

id: int¶

minimum: float¶

f_min: float¶

radius: float¶

bounds: ndarray¶

class stochkin.potentials.BasinNetwork1D(basins, s, U, labels)[source]¶

Bases: object

Collection of metastable basins on a 1D free-energy profile.

This is the 1D analogue of BasinNetwork. It is consumed by compute_ctmc_generator_fpe_1d() and the high-level workflow functions in stochkin.workflows.

Parameters:

basins (list)
s (ndarray)
U (ndarray)
labels (ndarray)

basins¶

List of detected Basin1D objects.

Type:: list[Basin1D]

s¶

1D coordinate grid.

Type:: ndarray

U¶

Free-energy values on the grid.

Type:: ndarray

labels¶

Basin id for each grid point, or −1 if unassigned.

Type:: ndarray of int

basins: list¶

s: ndarray¶

U: ndarray¶

labels: ndarray¶

property n_basins: int¶

which_basin(x)[source]¶: Return basin id (int) for position x, or None if outside grid or in an unlabeled region.

sample_point_in_basin(basin_id, rng=None)[source]¶: Sample a random x from grid points assigned to basin_id. Returns np.array([x]) or None if basin is empty.

stochkin.potentials.sample_potential_grid(potential, xlim=(-2.0, 2.0), ylim=(-2.0, 2.0), nx=200, ny=200)[source]¶

Sample a 2D potential on a regular grid.

For FES-based potentials this can be a hot path (basin detection). If the potential exposes a vectorized grid evaluator, use it.

Parameters:

potential (callable) – Potential callable that returns (U, F) for a single point.
xlim (tuple) – Bounds of the grid.
ylim (tuple) – Bounds of the grid.
nx (int) – Grid resolution.
ny (int) – Grid resolution.

Returns:

xs, ys (1D arrays)
U (2D array, shape (nx, ny))

stochkin.potentials.build_basin_network_from_fes_1d(s, F, max_basins=None, verbose=True)[source]¶

Build a BasinNetwork1D directly from a 1D FES on a grid.

Parameters:

s (1D array) – Grid points (e.g., CV values).
F (1D array) – Free energy (or potential) at each grid point.

stochkin.potentials.build_basin_network_from_potential_1d(potential, xlim=(-2.0, 2.0), ns=200, max_basins=None, verbose=True)[source]¶

Sample a 1D potential on a grid and build a BasinNetwork1D.

potentialcallable x->[U, F]: 1D potential taking x=[s] or x=(s,) and returning (U, F).

stochkin.potentials.detect_basins_for_mfpt_1d(potential, xlim=(-2.0, 2.0), ns=200, max_basins=None, verbose=True, core_fraction=None, core_cut=None)[source]¶

Convenience wrapper for MFPT analysis in 1D.

Samples the potential on a 1D grid, finds minima, and returns a BasinNetwork1D instance.

Parameters:

core_fraction (float | None)
core_cut (float | None)

stochkin.potentials.build_core_labels_from_full_labels(U_grid, labels_full, *, core_fraction=0.05, core_cut=None)[source]¶

Shrink full-partition basin labels to low-energy core sets.

Core sets are small regions around each basin minimum. They are used as Dirichlet boundary conditions in CTMC backward-equation solves (committors, exit times) so that the transition region between basins is part of the unknown domain.

For each basin b one of two rules applies:

If core_cut is provided (absolute energy units):

core = {cells with label == b and U ≤ U_min(b) + core_cut}

Otherwise (default):

core = lowest *core_fraction* energies within basin b

At least one cell (the minimum) is always included per basin.

Parameters:

U_grid (ndarray) – Free-energy values (1D or 2D, same shape as labels_full).
labels_full (ndarray of int) – Full-partition labels (≥ 0 inside a basin, −1 elsewhere).
core_fraction (float) – Quantile fraction (0, 1] of each basin to keep (default 0.05 = 5 %).
core_cut (float or None) – Absolute energy cutoff above the basin minimum. Overrides core_fraction when provided.

Returns:

core_labels – Same shape as labels_full; −1 outside cores.

Return type:

ndarray of int

stochkin.potentials.build_basin_network_from_potential(potential, xlim=(-2.0, 2.0), ylim=(-2.0, 2.0), nx=200, ny=200, max_basins=None, verbose=True, core_fraction=0.05, core_cut=None)[source]¶

Detect metastable basins on a 2D potential / FES.

The algorithm:

Samples U(x, y) on a regular nx × ny grid.
Finds local minima (strict 8-neighbour comparison).
Assigns each grid point to the nearest minimum (Voronoi-like).
Optionally shrinks basins to low-energy cores for CTMC analyses.

Parameters:

potential (callable) – potential([x, y]) -> (U, F) with F = −∇U.
xlim (tuple of float) – Grid bounds (lo, hi).
ylim (tuple of float) – Grid bounds (lo, hi).
nx (int) – Grid resolution.
ny (int) – Grid resolution.
max_basins (int or None) – If not None, keep only the max_basins deepest minima.
verbose (bool) – Print detection summary.
core_fraction (float) – Fraction of each basin (by energy ranking) kept as the basin core.
core_cut (float or None) – Absolute energy cutoff above the basin minimum (overrides core_fraction if provided).

Returns:

Network with basins, labels, labels_full, and core_labels attributes populated.

Return type:

BasinNetwork

stochkin.potentials.detect_basins_for_mfpt(potential, xlim=(-2.0, 2.0), ylim=(-2.0, 2.0), nx=200, ny=200, max_basins=None, verbose=True, core_fraction=0.05, core_cut=None)[source]¶

Convenience wrapper specifically for MFPT analysis:

Returns:: basin_network : BasinNetwork

Parameters:

core_fraction (float)
core_cut (float | None)