Point Processes — Primitives

Pure functions for every process family. Each family follows the same shape: an intensity \(\lambda(\cdot)\), a compensator \(\Lambda(\cdot)\) (the integrated intensity), a log-likelihood built from \(\sum_i \log \lambda(x_i) - \Lambda\), and a sampler. Everything is JAX — differentiable, jit-able, vmap-able — and PRNG keys are explicit arguments.

The operators and distributions are thin wrappers over these; call the primitives directly when you want a custom likelihood or a vmapped parameter sweep.

Homogeneous Poisson — temporal

Constant intensity \(\lambda(t) = \lambda\), so every quantity is closed-form: the compensator is \(\lambda t\), inter-event gaps are \(\mathrm{Exp}(\lambda)\), and memorylessness makes the mean residual life exactly \(1/\lambda\).

Pure-function primitives for temporal and spatial point processes.

Each submodule is free of class state and free of framework dependencies (no NumPyro — save that for the outer-layer wrappers — no equinox). These are the canonical reference implementations; the operator and distribution layers delegate to them.

hpp_intensity(t: Float[Array, ...], rate: Float[Array, ...]) -> Float[Array, ...]

Constant intensity :math:\lambda(t) = \lambda.

Broadcasts to the joint shape of t and rate so scalar t with batched rate (and vice versa) both work, matching the broadcasting semantics of the other HPP primitives.

Source code in src/xtremax/point_processes/primitives/hpp.py

def hpp_intensity(
    t: Float[Array, ...],
    rate: Float[Array, ...],
) -> Float[Array, ...]:
    """Constant intensity :math:`\\lambda(t) = \\lambda`.

    Broadcasts to the joint shape of ``t`` and ``rate`` so scalar ``t``
    with batched ``rate`` (and vice versa) both work, matching the
    broadcasting semantics of the other HPP primitives.
    """
    t = jnp.asarray(t)
    rate = jnp.asarray(rate)
    out_shape = jnp.broadcast_shapes(t.shape, rate.shape)
    return jnp.broadcast_to(rate, out_shape)

hpp_cumulative_intensity(t: Float[Array, ...], rate: Float[Array, ...]) -> Float[Array, ...]

Compensator :math:\Lambda(t) = \lambda t.

Source code in src/xtremax/point_processes/primitives/hpp.py

def hpp_cumulative_intensity(
    t: Float[Array, ...],
    rate: Float[Array, ...],
) -> Float[Array, ...]:
    """Compensator :math:`\\Lambda(t) = \\lambda t`."""
    return jnp.asarray(rate) * jnp.asarray(t)

hpp_log_prob(n_events: Int[Array, ...], rate: Float[Array, ...], T: Float[Array, ...]) -> Float[Array, ...]

Log-likelihood :math:\log L = n \log \lambda - \lambda T.

Parameters:

Name	Type	Description	Default
n_events	Int[Array, ...]	Observed event count.	required
rate	Float[Array, ...]	Intensity λ > 0.	required
T	Float[Array, ...]	Length of the observation window.	required

Returns:

Type	Description
Float[Array, ...]	Log-likelihood, broadcast over inputs.

Notes

The event times drop out because, given n, they are uniform on [0, T] and their density cancels with the Poisson count factors :math:n! and :math:T^n. What remains is :math:\lambda^n \exp(-\lambda T), which is what this returns in log-space. If you instead want the full Poisson-count log-probability, add :math:-\log(n!) - n \log T.

Source code in src/xtremax/point_processes/primitives/hpp.py

def hpp_log_prob(
    n_events: Int[Array, ...],
    rate: Float[Array, ...],
    T: Float[Array, ...],
) -> Float[Array, ...]:
    """Log-likelihood :math:`\\log L = n \\log \\lambda - \\lambda T`.

    Args:
        n_events: Observed event count.
        rate: Intensity ``λ > 0``.
        T: Length of the observation window.

    Returns:
        Log-likelihood, broadcast over inputs.

    Notes:
        The event times drop out because, given ``n``, they are uniform
        on ``[0, T]`` and their density cancels with the Poisson count
        factors :math:`n!` and :math:`T^n`. What remains is
        :math:`\\lambda^n \\exp(-\\lambda T)`, which is what this
        returns in log-space. If you instead want the full
        Poisson-count log-probability, add
        :math:`-\\log(n!) - n \\log T`.
    """
    n = jnp.asarray(n_events)
    rate = jnp.asarray(rate)
    T = jnp.asarray(T)
    return n * jnp.log(rate) - rate * T

hpp_sample(key: PRNGKeyArray, rate: Float[Array, ...], T: Float[Array, ...], max_events: int, sample_shape: tuple[int, ...] = ()) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Sample ordered event times from a homogeneous Poisson process.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	JAX PRNG key.	required
rate	Float[Array, ...]	Intensity λ > 0 (scalar or batched).	required
T	Float[Array, ...]	Observation window length (broadcasts with rate).	required
max_events	int	Static buffer size. The sampler truncates at this cap; choose comfortably above the expected count to avoid loss of realism (5 * λ T is a reasonable starting point).	required
sample_shape	tuple[int, ...]	Optional leading sample dimensions.	()

Returns:

Type	Description
Float[Array, ...]	Tuple (times, mask, n_events):
Bool[Array, ...]	times with shape sample_shape + batch_shape + (max_events,) — sorted, with padding positions set to T.
Int[Array, ...]	mask with the same shape as times — True where times is a real event, False for padding.
tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	n_events with shape sample_shape + batch_shape — the uncapped Poisson draw. If n_events > max_events the returned events have been truncated to the first max_events.

Notes

The algorithm draws n from :math:\mathrm{Poisson}(\lambda T) and then, unconditionally, generates max_events uniform candidates on [0, T]. Sorting and masking by rank lets the same padded buffer serve any batch element regardless of its count — and keeps the whole routine jit-compatible.

Source code in src/xtremax/point_processes/primitives/hpp.py

