datasets

This submodule provides a suite of one-dimensional potential energy landscapes and their corresponding biased-force variants for use in score-based and diffusion-based modeling experiments. Each dataset class encapsulates a specific analytic potential function and supporting machinery to generate samples, compute energies, and (where applicable) apply an external biasing force.

Classes: WPotential1D A symmetric double-well (W-shaped) potential in one dimension. BiasedForceWPotential1D The WPotential1D with an added constant biasing force term. ToyMembranePotential1D A simple membrane-like potential featuring a central barrier and flanking wells. BiasedForceToyMembranePotential1D The ToyMembranePotential1D augmented with an external biasing force. ToyMembrane2Potential1D An extended membrane potential with two barriers and three wells. BiasedForceToyMembrane2Potential1D The ToyMembrane2Potential1D with an additional biasing force. ToyMembrane3Potential1D A higher-order membrane potential featuring three barriers and four wells. BiasedForceToyMembrane3Potential1D The ToyMembrane3Potential1D augmented with an external biasing force.

All classes expose a consistent interface for: • Sampling data points from the potential's Boltzmann distribution. • Computing potential energies and (optional) biasing forces. • Integrating seamlessly with score-based learning workflows.

Usage example:

dataset = WPotential1D(num_samples=10000, temperature=1.0)
x, energy = dataset.sample()

BiasedForceWPotential1D(*, x=lambda: jnp.linspace(0, 1, 100)(), gamma=lambda x: 1.0, bias) dataclass

Bases: WPotential1D

Biased-force dataset for the 1D W-potential.

Extends WPotential1D by adding a bias force \(b(x,t)\).

Parameters:

• bias (callable) –

  Bias force function \(b(x,t)\).

Methods:

• sample –

  Simulate biased dynamics and return samples.

sample(key, dt=0.0001, n_steps=int(100000.0), n_samples=2048, beta=1.0)

Sample positions with bias force via Euler-Maruyama.

Source code in src/fpsl/datasets/datasets.py

def sample(
    self,
    key: JaxKey,
    dt: Float[ArrayLike, ''] = 1e-4,
    n_steps: Int[ArrayLike, ''] = int(1e5),
    n_samples: Int[ArrayLike, ''] = 2048,
    beta: Float[ArrayLike, ''] = 1.0,
) -> Float[ArrayLike, 'n_samples 1']:
    """Sample positions with bias force via Euler-Maruyama."""
    integrator = BiasedForceEulerMaruyamaIntegrator(
        bias_force=self.bias,
        potential=self.potential,
        n_dims=1,
        dt=dt,
        beta=beta,
        n_heatup=n_steps,
        gamma=self.gamma,
    )
    key1, key2 = jax.random.split(key)
    trajs, _, _ = integrator.integrate(
        key=key1,
        X=jax.random.uniform(key2, (n_samples, 1)),
        n_steps=0,
    )
    return trajs[-1] % 1

BiasedForceToyMembranePotential1D(*, x=lambda: jnp.linspace(0, 1, 100)(), gamma=lambda x: 1.0, bias) dataclass

Bases: ToyMembranePotential1D, BiasedForceWPotential1D

Biased-force dataset for the toy membrane potential.

BiasedForceToyMembrane2Potential1D(*, x=lambda: jnp.linspace(0, 1, 100)(), gamma=lambda x: 1.0, bias) dataclass

Bases: ToyMembrane2Potential1D, BiasedForceWPotential1D

Biased-force dataset for the second toy membrane potential.

BiasedForceToyMembrane3Potential1D(*, x=lambda: jnp.linspace(0, 1, 100)(), gamma=lambda x: 1.0, bias) dataclass

Bases: ToyMembrane3Potential1D, BiasedForceWPotential1D

Biased-force dataset for the third toy membrane potential.

WPotential1D(*, x=lambda: jnp.linspace(0, 1, 100)(), gamma=lambda x: 1.0) dataclass

Bases: DataSet

Dataset for the 1D W-potential.

