Analysis#

The analysis module provides functions for analyzing simulation results, including correlation analysis, kernel density estimation (KDE), entropy estimation, and mutual information computation.

Correlation Functions#

stochastix.analysis.autocorrelation #

autocorrelation(
    results: SimulationResults,
    n_points: int = 1000,
    species: str | tuple[str, ...] = '*',
) -> tuple[jnp.ndarray, jnp.ndarray]

Compute the autocorrelation of species trajectories.

This function first interpolates the simulation data onto a regular time grid. Then, it calculates the normalized autocorrelation for each specified species' trajectory. The core computation is JIT-compiled for performance.

Parameters:

results (SimulationResults) –

The SimulationResults object from a stochsimsolve simulation.
n_points (int, default: 1000 ) –

The number of points for the interpolation grid.
species (str | tuple[str, ...], default: '*' ) –

The species for which to compute the autocorrelation. Can be "*" for all species, or a string or tuple of strings for specific species.

Returns:

tuple[ndarray, ndarray] –
Tuple (lags, autocorrs) where:
- lags: The time lags for the autocorrelation, in the same units as the simulation time.
- autocorrs: A 2D array where autocorrs[:, i] is the autocorrelation of the i-th species.

stochastix.analysis.cross_correlation #

cross_correlation(
    results: SimulationResults,
    species1: str,
    species2: str,
    n_points: int = 1000,
) -> tuple[jnp.ndarray, jnp.ndarray]

Compute the cross-correlation between two species trajectories.

This function interpolates the simulation data onto a regular time grid and then computes the normalized cross-correlation between two specified species.

Parameters:

results (SimulationResults) –

The SimulationResults object from a stochsimsolve simulation.
species1 (str) –

The name of the first species.
species2 (str) –

The name of the second species.
n_points (int, default: 1000 ) –

The number of points for the interpolation grid.

Returns:

tuple[ndarray, ndarray] –
A tuple (lags, cross_corr) where:
- lags: The time lags for the cross-correlation.
- cross_corr: A 1D array of the cross-correlation values.

Kernel Density Estimation#

1D KDE#

stochastix.analysis.kde #

kde(
    x: ndarray,
    n_grid_points: int | None = None,
    min_max_vals: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    kernel: Literal = 'wendland_c2',
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray]

Kernel density estimation with a selectable 1D kernel.

This computes a JAX-compatible KDE centered at each sample. Supported kernels are 'triangular', 'exponential', 'gaussian', and 'wendland_c2'. Each sample's kernel is renormalized over the evaluation grid to avoid boundary mass loss on finite support.

Dirichlet smoothing

When density=True, applies add-α smoothing to the multinomial pmf implied by the soft counts before converting to a pdf: p_hat = (counts + α) / (N + α*K). When density=False, returns raw soft counts (no smoothing).

JIT-compatibility

For JIT-compatibility, provide concrete binning parameters. If n_grid_points or min_max_vals are None, bin parameters are derived from data outside of JIT.

Parameters:

x (ndarray) –

1D array of samples. If not 1D, it will be flattened.
n_grid_points (int | None, default: None ) –

Number of grid points. If None, inferred from the integer span [floor(min(x)), ceil(max(x))].
min_max_vals (tuple[float, float] | None, default: None ) –

Tuple (min_val, max_val) defining the bin range. If None, determined from data.
density (bool, default: True ) –

If True, returns a probability density function whose Riemann sum over the grid integrates to 1 (via normalization by sum * grid_step). If False, returns unnormalized counts/weights per grid point.
weights (ndarray | None, default: None ) –

Optional nonnegative weights per sample (same length as x). When provided, kernel contributions are multiplied by these weights.
bw_multiplier (float, default: 1.0 ) –

Kernel scale as a multiple of the bin width. Default is 1.0.
kernel (Literal, default: 'wendland_c2' ) –

Type of kernel to use. One of 'triangular', 'exponential', 'gaussian', or 'wendland_c2'. Default is 'wendland_c2'.
dirichlet_alpha (float | None, default: 0.1 ) –

Per-bin pseudo-count for Dirichlet smoothing. Default is 0.1. Note: dirichlet_kappa takes priority over this parameter if provided.
dirichlet_kappa (float | None, default: None ) –