def hpp_sample(
    key: PRNGKeyArray,
    rate: Float[Array, ...],
    T: Float[Array, ...],
    max_events: int,
    sample_shape: tuple[int, ...] = (),
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Sample ordered event times from a homogeneous Poisson process.

    Args:
        key: JAX PRNG key.
        rate: Intensity ``λ > 0`` (scalar or batched).
        T: Observation window length (broadcasts with ``rate``).
        max_events: Static buffer size. The sampler truncates at this
            cap; choose comfortably above the expected count to avoid
            loss of realism (``5 * λ T`` is a reasonable starting
            point).
        sample_shape: Optional leading sample dimensions.

    Returns:
        Tuple ``(times, mask, n_events)``:

        * ``times`` with shape ``sample_shape + batch_shape + (max_events,)``
            — sorted, with padding positions set to ``T``.
        * ``mask`` with the same shape as ``times`` — ``True`` where
            ``times`` is a real event, ``False`` for padding.
        * ``n_events`` with shape ``sample_shape + batch_shape`` — the
            uncapped Poisson draw. If ``n_events > max_events`` the
            returned events have been truncated to the first ``max_events``.

    Notes:
        The algorithm draws ``n`` from :math:`\\mathrm{Poisson}(\\lambda T)`
        and then, unconditionally, generates ``max_events`` uniform
        candidates on ``[0, T]``. Sorting and masking by rank lets the
        same padded buffer serve any batch element regardless of its
        count — and keeps the whole routine jit-compatible.
    """
    rate = jnp.asarray(rate)
    T = jnp.asarray(T)
    batch_shape = jnp.broadcast_shapes(rate.shape, T.shape)
    shape = tuple(sample_shape) + batch_shape

    key_n, key_times = random.split(key)

    expected = jnp.broadcast_to(rate * T, shape)
    n_events = random.poisson(key_n, expected, shape=shape)

    # Draw one uniform per buffer slot; mark the first ``n_events`` as
    # real events via ``mask`` and push padding positions to ``+inf``
    # before sorting so the real events occupy ranks ``0..n-1`` in the
    # sorted order without being biased to small values (which would
    # happen if we sorted first and then truncated to the first n).
    T_b = jnp.broadcast_to(T, shape)[..., None]
    ranks = jnp.arange(max_events)
    mask = ranks < n_events[..., None]
    raw = random.uniform(key_times, shape=(*shape, max_events)) * T_b
    raw_sortable = jnp.where(mask, raw, jnp.inf)
    times_sorted = jnp.sort(raw_sortable, axis=-1)
    # Real events are finite; padding positions are set to ``T`` (a
    # safe value inside the window) so downstream intensity evaluations
    # don't trip on ``+inf``.
    times = jnp.where(mask, times_sorted, T_b)
    return SampleResult(times, mask, n_events)

hpp_hazard(t: Float[Array, ...], rate: Float[Array, ...]) -> Float[Array, ...]

Constant hazard equal to rate.

Source code in src/xtremax/point_processes/primitives/hpp.py

def hpp_hazard(
    t: Float[Array, ...],
    rate: Float[Array, ...],
) -> Float[Array, ...]:
    """Constant hazard equal to ``rate``."""
    return hpp_intensity(t, rate)

hpp_survival(t: Float[Array, ...], given_time: Float[Array, ...], rate: Float[Array, ...]) -> Float[Array, ...]

:math:S(t|s) = \exp(-\lambda (t - s)) for t >= s.

Source code in src/xtremax/point_processes/primitives/hpp.py

def hpp_survival(
    t: Float[Array, ...],
    given_time: Float[Array, ...],
    rate: Float[Array, ...],
) -> Float[Array, ...]:
    """:math:`S(t|s) = \\exp(-\\lambda (t - s))` for ``t >= s``."""
    rate = jnp.asarray(rate)
    return jnp.exp(-rate * (jnp.asarray(t) - jnp.asarray(given_time)))

hpp_inter_event_log_prob(tau: Float[Array, ...], rate: Float[Array, ...]) -> Float[Array, ...]

Log density of τ ~ Exp(λ): :math:\log \lambda - \lambda \tau.

Source code in src/xtremax/point_processes/primitives/hpp.py

def hpp_inter_event_log_prob(
    tau: Float[Array, ...],
    rate: Float[Array, ...],
) -> Float[Array, ...]:
    """Log density of ``τ ~ Exp(λ)``: :math:`\\log \\lambda - \\lambda \\tau`."""
    tau = jnp.asarray(tau)
    rate = jnp.asarray(rate)
    return jnp.where(tau >= 0.0, jnp.log(rate) - rate * tau, -jnp.inf)

hpp_mean_residual_life(rate: Float[Array, ...], given_time: Float[Array, ...] | None = None) -> Float[Array, ...]

Memorylessness gives :math:\mathbb{E}[T - s \mid T > s] = 1/\lambda.

Parameters:

Name	Type	Description	Default
rate	Float[Array, ...]	Intensity.	required
given_time	Float[Array, ...] | None	Ignored — present for API symmetry with IPP, where mean residual life does depend on the conditioning time.	None

Returns:

Type	Description
Float[Array, ...]	1 / rate broadcast to the shape of rate.

Source code in src/xtremax/point_processes/primitives/hpp.py

def hpp_mean_residual_life(
    rate: Float[Array, ...],
    given_time: Float[Array, ...] | None = None,
) -> Float[Array, ...]:
    """Memorylessness gives :math:`\\mathbb{E}[T - s \\mid T > s] = 1/\\lambda`.

    Args:
        rate: Intensity.
        given_time: Ignored — present for API symmetry with IPP, where
            mean residual life does depend on the conditioning time.

    Returns:
        ``1 / rate`` broadcast to the shape of ``rate``.
    """
    del given_time
    return 1.0 / jnp.asarray(rate)

hpp_predict_count(start_time: Float[Array, ...], end_time: Float[Array, ...], rate: Float[Array, ...]) -> Float[Array, ...]

Expected event count on [start, end]: λ (end - start).

Source code in src/xtremax/point_processes/primitives/hpp.py

def hpp_predict_count(
    start_time: Float[Array, ...],
    end_time: Float[Array, ...],
    rate: Float[Array, ...],
) -> Float[Array, ...]:
    """Expected event count on ``[start, end]``: ``λ (end - start)``."""
    return jnp.asarray(rate) * (jnp.asarray(end_time) - jnp.asarray(start_time))

hpp_exceedance_log_prob(k: Int[Array, ...], start_time: Float[Array, ...], end_time: Float[Array, ...], rate: Float[Array, ...]) -> Float[Array, ...]

log P(N > k) over [start, end], computed via the Poisson CDF.

Returns log(1 - F_Poisson(k; λ (end - start))). Uses :func:jax.scipy.stats.poisson.cdf under the hood; for large k or small tails this can underflow — callers that care about very rare events should work in log-space end to end.

Source code in src/xtremax/point_processes/primitives/hpp.py

def hpp_exceedance_log_prob(
    k: Int[Array, ...],
    start_time: Float[Array, ...],
    end_time: Float[Array, ...],
    rate: Float[Array, ...],
) -> Float[Array, ...]:
    """log P(N > k) over ``[start, end]``, computed via the Poisson CDF.

    Returns ``log(1 - F_Poisson(k; λ (end - start)))``. Uses
    :func:`jax.scipy.stats.poisson.cdf` under the hood; for large ``k``
    or small tails this can underflow — callers that care about very
    rare events should work in log-space end to end.
    """
    lam = jnp.asarray(rate) * (jnp.asarray(end_time) - jnp.asarray(start_time))
    cdf_k = jax.scipy.stats.poisson.cdf(jnp.asarray(k), lam)
    return jnp.log(jnp.clip(1.0 - cdf_k, 0.0, 1.0))

hpp_return_period(probability: Float[Array, ...], rate: Float[Array, ...]) -> Float[Array, ...]

Time τ such that :math:P(\text{no event in } \tau) = 1 - p.

Equivalently :math:\tau = -\log(1 - p) / \lambda. The log1p form stays accurate as p → 0.

Source code in src/xtremax/point_processes/primitives/hpp.py

def hpp_return_period(
    probability: Float[Array, ...],
    rate: Float[Array, ...],
) -> Float[Array, ...]:
    """Time ``τ`` such that :math:`P(\\text{no event in } \\tau) = 1 - p`.

    Equivalently :math:`\\tau = -\\log(1 - p) / \\lambda`. The ``log1p``
    form stays accurate as ``p → 0``.
    """
    p = jnp.asarray(probability)
    return -jnp.log1p(-p) / jnp.asarray(rate)

Homogeneous Poisson — spatial

Complete spatial randomness (CSR) on a rectangular domain: the joint Janossy likelihood, the marginal Poisson count, and the CSR reference quantities that spatial diagnostics compare against.

Pure-function primitives for temporal and spatial point processes.

Each submodule is free of class state and free of framework dependencies (no NumPyro — save that for the outer-layer wrappers — no equinox). These are the canonical reference implementations; the operator and distribution layers delegate to them.

hpp_spatial_intensity(locations: Float[Array, ...], rate: Float[Array, ...]) -> Float[Array, ...]

Constant intensity λ(s) = λ, broadcast against locations.

Source code in src/xtremax/point_processes/primitives/hpp_spatial.py

def hpp_spatial_intensity(
    locations: Float[Array, ...],
    rate: Float[Array, ...],
) -> Float[Array, ...]:
    """Constant intensity ``λ(s) = λ``, broadcast against ``locations``."""
    locations = jnp.asarray(locations)
    rate = jnp.asarray(rate)
    return jnp.broadcast_to(rate, locations.shape[:-1])

hpp_spatial_log_prob(n_events: Int[Array, ...], rate: Float[Array, ...], domain_volume: Float[Array, ...]) -> Float[Array, ...]

Joint Janossy log-likelihood of a homogeneous Poisson pattern.

Parameters:

Name	Type	Description	Default
n_events	Int[Array, ...]	Observed event count.	required
rate	Float[Array, ...]	Intensity λ > 0.	required
domain_volume	Float[Array, ...]	Lebesgue measure |D|.	required

Returns:

Name	Type	Description
	Float[Array, ...]	Janossy log-likelihood n log λ − λ|D|. This is the
	Float[Array, ...]	constant-intensity specialisation of
	Float[Array, ...]	func:~xtremax.point_processes.primitives.ipp_spatial.ipp_spatial_log_prob,
	Float[Array, ...]	which evaluates :math:\sum_i \log \lambda(s_i) - \Lambda(D).
Setting	Float[Array, ...]	math:\lambda(s) = \lambda recovers exactly this
	Float[Array, ...]	expression.

Notes

The Janossy form drops both the :math:-\log(n!) Poisson-count normaliser and the :math:-n \log |D| iid-uniform location density. Those normalisers cancel out in likelihood ratios used for parameter estimation, so omitting them keeps the HPP log-prob directly comparable to IPP / Hawkes / renewal log-probs across the package. To recover the count-only log-pmf with the :math:\log(n!) term, use :func:hpp_spatial_count_log_prob.

Source code in src/xtremax/point_processes/primitives/hpp_spatial.py

def hpp_spatial_log_prob(
    n_events: Int[Array, ...],
    rate: Float[Array, ...],
    domain_volume: Float[Array, ...],
) -> Float[Array, ...]:
    """Joint Janossy log-likelihood of a homogeneous Poisson pattern.

    Args:
        n_events: Observed event count.
        rate: Intensity ``λ > 0``.
        domain_volume: Lebesgue measure ``|D|``.

    Returns:
        Janossy log-likelihood ``n log λ − λ|D|``. This is the
        constant-intensity specialisation of
        :func:`~xtremax.point_processes.primitives.ipp_spatial.ipp_spatial_log_prob`,
        which evaluates :math:`\\sum_i \\log \\lambda(s_i) - \\Lambda(D)`.
        Setting :math:`\\lambda(s) = \\lambda` recovers exactly this
        expression.

    Notes:
        The Janossy form drops both the :math:`-\\log(n!)` Poisson-count
        normaliser and the :math:`-n \\log |D|` iid-uniform location
        density. Those normalisers cancel out in likelihood ratios used
        for parameter estimation, so omitting them keeps the HPP log-prob
        directly comparable to IPP / Hawkes / renewal log-probs across
        the package. To recover the *count*-only log-pmf with the
        :math:`\\log(n!)` term, use :func:`hpp_spatial_count_log_prob`.
    """
    n = jnp.asarray(n_events)
    rate = jnp.asarray(rate)
    vol = jnp.asarray(domain_volume)
    return n * jnp.log(rate) - rate * vol

hpp_spatial_count_log_prob(n_events: Int[Array, ...], rate: Float[Array, ...], domain_volume: Float[Array, ...]) -> Float[Array, ...]

Marginal Poisson-count log-pmf log P(N = n).

Useful for predict-count workflows that only need the count distribution, ignoring conditional location density.

Source code in src/xtremax/point_processes/primitives/hpp_spatial.py

def hpp_spatial_count_log_prob(
    n_events: Int[Array, ...],
    rate: Float[Array, ...],
    domain_volume: Float[Array, ...],
) -> Float[Array, ...]:
    """Marginal Poisson-count log-pmf ``log P(N = n)``.

    Useful for predict-count workflows that only need the count
    distribution, ignoring conditional location density.
    """
    n = jnp.asarray(n_events)
    lam = jnp.asarray(rate) * jnp.asarray(domain_volume)
    return n * jnp.log(lam) - lam - jax.scipy.special.gammaln(n + 1)

hpp_spatial_sample(key: PRNGKeyArray, rate: Float[Array, ...], domain: RectangularDomain, max_events: int) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Sample a CSR realisation on a rectangular domain.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	JAX PRNG key.	required
rate	Float[Array, ...]	Intensity λ > 0 (scalar). Batched rates are not supported here — vmap externally.	required
domain	RectangularDomain	Spatial domain.	required
max_events	int	Static buffer cap. Choose generously above :math:\lambda |D|; common rule of thumb is 5 * λ |D|.	required

Returns:

Type	Description
Float[Array, ...]	Tuple (locations, mask, n_events):
Bool[Array, ...]	locations shape (max_events, d) — sorted along axis 0 by the first coordinate only (a deterministic order that pushes padding rows to the tail; ties between equal first coordinates keep their sampled order), with padding rows set to domain.lo (an arbitrary in-domain anchor that keeps downstream intensity calls inside the support).
Int[Array, ...]	mask shape (max_events,) — True at real events.
tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	n_events — uncapped Poisson draw; if greater than max_events the buffer truncated and the user should increase the cap.

Source code in src/xtremax/point_processes/primitives/hpp_spatial.py

def hpp_spatial_sample(
    key: PRNGKeyArray,
    rate: Float[Array, ...],
    domain: RectangularDomain,
    max_events: int,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Sample a CSR realisation on a rectangular domain.

    Args:
        key: JAX PRNG key.
        rate: Intensity ``λ > 0`` (scalar). Batched rates are not
            supported here — vmap externally.
        domain: Spatial domain.
        max_events: Static buffer cap. Choose generously above
            :math:`\\lambda |D|`; common rule of thumb is
            ``5 * λ |D|``.

    Returns:
        Tuple ``(locations, mask, n_events)``:

        * ``locations`` shape ``(max_events, d)`` — sorted along axis 0
            by the **first coordinate only** (a deterministic order that
            pushes padding rows to the tail; ties between equal first
            coordinates keep their sampled order), with padding rows set
            to ``domain.lo`` (an arbitrary in-domain anchor that keeps
            downstream intensity calls inside the support).
        * ``mask`` shape ``(max_events,)`` — ``True`` at real events.
        * ``n_events`` — uncapped Poisson draw; if greater than
            ``max_events`` the buffer truncated and the user should
            increase the cap.
    """
    rate = jnp.asarray(rate)
    vol = domain.volume()
    d = domain.n_dims

    key_n, key_locs = random.split(key)
    expected = rate * vol
    n_uncapped = random.poisson(key_n, expected)

    # Mark the first ``n_uncapped`` slots as real events. We put padding
    # at ``+inf`` along axis 0 and lex-sort, so the real events occupy
    # ranks ``0..n-1`` without order-statistic bias from the truncation.
    ranks = jnp.arange(max_events)
    mask = ranks < n_uncapped

    raw = domain.sample_uniform(key_locs, shape=(max_events,))
    # Lex-sort key: first coordinate is the primary key, second is the
    # secondary, etc. Push padding rows to +inf so they end up at the
    # tail of the sort and the real events fill the head.
    sort_key = jnp.where(mask[:, None], raw, jnp.inf)
    # Use the negative of the row-major encoding of the lex order: for
    # 2D and 3D this is just sorting by the first coordinate, which is
    # all we need for downstream consumers (the API does not promise
    # any particular spatial order, just a stable padding invariant).
    primary = sort_key[:, 0]
    order = jnp.argsort(primary)
    sorted_locs = raw[order]
    sorted_mask = mask[order]
    # Replace padding rows with ``domain.lo`` so log-intensity calls
    # never see ``+inf``.
    locations = jnp.where(sorted_mask[:, None], sorted_locs, domain.lo)
    # Sanity check on ``d``: our slicing ``[..., :d]`` implicitly trusts
    # the domain to have produced points of the right width. Since
    # ``domain.sample_uniform`` builds the trailing ``(d,)`` axis from
    # ``n_dims`` directly, we are safe by construction.
    del d
    return SpatialSampleResult(locations, sorted_mask, n_uncapped)

hpp_spatial_predict_count(rate: Float[Array, ...], subdomain_volume: Float[Array, ...]) -> tuple[Float[Array, ...], Float[Array, ...]]

Mean and variance of the count in a sub-domain (both λ |A|).

Source code in src/xtremax/point_processes/primitives/hpp_spatial.py

def hpp_spatial_predict_count(
    rate: Float[Array, ...],
    subdomain_volume: Float[Array, ...],
) -> tuple[Float[Array, ...], Float[Array, ...]]:
    """Mean and variance of the count in a sub-domain (both ``λ |A|``)."""
    mean = jnp.asarray(rate) * jnp.asarray(subdomain_volume)
    return mean, mean

hpp_spatial_nearest_neighbor_distance(rate: Float[Array, ...], n_dims: int) -> Float[Array, ...]

Expected nearest-neighbour distance for a CSR pattern in ℝᵈ.

For a homogeneous Poisson process,

.. math:: \mathbb{E}[D_1] = \frac{\Gamma(1 + 1/d)}{(\lambda \, V_d)^{1/d}},

with :math:V_d = \pi^{d/2} / \Gamma(d/2 + 1) the unit-ball volume.

Source code in src/xtremax/point_processes/primitives/hpp_spatial.py

def hpp_spatial_nearest_neighbor_distance(
    rate: Float[Array, ...],
    n_dims: int,
) -> Float[Array, ...]:
    """Expected nearest-neighbour distance for a CSR pattern in ``ℝᵈ``.

    For a homogeneous Poisson process,

    .. math::
        \\mathbb{E}[D_1] = \\frac{\\Gamma(1 + 1/d)}{(\\lambda \\, V_d)^{1/d}},

    with :math:`V_d = \\pi^{d/2} / \\Gamma(d/2 + 1)` the unit-ball volume.
    """
    d = float(n_dims)
    unit_ball = (jnp.pi ** (d / 2.0)) / jax.scipy.special.gamma(d / 2.0 + 1.0)
    gamma_term = jax.scipy.special.gamma(1.0 + 1.0 / d)
    return gamma_term / (jnp.asarray(rate) * unit_ball) ** (1.0 / d)

Homogeneous Poisson — spatiotemporal

The same constant-intensity process on \(D \times [t_0, t_1)\).

Pure-function primitives for temporal and spatial point processes.

Each submodule is free of class state and free of framework dependencies (no NumPyro — save that for the outer-layer wrappers — no equinox). These are the canonical reference implementations; the operator and distribution layers delegate to them.

hpp_spatiotemporal_intensity(locations: Float[Array, ...], times: Float[Array, ...], rate: Float[Array, ...]) -> Float[Array, ...]

Constant intensity λ(s, t) = λ, broadcast against times.

Source code in src/xtremax/point_processes/primitives/hpp_spatiotemporal.py

def hpp_spatiotemporal_intensity(
    locations: Float[Array, ...],
    times: Float[Array, ...],
    rate: Float[Array, ...],
) -> Float[Array, ...]:
    """Constant intensity ``λ(s, t) = λ``, broadcast against ``times``."""
    times = jnp.asarray(times)
    rate = jnp.asarray(rate)
    del locations  # the homogeneous intensity does not depend on locations
    return jnp.broadcast_to(rate, times.shape)

hpp_spatiotemporal_log_prob(n_events: Int[Array, ...], rate: Float[Array, ...], spatial_volume: Float[Array, ...], temporal_duration: Float[Array, ...]) -> Float[Array, ...]

Janossy log-likelihood n log λ − λ |D| (t1 − t0).

Same convention as :func:hpp_spatial_log_prob: drops the :math:-\log(n!) count normaliser and the iid-uniform location density so this stays directly comparable to the IPP-spatiotemporal log-prob (which uses the matching Janossy form).

Source code in src/xtremax/point_processes/primitives/hpp_spatiotemporal.py

def hpp_spatiotemporal_log_prob(
    n_events: Int[Array, ...],
    rate: Float[Array, ...],
    spatial_volume: Float[Array, ...],
    temporal_duration: Float[Array, ...],
) -> Float[Array, ...]:
    """Janossy log-likelihood ``n log λ − λ |D| (t1 − t0)``.

    Same convention as :func:`hpp_spatial_log_prob`: drops the
    :math:`-\\log(n!)` count normaliser and the iid-uniform location
    density so this stays directly comparable to the IPP-spatiotemporal
    log-prob (which uses the matching Janossy form).
    """
    n = jnp.asarray(n_events)
    rate = jnp.asarray(rate)
    vol = jnp.asarray(spatial_volume)
    dur = jnp.asarray(temporal_duration)
    return n * jnp.log(rate) - rate * vol * dur

hpp_spatiotemporal_count_log_prob(n_events: Int[Array, ...], rate: Float[Array, ...], spatial_volume: Float[Array, ...], temporal_duration: Float[Array, ...]) -> Float[Array, ...]

Marginal Poisson-count log-pmf log P(N = n) with mean λ |D| T.

Source code in src/xtremax/point_processes/primitives/hpp_spatiotemporal.py

def hpp_spatiotemporal_count_log_prob(
    n_events: Int[Array, ...],
    rate: Float[Array, ...],
    spatial_volume: Float[Array, ...],
    temporal_duration: Float[Array, ...],
) -> Float[Array, ...]:
    """Marginal Poisson-count log-pmf ``log P(N = n)`` with mean ``λ |D| T``."""
    n = jnp.asarray(n_events)
    lam = (
        jnp.asarray(rate) * jnp.asarray(spatial_volume) * jnp.asarray(temporal_duration)
    )
    return n * jnp.log(lam) - lam - jax.scipy.special.gammaln(n + 1)

hpp_spatiotemporal_sample(key: PRNGKeyArray, rate: Float[Array, ...], spatial: RectangularDomain, temporal: TemporalDomain, max_events: int) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Sample a homogeneous spatiotemporal Poisson realisation.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	JAX PRNG key.	required
rate	Float[Array, ...]	Intensity λ > 0 (scalar). Vmap externally for batched rates.	required
spatial	RectangularDomain	Spatial domain.	required
temporal	TemporalDomain	Temporal interval.	required
max_events	int	Static buffer cap; choose generously above :math:\lambda |D| T (rule of thumb 5·λ|D|T).	required

Returns:

Type	Description
Float[Array, ...]	Tuple (locations, times, mask, n_events):
Float[Array, ...]	locations shape (max_events, d) — sorted by time, padding rows set to spatial.lo.
Bool[Array, ...]	times shape (max_events,) — sorted ascending, padding rows set to temporal.t1 (so they live just outside the half-open interval but inside any quadrature support).
Int[Array, ...]	mask shape (max_events,) — True at real events.
tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	n_events — uncapped Poisson draw; if greater than max_events the buffer is truncated.

Source code in src/xtremax/point_processes/primitives/hpp_spatiotemporal.py

def hpp_spatiotemporal_sample(
    key: PRNGKeyArray,
    rate: Float[Array, ...],
    spatial: RectangularDomain,
    temporal: TemporalDomain,
    max_events: int,
) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Sample a homogeneous spatiotemporal Poisson realisation.

    Args:
        key: JAX PRNG key.
        rate: Intensity ``λ > 0`` (scalar). Vmap externally for batched
            rates.
        spatial: Spatial domain.
        temporal: Temporal interval.
        max_events: Static buffer cap; choose generously above
            :math:`\\lambda |D| T` (rule of thumb ``5·λ|D|T``).

    Returns:
        Tuple ``(locations, times, mask, n_events)``:

        * ``locations`` shape ``(max_events, d)`` — sorted by time,
            padding rows set to ``spatial.lo``.
        * ``times`` shape ``(max_events,)`` — sorted ascending, padding
            rows set to ``temporal.t1`` (so they live just outside the
            half-open interval but inside any quadrature support).
        * ``mask`` shape ``(max_events,)`` — ``True`` at real events.
        * ``n_events`` — uncapped Poisson draw; if greater than
            ``max_events`` the buffer is truncated.
    """
    rate = jnp.asarray(rate)
    vol = spatial.volume()
    dur = temporal.duration

    key_n, key_locs, key_times = random.split(key, 3)
    expected = rate * vol * dur
    n_uncapped = random.poisson(key_n, expected)

    ranks = jnp.arange(max_events)
    mask = ranks < n_uncapped

    raw_locs = spatial.sample_uniform(key_locs, shape=(max_events,))
    raw_times = temporal.sample_uniform(key_times, shape=(max_events,))

    # Sort by time. Push padding rows to +inf so the real events occupy
    # the head of the sort.
    sort_key = jnp.where(mask, raw_times, jnp.inf)
    order = jnp.argsort(sort_key)
    sorted_times = raw_times[order]
    sorted_locs = raw_locs[order]
    sorted_mask = mask[order]

    times = jnp.where(sorted_mask, sorted_times, temporal.t1)
    locations = jnp.where(sorted_mask[:, None], sorted_locs, spatial.lo)
    return SpatiotemporalSampleResult(locations, times, sorted_mask, n_uncapped)

hpp_spatiotemporal_predict_count(rate: Float[Array, ...], sub_volume: Float[Array, ...], sub_duration: Float[Array, ...]) -> tuple[Float[Array, ...], Float[Array, ...]]

Mean and variance of the count in a sub-slab (both λ |A| τ).

Source code in src/xtremax/point_processes/primitives/hpp_spatiotemporal.py

def hpp_spatiotemporal_predict_count(
    rate: Float[Array, ...],
    sub_volume: Float[Array, ...],
    sub_duration: Float[Array, ...],
) -> tuple[Float[Array, ...], Float[Array, ...]]:
    """Mean and variance of the count in a sub-slab (both ``λ |A| τ``)."""
    mean = jnp.asarray(rate) * jnp.asarray(sub_volume) * jnp.asarray(sub_duration)
    return mean, mean

Inhomogeneous Poisson — temporal

Intensity supplied as a log-intensity callable, so positivity is structural rather than a constraint. The compensator comes from composite trapezoid quadrature; sampling is either Lewis–Shedler thinning (general) or exact inversion when \(\Lambda^{-1}\) is available.

Pure-function primitives for temporal and spatial point processes.

Each submodule is free of class state and free of framework dependencies (no NumPyro — save that for the outer-layer wrappers — no equinox). These are the canonical reference implementations; the operator and distribution layers delegate to them.

ipp_intensity(t: Float[Array, ...], log_intensity_fn: Callable[[Array], Array]) -> Float[Array, ...]

Pointwise intensity λ(t) = exp(log_intensity_fn(t)).

Source code in src/xtremax/point_processes/primitives/ipp.py

def ipp_intensity(
    t: Float[Array, ...],
    log_intensity_fn: Callable[[Array], Array],
) -> Float[Array, ...]:
    """Pointwise intensity ``λ(t) = exp(log_intensity_fn(t))``."""
    return jnp.exp(log_intensity_fn(jnp.asarray(t)))

ipp_cumulative_intensity(t: Float[Array, ...], log_intensity_fn: Callable[[Array], Array], n_points: int = 100) -> Float[Array, ...]

Compensator :math:\Lambda(t) via composite trapezoid.

Source code in src/xtremax/point_processes/primitives/ipp.py

def ipp_cumulative_intensity(
    t: Float[Array, ...],
    log_intensity_fn: Callable[[Array], Array],
    n_points: int = 100,
) -> Float[Array, ...]:
    """Compensator :math:`\\Lambda(t)` via composite trapezoid."""
    return cumulative_log_intensity(log_intensity_fn, t, n_points=n_points)

ipp_log_prob(event_times: Float[Array, ...], mask: Bool[Array, ...], log_intensity_fn: Callable[[Array], Array], integrated_intensity: Float[Array, ...]) -> Float[Array, ...]

Log-likelihood of an observed IPP event sequence.

Parameters:

Name	Type	Description	Default
event_times	Float[Array, ...]	Padded array of event times. Padding positions are evaluated by log_intensity_fn but masked out of the sum — choose padding values inside the support of the log-intensity so evaluation does not raise.	required
mask	Bool[Array, ...]	Boolean mask with the same shape as event_times; True marks real events.	required
log_intensity_fn	Callable[[Array], Array]	Callable returning log λ(t) for an array of times. Must broadcast element-wise.	required
integrated_intensity	Float[Array, ...]	:math:\Lambda(T). The caller is responsible for computing this (via :func:integrate_log_intensity for numerical quadrature, or analytically when a closed form is available).	required

Returns:

Type	Description
Float[Array, ...]	Log-likelihood with shape equal to the leading (batch) dims of
Float[Array, ...]	event_times (i.e. mask.shape[:-1]).

Notes

Padding positions contribute 0 to the sum because mask zeros them out. Keeping the log-intensity call unmasked (rather than gathering only real events) preserves jit / vmap compatibility with static shapes.

Source code in src/xtremax/point_processes/primitives/ipp.py

def ipp_log_prob(
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    log_intensity_fn: Callable[[Array], Array],
    integrated_intensity: Float[Array, ...],
) -> Float[Array, ...]:
    """Log-likelihood of an observed IPP event sequence.

    Args:
        event_times: Padded array of event times. Padding positions are
            evaluated by ``log_intensity_fn`` but masked out of the sum
            — choose padding values inside the support of the
            log-intensity so evaluation does not raise.
        mask: Boolean mask with the same shape as ``event_times``;
            ``True`` marks real events.
        log_intensity_fn: Callable returning ``log λ(t)`` for an array
            of times. Must broadcast element-wise.
        integrated_intensity: :math:`\\Lambda(T)`. The caller is
            responsible for computing this (via
            :func:`integrate_log_intensity` for numerical quadrature,
            or analytically when a closed form is available).

    Returns:
        Log-likelihood with shape equal to the leading (batch) dims of
        ``event_times`` (i.e. ``mask.shape[:-1]``).

    Notes:
        Padding positions contribute ``0`` to the sum because ``mask``
        zeros them out. Keeping the log-intensity call unmasked (rather
        than gathering only real events) preserves ``jit`` / ``vmap``
        compatibility with static shapes.
    """
    log_intensities = log_intensity_fn(event_times)
    masked = jnp.where(mask, log_intensities, 0.0)
    return jnp.sum(masked, axis=-1) - jnp.asarray(integrated_intensity)

ipp_sample_thinning(key: PRNGKeyArray, log_intensity_fn: Callable[[Array], Array], T: Float[Array, ...], lambda_max: Float[Array, ...], max_candidates: int) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Sample via Lewis–Shedler thinning with a static buffer.

Algorithm (Lewis & Shedler, 1979):

1. Draw n candidates from a homogeneous Poisson process with rate λ_max on [0, T].
2. Accept candidate at time t with probability :math:\lambda(t) / \lambda_{\max}.

For JAX-friendliness the candidate buffer is of static size max_candidates: the Poisson draw is capped to this cap via a rank mask so the returned shape is deterministic.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	JAX PRNG key.	required
log_intensity_fn	Callable[[Array], Array]	Log-intensity callable.	required
T	Float[Array, ...]	Observation window length (scalar).	required
lambda_max	Float[Array, ...]	Upper bound on :math:\lambda(t) over [0, T]. Tighter bounds give higher acceptance rates; too loose a bound wastes work but is still exact.	required
max_candidates	int	Static cap on the candidate pool.	required

Returns:

Type	Description
Float[Array, ...]	Tuple (times, accepted_mask, n_candidates_uncapped):
Bool[Array, ...]	times with shape (max_candidates,) — the accepted events, sorted, compacted into a contiguous prefix; padding positions are set to T.
Int[Array, ...]	accepted_mask with shape (max_candidates,) — a contiguous prefix of True covering the accepted events. (Rejected candidates used to stay in the buffer as hole masks, which silently corrupted mask-contiguity consumers such as time_rescaling_residuals — the package invariant is now "mask is a contiguous prefix; padding time = T".)
tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	n_candidates_uncapped — the Poisson draw before capping at max_candidates; useful for diagnosing buffer over-runs.

Source code in src/xtremax/point_processes/primitives/ipp.py

def ipp_sample_thinning(
    key: PRNGKeyArray,
    log_intensity_fn: Callable[[Array], Array],
    T: Float[Array, ...],
    lambda_max: Float[Array, ...],
    max_candidates: int,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Sample via Lewis–Shedler thinning with a static buffer.

    Algorithm (Lewis & Shedler, 1979):

    1. Draw ``n`` candidates from a homogeneous Poisson process with
       rate ``λ_max`` on ``[0, T]``.
    2. Accept candidate at time ``t`` with probability
       :math:`\\lambda(t) / \\lambda_{\\max}`.

    For JAX-friendliness the candidate buffer is of static size
    ``max_candidates``: the Poisson draw is capped to this cap via a
    rank mask so the returned shape is deterministic.

    Args:
        key: JAX PRNG key.
        log_intensity_fn: Log-intensity callable.
        T: Observation window length (scalar).
        lambda_max: Upper bound on :math:`\\lambda(t)` over ``[0, T]``.
            Tighter bounds give higher acceptance rates; too loose a
            bound wastes work but is still exact.
        max_candidates: Static cap on the candidate pool.

    Returns:
        Tuple ``(times, accepted_mask, n_candidates_uncapped)``:

        * ``times`` with shape ``(max_candidates,)`` — the *accepted*
            events, sorted, compacted into a contiguous prefix; padding
            positions are set to ``T``.
        * ``accepted_mask`` with shape ``(max_candidates,)`` — a
            contiguous prefix of ``True`` covering the accepted events.
            (Rejected candidates used to stay in the buffer as hole
            masks, which silently corrupted mask-contiguity consumers
            such as ``time_rescaling_residuals`` — the package invariant
            is now "mask is a contiguous prefix; padding time = T".)
        * ``n_candidates_uncapped`` — the Poisson draw before capping
            at ``max_candidates``; useful for diagnosing buffer
            over-runs.
    """
    key_n, key_times, key_thin = random.split(key, 3)

    T = jnp.asarray(T)
    lambda_max = jnp.asarray(lambda_max)

    expected = lambda_max * T
    n_uncapped = random.poisson(key_n, expected)

    # Mark the first n_uncapped buffer positions as real candidates and
    # push padding to +inf before sorting so the real candidates occupy
    # sorted ranks 0..n-1 without order-statistic bias. (Sorting first
    # and then truncating to the first n would draw from the n smallest
    # of ``max_candidates`` uniforms, not from an actual HPP.)
    ranks = jnp.arange(max_candidates)
    candidate_mask = ranks < n_uncapped
    raw = random.uniform(key_times, shape=(max_candidates,)) * T
    raw_sortable = jnp.where(candidate_mask, raw, jnp.inf)
    candidate_times_sorted = jnp.sort(raw_sortable)
    # Replace padding positions with T so downstream log-intensity
    # calls evaluate inside the intended window.
    candidate_times = jnp.where(candidate_mask, candidate_times_sorted, T)

    log_intensities = log_intensity_fn(candidate_times)
    log_accept = log_intensities - jnp.log(lambda_max)
    u = random.uniform(key_thin, shape=(max_candidates,))
    thinning_accept = jnp.log(u) < log_accept
    accepted_raw = candidate_mask & thinning_accept

    # Compact the accepted events into a contiguous prefix (rejected
    # candidates used to remain in the buffer as hole masks). Pushing
    # rejected slots to +inf and re-sorting keeps the accepted times in
    # ascending order at ranks 0..n_accepted-1.
    accepted_sortable = jnp.where(accepted_raw, candidate_times, jnp.inf)
    compacted = jnp.sort(accepted_sortable)
    n_accepted = jnp.sum(accepted_raw)
    accepted_mask = ranks < n_accepted
    times = jnp.where(accepted_mask, compacted, T)
    return SampleResult(times, accepted_mask, n_uncapped)

ipp_sample_inversion(key: PRNGKeyArray, inverse_cumulative_intensity_fn: Callable[[Array], Array], Lambda_T: Float[Array, ...], max_events: int) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Exact inversion sampling when :math:\Lambda^{-1} is available.

When the inverse compensator is known in closed form, sampling is exact and does not require an upper bound:

1. Draw n from :math:\mathrm{Poisson}(\Lambda(T)).
2. Draw n iid points :math:u_i \sim \mathrm{Uniform}(0, \Lambda(T)).
3. Map through :math:\Lambda^{-1} and sort.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	JAX PRNG key.	required
inverse_cumulative_intensity_fn	Callable[[Array], Array]	Callable computing :math:\Lambda^{-1}(y) for y in [0, Λ(T)]. Must be vectorised over its input.	required
Lambda_T	Float[Array, ...]	:math:\Lambda(T) (scalar).	required
max_events	int	Static buffer size.	required

Returns:

Type	Description
Float[Array, ...]	Tuple (times, mask, n_events) with the same semantics as
Bool[Array, ...]	func:ipp_sample_thinning (no separate candidate/accepted
Int[Array, ...]	split — every position is either a real event or padding).

Source code in src/xtremax/point_processes/primitives/ipp.py

def ipp_sample_inversion(
    key: PRNGKeyArray,
    inverse_cumulative_intensity_fn: Callable[[Array], Array],
    Lambda_T: Float[Array, ...],
    max_events: int,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Exact inversion sampling when :math:`\\Lambda^{-1}` is available.

    When the inverse compensator is known in closed form, sampling is
    exact and does not require an upper bound:

    1. Draw ``n`` from :math:`\\mathrm{Poisson}(\\Lambda(T))`.
    2. Draw ``n`` iid points :math:`u_i \\sim \\mathrm{Uniform}(0, \\Lambda(T))`.
    3. Map through :math:`\\Lambda^{-1}` and sort.

    Args:
        key: JAX PRNG key.
        inverse_cumulative_intensity_fn: Callable computing
            :math:`\\Lambda^{-1}(y)` for ``y`` in ``[0, Λ(T)]``. Must
            be vectorised over its input.
        Lambda_T: :math:`\\Lambda(T)` (scalar).
        max_events: Static buffer size.

    Returns:
        Tuple ``(times, mask, n_events)`` with the same semantics as
        :func:`ipp_sample_thinning` (no separate candidate/accepted
        split — every position is either a real event or padding).
    """
    key_n, key_u = random.split(key)
    Lambda_T = jnp.asarray(Lambda_T)

    n_events = random.poisson(key_n, Lambda_T)

    # Mark the first ``n_events`` buffer slots as real and push padding
    # uniforms to ``+inf`` before sorting, so the real draws occupy
    # sorted ranks ``0..n-1`` without being biased to small values.
    # (Sorting all ``max_events`` uniforms and then truncating would
    # pick the ``n`` smallest — not ``n`` iid uniforms on ``[0, Λ(T)]``.)
    ranks = jnp.arange(max_events)
    mask = ranks < n_events
    u = random.uniform(key_u, shape=(max_events,)) * Lambda_T
    u_sortable = jnp.where(mask, u, jnp.inf)
    u_sorted = jnp.sort(u_sortable)
    # Replace padding positions with ``Λ(T)`` before applying the
    # inverse — keeps every buffer slot inside a representable domain
    # so ``Λ⁻¹`` never sees ``+inf``.
    u_clamped = jnp.where(mask, u_sorted, Lambda_T)
    times = inverse_cumulative_intensity_fn(u_clamped)

    # Final padding-position value: ``Λ⁻¹(Λ(T))`` = right edge of window.
    t_right = inverse_cumulative_intensity_fn(Lambda_T)
    times = jnp.where(mask, times, t_right)
    return SampleResult(times, mask, n_events)

ipp_hazard(t: Float[Array, ...], log_intensity_fn: Callable[[Array], Array]) -> Float[Array, ...]

Hazard equals the intensity for an IPP.

Source code in src/xtremax/point_processes/primitives/ipp.py

def ipp_hazard(
    t: Float[Array, ...],
    log_intensity_fn: Callable[[Array], Array],
) -> Float[Array, ...]:
    """Hazard equals the intensity for an IPP."""
    return ipp_intensity(t, log_intensity_fn)

ipp_cumulative_hazard(t: Float[Array, ...], given_time: Float[Array, ...], log_intensity_fn: Callable[[Array], Array], n_points: int = 100) -> Float[Array, ...]

:math:\Lambda(t) - \Lambda(s) via quadrature.

Source code in src/xtremax/point_processes/primitives/ipp.py

def ipp_cumulative_hazard(
    t: Float[Array, ...],
    given_time: Float[Array, ...],
    log_intensity_fn: Callable[[Array], Array],
    n_points: int = 100,
) -> Float[Array, ...]:
    """:math:`\\Lambda(t) - \\Lambda(s)` via quadrature."""
    return integrate_log_intensity(
        log_intensity_fn,
        jnp.asarray(given_time),
        jnp.asarray(t),
        n_points=n_points,
    )

ipp_survival(t: Float[Array, ...], given_time: Float[Array, ...], log_intensity_fn: Callable[[Array], Array], n_points: int = 100) -> Float[Array, ...]

:math:S(t|s) = \exp(-[\Lambda(t) - \Lambda(s)]) for t >= s.

Source code in src/xtremax/point_processes/primitives/ipp.py

def ipp_survival(
    t: Float[Array, ...],
    given_time: Float[Array, ...],
    log_intensity_fn: Callable[[Array], Array],
    n_points: int = 100,
) -> Float[Array, ...]:
    """:math:`S(t|s) = \\exp(-[\\Lambda(t) - \\Lambda(s)])` for ``t >= s``."""
    integrated = integrate_log_intensity(
        log_intensity_fn,
        jnp.asarray(given_time),
        jnp.asarray(t),
        n_points=n_points,
    )
    return jnp.exp(-integrated)

ipp_inter_event_log_prob(tau: Float[Array, ...], current_time: Float[Array, ...], log_intensity_fn: Callable[[Array], Array], n_points: int = 100) -> Float[Array, ...]

Log density of the next-event delay τ given last event at s.

For an IPP,

.. math:: f(\tau \mid s) = \lambda(s + \tau) \exp(-[\Lambda(s + \tau) - \Lambda(s)]).

Source code in src/xtremax/point_processes/primitives/ipp.py

def ipp_inter_event_log_prob(
    tau: Float[Array, ...],
    current_time: Float[Array, ...],
    log_intensity_fn: Callable[[Array], Array],
    n_points: int = 100,
) -> Float[Array, ...]:
    """Log density of the next-event delay ``τ`` given last event at ``s``.

    For an IPP,

    .. math::
        f(\\tau \\mid s) = \\lambda(s + \\tau)
        \\exp(-[\\Lambda(s + \\tau) - \\Lambda(s)]).
    """
    tau = jnp.asarray(tau)
    s = jnp.asarray(current_time)
    t = s + tau
    log_lambda = log_intensity_fn(t)
    cum = integrate_log_intensity(log_intensity_fn, s, t, n_points=n_points)
    return jnp.where(tau >= 0.0, log_lambda - cum, -jnp.inf)

ipp_predict_count(start_time: Float[Array, ...], end_time: Float[Array, ...], log_intensity_fn: Callable[[Array], Array], n_points: int = 100) -> Float[Array, ...]

Expected count :math:\Lambda(end) - \Lambda(start) on [start, end].

Source code in src/xtremax/point_processes/primitives/ipp.py

def ipp_predict_count(
    start_time: Float[Array, ...],
    end_time: Float[Array, ...],
    log_intensity_fn: Callable[[Array], Array],
    n_points: int = 100,
) -> Float[Array, ...]:
    r"""Expected count :math:`\Lambda(end) - \Lambda(start)` on ``[start, end]``."""
    return integrate_log_intensity(
        log_intensity_fn,
        jnp.asarray(start_time),
        jnp.asarray(end_time),
        n_points=n_points,
    )

Inhomogeneous Poisson — spatial

Pure-function primitives for temporal and spatial point processes.

Each submodule is free of class state and free of framework dependencies (no NumPyro — save that for the outer-layer wrappers — no equinox). These are the canonical reference implementations; the operator and distribution layers delegate to them.

ipp_spatial_intensity(locations: Float[Array, ...], log_intensity_fn: Callable[[Array], Array]) -> Float[Array, ...]

Pointwise intensity λ(s) = exp(log_intensity_fn(s)).

Source code in src/xtremax/point_processes/primitives/ipp_spatial.py

def ipp_spatial_intensity(
    locations: Float[Array, ...],
    log_intensity_fn: Callable[[Array], Array],
) -> Float[Array, ...]:
    """Pointwise intensity ``λ(s) = exp(log_intensity_fn(s))``."""
    return jnp.exp(log_intensity_fn(jnp.asarray(locations)))

ipp_spatial_log_prob(locations: Float[Array, ...], mask: Bool[Array, ...], log_intensity_fn: Callable[[Array], Array], integrated_intensity: Float[Array, ...]) -> Float[Array, ...]

Log-likelihood of an observed inhomogeneous spatial pattern.

Parameters:

Name	Type	Description	Default
locations	Float[Array, ...]	Padded locations (..., max_events, d). Padding rows are evaluated by log_intensity_fn for shape stability and masked out of the sum — choose padding inside the domain so log_intensity_fn does not raise.	required
mask	Bool[Array, ...]	Event mask shape (..., max_events).	required
log_intensity_fn	Callable[[Array], Array]	Callable mapping (..., n, d) locations to (..., n) log-intensities.	required
integrated_intensity	Float[Array, ...]	Λ(D) (computed by quadrature or analytically by the caller).	required

Returns:

Type	Description
Float[Array, ...]	Log-likelihood with shape mask.shape[:-1].

Source code in src/xtremax/point_processes/primitives/ipp_spatial.py

def ipp_spatial_log_prob(
    locations: Float[Array, ...],
    mask: Bool[Array, ...],
    log_intensity_fn: Callable[[Array], Array],
    integrated_intensity: Float[Array, ...],
) -> Float[Array, ...]:
    """Log-likelihood of an observed inhomogeneous spatial pattern.

    Args:
        locations: Padded locations ``(..., max_events, d)``. Padding
            rows are evaluated by ``log_intensity_fn`` for shape
            stability and masked out of the sum — choose padding
            inside the domain so ``log_intensity_fn`` does not raise.
        mask: Event mask shape ``(..., max_events)``.
        log_intensity_fn: Callable mapping ``(..., n, d)`` locations to
            ``(..., n)`` log-intensities.
        integrated_intensity: ``Λ(D)`` (computed by quadrature or
            analytically by the caller).

    Returns:
        Log-likelihood with shape ``mask.shape[:-1]``.
    """
    log_intensities = log_intensity_fn(locations)
    masked = jnp.where(mask, log_intensities, 0.0)
    return jnp.sum(masked, axis=-1) - jnp.asarray(integrated_intensity)

ipp_spatial_sample_thinning(key: PRNGKeyArray, log_intensity_fn: Callable[[Array], Array], domain: RectangularDomain, lambda_max: Float[Array, ...], max_candidates: int) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Lewis–Shedler thinning sampler with a static buffer.

Algorithm:

1. Generate n candidates from a homogeneous Poisson process at rate λ_max over D (uniform on the box).
2. Accept candidate s with probability :math:\lambda(s) / \lambda_{\max}.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	JAX PRNG key.	required
log_intensity_fn	Callable[[Array], Array]	Log-intensity callable.	required
domain	RectangularDomain	Spatial domain.	required
lambda_max	Float[Array, ...]	Upper bound on :math:\lambda(s) over D. Tighter bounds give higher acceptance.	required
max_candidates	int	Static cap on the candidate pool — choose comfortably above λ_max · |D|.	required

Returns:

Type	Description
Float[Array, ...]	Tuple (locations, accepted_mask, n_candidates_uncapped):
Bool[Array, ...]	locations shape (max_candidates, d) — the accepted points compacted into a contiguous prefix, ordered by the first coordinate; padding rows are set to domain.lo.
Int[Array, ...]	accepted_mask shape (max_candidates,) — a contiguous prefix of True covering the accepted points (the package invariant: masks are contiguous prefixes).
tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	n_candidates_uncapped — Poisson draw before capping; useful for diagnosing buffer over-runs.

Source code in src/xtremax/point_processes/primitives/ipp_spatial.py

def ipp_spatial_sample_thinning(
    key: PRNGKeyArray,
    log_intensity_fn: Callable[[Array], Array],
    domain: RectangularDomain,
    lambda_max: Float[Array, ...],
    max_candidates: int,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Lewis–Shedler thinning sampler with a static buffer.

    Algorithm:

    1. Generate ``n`` candidates from a homogeneous Poisson process at
       rate ``λ_max`` over ``D`` (uniform on the box).
    2. Accept candidate ``s`` with probability
       :math:`\\lambda(s) / \\lambda_{\\max}`.

    Args:
        key: JAX PRNG key.
        log_intensity_fn: Log-intensity callable.
        domain: Spatial domain.
        lambda_max: Upper bound on :math:`\\lambda(s)` over ``D``.
            Tighter bounds give higher acceptance.
        max_candidates: Static cap on the candidate pool — choose
            comfortably above ``λ_max · |D|``.

    Returns:
        Tuple ``(locations, accepted_mask, n_candidates_uncapped)``:

        * ``locations`` shape ``(max_candidates, d)`` — the *accepted*
            points compacted into a contiguous prefix, ordered by the
            first coordinate; padding rows are set to ``domain.lo``.
        * ``accepted_mask`` shape ``(max_candidates,)`` — a contiguous
            prefix of ``True`` covering the accepted points (the package
            invariant: masks are contiguous prefixes).
        * ``n_candidates_uncapped`` — Poisson draw before capping;
            useful for diagnosing buffer over-runs.
    """
    key_n, key_locs, key_thin = random.split(key, 3)

    lambda_max = jnp.asarray(lambda_max)
    expected = lambda_max * domain.volume()
    n_uncapped = random.poisson(key_n, expected)

    ranks = jnp.arange(max_candidates)
    candidate_mask = ranks < n_uncapped
    raw = domain.sample_uniform(key_locs, shape=(max_candidates,))

    # Sort by first coordinate so the buffer order is deterministic and
    # padding rows trail the real candidates (push padding to +inf).
    primary_for_sort = jnp.where(candidate_mask, raw[:, 0], jnp.inf)
    order = jnp.argsort(primary_for_sort)
    sorted_locs = raw[order]
    sorted_mask = candidate_mask[order]

    # Replace padding rows with ``domain.lo`` before evaluating the
    # log-intensity so the call stays inside the support.
    safe_locs = jnp.where(sorted_mask[:, None], sorted_locs, domain.lo)
    log_intensities = log_intensity_fn(safe_locs)

    log_accept = log_intensities - jnp.log(lambda_max)
    u = random.uniform(key_thin, shape=(max_candidates,))
    thinning_accept = jnp.log(u) < log_accept
    accepted_raw = sorted_mask & thinning_accept

    # Compact accepted points into a contiguous prefix (rejected
    # candidates used to remain as hole masks): stable argsort on the
    # rejection flag preserves the first-coordinate order within the
    # accepted block.
    order2 = jnp.argsort(~accepted_raw, stable=True)
    compacted = safe_locs[order2]
    n_accepted = jnp.sum(accepted_raw)
    accepted_mask = ranks < n_accepted
    locations = jnp.where(accepted_mask[:, None], compacted, domain.lo)
    return SpatialSampleResult(locations, accepted_mask, n_uncapped)

ipp_spatial_predict_count(log_intensity_fn: Callable[[Array], Array], subdomain: RectangularDomain, n_integration_points: int = 2048, method: str = 'qmc') -> Float[Array, ...]

Expected count over a sub-domain A.

:math:\mathbb{E}[N(A)] = \int_A \lambda(s) ds — same quadrature machinery as the :math:\Lambda(D) calculation.

Source code in src/xtremax/point_processes/primitives/ipp_spatial.py

def ipp_spatial_predict_count(
    log_intensity_fn: Callable[[Array], Array],
    subdomain: RectangularDomain,
    n_integration_points: int = 2048,
    method: str = "qmc",
) -> Float[Array, ...]:
    """Expected count over a sub-domain ``A``.

    :math:`\\mathbb{E}[N(A)] = \\int_A \\lambda(s) ds` — same quadrature
    machinery as the :math:`\\Lambda(D)` calculation.
    """
    # Local import to avoid a top-level cycle: the integration helper
    # itself depends on RectangularDomain, which we already imported.
    from xtremax.point_processes._integration_spatial import (
        integrate_log_intensity_spatial,
    )

    return integrate_log_intensity_spatial(
        log_intensity_fn,
        subdomain,
        n_points=n_integration_points,
        method=method,  # type: ignore[arg-type]
    )

Inhomogeneous Poisson — spatiotemporal

Beyond the core four, this family carries the marginalisation helpers you need for plotting — collapsing \(\lambda(s, t)\) onto space or time — and the residual-based goodness-of-fit tests that stand in for the temporal time-rescaling diagnostics.

Pure-function primitives for temporal and spatial point processes.

Each submodule is free of class state and free of framework dependencies (no NumPyro — save that for the outer-layer wrappers — no equinox). These are the canonical reference implementations; the operator and distribution layers delegate to them.

ipp_spatiotemporal_intensity(locations: Float[Array, ...], times: Float[Array, ...], log_intensity_fn: Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]]) -> Float[Array, ...]