This class integrates the 1D W-potential \(U(x)\) using the overdamped Langevin dynamics (Euler-Maruyama).

Parameters:

• x ((array_like, shape(dim1)), default: linspace(0,1,100) ) –

  Grid points for plotting the potential.

• gamma (callable, default: lambda x: 1.0 ) –

  Friction function \(\gamma(x)\), defaults to constant 1.

Methods:

• potential –

  Returns \(U(x)\) from w_potential_1d.

• plot_potential –

  Plot \(U(x)\) vs x.

• sample –

  Simulate and return samples modulo 1.

Returns:

• samples ( (ndarray, shape(n_samples, 1)) ) –

  Final positions of particles samples.

potential(x, t)

Evaluate the W-potential at x.

Parameters:

• x ((array_like, shape(1))) –

  Position in [0,1].

• t (float) –

  Time (ignored for static potential).

Returns:

• U ( float ) –

  Potential energy \(U(x)\).

Source code in src/fpsl/datasets/datasets.py

def potential(
    self,
    x: Float[ArrayLike, ' dim1'],
    t: Float[ArrayLike, ''],
) -> Float[ArrayLike, '']:
    r"""Evaluate the W-potential at x.

    Parameters
    ----------
    x : array_like, shape (1,)
        Position in [0,1].
    t : float
        Time (ignored for static potential).

    Returns
    -------
    U : float
        Potential energy $U(x)$.
    """
    return w_potential_1d(x[0])

plot_potential(x=None, title='Toy W-Potential')

Plot the potential energy curve.

Parameters:

• x (array_like, default: None ) –

  Grid points; defaults to self.x.

• title (str, default: 'Toy W-Potential' ) –

  Plot title.

Returns:

• ax ( Axes ) –

  Matplotlib Axes instance with the plot.

Source code in src/fpsl/datasets/datasets.py

def plot_potential(
    self,
    x: None | Float[ArrayLike, ' dim1'] = None,
    title: str = 'Toy W-Potential',
):
    """Plot the potential energy curve.

    Parameters
    ----------
    x : array_like, optional
        Grid points; defaults to self.x.
    title : str, optional
        Plot title.

    Returns
    -------
    ax : matplotlib.axes.Axes
        Matplotlib Axes instance with the plot.
    """
    if x is None:
        x = self.x
    vectorized_potential = jnp.vectorize(
        lambda xv: self.potential(jnp.array([xv]), 0.0)
    )
    ax = plt.gca()
    ax.plot(x, vectorized_potential(x), 'k--', label='ref')
    ax.set_xlabel('$x$')
    ax.set_ylabel('$U(x)$')
    ax.set_xlim(x[0], x[-1])
    ax.set_title(title)
    return ax

sample(key, dt=0.0001, n_steps=int(100000.0), n_samples=2048, beta=1.0)

Sample positions via Euler-Maruyama integration.

Parameters:

• key (JaxKey) –

  PRNG key.

• dt (float, default: 0.0001 ) –

  Time step size.

• n_steps (int, default: int(100000.0) ) –

  Number of heat-up steps.

• n_samples (int, default: 2048 ) –

  Number of independent trajectories.

• beta (float, default: 1.0 ) –

  Inverse temperature.

Returns:

• samples ( (ndarray, shape(n_samples, 1)) ) –

  Final positions modulo 1.

Source code in src/fpsl/datasets/datasets.py

def sample(
    self,
    key: JaxKey,
    dt: Float[ArrayLike, ''] = 1e-4,
    n_steps: Int[ArrayLike, ''] = int(1e5),
    n_samples: Int[ArrayLike, ''] = 2048,
    beta: Float[ArrayLike, ''] = 1.0,
) -> Float[ArrayLike, 'n_samples 1']:
    """Sample positions via Euler-Maruyama integration.

    Parameters
    ----------
    key : JaxKey
        PRNG key.
    dt : float
        Time step size.
    n_steps : int
        Number of heat-up steps.
    n_samples : int
        Number of independent trajectories.
    beta : float
        Inverse temperature.

    Returns
    -------
    samples : ndarray, shape (n_samples, 1)
        Final positions modulo 1.
    """
    integrator = EulerMaruyamaIntegrator(
        potential=self.potential,
        n_dims=1,
        dt=dt,
        beta=beta,
        n_heatup=n_steps,
        gamma=self.gamma,
    )
    key1, key2 = jax.random.split(key)
    trajs, _, _ = integrator.integrate(
        key=key1,
        X=jax.random.uniform(key2, (n_samples, 1)),
        n_steps=0,
    )
    return trajs[-1] % 1