Total pseudo-count for Dirichlet smoothing. If provided, takes priority over dirichlet_alpha and alpha = kappa / K where K is the number of grid points. If None, uses dirichlet_alpha instead.

Returns:

tuple[ndarray, ndarray] –
A tuple (grid, values) where:
- grid: 1D array of evaluation points (bin centers), shape (n_grid_points,).
- values: 1D array of KDE values at the grid points, shape (n_grid_points,). If density=True, these approximate a PDF.

Raises:

ValueError –

If weights length does not match x length.
ValueError –

If kernel is not a supported kernel name.

stochastix.analysis.kde_triangular #

kde_triangular(
    x: ndarray,
    n_grid_points: int | None = None,
    min_max_vals: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray]

Backward-compatible triangular specialization of kde.

stochastix.analysis.kde_exponential #

kde_exponential(
    x: ndarray,
    n_grid_points: int | None = None,
    min_max_vals: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray]

Backward-compatible exponential specialization of kde.

stochastix.analysis.kde_gaussian #

kde_gaussian(
    x: ndarray,
    n_grid_points: int | None = None,
    min_max_vals: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray]

Backward-compatible Gaussian specialization of kde.

stochastix.analysis.kde_wendland_c2 #

kde_wendland_c2(
    x: ndarray,
    n_grid_points: int | None = None,
    min_max_vals: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray]

Backward-compatible Wendland C2 specialization of kde.

Bandwidths below 1.0 approach a discrete pmf and usually lose useful gradients. Around 1.5 is a practical differentiation heuristic.

2D KDE#

stochastix.analysis.kde_2d #

Kernel density estimation functions for 2D data.

annotations `module-attribute` #

annotations = _Feature((3, 7, 0, 'beta', 1), None, 16777216)

kde_2d #