Evaluate λ(s, t) = exp(log λ(s, t)) on a batch of events.

Source code in src/xtremax/point_processes/primitives/ipp_spatiotemporal.py

def ipp_spatiotemporal_intensity(
    locations: Float[Array, ...],
    times: Float[Array, ...],
    log_intensity_fn: Callable[
        [Float[Array, ...], Float[Array, ...]], Float[Array, ...]
    ],
) -> Float[Array, ...]:
    """Evaluate ``λ(s, t) = exp(log λ(s, t))`` on a batch of events."""
    return jnp.exp(log_intensity_fn(locations, times))

ipp_spatiotemporal_log_prob(locations: Float[Array, ...], times: Float[Array, ...], mask: Bool[Array, ...], log_intensity_fn: Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]], integrated_intensity: Float[Array, ...]) -> Float[Array, ...]

Janossy log-likelihood :math:\sum_i \log \lambda(s_i, t_i) - \Lambda.

Parameters:

Name	Type	Description	Default
locations	Float[Array, ...]	(max_events, d_space) event locations.	required
times	Float[Array, ...]	(max_events,) event times.	required
mask	Bool[Array, ...]	(max_events,) boolean mask of real events.	required
log_intensity_fn	Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]]	Callable (s, t) → log λ(s, t).	required
integrated_intensity	Float[Array, ...]	Pre-computed :math:\Lambda = \int_{t_0}^{t_1}\!\!\int_D \lambda\, ds\, dt, typically from :func:xtremax.point_processes.integrate_log_intensity_spatiotemporal.	required

Returns:

Type	Description
Float[Array, ...]	Scalar log-likelihood.

Source code in src/xtremax/point_processes/primitives/ipp_spatiotemporal.py

def ipp_spatiotemporal_log_prob(
    locations: Float[Array, ...],
    times: Float[Array, ...],
    mask: Bool[Array, ...],
    log_intensity_fn: Callable[
        [Float[Array, ...], Float[Array, ...]], Float[Array, ...]
    ],
    integrated_intensity: Float[Array, ...],
) -> Float[Array, ...]:
    """Janossy log-likelihood :math:`\\sum_i \\log \\lambda(s_i, t_i) - \\Lambda`.

    Args:
        locations: ``(max_events, d_space)`` event locations.
        times: ``(max_events,)`` event times.
        mask: ``(max_events,)`` boolean mask of real events.
        log_intensity_fn: Callable ``(s, t) → log λ(s, t)``.
        integrated_intensity: Pre-computed
            :math:`\\Lambda = \\int_{t_0}^{t_1}\\!\\!\\int_D \\lambda\\, ds\\, dt`,
            typically from
            :func:`xtremax.point_processes.integrate_log_intensity_spatiotemporal`.

    Returns:
        Scalar log-likelihood.
    """
    log_lam = log_intensity_fn(locations, times)
    masked = jnp.where(mask, log_lam, 0.0)
    return jnp.sum(masked) - jnp.asarray(integrated_intensity)

ipp_spatiotemporal_sample_thinning(key: PRNGKeyArray, log_intensity_fn: Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]], spatial: RectangularDomain, temporal: TemporalDomain, lambda_max: Float[Array, ...], max_candidates: int) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Lewis–Shedler thinning on D × [t_0, t_1).

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key.	required
log_intensity_fn	Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]]	Callable (s, t) → log λ(s, t).	required
spatial	RectangularDomain	Spatial domain.	required
temporal	TemporalDomain	Temporal interval.	required
lambda_max	Float[Array, ...]	Upper bound on :math:\lambda(s, t) over the slab. Must be a true bound: a too-low lambda_max silently biases the sampler low in the intensity peaks.	required
max_candidates	int	Static cap on candidate-buffer size; choose generously above :math:\lambda_{\max} |D| T.	required

Returns:

Type	Description
Float[Array, ...]	(locations, times, mask, n_events) with the same shape
Float[Array, ...]	contract as :func:hpp_spatiotemporal_sample. Times are sorted
Bool[Array, ...]	ascending; locations follow the time order. n_events is the
Int[Array, ...]	uncapped candidate count (the Poisson draw before capping at
tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	max_candidates) — matching the temporal/spatial thinning
tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	siblings, so a value above max_candidates diagnoses a
tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	truncated candidate pool. The accepted-event count is
tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	mask.sum().

Source code in src/xtremax/point_processes/primitives/ipp_spatiotemporal.py