ToyMembranePotential1D(*, x=lambda: jnp.linspace(0, 1, 100)(), gamma=lambda x: 1.0) dataclass

Bases: WPotential1D

Dataset for a toy membrane potential in 1D.

Overrides potential with \(U(x)=\mathrm{toy\_membrane\_potential\_1d}(x)\). Inherits sampling and plotting behavior from WPotential1D.

potential(x, t)

Evaluate the toy membrane potential at x.

Source code in src/fpsl/datasets/datasets.py

def potential(
    self,
    x: Float[ArrayLike, ' dim1'],
    t: Float[ArrayLike, ''],
) -> Float[ArrayLike, '']:
    """Evaluate the toy membrane potential at x."""
    return toy_membrane_potential_1d(x[0])

plot_potential(x=None, title='Toy Membrane Potential')

Plot the toy membrane potential curve.

Source code in src/fpsl/datasets/datasets.py

def plot_potential(
    self,
    x: None | Float[ArrayLike, ' dim1'] = None,
    title: str = 'Toy Membrane Potential',
):
    """Plot the toy membrane potential curve."""
    return super().plot_potential(x, title)

ToyMembrane2Potential1D(*, x=lambda: jnp.linspace(0, 1, 100)(), gamma=lambda x: 1.0) dataclass

Bases: WPotential1D

Dataset for a second toy membrane potential in 1D.

Overrides potential with \(U(x)=\mathrm{toy\_membrane2\_potential\_1d}(x)\).

potential(x, t)

Evaluate the second toy membrane potential at x.

Source code in src/fpsl/datasets/datasets.py

def potential(
    self,
    x: Float[ArrayLike, ' dim1'],
    t: Float[ArrayLike, ''],
) -> Float[ArrayLike, '']:
    """Evaluate the second toy membrane potential at x."""
    return toy_membrane2_potential_1d(x[0])

plot_potential(x=None, title='Toy Membrane 2 Potential')

Plot the second toy membrane potential curve.

Source code in src/fpsl/datasets/datasets.py

def plot_potential(
    self,
    x: None | Float[ArrayLike, ' dim1'] = None,
    title: str = 'Toy Membrane 2 Potential',
):
    """Plot the second toy membrane potential curve."""
    return super().plot_potential(x, title)

ToyMembrane3Potential1D(*, x=lambda: jnp.linspace(0, 1, 100)(), gamma=lambda x: 1.0) dataclass

Bases: WPotential1D

Dataset for a third toy membrane potential in 1D.

Overrides potential with \(U(x)=\mathrm{toy\_membrane3\_potential\_1d}(x)\).

potential(x, t)

Evaluate the third toy membrane potential at x.

Source code in src/fpsl/datasets/datasets.py

def potential(
    self,
    x: Float[ArrayLike, ' dim1'],
    t: Float[ArrayLike, ''],
) -> Float[ArrayLike, '']:
    """Evaluate the third toy membrane potential at x."""
    return toy_membrane3_potential_1d(x[0])

plot_potential(x=None, title='Toy Membrane 3 Potential')

Plot the third toy membrane potential curve.

Source code in src/fpsl/datasets/datasets.py

def plot_potential(
    self,
    x: None | Float[ArrayLike, ' dim1'] = None,
    title: str = 'Toy Membrane 3 Potential',
):
    """Plot the third toy membrane potential curve."""
    return super().plot_potential(x, title)