stochkin.fpe¶

Fokker–Planck equation solvers and grid-based kinetic operators.

stochkin.fpe¶

Fokker–Planck equation (FPE) solvers and grid-based kinetic operators.

This module provides numerical tools for forward and backward Smoluchowski / Fokker–Planck equations on free-energy surfaces:

Forward FPE (probability evolution)

solve_fp_steady_state() – 2D time-stepping toward the steady-state density using FiPy (preferred) or an explicit NumPy fallback.
solve_fp_1d_from_fes() – 1D time-stepping on a free-energy profile \(F(s)\) with position-dependent diffusion \(D(s)\).

Discrete Fokker–Planck generator

build_fp_generator_from_fes() – build a detailed-balance-preserving rate matrix \(L\) on a 2D grid (SciPy sparse or dense).

Backward BVP solvers (1D, pure NumPy)

solve_committor_1d_from_fes() – committor \(q(s)\) from a tridiagonal backward equation.
solve_exit_time_1d_from_fes() – mean exit time \(\tau(s)\).
compute_ctmc_generator_fpe_1d() – multi-basin CTMC generator from backward 1D solves.
mfpt_1d_smolu_integral() – analytic Smoluchowski integral for \(\tau_{i \to j}\).

Backward BVP solvers (2D, FiPy)

solve_committor_fipy() – 2D committor on a FiPy mesh.
solve_exit_time_fipy() – 2D mean exit time.
compute_ctmc_generator_fpe_fipy() – multi-basin CTMC generator from backward 2D FiPy solves.

stochkin.fpe.solve_fp_steady_state(potential, xlim=(-2.0, 2.0), ylim=(-2.0, 2.0), nx=100, ny=100, D=1.0, beta=1.0, dt=0.001, n_steps=500, initial='uniform', normalize_each_step=True, viewer=False, plot_final=True)[source]¶

Solve the 2D forward Fokker–Planck equation toward steady state.

Evolves the probability density \(p(x,y,t)\) of the overdamped Langevin equation

\[dX_t = -D\beta\,\nabla U(X_t)\,dt + \sqrt{2D}\,dW_t\]

via the forward FPE in divergence form:

\[\partial_t p = \nabla\cdot(D\nabla p + D\beta\, p\,\nabla U)\]

Uses FiPy (implicit, preferred) when available; otherwise falls back to an explicit NumPy FTCS scheme.

Parameters:

potential (callable) – potential([x, y]) -> (U, F) with F = -\nabla U.
xlim ((float, float)) – Domain bounds.
ylim ((float, float)) – Domain bounds.
nx (int) – Grid points per axis.
ny (int) – Grid points per axis.
D (float) – Diffusion coefficient (constant).
beta (float) – Inverse temperature.
dt (float) – Time step.
n_steps (int) – Number of integration steps.
initial ({'uniform'} or callable) – Initial distribution.
normalize_each_step (bool) – Re-normalise \(p\) at every step.
viewer (bool) – Use FiPy’s live viewer (FiPy only).
plot_final (bool) – Plot the final density.

Returns:

Keys: 'p_grid', 'xs', 'ys', 'U_grid', 'Ux', 'Uy'.

Return type:

dict

stochkin.fpe.build_fp_generator_from_fes(xs, ys, U_grid, D, beta, sparse=True)[source]¶

Build a detailed-balance-preserving discrete FP generator on a regular grid.

Parameters:

xs (1D arrays) – Grid coordinates (monotonic, uniform spacing expected).
ys (1D arrays) – Grid coordinates (monotonic, uniform spacing expected).
U_grid ((nx, ny) array) – Potential / free energy on the grid.
D (float or (nx, ny) array) – Diffusion coefficient (scalar or field).
beta (float) – Inverse temperature.
sparse (bool) – If True, return a SciPy CSR matrix (requires SciPy). If False, return dense.

Returns:

L – Generator with off-diagonals >=0 and row sums 0. Indexing: k=i*ny + j.

Return type:

(N, N)

Notes

A common symmetric (“square-root”) discretization is used:

\[k_{i \to j} = \frac{D_{\text{face}}}{\Delta^2} \exp\!\bigl[-\tfrac{\beta}{2}(U_j - U_i)\bigr]\]

so that for constant D the stationary distribution satisfies \(\pi_i \propto \exp(-\beta U_i)\) up to discretization error.

stochkin.fpe.solve_fp_1d_from_fes(s, F, D, beta=1.0, dt=0.001, n_steps=1000, initial='boltzmann', normalize_each_step=True, viewer=False)[source]¶

1D forward Fokker–Planck solver on a free-energy profile.

Solves

\[\partial_t p(s,t) = \partial_s\bigl[D(s)\bigl( \partial_s p + \beta\, p\, F'(s)\bigr)\bigr]\]

with reflecting (no-flux) boundaries via FiPy.

Parameters:

s (array_like) – Uniform 1D grid.
F (array_like) – Free-energy values on s.
D (array_like) – Diffusion coefficient on s.
beta (float) – Inverse temperature.
dt (float) – Time step.
n_steps (int) – Number of integration steps.
initial ({'uniform', 'boltzmann'} or array_like) – Initial distribution.
normalize_each_step (bool) – Re-normalise at every step.
viewer (bool) – Use FiPy’s live viewer.

Returns:

Keys: 's', 'p', 'F', 'D', 'dF_ds'.

Return type:

dict

Raises:

ImportError – If FiPy is not installed.

stochkin.fpe.mfpt_1d_smolu_integral(s, F, D, beta, i_index, j_index)[source]¶

Analytic Smoluchowski integral for the 1D MFPT \(\tau_{i\to j}\).

Uses the classical formula

\[\tau_{i\to j} = \int_{s_i}^{s_j} \frac{e^{\beta F(z)}}{D(z)} \int_{s_{\min}}^{z} e^{-\beta F(y)}\,dy\;dz\]

evaluated via the trapezoidal rule.

Parameters:

s (array_like) – 1D grid (must be sorted).
F (array_like) – Free-energy values on s.
D (array_like) – Diffusion coefficient on s.
beta (float) – Inverse temperature.
i_index (int) – Grid indices of source and target.
j_index (int) – Grid indices of source and target.

Returns:

Mean first-passage time \(\tau_{i\to j}\).

Return type:

float

stochkin.fpe.solve_committor_1d_from_fes(s, F, D=1.0, beta=1.0, mask_q1=None, mask_q0=None)[source]¶

Solve the 1D committor BVP: d/ds(A q’) = 0 with internal Dirichlet sets.

Parameters:

s (1D arrays) – Grid and free energy (or potential). Must be (approximately) uniform.
F (1D arrays) – Grid and free energy (or potential). Must be (approximately) uniform.
D (float or 1D array) – Diffusion coefficient on the grid.
beta (float) – Inverse thermal energy in units compatible with F.
mask_q1 (boolean arrays) – Dirichlet sets for q=1 and q=0.
mask_q0 (boolean arrays) – Dirichlet sets for q=1 and q=0.

Returns:

q – Committor values.

Return type:

1D array

stochkin.fpe.solve_exit_time_1d_from_fes(s, F, D=1.0, beta=1.0, mask_absorb=None)[source]¶: Solve the 1D exit-time BVP: d/ds(A tau’) = -w with tau=0 on absorbing set.

stochkin.fpe.compute_ctmc_generator_fpe_1d(s, F, basin_network, D=1.0, beta=1.0, init_weight='boltzmann', verbose=True)[source]¶

Compute a CTMC generator between 1D basins using backward Smoluchowski BVPs.

This is the 1D analogue of compute_ctmc_generator_fpe_fipy(), but it uses a NumPy tridiagonal solver (no FiPy dependency).

Parameters:

s (1D arrays) – Uniform grid and free energy/potential.
F (1D arrays) – Uniform grid and free energy/potential.
basin_network (BasinNetwork1D) – Must provide labels defined on the same grid.
D (float or 1D array) – Diffusion coefficient.
beta (float) – Inverse thermal energy.
init_weight ({"boltzmann", "uniform"}) – Weight for basin-averages of exit times and committors.

Returns:

dict with keys

Return type:

K, exit_mean, k_out, p_branch, basin_ids, method.

stochkin.fpe.build_mesh_2d_from_edges(x_edges, y_edges, overlap=2)[source]¶: Build a FiPy 2D rectangular mesh from 1D edge arrays.

stochkin.fpe.make_backward_coefficient_A_face_from_cell_values(mesh, U_cell, D_cell, beta, name_U='U', name_D='D')[source]¶: Construct w=exp(-beta U) and A_face = (D*w).arithmeticFaceValue.

stochkin.fpe.solve_committor_fipy(mesh, A_face, mask_q1, mask_q0, large_value=None, enforce_reflecting=True, solver=None, sweep_tol=1e-10, max_sweeps=2000, verbose=False)[source]¶

Solve the 2D committor BVP on a FiPy mesh.

Solves \(\nabla\cdot(A\,\nabla q) = 0\) with internal Dirichlet constraints \(q=1\) (mask_q1) and \(q=0\) (mask_q0), enforced via large penalty terms.

Parameters:

mesh (FiPy mesh) – 2D mesh.
A_face (FiPy FaceVariable) – Backward coefficient \(A = D\,e^{-\beta U}\) on faces.
mask_q1 (array_like of bool) – Cell masks for the two Dirichlet sets.
mask_q0 (array_like of bool) – Cell masks for the two Dirichlet sets.
large_value (float or None) – Penalty magnitude (auto-scaled if None).
enforce_reflecting (bool) – Impose no-flux on exterior faces.
solver (optional) – FiPy solver parameters.
sweep_tol (optional) – FiPy solver parameters.
max_sweeps (optional) – FiPy solver parameters.
verbose (bool) – Print diagnostics.

Returns:

q (CellVariable) – Committor field.
info (dict) – {'residual', 'n_sweeps'}.

stochkin.fpe.solve_exit_time_fipy(mesh, A_face, w_var, mask_absorb, large_value=None, enforce_reflecting=True, solver=None, sweep_tol=1e-10, max_sweeps=4000, verbose=False)[source]¶

Solve the 2D mean exit-time BVP on a FiPy mesh.

Solves \(\nabla\cdot(A\,\nabla\tau) = -w\) with \(\tau=0\) on absorbing cells, enforced via penalty.

Parameters:

mesh (FiPy mesh) – 2D mesh.
A_face (FiPy FaceVariable) – Backward coefficient on faces.
w_var (CellVariable) – Boltzmann weight \(w = e^{-\beta(U-U_{\min})}\).
mask_absorb (array_like of bool) – Absorbing cell mask.
large_value (float or None) – Penalty magnitude.
enforce_reflecting (bool) – Impose no-flux on exterior faces.
solver (optional) – FiPy solver parameters.
sweep_tol (optional) – FiPy solver parameters.
max_sweeps (optional) – FiPy solver parameters.
verbose (bool) – Print diagnostics.

Returns:

tau (CellVariable) – Mean exit-time field.
info (dict) – {'residual', 'n_sweeps'}.

stochkin.fpe.compute_ctmc_generator_fpe_fipy(x_centers, y_centers, U_ij, basin_network, D=1.0, beta=1.0, x_edges=None, y_edges=None, init_weight='boltzmann', large_value=None, enforce_reflecting=True, solver=None, sweep_tol=1e-10, verbose=True)[source]¶

Compute a CTMC generator between basins using backward FPE FiPy solves.

For each basin i:

Solve exit time \(\tau_i\) to the union of other basins (absorbing). \(k_{\text{out}}(i) = 1 / \langle\tau_i\rangle\) (average over basin i).
For each \(j \neq i\), solve committor \(q_{ij}\) with \(q=1\) on basin j and \(q=0\) on all other basins. \(p_{ij} = \langle q_{ij}\rangle\) averaged over basin i.
Set \(K_{ij} = k_{\text{out}}(i)\, p_{ij}\) and \(K_{ii} = -\sum_{j \neq i} K_{ij}\).

Returns:: dict with keys
Return type:: K, exit_mean, k_out, p_branch, basin_ids, mesh, info, method.