def ipp_spatiotemporal_sample_thinning(
    key: PRNGKeyArray,
    log_intensity_fn: Callable[
        [Float[Array, ...], Float[Array, ...]], Float[Array, ...]
    ],
    spatial: RectangularDomain,
    temporal: TemporalDomain,
    lambda_max: Float[Array, ...],
    max_candidates: int,
) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Lewis–Shedler thinning on ``D × [t_0, t_1)``.

    Args:
        key: PRNG key.
        log_intensity_fn: Callable ``(s, t) → log λ(s, t)``.
        spatial: Spatial domain.
        temporal: Temporal interval.
        lambda_max: Upper bound on :math:`\\lambda(s, t)` over the slab.
            **Must** be a true bound: a too-low ``lambda_max`` silently
            biases the sampler low in the intensity peaks.
        max_candidates: Static cap on candidate-buffer size; choose
            generously above :math:`\\lambda_{\\max} |D| T`.

    Returns:
        ``(locations, times, mask, n_events)`` with the same shape
        contract as :func:`hpp_spatiotemporal_sample`. Times are sorted
        ascending; locations follow the time order. ``n_events`` is the
        **uncapped candidate count** (the Poisson draw before capping at
        ``max_candidates``) — matching the temporal/spatial thinning
        siblings, so a value above ``max_candidates`` diagnoses a
        truncated candidate pool. The accepted-event count is
        ``mask.sum()``.
    """
    lambda_max = jnp.asarray(lambda_max)
    vol = spatial.volume()
    dur = temporal.duration

    key_n, key_locs, key_times, key_thin = random.split(key, 4)
    expected = lambda_max * vol * dur
    n_candidates = random.poisson(key_n, expected)

    ranks = jnp.arange(max_candidates)
    candidate_mask = ranks < n_candidates

    raw_locs = spatial.sample_uniform(key_locs, shape=(max_candidates,))
    raw_times = temporal.sample_uniform(key_times, shape=(max_candidates,))
    log_lam = log_intensity_fn(raw_locs, raw_times)

    # Acceptance probability λ(s, t) / λ_max.
    log_accept = log_lam - jnp.log(lambda_max)
    u = random.uniform(key_thin, shape=(max_candidates,))
    accepted = jnp.log(u) < log_accept
    keep = candidate_mask & accepted

    # Sort kept events by time; push the rest to the tail.
    sort_key = jnp.where(keep, raw_times, jnp.inf)
    order = jnp.argsort(sort_key)
    sorted_times = raw_times[order]
    sorted_locs = raw_locs[order]
    sorted_keep = keep[order]

    # The post-thinning real count is at most ``max_candidates``; the
    # output mask records exactly which slots are real.
    final_mask = sorted_keep
    times = jnp.where(final_mask, sorted_times, temporal.t1)
    locations = jnp.where(final_mask[:, None], sorted_locs, spatial.lo)
    return SpatiotemporalSampleResult(locations, times, final_mask, n_candidates)

ipp_spatiotemporal_predict_count(integrated_intensity: Float[Array, ...]) -> tuple[Float[Array, ...], Float[Array, ...]]

Mean and variance of the Poisson count over the integration domain.

Source code in src/xtremax/point_processes/primitives/ipp_spatiotemporal.py

def ipp_spatiotemporal_predict_count(
    integrated_intensity: Float[Array, ...],
) -> tuple[Float[Array, ...], Float[Array, ...]]:
    """Mean and variance of the Poisson count over the integration domain."""
    lam = jnp.asarray(integrated_intensity)
    return lam, lam

ipp_spatiotemporal_intensity_surface_at_time(t: Float[Array, ...], log_intensity_fn: Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]], spatial: RectangularDomain, grid_size: int = 50) -> tuple[Float[Array, ...], Float[Array, ...]]

Evaluate λ(s, t) on a tensor-product spatial grid at fixed t.

Parameters:

Name	Type	Description	Default
t	Float[Array, ...]	Scalar query time.	required
log_intensity_fn	Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]]	Callable (s, t) → log λ(s, t).	required
spatial	RectangularDomain	Spatial domain.	required
grid_size	int	Static per-axis grid size.	50

Returns:

Type	Description
Float[Array, ...]	(grid, intensity) where grid has shape
Float[Array, ...]	(grid_size, ..., grid_size, d_space) and intensity has
tuple[Float[Array, ...], Float[Array, ...]]	shape (grid_size, ..., grid_size).

Source code in src/xtremax/point_processes/primitives/ipp_spatiotemporal.py

def ipp_spatiotemporal_intensity_surface_at_time(
    t: Float[Array, ...],
    log_intensity_fn: Callable[
        [Float[Array, ...], Float[Array, ...]], Float[Array, ...]
    ],
    spatial: RectangularDomain,
    grid_size: int = 50,
) -> tuple[Float[Array, ...], Float[Array, ...]]:
    """Evaluate ``λ(s, t)`` on a tensor-product spatial grid at fixed ``t``.

    Args:
        t: Scalar query time.
        log_intensity_fn: Callable ``(s, t) → log λ(s, t)``.
        spatial: Spatial domain.
        grid_size: Static per-axis grid size.

    Returns:
        ``(grid, intensity)`` where ``grid`` has shape
        ``(grid_size, ..., grid_size, d_space)`` and ``intensity`` has
        shape ``(grid_size, ..., grid_size)``.
    """
    d = spatial.n_dims
    axis_grids = [
        jnp.linspace(spatial.lo[i], spatial.hi[i], grid_size) for i in range(d)
    ]
    mesh = jnp.stack(jnp.meshgrid(*axis_grids, indexing="ij"), axis=-1)
    flat_s = mesh.reshape(-1, d)
    flat_t = jnp.full((flat_s.shape[0],), t)
    log_lam = log_intensity_fn(flat_s, flat_t)
    intensity = jnp.exp(log_lam).reshape((grid_size,) * d)
    return mesh, intensity

ipp_spatiotemporal_marginal_spatial_intensity(locations: Float[Array, ...], log_intensity_fn: Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]], temporal: TemporalDomain, n_time_points: int = 100) -> Float[Array, ...]

Compute :math:\lambda_S(s) = \int_{t_0}^{t_1} \lambda(s, t)\, dt.

Trapezoid quadrature on a n_time_points-uniform grid in time; vectorised over an arbitrary leading event axis.

Parameters:

Name	Type	Description	Default
locations	Float[Array, ...]	(N, d_space) query locations.	required
log_intensity_fn	Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]]	Callable (s, t) → log λ(s, t) accepting a shared spatial input broadcast against a time grid.	required
temporal	TemporalDomain	Temporal interval.	required
n_time_points	int	Static number of trapezoid nodes in time.	100

Returns:

Type	Description
Float[Array, ...]	(N,) array of marginal spatial intensities.

Source code in src/xtremax/point_processes/primitives/ipp_spatiotemporal.py

def ipp_spatiotemporal_marginal_spatial_intensity(
    locations: Float[Array, ...],
    log_intensity_fn: Callable[
        [Float[Array, ...], Float[Array, ...]], Float[Array, ...]
    ],
    temporal: TemporalDomain,
    n_time_points: int = 100,
) -> Float[Array, ...]:
    """Compute :math:`\\lambda_S(s) = \\int_{t_0}^{t_1} \\lambda(s, t)\\, dt`.

    Trapezoid quadrature on a ``n_time_points``-uniform grid in time;
    vectorised over an arbitrary leading event axis.

    Args:
        locations: ``(N, d_space)`` query locations.
        log_intensity_fn: Callable ``(s, t) → log λ(s, t)`` accepting a
            shared spatial input broadcast against a time grid.
        temporal: Temporal interval.
        n_time_points: Static number of trapezoid nodes in time.

    Returns:
        ``(N,)`` array of marginal spatial intensities.
    """
    locations = jnp.asarray(locations)
    n_query = locations.shape[0]
    t_grid = jnp.linspace(temporal.t0, temporal.t1, n_time_points)
    # Tile locations across the time axis: (N, T, d_space) and (N, T).
    s_tiled = jnp.repeat(locations[:, None, :], n_time_points, axis=1)
    t_tiled = jnp.broadcast_to(t_grid, (n_query, n_time_points))
    s_flat = s_tiled.reshape(-1, locations.shape[-1])
    t_flat = t_tiled.reshape(-1)
    intensities = jnp.exp(log_intensity_fn(s_flat, t_flat)).reshape(
        n_query, n_time_points
    )
    h = (temporal.t1 - temporal.t0) / (n_time_points - 1)
    weights = jnp.ones(n_time_points).at[0].set(0.5).at[-1].set(0.5)
    return jnp.sum(intensities * weights, axis=-1) * h

ipp_spatiotemporal_marginal_temporal_intensity(times: Float[Array, ...], log_intensity_fn: Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]], spatial: RectangularDomain, n_spatial_points: int = 256) -> Float[Array, ...]

Compute :math:\lambda_T(t) = \int_D \lambda(s, t)\, ds.

Halton-QMC over the spatial domain at each query time; cheap and deterministic, and converges as :math:O(\log^d N / N).

Parameters:

Name	Type	Description	Default
times	Float[Array, ...]	(M,) query times.	required
log_intensity_fn	Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]]	Callable (s, t) → log λ(s, t).	required
spatial	RectangularDomain	Spatial domain.	required
n_spatial_points	int	Static number of QMC nodes.	256

Returns:

Type	Description
Float[Array, ...]	(M,) array of marginal temporal intensities.

Source code in src/xtremax/point_processes/primitives/ipp_spatiotemporal.py

def ipp_spatiotemporal_marginal_temporal_intensity(
    times: Float[Array, ...],
    log_intensity_fn: Callable[
        [Float[Array, ...], Float[Array, ...]], Float[Array, ...]
    ],
    spatial: RectangularDomain,
    n_spatial_points: int = 256,
) -> Float[Array, ...]:
    """Compute :math:`\\lambda_T(t) = \\int_D \\lambda(s, t)\\, ds`.

    Halton-QMC over the spatial domain at each query time; cheap and
    deterministic, and converges as :math:`O(\\log^d N / N)`.

    Args:
        times: ``(M,)`` query times.
        log_intensity_fn: Callable ``(s, t) → log λ(s, t)``.
        spatial: Spatial domain.
        n_spatial_points: Static number of QMC nodes.

    Returns:
        ``(M,)`` array of marginal temporal intensities.
    """
    from xtremax.point_processes._integration_spatial import _halton_sequence

    times = jnp.asarray(times)
    n_query = times.shape[0]
    unit_cube = _halton_sequence(n_spatial_points, spatial.n_dims)
    spatial_grid = spatial.lo + unit_cube * spatial.side_lengths

    # (M, n_spatial_points, d_space) and (M, n_spatial_points)
    s_tiled = jnp.broadcast_to(
        spatial_grid[None], (n_query, n_spatial_points, spatial.n_dims)
    )
    t_tiled = jnp.repeat(times[:, None], n_spatial_points, axis=1)
    s_flat = s_tiled.reshape(-1, spatial.n_dims)
    t_flat = t_tiled.reshape(-1)
    intensities = jnp.exp(log_intensity_fn(s_flat, t_flat)).reshape(
        n_query, n_spatial_points
    )
    return jnp.mean(intensities, axis=-1) * spatial.volume()

ipp_spatiotemporal_pearson_residuals(locations: Float[Array, ...], times: Float[Array, ...], mask: Bool[Array, ...], log_intensity_fn: Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]], spatial: RectangularDomain, temporal: TemporalDomain, n_spatial_bins: int = 5, n_temporal_bins: int = 5, n_integration_points: int = 16) -> Float[Array, ...]

Pearson residuals on a tensor-product space-time grid.

For each cell A_k = R_k × [τ_k, τ_{k+1}),

.. math:: r_k = \frac{N(A_k) - \Lambda(A_k)}{\sqrt{\Lambda(A_k)}}.

Cells with :math:\Lambda(A_k) = 0 return 0 to keep the residual array finite.

Parameters:

Name	Type	Description	Default
locations	Float[Array, ...]	(max_events, d_space) event locations.	required
times	Float[Array, ...]	(max_events,) event times.	required
mask	Bool[Array, ...]	(max_events,) boolean mask of real events.	required
log_intensity_fn	Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]]	Callable (s, t) → log λ(s, t).	required
spatial	RectangularDomain	Spatial domain.	required
temporal	TemporalDomain	Temporal interval.	required
n_spatial_bins	int	Bins per spatial axis (so n_spatial_bins ** d spatial cells total).	5
n_temporal_bins	int	Bins along the time axis.	5
n_integration_points	int	Per-cell trapezoid nodes per axis when estimating :math:\Lambda(A_k).	16

Returns:

Type	Description
Float[Array, ...]	Flattened residual array of shape
Float[Array, ...]	(n_temporal_bins * n_spatial_bins ** d,).

Source code in src/xtremax/point_processes/primitives/ipp_spatiotemporal.py

def ipp_spatiotemporal_pearson_residuals(
    locations: Float[Array, ...],
    times: Float[Array, ...],
    mask: Bool[Array, ...],
    log_intensity_fn: Callable[
        [Float[Array, ...], Float[Array, ...]], Float[Array, ...]
    ],
    spatial: RectangularDomain,
    temporal: TemporalDomain,
    n_spatial_bins: int = 5,
    n_temporal_bins: int = 5,
    n_integration_points: int = 16,
) -> Float[Array, ...]:
    """Pearson residuals on a tensor-product space-time grid.

    For each cell ``A_k = R_k × [τ_k, τ_{k+1})``,

    .. math::
        r_k = \\frac{N(A_k) - \\Lambda(A_k)}{\\sqrt{\\Lambda(A_k)}}.

    Cells with :math:`\\Lambda(A_k) = 0` return ``0`` to keep the
    residual array finite.

    Args:
        locations: ``(max_events, d_space)`` event locations.
        times: ``(max_events,)`` event times.
        mask: ``(max_events,)`` boolean mask of real events.
        log_intensity_fn: Callable ``(s, t) → log λ(s, t)``.
        spatial: Spatial domain.
        temporal: Temporal interval.
        n_spatial_bins: Bins per spatial axis (so ``n_spatial_bins ** d``
            spatial cells total).
        n_temporal_bins: Bins along the time axis.
        n_integration_points: Per-cell trapezoid nodes per axis when
            estimating :math:`\\Lambda(A_k)`.

    Returns:
        Flattened residual array of shape
        ``(n_temporal_bins * n_spatial_bins ** d,)``.
    """
    d = spatial.n_dims
    s_edges = [
        jnp.linspace(spatial.lo[i], spatial.hi[i], n_spatial_bins + 1) for i in range(d)
    ]
    t_edges = jnp.linspace(temporal.t0, temporal.t1, n_temporal_bins + 1)

    # Bucket the events. Use jnp.digitize per axis with right-open bins.
    spatial_idx = jnp.stack(
        [
            jnp.clip(
                jnp.digitize(locations[:, i], s_edges[i]) - 1, 0, n_spatial_bins - 1
            )
            for i in range(d)
        ],
        axis=-1,
    )  # (N, d)
    temporal_idx = jnp.clip(jnp.digitize(times, t_edges) - 1, 0, n_temporal_bins - 1)

    # Encode (spatial_idx, temporal_idx) → flat cell index
    # row-major (temporal slowest, spatial axes packed in C order).
    spatial_flat = jnp.zeros_like(temporal_idx)
    stride = 1
    for axis in range(d - 1, -1, -1):
        spatial_flat = spatial_flat + spatial_idx[:, axis] * stride
        stride = stride * n_spatial_bins
    cell_idx = temporal_idx * (n_spatial_bins**d) + spatial_flat

    n_cells = n_temporal_bins * (n_spatial_bins**d)
    counts = jnp.zeros(n_cells)
    # Scatter-add a per-event weight: 1 for real events, 0 for padding.
    # Padding rows still scatter into a real cell (the digitize/clip
    # path collapses them to the boundary cell), but the 0 weight
    # leaves it untouched. The earlier sentinel-bucket trick was a
    # no-op under JAX's negative-index wrap, so drop it (Codex P1 in
    # PR #13 review).
    counts = counts.at[cell_idx].add(jnp.where(mask, 1.0, 0.0))

    # Per-cell integrated intensity by trapezoid quadrature on a small
    # tensor-product grid inside each cell.
    def cell_lambda(temporal_bin: int, spatial_bins: tuple[int, ...]) -> Array:
        sub_lo = jnp.array([s_edges[i][spatial_bins[i]] for i in range(d)])
        sub_hi = jnp.array([s_edges[i][spatial_bins[i] + 1] for i in range(d)])
        sub_t0 = t_edges[temporal_bin]
        sub_t1 = t_edges[temporal_bin + 1]

        n = n_integration_points
        spatial_grids = [jnp.linspace(sub_lo[i], sub_hi[i], n) for i in range(d)]
        time_grid = jnp.linspace(sub_t0, sub_t1, n)
        h_spatial = jnp.array([(sub_hi[i] - sub_lo[i]) / (n - 1) for i in range(d)])
        h_time = (sub_t1 - sub_t0) / (n - 1)

        base_w = jnp.ones(n).at[0].set(0.5).at[-1].set(0.5)
        mesh = jnp.meshgrid(*spatial_grids, time_grid, indexing="ij")
        flat_s = jnp.stack([axis.reshape(-1) for axis in mesh[:d]], axis=-1)
        flat_t = mesh[-1].reshape(-1)
        intensities = jnp.exp(log_intensity_fn(flat_s, flat_t)).reshape((n,) * (d + 1))

        weight = base_w
        for _ in range(d):  # d + 1 axes total → d more outer products
            weight = jnp.tensordot(weight, base_w, axes=0)
        cell_volume = jnp.prod(h_spatial) * h_time
        return jnp.sum(weight * intensities) * cell_volume

    # Build the per-cell λ list eagerly. The bin counts are static
    # Python ints so this is unrolled at trace time and fine for the
    # typical small ``n_*_bins`` choices we expect for residuals.
    spatial_index_grid = jnp.stack(
        jnp.meshgrid(*[jnp.arange(n_spatial_bins) for _ in range(d)], indexing="ij"),
        axis=-1,
    ).reshape(-1, d)
    expected = []
    for tb in range(n_temporal_bins):
        for sb in spatial_index_grid:
            expected.append(cell_lambda(tb, tuple(int(x) for x in sb)))
    expected_arr = jnp.stack(expected)

    safe_expected = jnp.where(expected_arr > 0.0, expected_arr, 1.0)
    residuals = (counts - expected_arr) / jnp.sqrt(safe_expected)
    return jnp.where(expected_arr > 0.0, residuals, 0.0)

ipp_spatiotemporal_chi_square_gof(locations: Float[Array, ...], times: Float[Array, ...], mask: Bool[Array, ...], log_intensity_fn: Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]], spatial: RectangularDomain, temporal: TemporalDomain, n_spatial_bins: int = 5, n_temporal_bins: int = 5, n_integration_points: int = 16) -> Float[Array, ...]

Pearson :math:\chi^2 statistic from binned residuals.

Returns :math:\sum_k r_k^2. Compare to a :math:\chi^2_{K - p} distribution with degrees of freedom equal to the number of cells minus the number of fitted parameters.

Source code in src/xtremax/point_processes/primitives/ipp_spatiotemporal.py

def ipp_spatiotemporal_chi_square_gof(
    locations: Float[Array, ...],
    times: Float[Array, ...],
    mask: Bool[Array, ...],
    log_intensity_fn: Callable[
        [Float[Array, ...], Float[Array, ...]], Float[Array, ...]
    ],
    spatial: RectangularDomain,
    temporal: TemporalDomain,
    n_spatial_bins: int = 5,
    n_temporal_bins: int = 5,
    n_integration_points: int = 16,
) -> Float[Array, ...]:
    """Pearson :math:`\\chi^2` statistic from binned residuals.

    Returns :math:`\\sum_k r_k^2`. Compare to a :math:`\\chi^2_{K - p}`
    distribution with degrees of freedom equal to the number of cells
    minus the number of fitted parameters.
    """
    residuals = ipp_spatiotemporal_pearson_residuals(
        locations,
        times,
        mask,
        log_intensity_fn,
        spatial,
        temporal,
        n_spatial_bins=n_spatial_bins,
        n_temporal_bins=n_temporal_bins,
        n_integration_points=n_integration_points,
    )
    return jnp.sum(residuals**2)

Renewal processes

Parameterised by the inter-event distribution rather than an intensity: the conditional intensity is the hazard evaluated since the last event, \(\lambda^*(t) = h(t - t_{N(t)})\). renewal_ogata_intensity_fn adapts a renewal process into the (t, history) -> λ*(t) closure that the thinning samplers expect.

Pure-function primitives for temporal and spatial point processes.

Each submodule is free of class state and free of framework dependencies (no NumPyro — save that for the outer-layer wrappers — no equinox). These are the canonical reference implementations; the operator and distribution layers delegate to them.

renewal_intensity(t: Float[Array, ...], event_times: Float[Array, ...], mask: Bool[Array, ...], inter_event_dist: dist.Distribution) -> Float[Array, ...]

Conditional intensity :math:\lambda^*(t) = h(t - t_{N(t)}).

Parameters:

Name	Type	Description	Default
t	Float[Array, ...]	Evaluation time.	required
event_times	Float[Array, ...]	Padded history.	required
mask	Bool[Array, ...]	Event mask.	required
inter_event_dist	Distribution	Inter-event distribution.	required

Returns:

Type	Description
Float[Array, ...]	h(τ) where τ is the time since the most recent event
Float[Array, ...]	strictly before t.

Source code in src/xtremax/point_processes/primitives/renewal.py

def renewal_intensity(
    t: Float[Array, ...],
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    inter_event_dist: dist.Distribution,
) -> Float[Array, ...]:
    r"""Conditional intensity :math:`\lambda^*(t) = h(t - t_{N(t)})`.

    Args:
        t: Evaluation time.
        event_times: Padded history.
        mask: Event mask.
        inter_event_dist: Inter-event distribution.

    Returns:
        ``h(τ)`` where ``τ`` is the time since the most recent event
        strictly before ``t``.
    """
    t = jnp.asarray(t)
    # Most recent event strictly before t (0 if none).
    before_t = mask & (event_times < t)
    last_time = jnp.max(
        jnp.where(before_t, event_times, jnp.zeros_like(event_times)), axis=-1
    )
    last_time = jnp.where(jnp.any(before_t, axis=-1), last_time, 0.0)
    tau = t - last_time
    return renewal_hazard(tau, inter_event_dist)

renewal_log_prob(event_times: Float[Array, ...], mask: Bool[Array, ...], T: Float[Array, ...], inter_event_dist: dist.Distribution) -> Float[Array, ...]

Log-likelihood of an observed renewal-process sequence.

Parameters:

Name	Type	Description	Default
event_times	Float[Array, ...]	Padded, sorted event times.	required
mask	Bool[Array, ...]	Boolean event mask.	required
T	Float[Array, ...]	Observation window length.	required
inter_event_dist	Distribution	NumPyro Distribution for inter-event times. Must expose log_prob and cdf.	required

Returns:

Type	Description
Float[Array, ...]	Scalar log-likelihood.

Notes

The final survival term uses log1p(-cdf(T - t_n)) with clipping to avoid log(0) when the observed gap has CDF very close to 1. The inter-event densities at padding positions are zeroed out by the mask.

Requires the package padded-buffer invariant (contiguous-prefix mask): gaps are taken between consecutive buffer slots, so a hole mask (a masked-out slot between two real events) silently corrupts the gap terms. Package samplers always emit contiguous-prefix masks.

Source code in src/xtremax/point_processes/primitives/renewal.py

def renewal_log_prob(
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    T: Float[Array, ...],
    inter_event_dist: dist.Distribution,
) -> Float[Array, ...]:
    """Log-likelihood of an observed renewal-process sequence.

    Args:
        event_times: Padded, sorted event times.
        mask: Boolean event mask.
        T: Observation window length.
        inter_event_dist: NumPyro Distribution for inter-event times.
            Must expose ``log_prob`` and ``cdf``.

    Returns:
        Scalar log-likelihood.

    Notes:
        The final survival term uses ``log1p(-cdf(T - t_n))`` with
        clipping to avoid ``log(0)`` when the observed gap has CDF
        very close to 1. The inter-event densities at padding positions
        are zeroed out by the mask.

        Requires the package padded-buffer invariant (contiguous-prefix
        mask): gaps are taken between consecutive *buffer slots*, so a
        hole mask (a masked-out slot between two real events) silently
        corrupts the gap terms. Package samplers always emit
        contiguous-prefix masks.
    """
    event_times = jnp.asarray(event_times)
    T_arr = jnp.asarray(T)

    # Inter-event gaps: Δt_i = t_i - t_{i-1} (with t_0 = 0).
    zero = jnp.zeros_like(event_times[..., :1])
    prev_times = jnp.concatenate([zero, event_times[..., :-1]], axis=-1)
    gaps = event_times - prev_times
    # Guard padding positions from producing NaN-producing log_probs.
    safe_gaps = jnp.where(mask, gaps, jnp.ones_like(gaps))
    log_f = inter_event_dist.log_prob(safe_gaps)
    log_density_sum = jnp.sum(jnp.where(mask, log_f, 0.0), axis=-1)

    # Tail survival from the last real event to the window edge.
    n_real = jnp.sum(mask, axis=-1)
    last_time = jnp.max(
        jnp.where(mask, event_times, jnp.zeros_like(event_times)), axis=-1
    )
    gap_after_last = T_arr - last_time
    # With no observed events the tail gap is the whole window.
    no_events = n_real == 0
    gap_after_last = jnp.where(no_events, T_arr, gap_after_last)
    cdf_tail = inter_event_dist.cdf(gap_after_last)
    log_survival = jnp.log(jnp.clip(1.0 - cdf_tail, 1e-30, 1.0))

    return log_density_sum + log_survival

renewal_sample(key: PRNGKeyArray, inter_event_dist: dist.Distribution, T: Float[Array, ...], max_events: int) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Sample a renewal-process event sequence by accumulating gaps.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	JAX PRNG key.	required
inter_event_dist	Distribution	Inter-event NumPyro Distribution. Must have a reparameterised sample that is jit friendly.	required
T	Float[Array, ...]	Observation window length.	required
max_events	int	Static buffer size.	required

Returns:

Type	Description
tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	Tuple (times, mask, n_events).

Notes

Draws max_events iid gaps unconditionally and cumulative- sums them; events past T are masked out. This keeps shape static while preserving the exact renewal distribution on [0, T]. If the buffer fills (sum of gaps < T) the returned sequence is truncated — choose max_events so that max_events * E[Δt] >> T.

Source code in src/xtremax/point_processes/primitives/renewal.py

def renewal_sample(
    key: PRNGKeyArray,
    inter_event_dist: dist.Distribution,
    T: Float[Array, ...],
    max_events: int,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Sample a renewal-process event sequence by accumulating gaps.

    Args:
        key: JAX PRNG key.
        inter_event_dist: Inter-event NumPyro Distribution. Must have
            a reparameterised ``sample`` that is ``jit`` friendly.
        T: Observation window length.
        max_events: Static buffer size.

    Returns:
        Tuple ``(times, mask, n_events)``.

    Notes:
        Draws ``max_events`` iid gaps unconditionally and cumulative-
        sums them; events past ``T`` are masked out. This keeps shape
        static while preserving the exact renewal distribution on
        ``[0, T]``. If the buffer fills (sum of gaps < ``T``) the
        returned sequence is truncated — choose ``max_events`` so that
        ``max_events * E[Δt] >> T``.
    """
    T_arr = jnp.asarray(T)
    gaps = inter_event_dist.sample(key, sample_shape=(max_events,))
    times = jnp.cumsum(gaps, axis=-1)
    mask = times < T_arr
    n_events = jnp.sum(mask, axis=-1).astype(jnp.int32)
    # Pad out-of-window positions to T so downstream intensity evaluations stay safe.
    times = jnp.where(mask, times, T_arr)
    return SampleResult(times, mask, n_events)

renewal_hazard(tau: Float[Array, ...], inter_event_dist: dist.Distribution) -> Float[Array, ...]

Hazard of the inter-event distribution: :math:h(\tau) = f(\tau)/S(\tau).

Implemented as exp(log_prob(τ) - log(1 - cdf(τ))) to stay numerically safe.

Source code in src/xtremax/point_processes/primitives/renewal.py

def renewal_hazard(
    tau: Float[Array, ...],
    inter_event_dist: dist.Distribution,
) -> Float[Array, ...]:
    r"""Hazard of the inter-event distribution: :math:`h(\tau) = f(\tau)/S(\tau)`.

    Implemented as ``exp(log_prob(τ) - log(1 - cdf(τ)))`` to stay
    numerically safe.
    """
    tau = jnp.asarray(tau)
    log_f = inter_event_dist.log_prob(tau)
    log_s = jnp.log(jnp.clip(1.0 - inter_event_dist.cdf(tau), 1e-30, 1.0))
    return jnp.exp(log_f - log_s)

renewal_cumulative_hazard(tau: Float[Array, ...], inter_event_dist: dist.Distribution) -> Float[Array, ...]

:math:\Lambda(\tau) = -\log S(\tau) = -\log(1 - F(\tau)).

Source code in src/xtremax/point_processes/primitives/renewal.py

def renewal_cumulative_hazard(
    tau: Float[Array, ...],
    inter_event_dist: dist.Distribution,
) -> Float[Array, ...]:
    r""":math:`\Lambda(\tau) = -\log S(\tau) = -\log(1 - F(\tau))`."""
    tau = jnp.asarray(tau)
    return -jnp.log(jnp.clip(1.0 - inter_event_dist.cdf(tau), 1e-30, 1.0))

renewal_survival(tau: Float[Array, ...], inter_event_dist: dist.Distribution) -> Float[Array, ...]

S(τ) = 1 − F(τ) clipped to be non-negative.

Source code in src/xtremax/point_processes/primitives/renewal.py

def renewal_survival(
    tau: Float[Array, ...],
    inter_event_dist: dist.Distribution,
) -> Float[Array, ...]:
    """``S(τ) = 1 − F(τ)`` clipped to be non-negative."""
    return jnp.clip(1.0 - inter_event_dist.cdf(jnp.asarray(tau)), 0.0, 1.0)

renewal_inter_event_log_prob(tau: Float[Array, ...], inter_event_dist: dist.Distribution) -> Float[Array, ...]

Log density of an inter-event gap: log f(τ).

Source code in src/xtremax/point_processes/primitives/renewal.py

def renewal_inter_event_log_prob(
    tau: Float[Array, ...],
    inter_event_dist: dist.Distribution,
) -> Float[Array, ...]:
    """Log density of an inter-event gap: ``log f(τ)``."""
    return inter_event_dist.log_prob(jnp.asarray(tau))

renewal_expected_count(T: Float[Array, ...], inter_event_dist: dist.Distribution, n_points: int = 200) -> Float[Array, ...]

Expected number of events on [0, T] via the renewal equation.

Solves :math:m(T) = F(T) + \\int_0^T m(T - s)\\, f(s)\\, ds on an n_points trapezoid grid by forward substitution: the discretized equation is lower-triangular in the grid values of m, so a single sweep over grid rows solves it exactly (up to quadrature error) at any horizon. The previous fixed 8-iteration Neumann series added one convolution power per iteration and therefore saturated near E[N] ≈ 9 regardless of T. For Exponential inter-events this recovers λT to quadrature tolerance.

Parameters:

Name	Type	Description	Default
T	Float[Array, ...]	Window length (scalar).	required
inter_event_dist	Distribution	Inter-event NumPyro Distribution.	required
n_points	int	Trapezoid grid size.	200

Returns:

Type	Description
Float[Array, ...]	Scalar expected count.