kde_2d(
    x1: ndarray,
    x2: ndarray,
    n_grid_points1: int | None = None,
    n_grid_points2: int | None = None,
    min_max_vals1: tuple[float, float] | None = None,
    min_max_vals2: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    kernel: Literal = 'wendland_c2',
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

Compute a 2D JAX-compatible KDE with selectable kernel.

This computes a JAX-compatible 2D KDE using a product of two 1D kernels centered at each sample, one per dimension. Supported kernels are 'triangular', 'exponential', 'gaussian', and 'wendland_c2'. Each sample's kernel is renormalized over the evaluation grid to avoid boundary mass loss on finite support.

Dirichlet smoothing

When density=True, applies add-α smoothing to the multinomial pmf implied by the soft counts before converting to a pdf: p_hat = (counts + α) / (N + α*K). When density=False, returns raw soft counts (no smoothing).

JIT-compatibility

For JIT-compatibility, provide concrete binning parameters. If n_grid_points1, n_grid_points2, min_max_vals1, or min_max_vals2 are None, bin parameters are derived from data outside of JIT.

Parameters:

x1 (ndarray) –

1D array of samples for the first dimension. If not 1D, it is flattened.
x2 (ndarray) –

1D array of samples for the second dimension. If not 1D, it is flattened. Must have the same length as x1.
n_grid_points1 (int | None, default: None ) –

Number of grid points for the first dimension. If None, inferred from the integer span [floor(min(x1)), ceil(max(x1))].
n_grid_points2 (int | None, default: None ) –

Number of grid points for the second dimension. If None, inferred from the integer span [floor(min(x2)), ceil(max(x2))].
min_max_vals1 (tuple[float, float] | None, default: None ) –

Tuple (min_val, max_val) defining the range for the first dimension. If None, determined from data.
min_max_vals2 (tuple[float, float] | None, default: None ) –

Tuple (min_val, max_val) defining the range for the second dimension. If None, determined from data.
density (bool, default: True ) –

If True, returns a probability density function whose Riemann sum over the grid integrates to 1 (via normalization by sum * grid_step1 * grid_step2). If False, returns unnormalized counts/weights per grid point.
weights (ndarray | None, default: None ) –

Optional nonnegative weights per sample (same length as x1 and x2). When provided, kernel contributions are multiplied by these weights.
bw_multiplier (float, default: 1.0 ) –

Kernel scale multiplier relative to the grid step in each dimension.
kernel (Literal, default: 'wendland_c2' ) –

Kernel family. Must be one of 'triangular', 'exponential', 'gaussian', or 'wendland_c2'. Default is 'wendland_c2'.
dirichlet_alpha (float | None, default: 0.1 ) –

Per-bin pseudo-count for Dirichlet smoothing. Default is 0.1. Note: dirichlet_kappa takes priority over this parameter if provided.
dirichlet_kappa (float | None, default: None ) –

Total pseudo-count for Dirichlet smoothing. If provided, takes priority over dirichlet_alpha and alpha = kappa / K where K is the total number of grid points (K1 * K2). If None, uses dirichlet_alpha instead.

Returns:

tuple[ndarray, ndarray, ndarray] –
A tuple (grid1, grid2, values) where:
- grid1: 1D array of evaluation points (bin centers) for the first dimension, shape (n_grid_points1,).
- grid2: 1D array of evaluation points (bin centers) for the second dimension, shape (n_grid_points2,).
- values: 2D array of KDE values at the grid points, shape (n_grid_points1, n_grid_points2). If density=True, these approximate a PDF.

Raises:

ValueError –

If x1 and x2 do not have the same length.
ValueError –

If weights length does not match x1/x2 length.
ValueError –

If kernel is not a supported kernel name.

kde_triangular_2d #

kde_triangular_2d(
    x1: ndarray,
    x2: ndarray,
    n_grid_points1: int | None = None,
    n_grid_points2: int | None = None,
    min_max_vals1: tuple[float, float] | None = None,
    min_max_vals2: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

Backward-compatible triangular specialization of kde_2d.

kde_exponential_2d #

kde_exponential_2d(
    x1: ndarray,
    x2: ndarray,
    n_grid_points1: int | None = None,
    n_grid_points2: int | None = None,
    min_max_vals1: tuple[float, float] | None = None,
    min_max_vals2: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

Backward-compatible exponential specialization of kde_2d.

kde_gaussian_2d #

kde_gaussian_2d(
    x1: ndarray,
    x2: ndarray,
    n_grid_points1: int | None = None,
    n_grid_points2: int | None = None,
    min_max_vals1: tuple[float, float] | None = None,
    min_max_vals2: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

Backward-compatible Gaussian specialization of kde_2d.

kde_wendland_c2_2d #

kde_wendland_c2_2d(
    x1: ndarray,
    x2: ndarray,
    n_grid_points1: int | None = None,
    n_grid_points2: int | None = None,
    min_max_vals1: tuple[float, float] | None = None,
    min_max_vals2: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

Backward-compatible Wendland C2 specialization of kde_2d.

stochastix.analysis.kde_triangular_2d #

kde_triangular_2d(
    x1: ndarray,
    x2: ndarray,
    n_grid_points1: int | None = None,
    n_grid_points2: int | None = None,
    min_max_vals1: tuple[float, float] | None = None,
    min_max_vals2: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

Backward-compatible triangular specialization of kde_2d.

stochastix.analysis.kde_exponential_2d #

kde_exponential_2d(
    x1: ndarray,
    x2: ndarray,
    n_grid_points1: int | None = None,
    n_grid_points2: int | None = None,
    min_max_vals1: tuple[float, float] | None = None,
    min_max_vals2: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

Backward-compatible exponential specialization of kde_2d.

stochastix.analysis.kde_gaussian_2d #

kde_gaussian_2d(
    x1: ndarray,
    x2: ndarray,
    n_grid_points1: int | None = None,
    n_grid_points2: int | None = None,
    min_max_vals1: tuple[float, float] | None = None,
    min_max_vals2: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

Backward-compatible Gaussian specialization of kde_2d.

stochastix.analysis.kde_wendland_c2_2d #

kde_wendland_c2_2d(
    x1: ndarray,
    x2: ndarray,
    n_grid_points1: int | None = None,
    n_grid_points2: int | None = None,
    min_max_vals1: tuple[float, float] | None = None,
    min_max_vals2: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]

Backward-compatible Wendland C2 specialization of kde_2d.

Simulation results#

stochastix.analysis.state_kde #

state_kde(
    results: SimulationResults,
    species: str | tuple[str, ...],
    n_grid_points: int | None = None,
    min_max_vals: tuple[float, float] | None = None,
    density: bool = True,
    t: Any = None,
    *,
    kde_type: str = 'wendland_c2',
    bw_multiplier: float = 1.0,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray]

Compute a kernel density estimate (KDE) of the state distribution.

This function is a convenience wrapper around KDE functions that specifically operates on the state of batched SimulationResults. It supports multiple kernel types for density estimation.

Parameters:

results (SimulationResults) –

The SimulationResults from a stochsimsolve simulation. This should contain a batch of simulation trajectories (e.g. from vmapping over stochsimsolve).
species (str | tuple[str, ...]) –

The species for which to compute the KDE. Can be a single species name or a tuple of names.
n_grid_points (int | None, default: None ) –

Number of grid points to use. If None, inferred from the integer span [floor(min(x)), ceil(max(x))].
min_max_vals (tuple[float, float] | None, default: None ) –

Tuple (min_val, max_val) defining the grid range. If None, determined from data.
density (bool, default: True ) –

If True, returns a probability density function whose Riemann sum over the grid integrates to 1. If False, returns unnormalized counts/weights per grid point.
t (Any, default: None ) –

Scalar physical time at which to compute the KDE. If None (default), uses the final recorded state.
kde_type (str, default: 'wendland_c2' ) –

Type of kernel to use. One of 'triangular', 'exponential', 'gaussian', or 'wendland_c2'. Default is 'wendland_c2'.
bw_multiplier (float, default: 1.0 ) –

Kernel bandwidth multiplier. Controls the width of the kernel relative to the grid step size. Default is 1.0.
dirichlet_alpha (float | None, default: 0.1 ) –

Per-bin pseudo-count for Dirichlet smoothing. Default is 0.1. Note: dirichlet_kappa takes priority over this parameter if provided.
dirichlet_kappa (float | None, default: None ) –

Total pseudo-count for Dirichlet smoothing. If provided, takes priority over dirichlet_alpha and alpha = kappa / K where K is the number of grid points. If None, uses dirichlet_alpha instead.

Returns:

tuple[ndarray, ndarray] –
A tuple (grid, values) where:
- grid: 1D array of evaluation points (grid centers), shape (n_grid_points,).
- values: 2D array where values[:, i] is the KDE values for the i-th species at the specified time point, shape (n_grid_points, n_species). If density=True, these approximate a PDF.

Entropy#

stochastix.analysis.entropy #

Entropy functions.

annotations `module-attribute` #

annotations = _Feature((3, 7, 0, 'beta', 1), None, 16777216)

entropy #

entropy(
    x: ndarray,
    n_grid_points: int | None = None,
    min_max_vals: tuple[float, float] | None = None,
    base: float = 2.0,
    *,
    kde_type: str = 'wendland_c2',
    bw_multiplier: float = 1.0,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> jnp.ndarray

Compute robust entropy of a 1D sample distribution.

This function computes entropy on grid cell masses induced by a 1D KDE soft histogram. It uses raw soft counts (density=False) and then applies optional Dirichlet smoothing to the implied pmf before entropy evaluation.

JIT-compatibility

For JIT-compatibility, provide concrete values for all grid parameters (n_grid_points, min_max_vals). If left as None, grid parameters are determined from the data (not JIT-able).

Taking entropy gradients

For gradient-based optimization, using Dirichlet smoothing (dirichlet_alpha > 0 or dirichlet_kappa > 0) is generally preferred and is the safest choice for stable autodiff. With no smoothing (both set to 0), entropy can remain finite, but gradients are piecewise and may be numerically fragile near zero-mass bins.

Parameters:

x (ndarray) –

1D sample array.
n_grid_points (int | None, default: None ) –

Number of grid points. If None, inferred from data.
min_max_vals (tuple[float, float] | None, default: None ) –

Tuple (min_val, max_val) for grid range. If None, inferred from data.
base (float, default: 2.0 ) –

Logarithmic base for entropy. Default is 2.0 (bits).
kde_type (str, default: 'wendland_c2' ) –

Type of kernel to use. One of 'triangular', 'exponential', 'gaussian', or 'wendland_c2'. Default is 'wendland_c2'.
bw_multiplier (float, default: 1.0 ) –

Kernel bandwidth multiplier. Default is 1.0.
dirichlet_alpha (float | None, default: 0.1 ) –

Per-bin pseudo-count for Dirichlet smoothing. Default is 0.1. Note: dirichlet_kappa takes priority over this parameter if provided.
dirichlet_kappa (float | None, default: None ) –

Total pseudo-count for Dirichlet smoothing. If provided, takes priority over dirichlet_alpha. If None, uses dirichlet_alpha instead. Default is None.

Returns:

ndarray –

Entropy of the smoothed grid cell-mass distribution in the requested
ndarray –

logarithmic base.

state_entropy #

state_entropy(
    results: SimulationResults,
    species_at_t: tuple[str, Any],
    n_grid_points: int | None = None,
    min_max_vals: tuple[float, float] | None = None,
    base: float = 2.0,
    *,
    kde_type: str = 'wendland_c2',
    bw_multiplier: float = 1.0,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> jnp.ndarray

Compute entropy of one species at a specific time point.

This function computes entropy of a species distribution at a selected physical time point from batched simulation results.

JIT-compatibility

For JIT-compatibility, provide concrete values for all grid parameters (n_grid_points, min_max_vals). If left as None, grid parameters are determined from the data (not JIT-able).

Taking entropy gradients

For gradient-based optimization, using Dirichlet smoothing (dirichlet_alpha > 0 or dirichlet_kappa > 0) is generally preferred and is the safest choice for stable autodiff. With no smoothing (both set to 0), entropy can remain finite, but gradients are piecewise and may be numerically fragile near zero-mass bins.

Parameters:

results (SimulationResults) –

Batched SimulationResults from simulation.
species_at_t (tuple[str, Any]) –

Tuple (species_name, t) where t is a scalar physical time.
n_grid_points (int | None, default: None ) –

Number of grid points. If None, inferred from data.
min_max_vals (tuple[float, float] | None, default: None ) –

Tuple (min_val, max_val) for grid range. If None, inferred from data.
base (float, default: 2.0 ) –

Logarithmic base for entropy. Default is 2.0 (bits).
kde_type (str, default: 'wendland_c2' ) –

Type of kernel to use. One of 'triangular', 'exponential', 'gaussian', or 'wendland_c2'. Default is 'wendland_c2'.
bw_multiplier (float, default: 1.0 ) –

Kernel bandwidth multiplier. Default is 1.0.
dirichlet_alpha (float | None, default: 0.1 ) –

Per-bin pseudo-count for Dirichlet smoothing. Default is 0.1. Note: dirichlet_kappa takes priority over this parameter if provided.
dirichlet_kappa (float | None, default: None ) –

Total pseudo-count for Dirichlet smoothing. If provided, takes priority over dirichlet_alpha. If None, uses dirichlet_alpha instead. Default is None.

Returns:

ndarray –

Entropy of the selected species at the requested time point.

kde #

kde(
    x: ndarray,
    n_grid_points: int | None = None,
    min_max_vals: tuple[float, float] | None = None,
    density: bool = True,
    weights: ndarray | None = None,
    bw_multiplier: float = 1.0,
    *,
    kernel: Literal = 'wendland_c2',
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> tuple[jnp.ndarray, jnp.ndarray]

Kernel density estimation with a selectable 1D kernel.

This computes a JAX-compatible KDE centered at each sample. Supported kernels are 'triangular', 'exponential', 'gaussian', and 'wendland_c2'. Each sample's kernel is renormalized over the evaluation grid to avoid boundary mass loss on finite support.

Dirichlet smoothing

When density=True, applies add-α smoothing to the multinomial pmf implied by the soft counts before converting to a pdf: p_hat = (counts + α) / (N + α*K). When density=False, returns raw soft counts (no smoothing).

JIT-compatibility

For JIT-compatibility, provide concrete binning parameters. If n_grid_points or min_max_vals are None, bin parameters are derived from data outside of JIT.

Parameters:

x (ndarray) –

1D array of samples. If not 1D, it will be flattened.
n_grid_points (int | None, default: None ) –

Number of grid points. If None, inferred from the integer span [floor(min(x)), ceil(max(x))].
min_max_vals (tuple[float, float] | None, default: None ) –

Tuple (min_val, max_val) defining the bin range. If None, determined from data.
density (bool, default: True ) –

If True, returns a probability density function whose Riemann sum over the grid integrates to 1 (via normalization by sum * grid_step). If False, returns unnormalized counts/weights per grid point.
weights (ndarray | None, default: None ) –

Optional nonnegative weights per sample (same length as x). When provided, kernel contributions are multiplied by these weights.
bw_multiplier (float, default: 1.0 ) –

Kernel scale as a multiple of the bin width. Default is 1.0.
kernel (Literal, default: 'wendland_c2' ) –

Type of kernel to use. One of 'triangular', 'exponential', 'gaussian', or 'wendland_c2'. Default is 'wendland_c2'.
dirichlet_alpha (float | None, default: 0.1 ) –

Per-bin pseudo-count for Dirichlet smoothing. Default is 0.1. Note: dirichlet_kappa takes priority over this parameter if provided.
dirichlet_kappa (float | None, default: None ) –

Total pseudo-count for Dirichlet smoothing. If provided, takes priority over dirichlet_alpha and alpha = kappa / K where K is the number of grid points. If None, uses dirichlet_alpha instead.

Returns:

tuple[ndarray, ndarray] –
A tuple (grid, values) where:
- grid: 1D array of evaluation points (bin centers), shape (n_grid_points,).
- values: 1D array of KDE values at the grid points, shape (n_grid_points,). If density=True, these approximate a PDF.

Raises:

ValueError –

If weights length does not match x length.
ValueError –

If kernel is not a supported kernel name.

normalize_time_scalar #

normalize_time_scalar(t, *, arg_name: str) -> jnp.ndarray

Normalize a scalar timestamp to float dtype without concretization.

pytree_to_state #

pytree_to_state(
    tree: Any, species: tuple[str, ...]
) -> jaxlib._jax.Array

Converts a PyTree or other initial state formats to a flat JAX array.

This function processes an initial state tree which can be a dictionary, an object with attributes (like a named tuple or an Equinox module), or an array-like object, and converts it into a flat JAX array of species counts. The order of species in the output array is determined by species.

Parameters:

tree (Any) –

The initial state. Can be a PyTree (dictionary, custom object) with leaves named after species, or an array-like object.
species (tuple[str, ...]) –

The species names, in the order they should appear in the output array.
dtype –

The data type for the output array.

Returns:

Array –

A 1D JAX array representing the state vector, ordered according to species.

stochastix.analysis.state_entropy #

state_entropy(
    results: SimulationResults,
    species_at_t: tuple[str, Any],
    n_grid_points: int | None = None,
    min_max_vals: tuple[float, float] | None = None,
    base: float = 2.0,
    *,
    kde_type: str = 'wendland_c2',
    bw_multiplier: float = 1.0,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> jnp.ndarray

Compute entropy of one species at a specific time point.

This function computes entropy of a species distribution at a selected physical time point from batched simulation results.

JIT-compatibility

For JIT-compatibility, provide concrete values for all grid parameters (n_grid_points, min_max_vals). If left as None, grid parameters are determined from the data (not JIT-able).

Taking entropy gradients

For gradient-based optimization, using Dirichlet smoothing (dirichlet_alpha > 0 or dirichlet_kappa > 0) is generally preferred and is the safest choice for stable autodiff. With no smoothing (both set to 0), entropy can remain finite, but gradients are piecewise and may be numerically fragile near zero-mass bins.

Parameters:

results (SimulationResults) –

Batched SimulationResults from simulation.
species_at_t (tuple[str, Any]) –

Tuple (species_name, t) where t is a scalar physical time.
n_grid_points (int | None, default: None ) –

Number of grid points. If None, inferred from data.
min_max_vals (tuple[float, float] | None, default: None ) –

Tuple (min_val, max_val) for grid range. If None, inferred from data.
base (float, default: 2.0 ) –

Logarithmic base for entropy. Default is 2.0 (bits).
kde_type (str, default: 'wendland_c2' ) –

Type of kernel to use. One of 'triangular', 'exponential', 'gaussian', or 'wendland_c2'. Default is 'wendland_c2'.
bw_multiplier (float, default: 1.0 ) –

Kernel bandwidth multiplier. Default is 1.0.
dirichlet_alpha (float | None, default: 0.1 ) –

Per-bin pseudo-count for Dirichlet smoothing. Default is 0.1. Note: dirichlet_kappa takes priority over this parameter if provided.
dirichlet_kappa (float | None, default: None ) –

Total pseudo-count for Dirichlet smoothing. If provided, takes priority over dirichlet_alpha. If None, uses dirichlet_alpha instead. Default is None.

Returns:

ndarray –

Entropy of the selected species at the requested time point.

Mutual Information#

stochastix.analysis.mutual_information #

mutual_information(
    x1: ndarray,
    x2: ndarray,
    n_grid_points1: int | None = None,
    n_grid_points2: int | None = None,
    min_max_vals1: tuple[float, float] | None = None,
    min_max_vals2: tuple[float, float] | None = None,
    base: float = 2.0,
    *,
    kde_type: str = 'wendland_c2',
    bw_multiplier: float = 1.0,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> jnp.ndarray

Compute the mutual information between two arrays.

This function uses KDE functions to compute the mutual information between two arrays x1 and x2. The mutual information is a measure of the mutual dependence between the two variables.

JIT-compatibility

For JIT-compatibility, provide concrete values for all grid parameters (n_grid_points1, n_grid_points2, min_max_vals1, min_max_vals2). If left as None, grid parameters are determined from the data (not JIT-able).

Taking MI gradients

For gradient-based optimization, using Dirichlet smoothing (dirichlet_alpha > 0 or dirichlet_kappa > 0) is generally preferred and is the safest choice for stable autodiff. With no smoothing (both set to 0), MI can remain finite, but gradients are piecewise and may be numerically fragile near zero-mass bins.

Parameters:

x1 (ndarray) –

1D array of the first input data.
x2 (ndarray) –

1D array of the second input data. Must have the same length as x1.
n_grid_points1 (int | None, default: None ) –

Number of grid points for the first dimension. If None, determined automatically.
n_grid_points2 (int | None, default: None ) –

Number of grid points for the second dimension. If None, determined automatically.
min_max_vals1 (tuple[float, float] | None, default: None ) –

Tuple (min_val, max_val) for the first dimension's grid range. If None, determined automatically.
min_max_vals2 (tuple[float, float] | None, default: None ) –

Tuple (min_val, max_val) for the second dimension's grid range. If None, determined automatically.
base (float, default: 2.0 ) –

The logarithmic base to use for the entropy calculation. Default is 2.0 (bits).
kde_type (str, default: 'wendland_c2' ) –

Type of kernel to use. One of 'triangular', 'exponential', 'gaussian', or 'wendland_c2'. Default is 'wendland_c2'.
bw_multiplier (float, default: 1.0 ) –

Kernel bandwidth multiplier. Controls the width of the kernel relative to the grid step size. Default is 1.0.
dirichlet_alpha (float | None, default: 0.1 ) –

Per-bin pseudo-count for Dirichlet smoothing. Default is 0.1. Note: dirichlet_kappa takes priority over this parameter if provided.
dirichlet_kappa (float | None, default: None ) –

Total pseudo-count for Dirichlet smoothing. If provided, takes priority over dirichlet_alpha. If None, uses dirichlet_alpha instead. Default is None.

Returns:

ndarray –

The mutual information between x1 and x2 in the specified base.

stochastix.analysis.state_mutual_info #

state_mutual_info(
    results: SimulationResults,
    species_at_t: Iterable[tuple[str, Any]],
    n_grid_points1: int | None = None,
    n_grid_points2: int | None = None,
    min_max_vals1: tuple[float, float] | None = None,
    min_max_vals2: tuple[float, float] | None = None,
    base: float = 2.0,
    *,
    kde_type: str = 'wendland_c2',
    bw_multiplier: float = 1.0,
    dirichlet_alpha: float | None = 0.1,
    dirichlet_kappa: float | None = None,
) -> jnp.ndarray

Compute the mutual information between two species at specific time points.

This function calculates the mutual information between the distributions of two species at two potentially different time points, t1 and t2, from batched simulation results. It uses KDE functions to ensure the entire computation is end-to-end differentiable, which is useful for gradient-based optimization of simulation parameters.

JIT-compatibility

For JIT-compatibility, provide concrete values for all grid parameters (n_grid_points1, n_grid_points2, min_max_vals1, min_max_vals2). If left as None, grid parameters are determined from the data (not JIT-able).

Taking MI gradients

For gradient-based optimization, using Dirichlet smoothing (dirichlet_alpha > 0 or dirichlet_kappa > 0) is generally preferred and is the safest choice for stable autodiff. With no smoothing (both set to 0), MI can remain finite, but gradients are piecewise and may be numerically fragile near zero-mass bins.

Parameters:

results (SimulationResults) –

The SimulationResults from a stochsimsolve simulation. This should contain a batch of simulation trajectories (e.g., from vmapping over stochsimsolve).
species_at_t (Iterable[tuple[str, Any]]) –

An iterable containing two tuples, where each tuple consists of a species name and a scalar physical time point, e.g., [('S1', t1), ('S2', t2)].
n_grid_points1 (int | None, default: None ) –

Number of grid points for the first species. If None, determined automatically (not JIT-compatible).
n_grid_points2 (int | None, default: None ) –

Number of grid points for the second species. If None, determined automatically (not JIT-compatible).
min_max_vals1 (tuple[float, float] | None, default: None ) –

Tuple (min_val, max_val) for the first species' grid range. If None, determined automatically (not JIT-compatible).
min_max_vals2 (tuple[float, float] | None, default: None ) –

Tuple (min_val, max_val) for the second species' grid range. If None, determined automatically (not JIT-compatible).
base (float, default: 2.0 ) –

The logarithmic base for the entropy calculation. Default is 2.0 (bits).
kde_type (str, default: 'wendland_c2' ) –

Type of kernel to use. One of 'triangular', 'exponential', 'gaussian', or 'wendland_c2'. Default is 'wendland_c2'.
bw_multiplier (float, default: 1.0 ) –

Kernel bandwidth multiplier. Controls the width of the kernel relative to the grid step size. Default is 1.0.
dirichlet_alpha (float | None, default: 0.1 ) –

Per-bin pseudo-count for Dirichlet smoothing. Default is 0.1. Note: dirichlet_kappa takes priority over this parameter if provided.
dirichlet_kappa (float | None, default: None ) –

Total pseudo-count for Dirichlet smoothing. If provided, takes priority over dirichlet_alpha. If None, uses dirichlet_alpha instead. Default is None.

Returns:

ndarray –

The mutual information between the distributions of the two specified
ndarray –

species at their respective time points.

Analysis#

Correlation Functions#

stochastix.analysis.autocorrelation #

stochastix.analysis.cross_correlation #

Kernel Density Estimation#

1D KDE#

stochastix.analysis.kde #

stochastix.analysis.kde_triangular #

stochastix.analysis.kde_exponential #

stochastix.analysis.kde_gaussian #

stochastix.analysis.kde_wendland_c2 #

2D KDE#

stochastix.analysis.kde_2d #

annotations module-attribute #

kde_2d #

kde_triangular_2d #

kde_exponential_2d #

kde_gaussian_2d #

kde_wendland_c2_2d #

stochastix.analysis.kde_triangular_2d #

stochastix.analysis.kde_exponential_2d #

stochastix.analysis.kde_gaussian_2d #

stochastix.analysis.kde_wendland_c2_2d #

Simulation results#

stochastix.analysis.state_kde #

Entropy#

stochastix.analysis.entropy #

annotations module-attribute #

entropy #

state_entropy #

kde #

normalize_time_scalar #

pytree_to_state #

stochastix.analysis.state_entropy #

Mutual Information#

stochastix.analysis.mutual_information #

stochastix.analysis.state_mutual_info #

annotations `module-attribute` #

annotations `module-attribute` #