Source code in src/xtremax/point_processes/primitives/renewal.py

def renewal_expected_count(
    T: Float[Array, ...],
    inter_event_dist: dist.Distribution,
    n_points: int = 200,
) -> Float[Array, ...]:
    r"""Expected number of events on ``[0, T]`` via the renewal equation.

    Solves :math:`m(T) = F(T) + \\int_0^T m(T - s)\\, f(s)\\, ds` on an
    ``n_points`` trapezoid grid by **forward substitution**: the
    discretized equation is lower-triangular in the grid values of
    ``m``, so a single sweep over grid rows solves it exactly (up to
    quadrature error) at *any* horizon. The previous fixed 8-iteration
    Neumann series added one convolution power per iteration and
    therefore saturated near ``E[N] ≈ 9`` regardless of ``T``. For
    ``Exponential`` inter-events this recovers ``λT`` to quadrature
    tolerance.

    Args:
        T: Window length (scalar).
        inter_event_dist: Inter-event NumPyro Distribution.
        n_points: Trapezoid grid size.

    Returns:
        Scalar expected count.
    """
    T_arr = jnp.asarray(T)
    grid = jnp.linspace(jnp.zeros_like(T_arr), T_arr, n_points)
    f_grid = jnp.exp(inter_event_dist.log_prob(grid))
    F_grid = inter_event_dist.cdf(grid)
    dt = grid[1] - grid[0]

    # Row i of the trapezoid-discretized equation:
    #   m_i = F_i + dt·[ f_0 m_i / 2 + Σ_{j=1}^{i-1} f_j m_{i-j} + f_i m_0 / 2 ]
    # The f_0 m_i / 2 term moves to the left-hand side; everything else
    # involves only already-computed rows, so a forward scan solves the
    # lower-triangular system in one pass.
    j = jnp.arange(n_points)
    denom = 1.0 - 0.5 * dt * f_grid[0]

    def row(m_cur: Array, i: Array) -> tuple[Array, None]:
        m_rev = m_cur[jnp.clip(i - j, 0, n_points - 1)]
        weights = jnp.where((j >= 1) & (j <= i - 1), 1.0, 0.0) + jnp.where(
            (j == i) & (i > 0), 0.5, 0.0
        )
        conv = jnp.sum(f_grid * m_rev * weights) * dt
        m_i = (F_grid[i] + conv) / denom
        return m_cur.at[i].set(m_i), None

    m_init = jnp.zeros_like(F_grid)
    m_final, _ = jax.lax.scan(row, m_init, jnp.arange(n_points))
    return m_final[-1]

renewal_ogata_intensity_fn(event_times: Float[Array, ...], mask: Bool[Array, ...], inter_event_dist: dist.Distribution)

Build a (t, history) -> λ*(t) closure usable by thinning samplers.

Convenience constructor for callers that want to mix renewal with other history-dependent constructions (e.g. a ThinningProcess wrapping a renewal base). The event_times / mask arguments are unused — they exist only for signature symmetry with the other *_intensity_fn builders; the returned closure evaluates against the live history the sampler passes in, so self-generated chains see their own events.

Source code in src/xtremax/point_processes/primitives/renewal.py

def renewal_ogata_intensity_fn(
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    inter_event_dist: dist.Distribution,
):
    """Build a ``(t, history) -> λ*(t)`` closure usable by thinning samplers.

    Convenience constructor for callers that want to mix renewal with
    other history-dependent constructions (e.g. a ThinningProcess
    wrapping a renewal base). The ``event_times`` / ``mask`` arguments
    are **unused** — they exist only for signature symmetry with the
    other ``*_intensity_fn`` builders; the returned closure evaluates
    against the *live* ``history`` the sampler passes in, so
    self-generated chains see their own events.
    """
    del event_times, mask  # unused; kept for signature symmetry

    def _fn(t: Array, history) -> Array:
        # Prefer the live history (passed in by the sampler) over the
        # captured times so self-generated chains see their own events.
        return renewal_intensity(t, history.times, history.mask, inter_event_dist)

    return _fn

Hawkes — temporal

Self-exciting processes where past events raise the intensity. The exponential-kernel variant has a closed-form compensator and a recursive intensity; the general-kernel variant takes any excitation function \(\phi\) at the cost of quadrature. Both sample via Ogata thinning, which needs the lambda_max upper bound.

Pure-function primitives for temporal and spatial point processes.

Each submodule is free of class state and free of framework dependencies (no NumPyro — save that for the outer-layer wrappers — no equinox). These are the canonical reference implementations; the operator and distribution layers delegate to them.

exp_hawkes_intensity(t: Float[Array, ...], event_times: Float[Array, ...], mask: Bool[Array, ...], mu: Float[Array, ...], alpha: Float[Array, ...], beta: Float[Array, ...]) -> Float[Array, ...]

Conditional intensity of the exponential-kernel Hawkes process.

:math:\lambda^*(t) = \mu + \alpha \sum_{t_i < t} e^{-\beta (t - t_i)}.

Parameters:

Name	Type	Description	Default
t	Float[Array, ...]	Evaluation time(s).	required
event_times	Float[Array, ...]	Padded, sorted past events.	required
mask	Bool[Array, ...]	Event mask.	required
mu	Float[Array, ...]	Baseline rate (positive).	required
alpha	Float[Array, ...]	Excitation amplitude (positive).	required
beta	Float[Array, ...]	Excitation decay rate (positive).	required

Returns:

Type	Description
Float[Array, ...]	math:\lambda^*(t) with shape equal to the broadcast of
Float[Array, ...]	t and the leading event-time batch dims.

Source code in src/xtremax/point_processes/primitives/hawkes.py

def exp_hawkes_intensity(
    t: Float[Array, ...],
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    mu: Float[Array, ...],
    alpha: Float[Array, ...],
    beta: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Conditional intensity of the exponential-kernel Hawkes process.

    :math:`\lambda^*(t) = \mu + \alpha \sum_{t_i < t} e^{-\beta (t - t_i)}`.

    Args:
        t: Evaluation time(s).
        event_times: Padded, sorted past events.
        mask: Event mask.
        mu: Baseline rate (positive).
        alpha: Excitation amplitude (positive).
        beta: Excitation decay rate (positive).

    Returns:
        :math:`\lambda^*(t)` with shape equal to the broadcast of
        ``t`` and the leading event-time batch dims.
    """
    t_arr = jnp.asarray(t)
    past = mask & (
        event_times < t_arr[..., None] if t_arr.ndim > 0 else event_times < t_arr
    )
    # Clip dt to [0, ∞) *before* the exp: padding/future events sit at
    # time T > t, so the raw dt is negative and exp(+β|dt|) overflows in
    # the discarded ``jnp.where`` branch — autodiff then propagates
    # 0·inf = NaN into the β gradient. The ``past`` mask still selects
    # the value. (Same fix as the spatiotemporal twin.)
    if t_arr.ndim == 0:
        dt_pos = jnp.clip(t_arr - event_times, 0.0, jnp.inf)
        contributions = jnp.where(past, alpha * jnp.exp(-beta * dt_pos), 0.0)
        return mu + jnp.sum(contributions, axis=-1)
    dt_pos = jnp.clip(t_arr[..., None] - event_times, 0.0, jnp.inf)
    contributions = jnp.where(past, alpha * jnp.exp(-beta * dt_pos), 0.0)
    return mu + jnp.sum(contributions, axis=-1)

exp_hawkes_cumulative_intensity(t: Float[Array, ...], event_times: Float[Array, ...], mask: Bool[Array, ...], mu: Float[Array, ...], alpha: Float[Array, ...], beta: Float[Array, ...]) -> Float[Array, ...]

Exponential-kernel Hawkes compensator.

:math:\Lambda(t) = \mu t + (\alpha / \beta) \sum_{t_i \leq t} (1 - e^{-\beta(t - t_i)}).

Source code in src/xtremax/point_processes/primitives/hawkes.py

def exp_hawkes_cumulative_intensity(
    t: Float[Array, ...],
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    mu: Float[Array, ...],
    alpha: Float[Array, ...],
    beta: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Exponential-kernel Hawkes compensator.

    :math:`\Lambda(t) = \mu t + (\alpha / \beta) \sum_{t_i \leq t}
    (1 - e^{-\beta(t - t_i)})`.
    """
    t_arr = jnp.asarray(t)
    # dt clipped to [0, ∞) before the exp so padding/future events cannot
    # overflow the discarded branch into a NaN β-gradient (see
    # exp_hawkes_intensity); the clip also makes the clipped per-event
    # term vanish exactly for future events.
    if t_arr.ndim == 0:
        past = mask & (event_times <= t_arr)
        dt_pos = jnp.clip(t_arr - event_times, 0.0, jnp.inf)
        per_event = (alpha / beta) * (1.0 - jnp.exp(-beta * dt_pos))
        return mu * t_arr + jnp.sum(jnp.where(past, per_event, 0.0), axis=-1)
    past = mask & (event_times <= t_arr[..., None])
    dt_pos = jnp.clip(t_arr[..., None] - event_times, 0.0, jnp.inf)
    per_event = (alpha / beta) * (1.0 - jnp.exp(-beta * dt_pos))
    return mu * t_arr + jnp.sum(jnp.where(past, per_event, 0.0), axis=-1)

exp_hawkes_log_prob(event_times: Float[Array, ...], mask: Bool[Array, ...], T: Float[Array, ...], mu: Float[Array, ...], alpha: Float[Array, ...], beta: Float[Array, ...]) -> Float[Array, ...]

Log-likelihood :math:\sum_i \log \lambda^*(t_i) - \Lambda(T) for exp Hawkes.

Uses the Ozaki recursion :math:R_i = e^{-\beta (t_i - t_{i-1})}(1 + R_{i-1}) so intensity evaluation at each observed event is O(1) amortised, giving O(n) total cost. The recursion consumes consecutive buffer gaps, so it requires the package padded-buffer invariant (contiguous-prefix mask, padding at T) that all package samplers guarantee.

Parameters:

Name	Type	Description	Default
event_times	Float[Array, ...]	Padded, sorted event times.	required
mask	Bool[Array, ...]	Event mask.	required
T	Float[Array, ...]	Observation window length.	required
mu	Float[Array, ...]	Baseline.	required
alpha	Float[Array, ...]	Excitation amplitude.	required
beta	Float[Array, ...]	Excitation decay.	required

Returns:

Type	Description
Float[Array, ...]	Scalar log-likelihood.

Source code in src/xtremax/point_processes/primitives/hawkes.py

def exp_hawkes_log_prob(
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    T: Float[Array, ...],
    mu: Float[Array, ...],
    alpha: Float[Array, ...],
    beta: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Log-likelihood :math:`\sum_i \log \lambda^*(t_i) - \Lambda(T)` for exp Hawkes.

    Uses the Ozaki recursion :math:`R_i = e^{-\beta (t_i - t_{i-1})}(1 + R_{i-1})`
    so intensity evaluation at each observed event is O(1) amortised,
    giving O(n) total cost. The recursion consumes consecutive buffer
    gaps, so it requires the package padded-buffer invariant
    (contiguous-prefix mask, padding at ``T``) that all package
    samplers guarantee.

    Args:
        event_times: Padded, sorted event times.
        mask: Event mask.
        T: Observation window length.
        mu: Baseline.
        alpha: Excitation amplitude.
        beta: Excitation decay.

    Returns:
        Scalar log-likelihood.
    """
    event_times = jnp.asarray(event_times)
    T_arr = jnp.asarray(T)

    # Ozaki recursion for intensities at observed event times.
    # Using the A-variable A_i = 1 + Σ_{j<i} exp(-β(t_i - t_j)) = 1 + R_i
    # which satisfies A_i = 1 + exp(-β·gap_i) · A_{i-1}, A_{-1} := 0,
    # so R_i = A_i − 1 and the recursion is numerically clean.
    zero = jnp.zeros_like(event_times[..., :1])
    prev_times = jnp.concatenate([zero, event_times[..., :-1]], axis=-1)
    gaps = event_times - prev_times

    def step(A_prev: Array, inp: tuple[Array, Array]) -> tuple[Array, Array]:
        gap, valid = inp
        A_new = jnp.where(
            valid,
            1.0 + jnp.exp(-beta * gap) * A_prev,
            A_prev,
        )
        # Output R_i = A_new - 1 when the current event is valid; at
        # padding positions emit 0 (masked out of the sum anyway).
        R_i = jnp.where(valid, A_new - 1.0, 0.0)
        return A_new, R_i

    _, R_strict = jax.lax.scan(
        step,
        jnp.zeros(event_times.shape[:-1], dtype=event_times.dtype),
        (jnp.moveaxis(gaps, -1, 0), jnp.moveaxis(mask, -1, 0)),
    )
    R_strict = jnp.moveaxis(R_strict, 0, -1)
    lam_at_events = mu + alpha * R_strict
    # Guard the log at padding positions (lam_at_events = μ > 0 there anyway).
    log_lam = jnp.log(jnp.clip(lam_at_events, 1e-30, jnp.inf))
    sum_log_lam = jnp.sum(jnp.where(mask, log_lam, 0.0), axis=-1)

    Lambda_T = exp_hawkes_cumulative_intensity(
        T_arr, event_times, mask, mu, alpha, beta
    )
    return sum_log_lam - Lambda_T

exp_hawkes_sample(key: PRNGKeyArray, T: Float[Array, ...], mu: Float[Array, ...], alpha: Float[Array, ...], beta: Float[Array, ...], max_events: int, max_candidates: int | None = None) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Ogata thinning sample of an exponential-kernel Hawkes process.

Uses the adaptive envelope :math:\bar\lambda_k = \mu + \alpha k valid for :math:t \geq t_k. With a decaying excitation this bound is tight just after each event and loosens in between — efficient enough for :math:\alpha / \beta < 1 (sub-critical).

Source code in src/xtremax/point_processes/primitives/hawkes.py

def exp_hawkes_sample(
    key: PRNGKeyArray,
    T: Float[Array, ...],
    mu: Float[Array, ...],
    alpha: Float[Array, ...],
    beta: Float[Array, ...],
    max_events: int,
    max_candidates: int | None = None,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Ogata thinning sample of an exponential-kernel Hawkes process.

    Uses the adaptive envelope :math:`\\bar\\lambda_k = \\mu + \\alpha k`
    valid for :math:`t \\geq t_k`. With a decaying excitation this
    bound is tight just after each event and loosens in between —
    efficient enough for :math:`\\alpha / \\beta < 1` (sub-critical).
    """
    mu_arr = jnp.asarray(mu)
    alpha_arr = jnp.asarray(alpha)
    beta_arr = jnp.asarray(beta)

    def intensity(t: Array, history: EventHistory) -> Array:
        return exp_hawkes_intensity(
            t, history.times, history.mask, mu_arr, alpha_arr, beta_arr
        )

    def lambda_max(history: EventHistory) -> Array:
        return exp_hawkes_lambda_max(history, mu_arr, alpha_arr)

    return thinning_sample(
        key=key,
        conditional_intensity_fn=intensity,
        lambda_max_fn=lambda_max,
        T=T,
        max_events=max_events,
        max_candidates=max_candidates,
    )

exp_hawkes_lambda_max(history: EventHistory, mu: Float[Array, ...], alpha: Float[Array, ...]) -> Float[Array, ...]

Upper bound :math:\mu + \alpha \cdot N(t) valid from the last event onward.

Source code in src/xtremax/point_processes/primitives/hawkes.py

def exp_hawkes_lambda_max(
    history: EventHistory,
    mu: Float[Array, ...],
    alpha: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Upper bound :math:`\mu + \alpha \cdot N(t)` valid from the last event onward."""
    n = jnp.sum(history.mask).astype(mu.dtype)
    return mu + alpha * n

general_hawkes_intensity(t: Float[Array, ...], event_times: Float[Array, ...], mask: Bool[Array, ...], mu: Float[Array, ...], kernel_fn: Callable[[Array], Array]) -> Float[Array, ...]

Conditional intensity for a general kernel.

:math:\lambda^*(t) = \mu + \sum_{t_i < t} \phi(t - t_i).

Source code in src/xtremax/point_processes/primitives/hawkes.py

def general_hawkes_intensity(
    t: Float[Array, ...],
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    mu: Float[Array, ...],
    kernel_fn: Callable[[Array], Array],
) -> Float[Array, ...]:
    r"""Conditional intensity for a general kernel.

    :math:`\lambda^*(t) = \mu + \sum_{t_i < t} \phi(t - t_i)`.
    """
    t_arr = jnp.asarray(t)
    # dt clamped to [0, ∞) before the user kernel: padding/future events
    # produce negative dt, where an arbitrary kernel may be non-finite or
    # have non-finite gradients — the discarded ``jnp.where`` branch then
    # poisons autodiff with 0·inf = NaN. Kernels are only ever evaluated
    # on their supported domain τ ≥ 0.
    if t_arr.ndim == 0:
        past = mask & (event_times < t_arr)
        dt_pos = jnp.clip(t_arr - event_times, 0.0, jnp.inf)
        kernel_vals = kernel_fn(dt_pos)
        return mu + jnp.sum(jnp.where(past, kernel_vals, 0.0), axis=-1)
    past = mask & (event_times < t_arr[..., None])
    dt_pos = jnp.clip(t_arr[..., None] - event_times, 0.0, jnp.inf)
    kernel_vals = kernel_fn(dt_pos)
    return mu + jnp.sum(jnp.where(past, kernel_vals, 0.0), axis=-1)

general_hawkes_cumulative_intensity(T: Float[Array, ...], event_times: Float[Array, ...], mask: Bool[Array, ...], mu: Float[Array, ...], kernel_integral_fn: Callable[[Array, Array], Array]) -> Float[Array, ...]

Compensator :math:\mu T + \sum_i \int_{t_i}^{T} \phi(s - t_i) ds.

Supports scalar and array-valued T. Array T takes a pairwise branch — upper[k, i] = clip(T[k] - t_i) — mirroring the exponential-kernel primitive; the previous elementwise subtraction made cumulative_intensity(t=event_times) return upper ≡ 0, so every kernel integral vanished and GOF diagnostics silently degenerated to a homogeneous Poisson with rate μ.

Source code in src/xtremax/point_processes/primitives/hawkes.py

def general_hawkes_cumulative_intensity(
    T: Float[Array, ...],
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    mu: Float[Array, ...],
    kernel_integral_fn: Callable[[Array, Array], Array],
) -> Float[Array, ...]:
    r"""Compensator :math:`\mu T + \sum_i \int_{t_i}^{T} \phi(s - t_i) ds`.

    Supports scalar and array-valued ``T``. Array ``T`` takes a pairwise
    branch — ``upper[k, i] = clip(T[k] - t_i)`` — mirroring the
    exponential-kernel primitive; the previous elementwise subtraction
    made ``cumulative_intensity(t=event_times)`` return ``upper ≡ 0``,
    so every kernel integral vanished and GOF diagnostics silently
    degenerated to a homogeneous Poisson with rate ``μ``.
    """
    T_arr = jnp.asarray(T)
    if T_arr.ndim == 0:
        lower = jnp.zeros_like(event_times)
        upper = jnp.clip(T_arr - event_times, 0.0, jnp.inf)
        per_event = kernel_integral_fn(lower, upper)
        return mu * T_arr + jnp.sum(jnp.where(mask, per_event, 0.0), axis=-1)
    # ``T[..., None]`` adds one axis that broadcasts against the events
    # axis — pairwise for an unbatched vector T against (n,) events, and
    # elementwise (one T per history) for batched aligned inputs, exactly
    # like the exponential-kernel primitive. Masking through the same
    # broadcast (rather than an unconditional ``mask[..., None, :]``,
    # which mixed masks across batch rows) keeps both cases correct.
    upper = jnp.clip(T_arr[..., None] - event_times, 0.0, jnp.inf)
    lower = jnp.zeros_like(upper)
    per_event = kernel_integral_fn(lower, upper)
    past = mask & (event_times <= T_arr[..., None])
    return mu * T_arr + jnp.sum(jnp.where(past, per_event, 0.0), axis=-1)

general_hawkes_log_prob(event_times: Float[Array, ...], mask: Bool[Array, ...], T: Float[Array, ...], mu: Float[Array, ...], kernel_fn: Callable[[Array], Array], kernel_integral_fn: Callable[[Array, Array], Array]) -> Float[Array, ...]

Log-likelihood for a general-kernel Hawkes process.

Parameters:

Name	Type	Description	Default
event_times	Float[Array, ...]	Padded event times.	required
mask	Bool[Array, ...]	Event mask.	required
T	Float[Array, ...]	Window length.	required
mu	Float[Array, ...]	Baseline rate.	required
kernel_fn	Callable[[Array], Array]	φ(τ) -> Array broadcast-friendly.	required
kernel_integral_fn	Callable[[Array, Array], Array]	(a, b) -> ∫_a^b φ(s) ds broadcast-friendly. Used for the compensator term; closed-form forms give exact gradients, quadrature fallbacks (see the operator layer) trade some accuracy for generality.	required

Returns:

Type	Description
Float[Array, ...]	Scalar log-likelihood.

Source code in src/xtremax/point_processes/primitives/hawkes.py

def general_hawkes_log_prob(
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    T: Float[Array, ...],
    mu: Float[Array, ...],
    kernel_fn: Callable[[Array], Array],
    kernel_integral_fn: Callable[[Array, Array], Array],
) -> Float[Array, ...]:
    r"""Log-likelihood for a general-kernel Hawkes process.

    Args:
        event_times: Padded event times.
        mask: Event mask.
        T: Window length.
        mu: Baseline rate.
        kernel_fn: ``φ(τ) -> Array`` broadcast-friendly.
        kernel_integral_fn: ``(a, b) -> ∫_a^b φ(s) ds`` broadcast-friendly.
            Used for the compensator term; closed-form forms give exact
            gradients, quadrature fallbacks (see the operator layer)
            trade some accuracy for generality.

    Returns:
        Scalar log-likelihood.
    """
    # λ*(t_i) per event via O(n²) pairwise construction (padded-mask-safe).
    # Precedence is decided by TIMESTAMP, not buffer position — index-based
    # comparison made log_prob change under a row permutation of the same
    # event set (same fix as the spatiotemporal twin). dt is clamped to
    # [0, ∞) before the user kernel so discarded (non-causal) pairs cannot
    # feed a negative argument whose value/gradient may be non-finite.
    dt = event_times[..., :, None] - event_times[..., None, :]
    dt_pos = jnp.clip(dt, 0.0, jnp.inf)
    kernel_vals = kernel_fn(dt_pos)
    strictly_before = event_times[..., None, :] < event_times[..., :, None]
    pair_valid = strictly_before & mask[..., None, :] & mask[..., :, None]
    per_i = jnp.sum(jnp.where(pair_valid, kernel_vals, 0.0), axis=-1)
    lam_at_events = mu + per_i
    log_lam = jnp.log(jnp.clip(lam_at_events, 1e-30, jnp.inf))
    sum_log_lam = jnp.sum(jnp.where(mask, log_lam, 0.0), axis=-1)

    Lambda_T = general_hawkes_cumulative_intensity(
        T, event_times, mask, mu, kernel_integral_fn
    )
    return sum_log_lam - Lambda_T

general_hawkes_sample(key: PRNGKeyArray, T: Float[Array, ...], mu: Float[Array, ...], kernel_fn: Callable[[Array], Array], kernel_max_fn: Callable[[EventHistory], Array], max_events: int, max_candidates: int | None = None) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Ogata thinning for a general-kernel Hawkes process.

The caller supplies kernel_max_fn(history) -> λ_bar — the history-dependent upper bound. For monotone-decreasing kernels with maximum at :math:\tau = 0, a safe choice is :math:\mu + \phi(0) \cdot N(t); tighter bounds give higher acceptance.

Source code in src/xtremax/point_processes/primitives/hawkes.py

def general_hawkes_sample(
    key: PRNGKeyArray,
    T: Float[Array, ...],
    mu: Float[Array, ...],
    kernel_fn: Callable[[Array], Array],
    kernel_max_fn: Callable[[EventHistory], Array],
    max_events: int,
    max_candidates: int | None = None,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    r"""Ogata thinning for a general-kernel Hawkes process.

    The caller supplies ``kernel_max_fn(history) -> λ_bar`` — the
    history-dependent upper bound. For monotone-decreasing kernels
    with maximum at :math:`\tau = 0`, a safe choice is
    :math:`\mu + \phi(0) \cdot N(t)`; tighter bounds give higher
    acceptance.
    """
    mu_arr = jnp.asarray(mu)

    def intensity(t: Array, history: EventHistory) -> Array:
        return general_hawkes_intensity(
            t, history.times, history.mask, mu_arr, kernel_fn
        )

    return thinning_sample(
        key=key,
        conditional_intensity_fn=intensity,
        lambda_max_fn=kernel_max_fn,
        T=T,
        max_events=max_events,
        max_candidates=max_candidates,
    )

Hawkes — spatiotemporal (ETAS)

Separable exp-Gaussian excitation: an exponential temporal decay times a Gaussian spatial kernel, the ETAS form used for aftershock and contagion modelling.

Pure-function primitives for temporal and spatial point processes.

Each submodule is free of class state and free of framework dependencies (no NumPyro — save that for the outer-layer wrappers — no equinox). These are the canonical reference implementations; the operator and distribution layers delegate to them.

stpp_hawkes_intensity(s: Float[Array, ...], t: Float[Array, ...], event_locations: Float[Array, ...], event_times: Float[Array, ...], mask: Bool[Array, ...], mu: Float[Array, ...], alpha: Float[Array, ...], beta: Float[Array, ...], sigma: Float[Array, ...]) -> Float[Array, ...]

Conditional intensity :math:\lambda^*(s, t) of the ETAS Hawkes process.

Parameters:

Name	Type	Description	Default
s	Float[Array, ...]	(d,) query location (scalar event) or (N, d) batch.	required
t	Float[Array, ...]	Scalar or (N,) query time.	required
event_locations	Float[Array, ...]	(max_events, d) past event locations.	required
event_times	Float[Array, ...]	(max_events,) past event times.	required
mask	Bool[Array, ...]	(max_events,) boolean event mask.	required
mu	Float[Array, ...]	Baseline rate (positive scalar).	required
alpha	Float[Array, ...]	Excitation amplitude (must satisfy alpha < 1 for sub-criticality).	required
beta	Float[Array, ...]	Temporal decay rate (positive).	required
sigma	Float[Array, ...]	Spatial bandwidth (positive).	required

Returns:

Type	Description
Float[Array, ...]	Intensity at the query, with shape matching the broadcast of
Float[Array, ...]	t and s's leading axes.

Source code in src/xtremax/point_processes/primitives/hawkes_spatiotemporal.py

def stpp_hawkes_intensity(
    s: Float[Array, ...],
    t: Float[Array, ...],
    event_locations: Float[Array, ...],
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    mu: Float[Array, ...],
    alpha: Float[Array, ...],
    beta: Float[Array, ...],
    sigma: Float[Array, ...],
) -> Float[Array, ...]:
    r"""Conditional intensity :math:`\lambda^*(s, t)` of the ETAS Hawkes process.

    Args:
        s: ``(d,)`` query location (scalar event) or ``(N, d)`` batch.
        t: Scalar or ``(N,)`` query time.
        event_locations: ``(max_events, d)`` past event locations.
        event_times: ``(max_events,)`` past event times.
        mask: ``(max_events,)`` boolean event mask.
        mu: Baseline rate (positive scalar).
        alpha: Excitation amplitude (must satisfy ``alpha < 1`` for
            sub-criticality).
        beta: Temporal decay rate (positive).
        sigma: Spatial bandwidth (positive).

    Returns:
        Intensity at the query, with shape matching the broadcast of
        ``t`` and ``s``'s leading axes.
    """
    s = jnp.asarray(s)
    t_arr = jnp.asarray(t)

    if t_arr.ndim == 0:
        past = mask & (event_times < t_arr)
        # Clip ``dt`` to non-negative before the exponential. With
        # negative ``dt`` (future / padding events) ``exp(-β·dt)`` can
        # overflow to ``+inf``, and ``inf · 0`` from the masked-out
        # branch then poisons the result with ``nan`` even though the
        # term should contribute zero. The mask is still applied below
        # to zero out the future events; ``dt_pos`` only keeps
        # intermediates finite.
        dt_pos = jnp.clip(t_arr - event_times, 0.0, jnp.inf)
        delta = s[..., None, :] - event_locations
        spatial = _gaussian_kernel(delta, sigma)
        temporal = beta * jnp.exp(-beta * dt_pos)
        contribs = jnp.where(past, alpha * temporal * spatial, 0.0)
        return mu + jnp.sum(contribs, axis=-1)

    past = mask & (event_times < t_arr[..., None])
    dt_pos = jnp.clip(t_arr[..., None] - event_times, 0.0, jnp.inf)
    delta = s[..., None, :] - event_locations
    spatial = _gaussian_kernel(delta, sigma)
    temporal = beta * jnp.exp(-beta * dt_pos)
    contribs = jnp.where(past, alpha * temporal * spatial, 0.0)
    return mu + jnp.sum(contribs, axis=-1)

stpp_hawkes_compensator(event_locations: Float[Array, ...], event_times: Float[Array, ...], mask: Bool[Array, ...], mu: Float[Array, ...], alpha: Float[Array, ...], beta: Float[Array, ...], sigma: Float[Array, ...], spatial: RectangularDomain, temporal: TemporalDomain, boundary_correction: bool = True) -> Float[Array, ...]

Closed-form compensator :math:\Lambda of the ETAS Hawkes process.

.. math:: \Lambda = \mu |D| (t_1 - t_0) + \alpha \sum_i (1 - e^{-\beta (t_1 - t_i)}) \kappa_i,

where :math:\kappa_i = \int_D k_s(s - s_i)\, ds is exactly :math:1 if boundary_correction=False (infinite-domain approximation, valid when :math:\sigma is small relative to the box and events are interior) or computed via axis-wise erf when boundary_correction=True.

Source code in src/xtremax/point_processes/primitives/hawkes_spatiotemporal.py

def stpp_hawkes_compensator(
    event_locations: Float[Array, ...],
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    mu: Float[Array, ...],
    alpha: Float[Array, ...],
    beta: Float[Array, ...],
    sigma: Float[Array, ...],
    spatial: RectangularDomain,
    temporal: TemporalDomain,
    boundary_correction: bool = True,
) -> Float[Array, ...]:
    r"""Closed-form compensator :math:`\Lambda` of the ETAS Hawkes process.

    .. math::
        \Lambda = \mu |D| (t_1 - t_0) + \alpha \sum_i
                  (1 - e^{-\beta (t_1 - t_i)}) \kappa_i,

    where :math:`\kappa_i = \int_D k_s(s - s_i)\, ds` is exactly
    :math:`1` if ``boundary_correction=False`` (infinite-domain
    approximation, valid when :math:`\sigma` is small relative to the
    box and events are interior) or computed via axis-wise erf when
    ``boundary_correction=True``.
    """
    dt = jnp.clip(temporal.t1 - event_times, 0.0, jnp.inf)
    temporal_factor = 1.0 - jnp.exp(-beta * dt)

    if boundary_correction:
        spatial_factor = jax.vmap(_gaussian_box_mass, in_axes=(0, None, None))(
            event_locations, spatial, sigma
        )
    else:
        spatial_factor = jnp.ones_like(temporal_factor)

    excitation = alpha * jnp.sum(
        jnp.where(mask, temporal_factor * spatial_factor, 0.0), axis=-1
    )
    baseline = mu * spatial.volume() * temporal.duration
    return baseline + excitation

stpp_hawkes_log_prob(locations: Float[Array, ...], times: Float[Array, ...], mask: Bool[Array, ...], mu: Float[Array, ...], alpha: Float[Array, ...], beta: Float[Array, ...], sigma: Float[Array, ...], spatial: RectangularDomain, temporal: TemporalDomain, boundary_correction: bool = True) -> Float[Array, ...]

Log-likelihood :math:\sum_i \log \lambda^*(s_i, t_i) - \Lambda.

Computed by O(n²) pairwise construction of the per-event intensities (matches the temporal :func:general_hawkes_log_prob pattern); for typical event counts under a few thousand this is competitive with the recursive Ozaki form and far simpler to extend to non-exponential temporal kernels in a follow-up.

Source code in src/xtremax/point_processes/primitives/hawkes_spatiotemporal.py

def stpp_hawkes_log_prob(
    locations: Float[Array, ...],
    times: Float[Array, ...],
    mask: Bool[Array, ...],
    mu: Float[Array, ...],
    alpha: Float[Array, ...],
    beta: Float[Array, ...],
    sigma: Float[Array, ...],
    spatial: RectangularDomain,
    temporal: TemporalDomain,
    boundary_correction: bool = True,
) -> Float[Array, ...]:
    r"""Log-likelihood :math:`\sum_i \log \lambda^*(s_i, t_i) - \Lambda`.

    Computed by O(n²) pairwise construction of the per-event
    intensities (matches the temporal :func:`general_hawkes_log_prob`
    pattern); for typical event counts under a few thousand this is
    competitive with the recursive Ozaki form and far simpler to
    extend to non-exponential temporal kernels in a follow-up.
    """
    locations = jnp.asarray(locations)
    times = jnp.asarray(times)

    dt = times[..., :, None] - times[..., None, :]
    delta = locations[..., :, None, :] - locations[..., None, :, :]
    spatial_vals = _gaussian_kernel(delta, sigma)
    temporal_vals = beta * jnp.exp(-beta * jnp.clip(dt, 0.0, jnp.inf))

    # Define causal precedence by *timestamp*, not buffer position.
    # Position-based masking (``idx_j < idx_i``) silently treats
    # whichever rows the caller wrote first as "earlier" — for an
    # unsorted buffer this turns future events into spurious parents,
    # and the ``clip(dt, 0, inf)`` then injects them at the maximum
    # temporal-kernel value. Timestamp comparison is correct
    # regardless of the row order chosen by the caller (Codex P2 in
    # PR #13 review).
    strictly_before = times[..., None, :] < times[..., :, None]
    pair_valid = strictly_before & mask[..., None, :]
    contribs = alpha * temporal_vals * spatial_vals
    per_event = jnp.sum(jnp.where(pair_valid, contribs, 0.0), axis=-1)
    lam_at_events = mu + per_event
    log_lam = jnp.log(jnp.clip(lam_at_events, 1e-30, jnp.inf))
    sum_log_lam = jnp.sum(jnp.where(mask, log_lam, 0.0), axis=-1)

    Lambda = stpp_hawkes_compensator(
        locations,
        times,
        mask,
        mu,
        alpha,
        beta,
        sigma,
        spatial,
        temporal,
        boundary_correction=boundary_correction,
    )
    return sum_log_lam - Lambda

stpp_hawkes_sample(key: PRNGKeyArray, mu: Float[Array, ...], alpha: Float[Array, ...], beta: Float[Array, ...], sigma: Float[Array, ...], spatial: RectangularDomain, temporal: TemporalDomain, max_events: int) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Ogata-style spatiotemporal thinning sample.

Algorithm (per Daley & Vere-Jones, §7.3):

1. Maintain a running history (locs, times, mask).
2. At each step, draw a candidate inter-event time from Exp(λ̄) with adaptive bound :math:\bar\lambda from :func:stpp_hawkes_lambda_max.
3. Draw a candidate location uniformly on D.
4. Accept with probability :math:\lambda^*(s, t)/\bar\lambda; otherwise discard and try again.

Loops via :func:jax.lax.scan over a fixed budget of 5 * max_events candidates so the function is jit-compatible.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	PRNG key.	required
mu	Float[Array, ...]	Baseline rate.	required
alpha	Float[Array, ...]	Excitation amplitude (sub-critical: alpha < 1).	required
beta	Float[Array, ...]	Temporal decay.	required
sigma	Float[Array, ...]	Spatial bandwidth.	required
spatial	RectangularDomain	Spatial domain.	required
temporal	TemporalDomain	Temporal interval.	required
max_events	int	Static buffer cap.	required

Returns:

Type	Description
Float[Array, ...]	(locations, times, mask, n_events) with times sorted
Float[Array, ...]	ascending. n_events counts accepted events and can never
Bool[Array, ...]	exceed max_events — acceptance stops once the buffer is
Int[Array, ...]	full, so a saturated buffer (n_events == max_events, equiv.
tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	mask.sum() == max_events) signals probable truncation and
tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	the user should increase max_events. There is no separate
tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	uncapped attempt count.

Source code in src/xtremax/point_processes/primitives/hawkes_spatiotemporal.py

def stpp_hawkes_sample(
    key: PRNGKeyArray,
    mu: Float[Array, ...],
    alpha: Float[Array, ...],
    beta: Float[Array, ...],
    sigma: Float[Array, ...],
    spatial: RectangularDomain,
    temporal: TemporalDomain,
    max_events: int,
) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    r"""Ogata-style spatiotemporal thinning sample.

    Algorithm (per Daley & Vere-Jones, §7.3):

    1. Maintain a running history ``(locs, times, mask)``.
    2. At each step, draw a candidate inter-event time from
       ``Exp(λ̄)`` with adaptive bound :math:`\bar\lambda` from
       :func:`stpp_hawkes_lambda_max`.
    3. Draw a candidate location uniformly on ``D``.
    4. Accept with probability
       :math:`\lambda^*(s, t)/\bar\lambda`; otherwise discard and try
       again.

    Loops via :func:`jax.lax.scan` over a fixed budget of
    ``5 * max_events`` candidates so the function is jit-compatible.

    Args:
        key: PRNG key.
        mu: Baseline rate.
        alpha: Excitation amplitude (sub-critical: ``alpha < 1``).
        beta: Temporal decay.
        sigma: Spatial bandwidth.
        spatial: Spatial domain.
        temporal: Temporal interval.
        max_events: Static buffer cap.

    Returns:
        ``(locations, times, mask, n_events)`` with ``times`` sorted
        ascending. ``n_events`` counts *accepted* events and can never
        exceed ``max_events`` — acceptance stops once the buffer is
        full, so a saturated buffer (``n_events == max_events``, equiv.
        ``mask.sum() == max_events``) signals probable truncation and
        the user should increase ``max_events``. There is no separate
        uncapped attempt count.
    """
    mu_arr = jnp.asarray(mu)
    alpha_arr = jnp.asarray(alpha)
    beta_arr = jnp.asarray(beta)
    sigma_arr = jnp.asarray(sigma)

    d = spatial.n_dims
    n_max = max_events
    candidate_budget = 5 * n_max

    init_locs = jnp.zeros((n_max, d), dtype=mu_arr.dtype)
    init_times = jnp.full((n_max,), temporal.t1, dtype=mu_arr.dtype)
    init_mask = jnp.zeros((n_max,), dtype=jnp.bool_)
    init_count = jnp.asarray(0, dtype=jnp.int32)
    init_t = jnp.asarray(temporal.t0, dtype=mu_arr.dtype)

    keys = random.split(key, candidate_budget)

    box_vol = spatial.volume()

    def step(carry, k_step):
        locs, times, mask, count, t_curr = carry
        k_dt, k_loc, k_acc, k_thin = random.split(k_step, 4)

        # Per-unit-volume peak intensity, valid as an upper bound on
        # ``λ*(s, t)`` for any ``s ∈ D`` from the current step onward.
        # Floor with ``eps`` so that a μ=0, count=0 history does not
        # divide by zero in either ``dt`` or the acceptance ratio
        # (Codex P1 in PR #13 review). The clamp matches the pattern in
        # :func:`thinning_sample`.
        eps = jnp.asarray(1e-30, dtype=mu_arr.dtype)
        lam_bar_per_vol = jnp.maximum(
            stpp_hawkes_lambda_max(count, mu_arr, alpha_arr, beta_arr, sigma_arr, d),
            eps,
        )
        # The inter-event waiting time over the whole domain is
        # exponential with rate ``λ̄·|D|``: integrating the per-volume
        # bound across ``D`` gives the rate at which space-time events
        # arrive anywhere in the box.
        lam_bar_total = jnp.maximum(lam_bar_per_vol * box_vol, eps)
        u = random.uniform(k_dt, minval=1e-30, maxval=1.0)
        dt = -jnp.log(u) / lam_bar_total
        t_next = t_curr + dt

        s_cand = spatial.sample_uniform(k_loc, shape=())
        # Intensity uses the *current* history (events before t_next).
        lam = stpp_hawkes_intensity(
            s_cand,
            t_next,
            locs,
            times,
            mask,
            mu_arr,
            alpha_arr,
            beta_arr,
            sigma_arr,
        )
        u_acc = random.uniform(k_acc, minval=1e-30, maxval=1.0)
        # Acceptance is against the per-volume bound — the candidate
        # location is uniform on ``D`` so the proposal already accounts
        # for the spatial integration in ``λ_bar_total``.
        accept_in_box = u_acc < (lam / lam_bar_per_vol)
        del k_thin

        in_window = t_next < temporal.t1
        room = count < n_max
        accept = accept_in_box & in_window & room

        # Clip the write index to ``[0, n_max - 1]``. Using ``n_max`` as
        # a sentinel for rejected candidates is undefined: in JAX the
        # default ``promise_in_bounds`` semantics make scatter/gather at
        # that index backend-dependent and can corrupt buffer state on
        # rejection-heavy runs (Copilot/Codex P1 in PR #13 review). The
        # ``accept`` flag in ``jnp.where`` already keeps rejected
        # candidates from changing the slot — the clip just keeps the
        # index in range for the no-op read-modify-write.
        safe_idx = jnp.clip(count, 0, n_max - 1)
        locs = locs.at[safe_idx].set(jnp.where(accept, s_cand, locs[safe_idx]))
        times = times.at[safe_idx].set(jnp.where(accept, t_next, times[safe_idx]))
        mask = mask.at[safe_idx].set(jnp.where(accept, True, mask[safe_idx]))
        count = count + accept.astype(jnp.int32)
        # Always advance t_curr by dt — even rejected candidates consume
        # the exponential clock per Ogata's argument.
        t_curr = jnp.minimum(t_next, temporal.t1)
        return (locs, times, mask, count, t_curr), None

    (locs, times, mask, count, _), _ = jax.lax.scan(
        step,
        (init_locs, init_times, init_mask, init_count, init_t),
        keys,
    )

    # Buffer is already filled in event-arrival order; events are
    # naturally time-sorted because the Ogata sampler walks ``t``
    # monotonically forward.
    locations = jnp.where(mask[:, None], locs, spatial.lo)
    times = jnp.where(mask, times, temporal.t1)
    return SpatiotemporalSampleResult(locations, times, mask, count)

stpp_hawkes_lambda_max(n_events: Int[Array, ...], mu: Float[Array, ...], alpha: Float[Array, ...], beta: Float[Array, ...], sigma: Float[Array, ...], n_dims: int) -> Float[Array, ...]

Adaptive upper bound :math:\bar\lambda = \mu + N \alpha \beta k_s(0).

The temporal kernel peaks at :math:\beta and the spatial kernel peaks at :math:(2\pi)^{-d/2} / \sigma^d. For a sub-critical process this bound is finite and tight just after the most recent event; it loosens between events but stays a true upper bound.

Source code in src/xtremax/point_processes/primitives/hawkes_spatiotemporal.py

def stpp_hawkes_lambda_max(
    n_events: Int[Array, ...],
    mu: Float[Array, ...],
    alpha: Float[Array, ...],
    beta: Float[Array, ...],
    sigma: Float[Array, ...],
    n_dims: int,
) -> Float[Array, ...]:
    r"""Adaptive upper bound :math:`\bar\lambda = \mu + N \alpha \beta k_s(0)`.

    The temporal kernel peaks at :math:`\beta` and the spatial kernel
    peaks at :math:`(2\pi)^{-d/2} / \sigma^d`. For a sub-critical
    process this bound is finite and tight just after the most recent
    event; it loosens between events but stays a true upper bound.
    """
    n = jnp.asarray(n_events).astype(jnp.asarray(mu).dtype)
    peak_spatial = (2.0 * jnp.pi) ** (-0.5 * n_dims) / (sigma**n_dims)
    return mu + n * alpha * beta * peak_spatial

Marks

Marked processes factorise as ground process × mark distribution, so the mark contribution to the likelihood is a separate additive term. These functions supply that term and the matching samplers for each event geometry.

Pure-function primitives for temporal and spatial point processes.

Each submodule is free of class state and free of framework dependencies (no NumPyro — save that for the outer-layer wrappers — no equinox). These are the canonical reference implementations; the operator and distribution layers delegate to them.

marks_log_prob(event_times: Float[Array, ...], marks: Float[Array, ...], mask: Bool[Array, ...], mark_distribution_fn: Callable[[Array, EventHistory], dist.Distribution], history_at_each_event: bool = True) -> Float[Array, ...]

Sum of log f(mᵢ | tᵢ, Hᵢ) over real events.

Parameters:

Name	Type	Description	Default
event_times	Float[Array, ...]	Padded event times (max_events,).	required
marks	Float[Array, ...]	Padded marks aligned with event_times. Shape (max_events,) for scalar marks, (max_events, d) for d-dimensional marks.	required
mask	Bool[Array, ...]	Event mask.	required
mark_distribution_fn	Callable[[Array, EventHistory], Distribution]	Callable (t, history) -> Distribution.	required
history_at_each_event	bool	When True (the default) the history passed to mark_distribution_fn(tᵢ, Hᵢ) contains events {t_j : j < i}; when False it is the full observed history — this saves the per-event history construction but is only correct for history-independent mark distributions.	True

Returns:

Type	Description
Float[Array, ...]	Scalar log-likelihood contribution from the marks.

Notes

At padding positions mark_distribution_fn is still called (to keep shapes static) but the padding mark is first replaced with a value inside the distribution's support (:func:safe_padding_mark) and the log-prob is masked out. Without the substitution, bounded-support mark laws (Gamma, LogNormal, ...) return -inf at zero padding and autodiff propagates NaN gradients through the discarded jnp.where branch — the same guard the spatial variants already had.

Source code in src/xtremax/point_processes/primitives/marked.py

def marks_log_prob(
    event_times: Float[Array, ...],
    marks: Float[Array, ...],
    mask: Bool[Array, ...],
    mark_distribution_fn: Callable[[Array, EventHistory], dist.Distribution],
    history_at_each_event: bool = True,
) -> Float[Array, ...]:
    r"""Sum of ``log f(mᵢ | tᵢ, Hᵢ)`` over real events.

    Args:
        event_times: Padded event times ``(max_events,)``.
        marks: Padded marks aligned with ``event_times``. Shape
            ``(max_events,)`` for scalar marks, ``(max_events, d)`` for
            ``d``-dimensional marks.
        mask: Event mask.
        mark_distribution_fn: Callable ``(t, history) -> Distribution``.
        history_at_each_event: When ``True`` (the default) the history
            passed to ``mark_distribution_fn(tᵢ, Hᵢ)`` contains events
            ``{t_j : j < i}``; when ``False`` it is the full observed
            history — this saves the per-event history construction but
            is only correct for history-independent mark distributions.

    Returns:
        Scalar log-likelihood contribution from the marks.

    Notes:
        At padding positions ``mark_distribution_fn`` is still called
        (to keep shapes static) but the padding mark is first replaced
        with a value inside the distribution's support
        (:func:`safe_padding_mark`) and the log-prob is masked out.
        Without the substitution, bounded-support mark laws (Gamma,
        LogNormal, ...) return ``-inf`` at zero padding and autodiff
        propagates NaN gradients through the discarded ``jnp.where``
        branch — the same guard the spatial variants already had.
    """
    event_times = jnp.asarray(event_times)
    marks = jnp.asarray(marks)

    if not history_at_each_event:
        full_history = EventHistory(
            times=event_times, mask=mask, marks=marks if marks.ndim == 2 else None
        )

        def per_event(i: Int[Array, ...]) -> Float[Array, ...]:
            t_i = event_times[i]
            d = mark_distribution_fn(t_i, full_history)
            # Padding marks may sit outside d's support (zero padding vs
            # a positive-support law); substitute an in-support value so
            # the masked-out branch cannot poison gradients with NaN.
            safe_m = jnp.where(mask[i], marks[i], safe_padding_mark(d, marks[i]))
            return d.log_prob(safe_m)

    else:
        n_max = event_times.shape[-1]
        idx_grid = jnp.arange(n_max)
        # A per-event history uses mask[..., j] & (j < i).
        marks_2d = marks if marks.ndim == 2 else marks[..., None]

        def per_event(i: Int[Array, ...]) -> Float[Array, ...]:
            before_i = (idx_grid < i) & mask
            h_i = EventHistory(
                times=jnp.where(before_i, event_times, 0.0),
                mask=before_i,
                marks=jnp.where(before_i[..., None], marks_2d, 0.0),
            )
            d = mark_distribution_fn(event_times[i], h_i)
            # Same in-support substitution as the flat-history branch.
            safe_m = jnp.where(mask[i], marks[i], safe_padding_mark(d, marks[i]))
            return d.log_prob(safe_m)

    idx = jnp.arange(event_times.shape[-1])
    per_event_log_prob = jax.vmap(per_event)(idx)
    # Mark distributions on multi-dim marks return per-draw scalars, on
    # batched marks they return per-component values — sum any trailing dims.
    while per_event_log_prob.ndim > mask.ndim:
        per_event_log_prob = jnp.sum(per_event_log_prob, axis=-1)
    return jnp.sum(jnp.where(mask, per_event_log_prob, 0.0), axis=-1)

sample_marks_at_times(key: PRNGKeyArray, event_times: Float[Array, ...], mask: Bool[Array, ...], mark_distribution_fn: Callable[[Array, EventHistory], dist.Distribution], history_at_each_event: bool = True) -> Float[Array, ...]

Draw marks mᵢ ~ f(· | tᵢ, Hᵢ) at each observed time.

When history_at_each_event=True sampling is sequential: each draw's history pytree includes the previously drawn marks, so mark distributions can legitimately depend on both past times and past marks. When False a single vmap fans the draws out in parallel and the mark distribution is evaluated against the full observed history with marks=None — cheaper, but only correct for mark laws that do not look at prior marks.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	JAX PRNG key.	required
event_times	Float[Array, ...]	Padded event times.	required
mask	Bool[Array, ...]	Event mask.	required
mark_distribution_fn	Callable[[Array, EventHistory], Distribution]	Mark distribution callable.	required
history_at_each_event	bool	See above.	True

Returns:

Type	Description
Float[Array, ...]	Array of marks shaped event_times.shape for scalar marks or
Float[Array, ...]	event_times.shape + (mark_dim,) for vector marks, with
Float[Array, ...]	padding positions zeroed out.

Source code in src/xtremax/point_processes/primitives/marked.py

def sample_marks_at_times(
    key: PRNGKeyArray,
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    mark_distribution_fn: Callable[[Array, EventHistory], dist.Distribution],
    history_at_each_event: bool = True,
) -> Float[Array, ...]:
    """Draw marks ``mᵢ ~ f(· | tᵢ, Hᵢ)`` at each observed time.

    When ``history_at_each_event=True`` sampling is sequential: each
    draw's history pytree includes the previously drawn marks, so mark
    distributions can legitimately depend on both past times and past
    marks. When ``False`` a single ``vmap`` fans the draws out in
    parallel and the mark distribution is evaluated against the full
    observed history with ``marks=None`` — cheaper, but only correct
    for mark laws that do not look at prior marks.

    Args:
        key: JAX PRNG key.
        event_times: Padded event times.
        mask: Event mask.
        mark_distribution_fn: Mark distribution callable.
        history_at_each_event: See above.

    Returns:
        Array of marks shaped ``event_times.shape`` for scalar marks or
        ``event_times.shape + (mark_dim,)`` for vector marks, with
        padding positions zeroed out.
    """
    n_max = event_times.shape[-1]
    keys = random.split(key, n_max)

    if not history_at_each_event:
        full_history = EventHistory(times=event_times, mask=mask, marks=None)

        def per_event(i: Int[Array, ...], k: PRNGKeyArray) -> Array:
            d = mark_distribution_fn(event_times[i], full_history)
            return d.sample(k)

        idx = jnp.arange(n_max)
        marks = jax.vmap(per_event)(idx, keys)
        if marks.ndim == mask.ndim:
            return jnp.where(mask, marks, 0.0)
        return jnp.where(mask[..., None], marks, 0.0)

    # Sequential path: probe the mark dist once on an all-padding history
    # so we know the mark shape/dtype, preallocate a 2-D marks buffer on
    # the carry (scalar marks get a trailing ``1`` to keep the buffer
    # uniformly 2-D), and scan. Each step appends the drawn mark to the
    # carry so subsequent steps can condition on ``history.marks``.
    #
    # The probe history carries a 2-D marks BUFFER — the same pytree
    # structure the scan carry has — rather than ``marks=None``: a
    # history-dependent mark law that reads ``history.marks`` would
    # otherwise see two different structures between probe and scan.
    # The buffer is probed at dimension 1 first and re-probed once the
    # true mark dimension is known (vector-mark laws that index columns
    # of ``history.marks`` fail loudly at the probe, as before).
    def _probe(marks_buffer: Array) -> jax.ShapeDtypeStruct:
        hist = EventHistory(
            times=jnp.zeros_like(event_times),
            mask=jnp.zeros_like(mask),
            marks=marks_buffer,
        )
        return jax.eval_shape(
            lambda et, k: mark_distribution_fn(et, hist).sample(k),
            event_times[0],
            keys[0],
        )

    probe_spec = _probe(jnp.zeros((n_max, 1), dtype=event_times.dtype))
    mark_event_shape = probe_spec.shape
    mark_dtype = probe_spec.dtype
    is_scalar_mark = len(mark_event_shape) == 0
    mark_dim = 1 if is_scalar_mark else mark_event_shape[-1]
    if mark_dim != 1:
        probe_spec = _probe(jnp.zeros((n_max, mark_dim), dtype=mark_dtype))
        mark_dtype = probe_spec.dtype
    initial_carry = EventHistory(
        times=jnp.zeros_like(event_times),
        mask=jnp.zeros_like(mask),
        marks=jnp.zeros((n_max, mark_dim), dtype=mark_dtype),
    )

    def step(
        carry: EventHistory, inp: tuple[Int[Array, ...], PRNGKeyArray]
    ) -> tuple[EventHistory, Array]:
        i, k = inp
        d = mark_distribution_fn(event_times[i], carry)
        m = d.sample(k)
        m_row = jnp.expand_dims(m, axis=-1) if is_scalar_mark else m

        valid_i = mask[i]
        new_times = carry.times.at[i].set(
            jnp.where(valid_i, event_times[i], carry.times[i])
        )
        new_mask = carry.mask.at[i].set(
            jnp.where(valid_i, jnp.bool_(True), carry.mask[i])
        )
        zero_row = jnp.zeros_like(m_row)
        new_marks_row = jnp.where(valid_i, m_row, zero_row)
        new_marks = carry.marks.at[i].set(new_marks_row)
        new_carry = EventHistory(times=new_times, mask=new_mask, marks=new_marks)

        # Emitted per-step mark (in the user's original scalar/vector shape).
        emitted = jnp.where(valid_i, m, jnp.zeros_like(m))
        return new_carry, emitted

    idx = jnp.arange(n_max)
    _, marks = jax.lax.scan(step, initial_carry, (idx, keys))
    return marks

spatial_marks_log_prob(locations: Float[Array, ...], marks: Float[Array, ...], mask: Bool[Array, ...], mark_distribution_fn: Callable[[Array], dist.Distribution]) -> Float[Array, ...]

Sum of log f(mᵢ | sᵢ) over real events.

Parameters:

Name	Type	Description	Default
locations	Float[Array, ...]	Padded locations (max_events, d).	required
marks	Float[Array, ...]	Padded marks. Shape (max_events,) for scalar marks or (max_events, mark_dim) for vector marks.	required
mask	Bool[Array, ...]	Event mask shape (max_events,).	required
mark_distribution_fn	Callable[[Array], Distribution]	Callable s -> Distribution. The distribution must have event_shape == () for scalar marks or event_shape == (mark_dim,) for vector marks.	required

Returns:

Type	Description
Float[Array, ...]	Scalar log-likelihood contribution from marks.

Notes

log_prob is still called at padding positions for shape stability, but padding marks are first replaced with a value inside the distribution's support (:func:_safe_padding_mark). Without this, distributions with bounded support produce -inf log-probs and NaN gradients even though the result is masked out — JAX's autodiff threads through both branches of jnp.where.

Source code in src/xtremax/point_processes/primitives/marked_spatial.py

def spatial_marks_log_prob(
    locations: Float[Array, ...],
    marks: Float[Array, ...],
    mask: Bool[Array, ...],
    mark_distribution_fn: Callable[[Array], dist.Distribution],
) -> Float[Array, ...]:
    """Sum of ``log f(mᵢ | sᵢ)`` over real events.

    Args:
        locations: Padded locations ``(max_events, d)``.
        marks: Padded marks. Shape ``(max_events,)`` for scalar
            marks or ``(max_events, mark_dim)`` for vector marks.
        mask: Event mask shape ``(max_events,)``.
        mark_distribution_fn: Callable ``s -> Distribution``. The
            distribution must have ``event_shape == ()`` for scalar
            marks or ``event_shape == (mark_dim,)`` for vector marks.

    Returns:
        Scalar log-likelihood contribution from marks.

    Notes:
        ``log_prob`` is still called at padding positions for shape
        stability, but padding marks are first replaced with a value
        inside the distribution's support
        (:func:`_safe_padding_mark`). Without this, distributions with
        bounded support produce ``-inf`` log-probs and NaN gradients
        even though the result is masked out — JAX's autodiff threads
        through both branches of ``jnp.where``.
    """
    locations = jnp.asarray(locations)
    marks = jnp.asarray(marks)

    def per_event(s_i: Array, m_i: Array, valid: Array) -> Array:
        d = mark_distribution_fn(s_i)
        # Replace padding mark with a value inside ``d``'s support so
        # ``log_prob`` returns finite values (and finite gradients);
        # the result is then masked out below.
        safe_m = jnp.where(valid, m_i, _safe_padding_mark(d, m_i))
        return d.log_prob(safe_m)

    # Vmap over the leading event axis (the only batched axis of these
    # 1-D / 2-D arrays per the shape contract documented at the
    # module level).
    per_event_log_prob = jax.vmap(per_event)(locations, marks, mask)
    while per_event_log_prob.ndim > mask.ndim:
        per_event_log_prob = jnp.sum(per_event_log_prob, axis=-1)
    return jnp.sum(jnp.where(mask, per_event_log_prob, 0.0), axis=-1)

sample_spatial_marks_at_locations(key: PRNGKeyArray, locations: Float[Array, ...], mask: Bool[Array, ...], mark_distribution_fn: Callable[[Array], dist.Distribution]) -> Float[Array, ...]

Draw marks mᵢ ~ f(· | sᵢ) at every location.

Independent across events, so a single vmap over PRNG splits is sufficient — no scan is needed (unlike the temporal marked PP, which threads history through each step).

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	JAX PRNG key.	required
locations	Float[Array, ...]	Padded locations (max_events, d).	required
mask	Bool[Array, ...]	Event mask (max_events,).	required
mark_distribution_fn	Callable[[Array], Distribution]	Mark distribution callable.	required

Returns:

Type	Description
Float[Array, ...]	Marks aligned with locations. Padding positions are filled
Float[Array, ...]	with a dtype-preserving zero (zeros_like(marks)) so discrete
Float[Array, ...]	/ integer marks keep their dtype. Scalar marks have shape
Float[Array, ...]	(max_events,); vector marks have shape (max_events, mark_dim).

Source code in src/xtremax/point_processes/primitives/marked_spatial.py

def sample_spatial_marks_at_locations(
    key: PRNGKeyArray,
    locations: Float[Array, ...],
    mask: Bool[Array, ...],
    mark_distribution_fn: Callable[[Array], dist.Distribution],
) -> Float[Array, ...]:
    """Draw marks ``mᵢ ~ f(· | sᵢ)`` at every location.

    Independent across events, so a single ``vmap`` over PRNG splits is
    sufficient — no scan is needed (unlike the temporal marked PP,
    which threads history through each step).

    Args:
        key: JAX PRNG key.
        locations: Padded locations ``(max_events, d)``.
        mask: Event mask ``(max_events,)``.
        mark_distribution_fn: Mark distribution callable.

    Returns:
        Marks aligned with ``locations``. Padding positions are filled
        with a dtype-preserving zero (`zeros_like(marks)`) so discrete
        / integer marks keep their dtype. Scalar marks have shape
        ``(max_events,)``; vector marks have shape ``(max_events, mark_dim)``.
    """
    n_max = locations.shape[-2]
    keys = random.split(key, n_max)

    def per_event(s_i: Array, k: PRNGKeyArray) -> Array:
        d = mark_distribution_fn(s_i)
        return d.sample(k)

    marks = jax.vmap(per_event)(locations, keys)
    # ``zeros_like(marks)`` preserves the mark dtype (important for
    # discrete distributions like Categorical / Poisson that produce
    # integer arrays, which a literal ``0.0`` would silently upcast).
    zero = jnp.zeros_like(marks)
    if marks.ndim == mask.ndim:
        return jnp.where(mask, marks, zero)
    return jnp.where(mask[..., None], marks, zero)

spatiotemporal_marks_log_prob(locations: Float[Array, ...], times: Float[Array, ...], marks: Float[Array, ...], mask: Bool[Array, ...], mark_distribution_fn: Callable[[Array, Array], dist.Distribution]) -> Float[Array, ...]

Sum of :math:\log f(m_i \mid s_i, t_i) over real events.

Parameters:

Name	Type	Description	Default
locations	Float[Array, ...]	(max_events, d_space) event locations.	required
times	Float[Array, ...]	(max_events,) event times.	required
marks	Float[Array, ...]	(max_events,) or (max_events, mark_dim) marks.	required
mask	Bool[Array, ...]	(max_events,) boolean event mask.	required
mark_distribution_fn	Callable[[Array, Array], Distribution]	Callable (s, t) -> Distribution.	required

Returns:

Type	Description
Float[Array, ...]	Scalar log-likelihood contribution from marks. Padding rows are
Float[Array, ...]	substituted with a value inside the distribution's support
Float[Array, ...]	before log_prob is called, so distributions with bounded
Float[Array, ...]	support do not produce NaN gradients.

Source code in src/xtremax/point_processes/primitives/marked_spatiotemporal.py

def spatiotemporal_marks_log_prob(
    locations: Float[Array, ...],
    times: Float[Array, ...],
    marks: Float[Array, ...],
    mask: Bool[Array, ...],
    mark_distribution_fn: Callable[[Array, Array], dist.Distribution],
) -> Float[Array, ...]:
    """Sum of :math:`\\log f(m_i \\mid s_i, t_i)` over real events.

    Args:
        locations: ``(max_events, d_space)`` event locations.
        times: ``(max_events,)`` event times.
        marks: ``(max_events,)`` or ``(max_events, mark_dim)`` marks.
        mask: ``(max_events,)`` boolean event mask.
        mark_distribution_fn: Callable ``(s, t) -> Distribution``.

    Returns:
        Scalar log-likelihood contribution from marks. Padding rows are
        substituted with a value inside the distribution's support
        before ``log_prob`` is called, so distributions with bounded
        support do not produce NaN gradients.
    """
    locations = jnp.asarray(locations)
    times = jnp.asarray(times)
    marks = jnp.asarray(marks)

    def per_event(s_i: Array, t_i: Array, m_i: Array, valid: Array) -> Array:
        d = mark_distribution_fn(s_i, t_i)
        safe_m = jnp.where(valid, m_i, _safe_padding_mark(d, m_i))
        return d.log_prob(safe_m)

    per_event_log_prob = jax.vmap(per_event)(locations, times, marks, mask)
    while per_event_log_prob.ndim > mask.ndim:
        per_event_log_prob = jnp.sum(per_event_log_prob, axis=-1)
    return jnp.sum(jnp.where(mask, per_event_log_prob, 0.0), axis=-1)

sample_spatiotemporal_marks(key: PRNGKeyArray, locations: Float[Array, ...], times: Float[Array, ...], mask: Bool[Array, ...], mark_distribution_fn: Callable[[Array, Array], dist.Distribution]) -> Float[Array, ...]

Draw marks :math:m_i \sim f(\cdot \mid s_i, t_i) at every event.

Independent across events, so a single vmap over PRNG splits suffices.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	JAX PRNG key.	required
locations	Float[Array, ...]	(max_events, d_space) locations.	required
times	Float[Array, ...]	(max_events,) times.	required
mask	Bool[Array, ...]	(max_events,) event mask.	required
mark_distribution_fn	Callable[[Array, Array], Distribution]	Callable (s, t) -> Distribution.	required

Returns:

Type	Description
Float[Array, ...]	Marks with the same leading shape as times. Padding rows
Float[Array, ...]	are filled with zeros_like(marks) to preserve the mark
Float[Array, ...]	dtype (important for discrete distributions).

Source code in src/xtremax/point_processes/primitives/marked_spatiotemporal.py

def sample_spatiotemporal_marks(
    key: PRNGKeyArray,
    locations: Float[Array, ...],
    times: Float[Array, ...],
    mask: Bool[Array, ...],
    mark_distribution_fn: Callable[[Array, Array], dist.Distribution],
) -> Float[Array, ...]:
    """Draw marks :math:`m_i \\sim f(\\cdot \\mid s_i, t_i)` at every event.

    Independent across events, so a single ``vmap`` over PRNG splits
    suffices.

    Args:
        key: JAX PRNG key.
        locations: ``(max_events, d_space)`` locations.
        times: ``(max_events,)`` times.
        mask: ``(max_events,)`` event mask.
        mark_distribution_fn: Callable ``(s, t) -> Distribution``.

    Returns:
        Marks with the same leading shape as ``times``. Padding rows
        are filled with ``zeros_like(marks)`` to preserve the mark
        dtype (important for discrete distributions).
    """
    n_max = times.shape[-1]
    keys = random.split(key, n_max)

    def per_event(s_i: Array, t_i: Array, k: PRNGKeyArray) -> Array:
        d = mark_distribution_fn(s_i, t_i)
        return d.sample(k)

    marks = jax.vmap(per_event)(locations, times, keys)
    zero = jnp.zeros_like(marks)
    if marks.ndim == mask.ndim:
        return jnp.where(mask, marks, zero)
    return jnp.where(mask[..., None], marks, zero)

Thinning & observation operators

Thinning models partial observation: a base process generates events and a retention probability decides which are recorded. Useful as an observation operator on top of any ground process.

Pure-function primitives for temporal and spatial point processes.

Each submodule is free of class state and free of framework dependencies (no NumPyro — save that for the outer-layer wrappers — no equinox). These are the canonical reference implementations; the operator and distribution layers delegate to them.

thinning_sample(key: PRNGKeyArray, conditional_intensity_fn: Callable[[Array, EventHistory], Array], lambda_max_fn: Callable[[EventHistory], Array], T: Float[Array, ...], max_events: int, max_candidates: int | None = None, initial_history: EventHistory | None = None) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Sample a history-dependent temporal point process via Ogata thinning.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	JAX PRNG key.	required
conditional_intensity_fn	Callable[[Array, EventHistory], Array]	Callable (t, history) -> λ*(t | H). history is the :class:EventHistory accumulated before t; callers must not peek at anything at or after t.	required
lambda_max_fn	Callable[[EventHistory], Array]	Callable history -> λ_bar. Must return a scalar >= the supremum of conditional_intensity_fn over the remaining window [last_time, T] given the current history. Over-estimates are safe (lose efficiency); under-estimates break the sampler.	required
T	Float[Array, ...]	Observation window length (scalar).	required
max_events	int	Static cap on the number of accepted events.	required
max_candidates	int | None	Static cap on the number of proposals. Defaults to 4 * max_events if None. Choose comfortably above the expected count to avoid truncation under low acceptance.	None
initial_history	EventHistory | None	Optional starting history. Defaults to an empty buffer of size max_events.	None

Returns:

Type	Description
Float[Array, ...]	Tuple (times, mask, n_proposals):
Bool[Array, ...]	times with shape (max_events,) — sorted, padding positions filled with T.
Int[Array, ...]	mask — True at accepted positions.
tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]	n_proposals — number of thinning proposals actually consumed (useful for diagnosing buffer over-runs).

Source code in src/xtremax/point_processes/primitives/thinning.py

def thinning_sample(
    key: PRNGKeyArray,
    conditional_intensity_fn: Callable[[Array, EventHistory], Array],
    lambda_max_fn: Callable[[EventHistory], Array],
    T: Float[Array, ...],
    max_events: int,
    max_candidates: int | None = None,
    initial_history: EventHistory | None = None,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Sample a history-dependent temporal point process via Ogata thinning.

    Args:
        key: JAX PRNG key.
        conditional_intensity_fn: Callable ``(t, history) -> λ*(t | H)``.
            ``history`` is the :class:`EventHistory` accumulated *before*
            ``t``; callers must not peek at anything at or after ``t``.
        lambda_max_fn: Callable ``history -> λ_bar``. Must return a
            scalar ``>=`` the supremum of ``conditional_intensity_fn``
            over the remaining window ``[last_time, T]`` given the
            current history. Over-estimates are safe (lose efficiency);
            under-estimates break the sampler.
        T: Observation window length (scalar).
        max_events: Static cap on the number of accepted events.
        max_candidates: Static cap on the number of proposals. Defaults
            to ``4 * max_events`` if ``None``. Choose comfortably above
            the expected count to avoid truncation under low acceptance.
        initial_history: Optional starting history. Defaults to an
            empty buffer of size ``max_events``.

    Returns:
        Tuple ``(times, mask, n_proposals)``:

        * ``times`` with shape ``(max_events,)`` — sorted, padding
            positions filled with ``T``.
        * ``mask`` — ``True`` at accepted positions.
        * ``n_proposals`` — number of thinning proposals actually
            consumed (useful for diagnosing buffer over-runs).
    """
    T_arr = jnp.asarray(T)
    if max_candidates is None:
        max_candidates = 4 * max_events

    if initial_history is None:
        initial_history = EventHistory.empty(max_events=max_events, dtype=T_arr.dtype)

    keys = random.split(key, max_candidates)

    def step(
        carry: tuple[Float[Array, ...], EventHistory, Int[Array, ...]],
        key_i: PRNGKeyArray,
    ) -> tuple[tuple[Float[Array, ...], EventHistory, Int[Array, ...]], None]:
        t_prev, history, consumed = carry
        key_dt, key_u = random.split(key_i)

        # Guard λ_bar against 0 so ``dt`` can never be ``+inf``; a zero
        # upper bound is only reachable when the user supplies an
        # identically-zero process, in which case a tiny effective rate
        # still overshoots ``T`` on the first step and terminates.
        lambda_bar = jnp.maximum(lambda_max_fn(history), 1e-30)
        # Exp(λ_bar) draw with a log1p-safe form (``u`` is in (0, 1)).
        # Clip both ends away from 0 / 1: ``u == 0`` would give
        # ``dt == 0`` (duplicate times), and ``u == 1`` gives ``-inf``.
        u1 = random.uniform(key_dt, dtype=T_arr.dtype)
        u1 = jnp.clip(u1, 1e-7, 1.0 - 1e-7)
        dt = -jnp.log1p(-u1) / lambda_bar
        t_new = t_prev + dt

        in_window = t_new < T_arr
        lambda_new = conditional_intensity_fn(t_new, history)
        u2 = random.uniform(key_u, dtype=T_arr.dtype)
        accepted = in_window & (u2 * lambda_bar < lambda_new)

        new_history = history.append(t_new, accepted=accepted)
        # Always advance to the new proposal; once past ``T`` subsequent
        # proposals stay outside the window and ``in_window`` stays
        # ``False``, naturally stopping acceptance. Resetting to
        # ``t_prev`` on overshoot would let the sampler re-propose into
        # the window and double-count events near the right edge.
        t_out = t_new
        # Consumed counts proposals actually processed (i.e. within [0, T]).
        new_consumed = consumed + in_window.astype(consumed.dtype)
        return (t_out, new_history, new_consumed), None

    initial = (
        jnp.asarray(0.0, dtype=T_arr.dtype),
        initial_history,
        jnp.asarray(0, dtype=jnp.int32),
    )
    (_, final_history, n_proposals), _ = jax.lax.scan(step, initial, keys)

    # Final pass: pad padding positions with T so downstream intensity
    # evaluations stay inside the window.
    times = jnp.where(final_history.mask, final_history.times, T_arr)
    return SampleResult(times, final_history.mask, n_proposals)

thinning_retention_log_prob(retention_fn: Callable[[Array, EventHistory, Array | None], Array], event_times: Float[Array, ...], mask: Bool[Array, ...], history: EventHistory, marks: Float[Array, ...] | None = None) -> Float[Array, ...]

Sum of log p(tᵢ | Hᵢ, mᵢ) over retained events.

At each event index i the retention callable sees only events {j : j < i} in the history — so the contract :math:p(t_i \mid \mathcal{H}_i, m_i) is honoured exactly, even for history-dependent retention.

Parameters:

Name	Type	Description	Default
retention_fn	Callable[[Array, EventHistory, Array | None], Array]	Retention callable.	required
event_times	Float[Array, ...]	Observed times (padded).	required
mask	Bool[Array, ...]	Observed mask.	required
history	EventHistory	The full observed history. Restricted per-event to the prefix {j : j < i} inside this routine.	required
marks	Float[Array, ...] | None	Optional marks aligned with event_times; passed to retention_fn as proposed_mark at each position.	None

Returns:

Type	Description
Float[Array, ...]	Scalar sum; padding positions contribute 0 via the mask.

Source code in src/xtremax/point_processes/primitives/thinning.py

def thinning_retention_log_prob(
    retention_fn: Callable[[Array, EventHistory, Array | None], Array],
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    history: EventHistory,
    marks: Float[Array, ...] | None = None,
) -> Float[Array, ...]:
    """Sum of ``log p(tᵢ | Hᵢ, mᵢ)`` over retained events.

    At each event index ``i`` the retention callable sees only events
    ``{j : j < i}`` in the history — so the contract
    :math:`p(t_i \\mid \\mathcal{H}_i, m_i)` is honoured exactly, even
    for history-dependent retention.

    Args:
        retention_fn: Retention callable.
        event_times: Observed times (padded).
        mask: Observed mask.
        history: The full observed history. Restricted per-event to the
            prefix ``{j : j < i}`` inside this routine.
        marks: Optional marks aligned with ``event_times``; passed to
            ``retention_fn`` as ``proposed_mark`` at each position.

    Returns:
        Scalar sum; padding positions contribute ``0`` via the mask.
    """
    n = event_times.shape[-1]
    index_grid = jnp.arange(n)

    def eval_at(i: Int[Array, ...]) -> Float[Array, ...]:
        prefix = _prefix_history(history, index_grid < i)
        t_i = event_times[i]
        mark_i = None if marks is None else marks[i]
        return retention_fn(t_i, prefix, mark_i)

    probs = jax.vmap(eval_at)(index_grid)
    log_p = jnp.log(jnp.clip(probs, 1e-30, 1.0))
    return jnp.sum(jnp.where(mask, log_p, 0.0), axis=-1)

retention_compensator(retention_fn: Callable[[Array, EventHistory, Array | None], Array], conditional_intensity_fn: Callable[[Array, EventHistory], Array], history: EventHistory, T: Float[Array, ...], n_points: int = 100, mark_sampler_fn: Callable[[Array, EventHistory, PRNGKeyArray], Array] | None = None, mark_sample_key: PRNGKeyArray | None = None, n_mark_samples: int = 8) -> Float[Array, ...]

Compute the retention compensator via trapezoid quadrature.

:math:\int_0^T (1 - \bar p(t|H(t)))\,\lambda_g(t|H(t))\, dt.

Used by :class:~xtremax.point_processes.operators.ThinningProcess to correct the thinned-process log-likelihood by the expected compensator of the discarded events. At each quadrature node t the retention and intensity callables see only the observed events with time < t, so a history-dependent retention never "peeks" at events in the future of its query point.

For mark-independent retention (the common case), retention_fn is called as retention_fn(t, prefix, None) and :math:\bar p = p. For mark-dependent retention, pass a mark_sampler_fn that draws one mark from the base's mark distribution given (t, prefix, key). The retention is then averaged over n_mark_samples such marks per grid point, yielding a Monte-Carlo estimate of :math:\bar p(t|H) = \mathbb{E}_{m \sim f(\cdot|t,H)}[p(t|H,m)]. Without mark_sampler_fn, mark-dependent retention silently receives None; callers who need the mark-averaged correction must supply the sampler explicitly.

Parameters:

Name	Type	Description	Default
retention_fn	Callable[[Array, EventHistory, Array | None], Array]	Retention callable.	required
conditional_intensity_fn	Callable[[Array, EventHistory], Array]	Intensity callable (the ground intensity when the base is a marked process — see the operator layer).	required
history	EventHistory	Observed (retained) history. Filtered per-t to the prefix {i : times[i] < t} inside this routine.	required
T	Float[Array, ...]	Window upper limit.	required
n_points	int	Trapezoid grid size (static Python int).	100
mark_sampler_fn	Callable[[Array, EventHistory, PRNGKeyArray], Array] | None	Optional (t, prefix, key) -> mark for MC marginalisation over marks in mark-dependent retention.	None
mark_sample_key	PRNGKeyArray | None	JAX key used for mark MC; required when mark_sampler_fn is set.	None
n_mark_samples	int	Number of mark MC samples per grid point.	8

Returns:

Type	Description
Float[Array, ...]	Scalar integral.

Source code in src/xtremax/point_processes/primitives/thinning.py

def retention_compensator(
    retention_fn: Callable[[Array, EventHistory, Array | None], Array],
    conditional_intensity_fn: Callable[[Array, EventHistory], Array],
    history: EventHistory,
    T: Float[Array, ...],
    n_points: int = 100,
    mark_sampler_fn: Callable[[Array, EventHistory, PRNGKeyArray], Array] | None = None,
    mark_sample_key: PRNGKeyArray | None = None,
    n_mark_samples: int = 8,
) -> Float[Array, ...]:
    r"""Compute the retention compensator via trapezoid quadrature.

    :math:`\int_0^T (1 - \bar p(t|H(t)))\,\lambda_g(t|H(t))\, dt`.

    Used by :class:`~xtremax.point_processes.operators.ThinningProcess`
    to correct the thinned-process log-likelihood by the expected
    compensator of the *discarded* events. At each quadrature node
    ``t`` the retention and intensity callables see only the observed
    events with ``time < t``, so a history-dependent retention never
    "peeks" at events in the future of its query point.

    For **mark-independent** retention (the common case), ``retention_fn``
    is called as ``retention_fn(t, prefix, None)`` and :math:`\bar p = p`.
    For **mark-dependent** retention, pass a ``mark_sampler_fn`` that
    draws one mark from the base's mark distribution given
    ``(t, prefix, key)``. The retention is then averaged over
    ``n_mark_samples`` such marks per grid point, yielding a Monte-Carlo
    estimate of
    :math:`\bar p(t|H) = \mathbb{E}_{m \sim f(\cdot|t,H)}[p(t|H,m)]`.
    Without ``mark_sampler_fn``, mark-dependent retention silently
    receives ``None``; callers who need the mark-averaged correction
    must supply the sampler explicitly.

    Args:
        retention_fn: Retention callable.
        conditional_intensity_fn: Intensity callable (the **ground**
            intensity when the base is a marked process — see the
            operator layer).
        history: Observed (retained) history. Filtered per-``t`` to the
            prefix ``{i : times[i] < t}`` inside this routine.
        T: Window upper limit.
        n_points: Trapezoid grid size (static Python int).
        mark_sampler_fn: Optional ``(t, prefix, key) -> mark`` for MC
            marginalisation over marks in mark-dependent retention.
        mark_sample_key: JAX key used for mark MC; required when
            ``mark_sampler_fn`` is set.
        n_mark_samples: Number of mark MC samples per grid point.

    Returns:
        Scalar integral.
    """
    T_arr = jnp.asarray(T)
    grid = jnp.linspace(jnp.zeros_like(T_arr), T_arr, n_points)

    if mark_sampler_fn is None:

        def integrand(t: Array) -> Array:
            prefix = _prefix_history(history, history.times < t)
            # Clip ``p`` into [0, 1] so a badly-behaved retention (e.g.
            # numerical over/undershoot at grid points) cannot flip the
            # compensator sign.
            p = jnp.clip(retention_fn(t, prefix, None), 0.0, 1.0)
            lam = conditional_intensity_fn(t, prefix)
            return (1.0 - p) * lam

        values = jax.vmap(integrand)(grid)
        return jnp.trapezoid(values, grid)

    if mark_sample_key is None:
        raise ValueError(
            "retention_compensator: mark_sampler_fn requires mark_sample_key."
        )
    per_t_keys = random.split(mark_sample_key, n_points)

    def integrand_mc(t: Array, key_t: PRNGKeyArray) -> Array:
        prefix = _prefix_history(history, history.times < t)
        mark_keys = random.split(key_t, n_mark_samples)

        def one_sample(k: PRNGKeyArray) -> Array:
            m = mark_sampler_fn(t, prefix, k)
            return jnp.clip(retention_fn(t, prefix, m), 0.0, 1.0)

        p_bar = jnp.mean(jax.vmap(one_sample)(mark_keys))
        lam = conditional_intensity_fn(t, prefix)
        return (1.0 - p_bar) * lam

    values = jax.vmap(integrand_mc)(grid, per_t_keys)
    return jnp.trapezoid(values, grid)

Diagnostics

Temporal diagnostics rest on the time-rescaling theorem: under the correct model, rescaled inter-event times are \(\mathrm{Exp}(1)\), which turns model checking into a KS test or a QQ plot. Spatial diagnostics instead compare against CSR through Ripley's \(K\) and its variance-stabilised relatives.

Pure-function primitives for temporal and spatial point processes.

Each submodule is free of class state and free of framework dependencies (no NumPyro — save that for the outer-layer wrappers — no equinox). These are the canonical reference implementations; the operator and distribution layers delegate to them.

time_rescaling_residuals(event_times: Float[Array, ...], mask: Bool[Array, ...], cumulative_intensity_fn: Callable[[Array], Array]) -> tuple[Float[Array, ...], Bool[Array, ...]]

Rescaled inter-event times under the compensator.

Parameters:

Name	Type	Description	Default
event_times	Float[Array, ...]	Padded, sorted event times.	required
mask	Bool[Array, ...]	True for real events.	required
cumulative_intensity_fn	Callable[[Array], Array]	Callable computing :math:\Lambda(t) pointwise. For IPP this will typically be a closure around :func:cumulative_log_intensity; for HPP it is simply λ t.	required

Returns:

Type	Description
Float[Array, ...]	Tuple (residuals, residual_mask):
Bool[Array, ...]	residuals with shape equal to event_times — where position i holds :math:\Lambda(t_i) - \Lambda(t_{i-1}) with :math:t_0 = 0. Padding positions are 0.
tuple[Float[Array, ...], Bool[Array, ...]]	residual_mask — True for positions that correspond to real events. Same shape as mask.

Notes

Residuals at padding positions are deliberately returned as 0 (rather than NaN) so downstream reductions that forget to mask do not propagate NaNs — they will instead see a biased statistic. All consumers in this module apply the mask.

The previous-event compensator is taken from the previous valid event, not the previous buffer slot, so the function is robust to hole masks (a masked-out slot between two real events no longer shrinks the following residual). Samplers in this package emit contiguous-prefix masks (see the package docstring), but diagnostics stay correct either way.

Source code in src/xtremax/point_processes/primitives/diagnostics.py

def time_rescaling_residuals(
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    cumulative_intensity_fn: Callable[[Array], Array],
) -> tuple[Float[Array, ...], Bool[Array, ...]]:
    """Rescaled inter-event times under the compensator.

    Args:
        event_times: Padded, sorted event times.
        mask: ``True`` for real events.
        cumulative_intensity_fn: Callable computing :math:`\\Lambda(t)`
            pointwise. For IPP this will typically be a closure around
            :func:`cumulative_log_intensity`; for HPP it is simply
            ``λ t``.

    Returns:
        Tuple ``(residuals, residual_mask)``:

        * ``residuals`` with shape equal to ``event_times`` — where
            position ``i`` holds :math:`\\Lambda(t_i) - \\Lambda(t_{i-1})`
            with :math:`t_0 = 0`. Padding positions are ``0``.
        * ``residual_mask`` — ``True`` for positions that correspond to
            real events. Same shape as ``mask``.

    Notes:
        Residuals at padding positions are deliberately returned as
        ``0`` (rather than ``NaN``) so downstream reductions that
        forget to mask do not propagate NaNs — they will instead see a
        biased statistic. All consumers in this module apply the mask.

        The previous-event compensator is taken from the previous
        *valid* event, not the previous buffer slot, so the function is
        robust to hole masks (a masked-out slot between two real events
        no longer shrinks the following residual). Samplers in this
        package emit contiguous-prefix masks (see the package
        docstring), but diagnostics stay correct either way.
    """
    Lambda_t = cumulative_intensity_fn(event_times)
    # Λ at the previous VALID event: running max of the masked Λ values
    # (times are sorted, so Λ is monotone over valid slots and the
    # cumulative max equals the last valid Λ), shifted one slot right,
    # floored at Λ(0) = 0.
    Lambda_valid = jnp.where(mask, Lambda_t, -jnp.inf)
    # lax.cummax requires a non-negative axis index.
    cummax = jax.lax.cummax(Lambda_valid, axis=Lambda_valid.ndim - 1)
    neg_inf = jnp.full_like(cummax[..., :1], -jnp.inf)
    shifted = jnp.concatenate([neg_inf, cummax[..., :-1]], axis=-1)
    Lambda_prev = jnp.maximum(shifted, 0.0)
    residuals = Lambda_t - Lambda_prev
    residuals = jnp.where(mask, residuals, 0.0)
    return residuals, mask

compensator_curve(event_times: Float[Array, ...], mask: Bool[Array, ...], cumulative_intensity_fn: Callable[[Array], Array]) -> tuple[Float[Array, ...], Float[Array, ...]]

Pairs (t_i, Λ(t_i)) for a compensator plot.

Under the null the points lie along the line :math:y = \Lambda(t_i) with slope 1 when plotted against the event index i — the purpose of the plot is to see deviations from that line.

Parameters:

Name	Type	Description	Default
event_times	Float[Array, ...]	Padded event times.	required
mask	Bool[Array, ...]	Event mask.	required
cumulative_intensity_fn	Callable[[Array], Array]	Compensator callable.	required

Returns:

Type	Description
Float[Array, ...]	Tuple (times, compensator) where padding positions are
Float[Array, ...]	replaced with NaN for plot-friendliness.

Source code in src/xtremax/point_processes/primitives/diagnostics.py

def compensator_curve(
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    cumulative_intensity_fn: Callable[[Array], Array],
) -> tuple[Float[Array, ...], Float[Array, ...]]:
    """Pairs ``(t_i, Λ(t_i))`` for a compensator plot.

    Under the null the points lie along the line
    :math:`y = \\Lambda(t_i)` with slope 1 when plotted against the
    event index ``i`` — the purpose of the plot is to see deviations
    from that line.

    Args:
        event_times: Padded event times.
        mask: Event mask.
        cumulative_intensity_fn: Compensator callable.

    Returns:
        Tuple ``(times, compensator)`` where padding positions are
        replaced with ``NaN`` for plot-friendliness.
    """
    Lambda_t = cumulative_intensity_fn(event_times)
    times = jnp.where(mask, event_times, jnp.nan)
    cum = jnp.where(mask, Lambda_t, jnp.nan)
    return times, cum

ks_statistic_exp1(residuals: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Two-sided Kolmogorov–Smirnov statistic against Exp(1).

Parameters:

Name	Type	Description	Default
residuals	Float[Array, ...]	Time-rescaled residuals.	required
mask	Bool[Array, ...]	True at real-event positions.	required

Returns:

Type	Description
Float[Array, ...]	math:\max_i |F_n(\tau_{(i)}) - (1 - e^{-\tau_{(i)}})| where
Float[Array, ...]	math:F_n is the empirical CDF computed from the unmasked
Float[Array, ...]	residuals only.

Notes

With padding present this requires a little care: we sort residuals with padding positions pushed to +inf so they land at the top of the sorted order, and compute empirical ranks using the per-batch real count. Padding contributes no mass to the empirical CDF.

Source code in src/xtremax/point_processes/primitives/diagnostics.py

def ks_statistic_exp1(
    residuals: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Two-sided Kolmogorov–Smirnov statistic against ``Exp(1)``.

    Args:
        residuals: Time-rescaled residuals.
        mask: ``True`` at real-event positions.

    Returns:
        :math:`\\max_i |F_n(\\tau_{(i)}) - (1 - e^{-\\tau_{(i)}})|` where
        :math:`F_n` is the empirical CDF computed from the unmasked
        residuals only.

    Notes:
        With padding present this requires a little care: we sort
        residuals with padding positions pushed to ``+inf`` so they
        land at the top of the sorted order, and compute empirical
        ranks using the per-batch real count. Padding contributes no
        mass to the empirical CDF.
    """
    n_real = jnp.sum(mask, axis=-1)
    # Push padding to +inf so real residuals occupy ranks 0..n_real-1.
    sort_keys = jnp.where(mask, residuals, jnp.inf)
    sorted_res = jnp.sort(sort_keys, axis=-1)

    max_len = residuals.shape[-1]
    ranks = jnp.arange(1, max_len + 1, dtype=residuals.dtype)
    n_broadcast = jnp.expand_dims(n_real.astype(residuals.dtype), axis=-1)
    inv_n = 1.0 / jnp.maximum(n_broadcast, 1.0)

    theoretical = 1.0 - jnp.exp(-sorted_res)
    # Two-sided one-sample KS: at rank i the empirical CDF jumps from
    # (i-1)/n to i/n, so the supremum of |F_n(x) - F(x)| can land on
    # either side of the jump. Taking only ``|i/n - F|`` (as the first
    # version did) misses the ``F - (i-1)/n`` branch and under-reports
    # the distance when the theoretical CDF lands just below a step.
    upper = ranks * inv_n - theoretical
    lower = theoretical - (ranks - 1.0) * inv_n
    diff = jnp.maximum(upper, lower)
    valid = ranks <= n_broadcast
    return jnp.max(jnp.where(valid, diff, 0.0), axis=-1)

qq_exp1_quantiles(residuals: Float[Array, ...], mask: Bool[Array, ...]) -> tuple[Float[Array, ...], Float[Array, ...]]

Theoretical vs empirical quantiles against Exp(1) for a QQ plot.

Parameters:

Name	Type	Description	Default
residuals	Float[Array, ...]	Time-rescaling residuals.	required
mask	Bool[Array, ...]	True at real events.	required

Returns:

Type	Description
Float[Array, ...]	Tuple (theoretical, empirical) of shape equal to
Float[Array, ...]	residuals. Positions with rank > n_real are filled
tuple[Float[Array, ...], Float[Array, ...]]	with NaN so a plotter can detect and drop them.

Source code in src/xtremax/point_processes/primitives/diagnostics.py

def qq_exp1_quantiles(
    residuals: Float[Array, ...],
    mask: Bool[Array, ...],
) -> tuple[Float[Array, ...], Float[Array, ...]]:
    """Theoretical vs empirical quantiles against ``Exp(1)`` for a QQ plot.

    Args:
        residuals: Time-rescaling residuals.
        mask: ``True`` at real events.

    Returns:
        Tuple ``(theoretical, empirical)`` of shape equal to
        ``residuals``. Positions with ``rank > n_real`` are filled
        with ``NaN`` so a plotter can detect and drop them.
    """
    n_real = jnp.sum(mask, axis=-1)
    sort_keys = jnp.where(mask, residuals, jnp.inf)
    sorted_res = jnp.sort(sort_keys, axis=-1)

    max_len = residuals.shape[-1]
    ranks = jnp.arange(1, max_len + 1, dtype=residuals.dtype)
    n_broadcast = jnp.expand_dims(n_real.astype(residuals.dtype), axis=-1)
    # Plotting positions (i - 0.5) / n avoid the F⁻¹(1) = ∞ edge case.
    plot_positions = (ranks - 0.5) / jnp.maximum(n_broadcast, 1.0)
    theoretical = -jnp.log1p(-plot_positions)

    valid = ranks <= n_broadcast
    theoretical = jnp.where(valid, theoretical, jnp.nan)
    empirical = jnp.where(valid, sorted_res, jnp.nan)
    return theoretical, empirical

csr_ripleys_k(r: Float[Array, ...], n_dims: int) -> Float[Array, ...]

Theoretical Ripley's K-function under CSR, :math:K(r) = V_d r^d.

Special-cased for the common dimensions to keep the formulas legible: :math:2r (1D), :math:\pi r^2 (2D), :math:\frac{4}{3}\pi r^3 (3D). Higher dimensions fall through to the generic V_d r^d form.

Source code in src/xtremax/point_processes/primitives/diagnostics_spatial.py

def csr_ripleys_k(
    r: Float[Array, ...],
    n_dims: int,
) -> Float[Array, ...]:
    """Theoretical Ripley's K-function under CSR, :math:`K(r) = V_d r^d`.

    Special-cased for the common dimensions to keep the formulas
    legible: :math:`2r` (1D), :math:`\\pi r^2` (2D),
    :math:`\\frac{4}{3}\\pi r^3` (3D). Higher dimensions fall through
    to the generic ``V_d r^d`` form.
    """
    r = jnp.asarray(r)
    d = n_dims
    if d == 1:
        return 2.0 * r
    if d == 2:
        return jnp.pi * r**2
    if d == 3:
        return (4.0 / 3.0) * jnp.pi * r**3
    return unit_ball_volume(d) * r**d

csr_l_function(r: Float[Array, ...], n_dims: int) -> Float[Array, ...]

Besag's L-function, the variance-stabilised K.

:math:L(r) = (K(r) / V_d)^{1/d}. Under CSR, L(r) = r identically, so any departure of the empirical estimator from the diagonal is a sign of clustering or regularity.

Source code in src/xtremax/point_processes/primitives/diagnostics_spatial.py

def csr_l_function(
    r: Float[Array, ...],
    n_dims: int,
) -> Float[Array, ...]:
    """Besag's L-function, the variance-stabilised K.

    :math:`L(r) = (K(r) / V_d)^{1/d}`. Under CSR, ``L(r) = r``
    identically, so any departure of the empirical estimator from the
    diagonal is a sign of clustering or regularity.
    """
    K = csr_ripleys_k(r, n_dims)
    v_d = unit_ball_volume(n_dims)
    return (K / v_d) ** (1.0 / float(n_dims))

csr_pair_correlation(r: Float[Array, ...], n_dims: int) -> Float[Array, ...]

Pair correlation function under CSR: g(r) ≡ 1.

Source code in src/xtremax/point_processes/primitives/diagnostics_spatial.py

def csr_pair_correlation(
    r: Float[Array, ...],
    n_dims: int,
) -> Float[Array, ...]:
    """Pair correlation function under CSR: ``g(r) ≡ 1``."""
    del n_dims
    r = jnp.asarray(r)
    return jnp.ones_like(r)

unit_ball_volume(n_dims: int) -> float

Volume of the unit ℝᵈ ball: :math:V_d = \pi^{d/2}/\Gamma(d/2+1).

Source code in src/xtremax/point_processes/primitives/diagnostics_spatial.py

def unit_ball_volume(n_dims: int) -> float:
    """Volume of the unit ``ℝᵈ`` ball: :math:`V_d = \\pi^{d/2}/\\Gamma(d/2+1)`."""
    d = float(n_dims)
    return float((jnp.pi ** (d / 2.0)) / jax.scipy.special.gamma(d / 2.0 + 1.0))