Point Processes — Operators

Immutable equinox.Module process objects that bundle an intensity model with sampling. Operators are pytrees, safe under jit / grad / vmap, and delegate all math to the primitives — so an operator is a convenience wrapper, never a separate implementation.

Use these when you want a process as a reusable object (fit it, sample from it, check it); use the primitives when you want a single function inside a larger model.

Temporal processes

The full temporal family, from constant-rate Poisson through self-exciting Hawkes. These are the operators that inherit the goodness-of-fit mixin, so each exposes residuals, goodness_of_fit, and compensator_curve built on the time-rescaling theorem.

Equinox-based temporal and spatial point process operators.

HomogeneousPoissonProcess(rate: ArrayLike, observation_window: ArrayLike)

Bases: GoodnessOfFitMixin, Module

Homogeneous Poisson process on [0, T].

Parameters:

Name	Type	Description	Default
rate	ArrayLike	Intensity λ > 0. Stored as a JAX array.	required
observation_window	ArrayLike	Window length T > 0.	required

Source code in src/xtremax/point_processes/operators/temporal.py

def __init__(
    self,
    rate: ArrayLike,
    observation_window: ArrayLike,
) -> None:
    self.rate = jnp.asarray(rate)
    self.observation_window = jnp.asarray(observation_window)

log_prob(n_events: Int[Array, ...]) -> Float[Array, ...]

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

Source code in src/xtremax/point_processes/operators/temporal.py

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

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

Draw a padded event sequence.

Returns the same (times, mask, n_events) triple as :func:hpp_sample.

Source code in src/xtremax/point_processes/operators/temporal.py

def sample(
    self,
    key: PRNGKeyArray,
    max_events: int,
    sample_shape: tuple[int, ...] = (),
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Draw a padded event sequence.

    Returns the same ``(times, mask, n_events)`` triple as
    :func:`hpp_sample`.
    """
    return hpp_sample(
        key, self.rate, self.observation_window, max_events, sample_shape
    )

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

Pointwise intensity (constant, equal to rate).

Source code in src/xtremax/point_processes/operators/temporal.py

def intensity(self, t: Float[Array, ...]) -> Float[Array, ...]:
    """Pointwise intensity (constant, equal to ``rate``)."""
    return hpp_intensity(t, self.rate)

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

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

Source code in src/xtremax/point_processes/operators/temporal.py

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

survival(t: Float[Array, ...], given_time: Float[Array, ...] | float = 0.0) -> Float[Array, ...]

:math:S(t|s) = e^{-\lambda(t - s)}.

Source code in src/xtremax/point_processes/operators/temporal.py

def survival(
    self,
    t: Float[Array, ...],
    given_time: Float[Array, ...] | float = 0.0,
) -> Float[Array, ...]:
    """:math:`S(t|s) = e^{-\\lambda(t - s)}`."""
    return hpp_survival(t, given_time, self.rate)

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

Constant hazard equal to rate.

Source code in src/xtremax/point_processes/operators/temporal.py

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

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

Log density of inter-event delays τ ~ Exp(λ).

Source code in src/xtremax/point_processes/operators/temporal.py

def inter_event_log_prob(
    self,
    tau: Float[Array, ...],
) -> Float[Array, ...]:
    """Log density of inter-event delays ``τ ~ Exp(λ)``."""
    return hpp_inter_event_log_prob(tau, self.rate)

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

Memorylessness: :math:1/\lambda regardless of given_time.

Source code in src/xtremax/point_processes/operators/temporal.py

def mean_residual_life(
    self,
    given_time: Float[Array, ...] | None = None,
) -> Float[Array, ...]:
    """Memorylessness: :math:`1/\\lambda` regardless of ``given_time``."""
    return hpp_mean_residual_life(self.rate, given_time)

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

Expected count on [start, end].

Source code in src/xtremax/point_processes/operators/temporal.py

def predict_count(
    self,
    start_time: Float[Array, ...],
    end_time: Float[Array, ...],
) -> Float[Array, ...]:
    """Expected count on ``[start, end]``."""
    return hpp_predict_count(start_time, end_time, self.rate)

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

Time :math:\tau with :math:P(\text{event within } \tau) = p.

Source code in src/xtremax/point_processes/operators/temporal.py

def return_period(
    self,
    probability: Float[Array, ...],
) -> Float[Array, ...]:
    """Time :math:`\\tau` with :math:`P(\\text{event within } \\tau) = p`."""
    return hpp_return_period(probability, self.rate)

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

Log-probability of more than k events on [start, end].

Source code in src/xtremax/point_processes/operators/temporal.py

def exceedance_log_prob(
    self,
    k: Int[Array, ...],
    start_time: Float[Array, ...],
    end_time: Float[Array, ...],
) -> Float[Array, ...]:
    """Log-probability of more than ``k`` events on ``[start, end]``."""
    return hpp_exceedance_log_prob(k, start_time, end_time, self.rate)

InhomogeneousPoissonProcess(log_intensity_fn: Callable[[Array], Array], observation_window: ArrayLike, integrated_intensity: ArrayLike | None = None, lambda_max: ArrayLike | None = None, n_integration_points: int = 100)

Bases: GoodnessOfFitMixin, LiveIntensityMixin, Module

Inhomogeneous Poisson process on [0, T].

Parameters:

Name	Type	Description	Default
log_intensity_fn	Callable[[Array], Array]	Callable mapping an array of times to an array of log-intensities (same shape). If this is itself an eqx.Module its parameters are part of the PyTree.	required
observation_window	ArrayLike	Window length T > 0.	required
integrated_intensity	ArrayLike | None	:math:\Lambda(T). Pass a fixed value only if the intensity is static; during training this cache would go stale. Otherwise leave as None and the operator will compute Λ(T) from the current log_intensity_fn on every call — exactly when the intensity module exposes .integrate(a, b) (:class:PiecewiseConstantLogIntensity does), via quadrature otherwise.	None
lambda_max	ArrayLike | None	Upper bound on :math:\lambda(t). Same staleness caveat as integrated_intensity: pass a fixed bound only if you will not mutate the intensity. When None, the operator reads the live bound from log_intensity_fn.max_intensity() at sample time if the module exposes it; otherwise :meth:sample raises.	None
n_integration_points	int	Static grid size used by every method that falls back to quadrature. Override once here rather than plumbing kwargs through every call.	100

Note

log_intensity_fn is stored as a regular PyTree leaf — not a static field. This way, when the function is itself an eqx.Module (e.g. an MLP), its array leaves participate in the PyTree and are differentiable through :meth:log_prob. A plain Python callable stored as a leaf makes the operator incompatible with raw jax.jit / jax.vmap (functions are not valid traced leaves) — use eqx.filter_jit / eqx.filter_vmap or wrap the callable in an eqx.Module.

Source code in src/xtremax/point_processes/operators/temporal.py

def __init__(
    self,
    log_intensity_fn: Callable[[Array], Array],
    observation_window: ArrayLike,
    integrated_intensity: ArrayLike | None = None,
    lambda_max: ArrayLike | None = None,
    n_integration_points: int = 100,
) -> None:
    self.log_intensity_fn = log_intensity_fn
    self.observation_window = jnp.asarray(observation_window)
    self.integrated_intensity = (
        None if integrated_intensity is None else jnp.asarray(integrated_intensity)
    )
    self.lambda_max = None if lambda_max is None else jnp.asarray(lambda_max)
    self.n_integration_points = n_integration_points

from_fn(log_intensity_fn: Callable[[Array], Array], observation_window: Float[Array, ...], **kwargs) -> InhomogeneousPoissonProcess classmethod

Constructor alias for clarity at call sites.

Source code in src/xtremax/point_processes/operators/temporal.py

@classmethod
def from_fn(
    cls,
    log_intensity_fn: Callable[[Array], Array],
    observation_window: Float[Array, ...],
    **kwargs,
) -> InhomogeneousPoissonProcess:
    """Constructor alias for clarity at call sites."""
    return cls(log_intensity_fn, observation_window, **kwargs)

from_piecewise_constant(bin_edges: ArrayLike, rates: ArrayLike, **kwargs: Any) -> InhomogeneousPoissonProcess classmethod

Build an IPP with piecewise-constant intensity.

Parameters:

Name	Type	Description	Default
bin_edges	ArrayLike	Sorted array of length E with bin_edges[0] = 0 and bin_edges[-1] = T.	required
rates	ArrayLike	Array of length E - 1 giving the rate on each bin. Must be strictly positive.	required
**kwargs	Any	Forwarded to :meth:__init__ (e.g. n_integration_points).	{}

Returns:

Type	Description
InhomogeneousPoissonProcess	Operator whose log_intensity_fn is a
InhomogeneousPoissonProcess	class:PiecewiseConstantLogIntensity module — its
InhomogeneousPoissonProcess	bin_edges and log_rates are real PyTree leaves,
InhomogeneousPoissonProcess	so you can differentiate through them with
InhomogeneousPoissonProcess	eqx.filter_grad. integrated_intensity is the exact
InhomogeneousPoissonProcess	sum of rates × bin_widths, and lambda_max is
InhomogeneousPoissonProcess	max(rates).

Source code in src/xtremax/point_processes/operators/temporal.py

@classmethod
def from_piecewise_constant(
    cls,
    bin_edges: ArrayLike,
    rates: ArrayLike,
    **kwargs: Any,
) -> InhomogeneousPoissonProcess:
    """Build an IPP with piecewise-constant intensity.

    Args:
        bin_edges: Sorted array of length ``E`` with
            ``bin_edges[0] = 0`` and ``bin_edges[-1] = T``.
        rates: Array of length ``E - 1`` giving the rate on each
            bin. Must be strictly positive.
        **kwargs: Forwarded to :meth:`__init__` (e.g.
            ``n_integration_points``).

    Returns:
        Operator whose ``log_intensity_fn`` is a
        :class:`PiecewiseConstantLogIntensity` module — its
        ``bin_edges`` and ``log_rates`` are real PyTree leaves,
        so you can differentiate through them with
        ``eqx.filter_grad``. ``integrated_intensity`` is the exact
        sum of ``rates × bin_widths``, and ``lambda_max`` is
        ``max(rates)``.
    """
    bin_edges_arr = jnp.asarray(bin_edges)
    rates_arr = jnp.asarray(rates)

    log_intensity_fn = PiecewiseConstantLogIntensity(
        bin_edges=bin_edges_arr, rates=rates_arr
    )

    # Leave ``integrated_intensity`` / ``lambda_max`` as ``None`` so
    # the operator queries the live values from the intensity
    # module on every call. This is what lets gradient-based
    # training of ``log_rates`` see a correct ``-Λ(T)`` term and
    # keeps the thinning envelope tight after updates.
    return cls(
        log_intensity_fn,
        observation_window=bin_edges_arr[-1],
        integrated_intensity=None,
        lambda_max=None,
        **kwargs,
    )

effective_integrated_intensity(a: ArrayLike | None = None, b: ArrayLike | None = None) -> Float[Array, ...]

Return :math:\Lambda(b) - \Lambda(a) using live leaves.

Precedence: (1) use :attr:integrated_intensity if the caller pinned a fixed endpoint-to-endpoint value and is asking for the default [0, T] interval; (2) otherwise use log_intensity_fn.integrate(a, b) when the module exposes it (exact closed form, differentiable); (3) otherwise fall back to quadrature with n_integration_points nodes. Never reads a cached value for non-default intervals, since the cache holds only :math:\Lambda(T).

Parameters:

Name	Type	Description	Default
a	ArrayLike | None	Lower limit (default 0).	None
b	ArrayLike | None	Upper limit (default :attr:observation_window).	None

Source code in src/xtremax/point_processes/operators/temporal.py

def effective_integrated_intensity(
    self,
    a: ArrayLike | None = None,
    b: ArrayLike | None = None,
) -> Float[Array, ...]:
    """Return :math:`\\Lambda(b) - \\Lambda(a)` using live leaves.

    Precedence: (1) use :attr:`integrated_intensity` if the caller
    pinned a fixed endpoint-to-endpoint value and is asking for the
    default ``[0, T]`` interval; (2) otherwise use
    ``log_intensity_fn.integrate(a, b)`` when the module exposes it
    (exact closed form, differentiable); (3) otherwise fall back to
    quadrature with ``n_integration_points`` nodes. Never reads a
    cached value for non-default intervals, since the cache holds
    only :math:`\\Lambda(T)`.

    Args:
        a: Lower limit (default ``0``).
        b: Upper limit (default :attr:`observation_window`).
    """
    a_arr = jnp.zeros_like(self.observation_window) if a is None else jnp.asarray(a)
    b_arr = self.observation_window if b is None else jnp.asarray(b)
    default_interval = a is None and b is None

    if default_interval and self.integrated_intensity is not None:
        return self.integrated_intensity

    integrate = getattr(self.log_intensity_fn, "integrate", None)
    if integrate is not None:
        live = integrate(a_arr, b_arr)
    else:
        live = ipp_predict_count(
            a_arr,
            b_arr,
            self.log_intensity_fn,
            n_points=self.n_integration_points,
        )
    if self.integrated_intensity is not None:
        # Honour the pin for explicit full-window endpoints (e.g.
        # ``cumulative_hazard(op.observation_window)``) so those
        # queries agree with the Λ(T) used by log_prob. The
        # selection is an elementwise ``where`` — jit/vmap-safe and
        # shape-preserving for batched endpoints, where a Python
        # bool guard would either trace-fail or collapse the batch
        # to the scalar pin.
        is_full_window = (a_arr == 0.0) & (b_arr == self.observation_window)
        return jnp.where(is_full_window, self.integrated_intensity, live)
    return live

log_prob(event_times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Log-likelihood :math:\sum_i \log \lambda(t_i) - \Lambda(T).

Uses :meth:effective_integrated_intensity so that gradients through the intensity module's parameters flow into both the sum-of-log-intensities and the :math:-\Lambda(T) term.

Source code in src/xtremax/point_processes/operators/temporal.py

def log_prob(
    self,
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Log-likelihood :math:`\\sum_i \\log \\lambda(t_i) - \\Lambda(T)`.

    Uses :meth:`effective_integrated_intensity` so that gradients
    through the intensity module's parameters flow into both the
    sum-of-log-intensities and the :math:`-\\Lambda(T)` term.
    """
    Lambda_T = self.effective_integrated_intensity()
    return ipp_log_prob(event_times, mask, self.log_intensity_fn, Lambda_T)

sample(key: PRNGKeyArray, max_candidates: int) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Thinning-based sampler.

Uses :meth:effective_lambda_max so that, after an optimiser update to the intensity module's parameters, the thinning envelope still exceeds the live λ(t) and the sampler remains valid.

Source code in src/xtremax/point_processes/operators/temporal.py

def sample(
    self,
    key: PRNGKeyArray,
    max_candidates: int,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Thinning-based sampler.

    Uses :meth:`effective_lambda_max` so that, after an optimiser
    update to the intensity module's parameters, the thinning
    envelope still exceeds the live ``λ(t)`` and the sampler
    remains valid.
    """
    lam_max = self.effective_lambda_max()
    return ipp_sample_thinning(
        key,
        self.log_intensity_fn,
        self.observation_window,
        lam_max,
        max_candidates,
    )

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

Exact sampler when the inverse compensator is available.

Source code in src/xtremax/point_processes/operators/temporal.py

def sample_inversion(
    self,
    key: PRNGKeyArray,
    inverse_cumulative_intensity_fn: Callable[[Array], Array],
    max_events: int,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Exact sampler when the inverse compensator is available."""
    Lambda_T = self.effective_integrated_intensity()
    return ipp_sample_inversion(
        key, inverse_cumulative_intensity_fn, Lambda_T, max_events
    )

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

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

Source code in src/xtremax/point_processes/operators/temporal.py

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

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

Compensator :math:\Lambda(t) = \int_0^t \lambda(s) ds.

Prefers the intensity module's closed-form .integrate when available, falling back to quadrature otherwise. This keeps diagnostics (residuals, compensator curve) consistent with :meth:log_prob after the intensity parameters have been updated.

Source code in src/xtremax/point_processes/operators/temporal.py

def cumulative_intensity(self, t: Float[Array, ...]) -> Float[Array, ...]:
    """Compensator :math:`\\Lambda(t) = \\int_0^t \\lambda(s) ds`.

    Prefers the intensity module's closed-form ``.integrate`` when
    available, falling back to quadrature otherwise. This keeps
    diagnostics (residuals, compensator curve) consistent with
    :meth:`log_prob` after the intensity parameters have been
    updated.
    """
    t = jnp.asarray(t)
    integrate = getattr(self.log_intensity_fn, "integrate", None)
    if integrate is not None:
        zeros = jnp.zeros_like(t)
        # Vectorise the module's ``integrate(a, b)`` over ``t``
        # whenever it is a non-scalar — the closed form accepts
        # scalars in its native form.
        if t.ndim == 0:
            return integrate(zeros, t)
        flat_t = t.reshape(-1)
        flat_zeros = jnp.zeros_like(flat_t)
        vals = jax.vmap(integrate)(flat_zeros, flat_t)
        return vals.reshape(t.shape)
    return ipp_cumulative_intensity(
        t, self.log_intensity_fn, n_points=self.n_integration_points
    )

survival(t: Float[Array, ...], given_time: Float[Array, ...] | float = 0.0) -> Float[Array, ...]

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

Routed through :meth:effective_integrated_intensity so the result agrees with :meth:log_prob (closed-form .integrate when available; quadrature otherwise).

Source code in src/xtremax/point_processes/operators/temporal.py

def survival(
    self,
    t: Float[Array, ...],
    given_time: Float[Array, ...] | float = 0.0,
) -> Float[Array, ...]:
    """:math:`S(t|s) = \\exp(-[\\Lambda(t) - \\Lambda(s)])`.

    Routed through :meth:`effective_integrated_intensity` so the
    result agrees with :meth:`log_prob` (closed-form ``.integrate``
    when available; quadrature otherwise).
    """
    return jnp.exp(-self.cumulative_hazard(t, given_time))

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

Hazard equals intensity for an IPP.

Source code in src/xtremax/point_processes/operators/temporal.py

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

cumulative_hazard(t: Float[Array, ...], given_time: Float[Array, ...] | float = 0.0) -> Float[Array, ...]

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

Routed through :meth:effective_integrated_intensity (pinned → closed-form .integrate → quadrature) so every integral- consuming method uses the same integrator as :meth:log_prob.

Source code in src/xtremax/point_processes/operators/temporal.py

def cumulative_hazard(
    self,
    t: Float[Array, ...],
    given_time: Float[Array, ...] | float = 0.0,
) -> Float[Array, ...]:
    """:math:`\\Lambda(t) - \\Lambda(s)`.

    Routed through :meth:`effective_integrated_intensity` (pinned →
    closed-form ``.integrate`` → quadrature) so every integral-
    consuming method uses the same integrator as :meth:`log_prob`.
    """
    return self.effective_integrated_intensity(given_time, t)

inter_event_log_prob(tau: Float[Array, ...], current_time: Float[Array, ...] | float = 0.0) -> Float[Array, ...]

Log density of next-event delay given last event at current_time.

The compensator term is routed through :meth:effective_integrated_intensity for consistency with :meth:log_prob.

Source code in src/xtremax/point_processes/operators/temporal.py

def inter_event_log_prob(
    self,
    tau: Float[Array, ...],
    current_time: Float[Array, ...] | float = 0.0,
) -> Float[Array, ...]:
    """Log density of next-event delay given last event at ``current_time``.

    The compensator term is routed through
    :meth:`effective_integrated_intensity` for consistency with
    :meth:`log_prob`.
    """
    tau = jnp.asarray(tau)
    s = jnp.asarray(current_time)
    t_next = s + tau
    log_lambda = self.log_intensity_fn(t_next)
    cum = self.effective_integrated_intensity(s, t_next)
    return jnp.where(tau >= 0.0, log_lambda - cum, -jnp.inf)

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

Expected count :math:\Lambda(\text{end}) - \Lambda(\text{start}).

Routed through :meth:effective_integrated_intensity so it agrees exactly with :meth:log_prob's compensator (closed form when the intensity module exposes .integrate).

Source code in src/xtremax/point_processes/operators/temporal.py

def predict_count(
    self,
    start_time: Float[Array, ...],
    end_time: Float[Array, ...],
) -> Float[Array, ...]:
    """Expected count :math:`\\Lambda(\\text{end}) - \\Lambda(\\text{start})`.

    Routed through :meth:`effective_integrated_intensity` so it
    agrees exactly with :meth:`log_prob`'s compensator (closed form
    when the intensity module exposes ``.integrate``).
    """
    return self.effective_integrated_intensity(start_time, end_time)

RenewalProcess(inter_event_dist: dist.Distribution, observation_window: ArrayLike, n_integration_points: int = 100)

Bases: GoodnessOfFitMixin, Module

Renewal process on [0, T] with a user-supplied inter-event law.

Parameters:

Name	Type	Description	Default
inter_event_dist	Distribution	NumPyro Distribution for inter-event times. Must implement sample, log_prob, and cdf. Any Distribution parameters that are JAX arrays remain live PyTree leaves.	required
observation_window	ArrayLike	Window length T > 0.	required
n_integration_points	int	Grid size used by compensator-based diagnostics that cannot be reduced to the hazard closed form.	100

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

def __init__(
    self,
    inter_event_dist: dist.Distribution,
    observation_window: ArrayLike,
    n_integration_points: int = 100,
) -> None:
    self.inter_event_dist = inter_event_dist
    self.observation_window = jnp.asarray(observation_window)
    self.n_integration_points = n_integration_points

log_prob(event_times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Joint event-time log-likelihood.

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

def log_prob(
    self,
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Joint event-time log-likelihood."""
    return renewal_log_prob(
        event_times, mask, self.observation_window, self.inter_event_dist
    )

sample(key: PRNGKeyArray, max_events: int) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Sample event times by accumulating iid gaps.

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

def sample(
    self,
    key: PRNGKeyArray,
    max_events: int,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Sample event times by accumulating iid gaps."""
    return renewal_sample(
        key, self.inter_event_dist, self.observation_window, max_events
    )

intensity(t: Float[Array, ...], event_times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Conditional intensity λ*(t) = h(t − t_{N(t)}).

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

def intensity(
    self,
    t: Float[Array, ...],
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Conditional intensity ``λ*(t) = h(t − t_{N(t)})``."""
    return renewal_intensity(t, event_times, mask, self.inter_event_dist)

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

Inter-event hazard h(τ) = f(τ)/S(τ).

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

def hazard(self, tau: Float[Array, ...]) -> Float[Array, ...]:
    """Inter-event hazard ``h(τ) = f(τ)/S(τ)``."""
    return renewal_hazard(tau, self.inter_event_dist)

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

Λ(τ) = −log S(τ).

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

def cumulative_hazard(self, tau: Float[Array, ...]) -> Float[Array, ...]:
    """``Λ(τ) = −log S(τ)``."""
    return renewal_cumulative_hazard(tau, self.inter_event_dist)

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

Inter-event survival S(τ) = 1 − F(τ).

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

def survival(self, tau: Float[Array, ...]) -> Float[Array, ...]:
    """Inter-event survival ``S(τ) = 1 − F(τ)``."""
    return renewal_survival(tau, self.inter_event_dist)

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

Log density of a single inter-event gap.

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

def inter_event_log_prob(self, tau: Float[Array, ...]) -> Float[Array, ...]:
    """Log density of a single inter-event gap."""
    return renewal_inter_event_log_prob(tau, self.inter_event_dist)

predict_count(start_time: Float[Array, ...] | float = 0.0, end_time: Float[Array, ...] | None = None) -> Float[Array, ...]

Expected event count on [start, end] via renewal equation.

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

def predict_count(
    self,
    start_time: Float[Array, ...] | float = 0.0,
    end_time: Float[Array, ...] | None = None,
) -> Float[Array, ...]:
    r"""Expected event count on ``[start, end]`` via renewal equation."""
    if end_time is None:
        end_time = self.observation_window
    T = jnp.asarray(end_time) - jnp.asarray(start_time)
    return renewal_expected_count(
        T, self.inter_event_dist, n_points=self.n_integration_points
    )

ExponentialHawkes(mu: ArrayLike, alpha: ArrayLike, beta: ArrayLike, observation_window: ArrayLike)

Bases: GoodnessOfFitMixin, Module

Hawkes process with exponential excitation on [0, T].

Parameters:

Name	Type	Description	Default
mu	ArrayLike	Baseline rate μ > 0.	required
alpha	ArrayLike	Excitation amplitude α >= 0.	required
beta	ArrayLike	Excitation decay β > 0.	required
observation_window	ArrayLike	Window length T > 0.	required

Notes

The process is sub-critical (stationary) iff :math:\alpha / \beta < 1; the sampler and log-likelihood are well-defined regardless, but super-critical branching ratios make the adaptive thinning envelope blow up and the sampler will often truncate at max_events.

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

def __init__(
    self,
    mu: ArrayLike,
    alpha: ArrayLike,
    beta: ArrayLike,
    observation_window: ArrayLike,
) -> None:
    self.mu = jnp.asarray(mu)
    self.alpha = jnp.asarray(alpha)
    self.beta = jnp.asarray(beta)
    self.observation_window = jnp.asarray(observation_window)

log_prob(event_times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Log-likelihood using Ozaki's O(n) recursion.

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

def log_prob(
    self,
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Log-likelihood using Ozaki's O(n) recursion."""
    return exp_hawkes_log_prob(
        event_times, mask, self.observation_window, self.mu, self.alpha, self.beta
    )

sample(key: PRNGKeyArray, max_events: int, max_candidates: int | None = None) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Adaptive thinning sample; envelope μ + α N(t).

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

def sample(
    self,
    key: PRNGKeyArray,
    max_events: int,
    max_candidates: int | None = None,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Adaptive thinning sample; envelope ``μ + α N(t)``."""
    return exp_hawkes_sample(
        key,
        self.observation_window,
        self.mu,
        self.alpha,
        self.beta,
        max_events,
        max_candidates=max_candidates,
    )

intensity(t: Float[Array, ...], event_times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Conditional intensity λ*(t | H).

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

def intensity(
    self,
    t: Float[Array, ...],
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Conditional intensity ``λ*(t | H)``."""
    return exp_hawkes_intensity(
        t, event_times, mask, self.mu, self.alpha, self.beta
    )

cumulative_intensity(t: Float[Array, ...], event_times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Compensator Λ(t | H).

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

def cumulative_intensity(
    self,
    t: Float[Array, ...],
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Compensator ``Λ(t | H)``."""
    return exp_hawkes_cumulative_intensity(
        t, event_times, mask, self.mu, self.alpha, self.beta
    )

hazard(t: Float[Array, ...], event_times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Hazard equals intensity for a simple point process.

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

def hazard(
    self,
    t: Float[Array, ...],
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Hazard equals intensity for a simple point process."""
    return self.intensity(t, event_times, mask)

effective_lambda_max(history: EventHistory) -> Float[Array, ...]

Thinning envelope μ + α N(t) read from live parameters.

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

def effective_lambda_max(self, history: EventHistory) -> Float[Array, ...]:
    """Thinning envelope ``μ + α N(t)`` read from live parameters."""
    n = jnp.sum(history.mask).astype(self.mu.dtype)
    return self.mu + self.alpha * n

effective_integrated_intensity(event_times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Compensator Λ(T | H) evaluated at the window edge.

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

def effective_integrated_intensity(
    self,
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Compensator ``Λ(T | H)`` evaluated at the window edge."""
    return self.cumulative_intensity(self.observation_window, event_times, mask)

GeneralHawkesProcess(mu: ArrayLike, kernel: eqx.Module, observation_window: ArrayLike, n_integration_points: int = 100, max_kernel_value: ArrayLike | None = None)

Bases: GoodnessOfFitMixin, Module

Hawkes process with a user-supplied kernel module.

Parameters:

Name	Type	Description	Default
mu	ArrayLike	Baseline rate μ > 0.	required
kernel	Module	eqx.Module exposing .kernel(dt) -> Array. When it also exposes .kernel_integral(a, b) -> Array the compensator uses the closed form (exact, differentiable); otherwise the operator falls back to trapezoid quadrature with n_integration_points nodes.	required
observation_window	ArrayLike	Window length.	required
n_integration_points	int	Quadrature grid size for the fallback compensator and for diagnostics.	100
max_kernel_value	ArrayLike | None	Optional static upper bound on :math:\phi(\tau) over τ >= 0. When None the operator tries kernel.max_kernel() and raises if absent. Used only by :meth:sample to build the thinning envelope μ + kernel_max · N(t).	None

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

def __init__(
    self,
    mu: ArrayLike,
    kernel: eqx.Module,
    observation_window: ArrayLike,
    n_integration_points: int = 100,
    max_kernel_value: ArrayLike | None = None,
) -> None:
    self.mu = jnp.asarray(mu)
    self.kernel = kernel
    self.observation_window = jnp.asarray(observation_window)
    self.n_integration_points = n_integration_points
    self.max_kernel_value = (
        None if max_kernel_value is None else jnp.asarray(max_kernel_value)
    )

MarkedTemporalPointProcess(ground: eqx.Module, mark_distribution_fn: Callable[[Array, EventHistory], dist.Distribution], mark_dim: int | None = None, history_at_each_event: bool = True)

Bases: SeparableMarkedPP

Separable marked TPP: ground process + time-/history-conditioned marks.

Parameters:

Name	Type	Description	Default
ground	Module	Any temporal point process operator exposing log_prob(times, mask) and sample(key, max_events) -> (times, mask, n).	required
mark_distribution_fn	Callable[[Array, EventHistory], Distribution]	Callable (t, history) -> Distribution. History is an :class:EventHistory containing events before t (empty at the first event). Any NumPyro Distribution works; scalar, multivariate, and discrete marks all follow the same protocol.	required
mark_dim	int | None	Dimensionality of each mark. Use None for scalar marks (mark shape ()) and an int for vector marks (mark shape (mark_dim,)).	None
history_at_each_event	bool	Pass per-event histories to mark_distribution_fn. When False the full observed history is used instead — cheaper but only correct for history-independent mark distributions.	True

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

def __init__(
    self,
    ground: eqx.Module,
    mark_distribution_fn: Callable[[Array, EventHistory], dist.Distribution],
    mark_dim: int | None = None,
    history_at_each_event: bool = True,
) -> None:
    self.ground = ground
    self.mark_distribution_fn = mark_distribution_fn
    self.mark_dim = mark_dim
    self.history_at_each_event = history_at_each_event

observation_window: Float[Array, ...] property

Forward to the ground operator.

log_prob(event_times: Float[Array, ...], mask: Bool[Array, ...], marks: Float[Array, ...]) -> Float[Array, ...]

Joint log-likelihood: ground + marks.

The ground-process log-prob signature varies across families — Hawkes needs (times, mask), HPP needs count, etc. This wrapper passes (times, mask) when the ground operator accepts that signature and falls back to just the mask-derived count for HPP-like families.

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

def log_prob(
    self,
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    marks: Float[Array, ...],
) -> Float[Array, ...]:
    """Joint log-likelihood: ground + marks.

    The ground-process log-prob signature varies across families —
    Hawkes needs ``(times, mask)``, HPP needs count, etc. This
    wrapper passes ``(times, mask)`` when the ground operator
    accepts that signature and falls back to just the mask-derived
    count for HPP-like families.
    """
    ground_log_prob = self._call_ground_log_prob(event_times, mask)
    mark_log_prob = marks_log_prob(
        event_times,
        marks,
        mask,
        self.mark_distribution_fn,
        history_at_each_event=self.history_at_each_event,
    )
    return ground_log_prob + mark_log_prob

sample(key: PRNGKeyArray, max_events: int, **ground_kwargs) -> MarkedSampleResult

Sample ground event times, then draw marks at each time.

Returns:

Type	Description
MarkedSampleResult	class:MarkedSampleResult (times, mask, marks) —
MarkedSampleResult	unpacks like the tuple it replaces. marks has shape
MarkedSampleResult	(max_events,) for scalar marks or
MarkedSampleResult	(max_events, mark_dim) for vector marks; padding
MarkedSampleResult	positions are filled with zeros. Any operator-specific
MarkedSampleResult	sample kwargs (e.g. max_candidates for a thinning-based
MarkedSampleResult	ground sampler) are forwarded to the ground.

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

def sample(
    self,
    key: PRNGKeyArray,
    max_events: int,
    **ground_kwargs,
) -> MarkedSampleResult:
    """Sample ground event times, then draw marks at each time.

    Returns:
        :class:`MarkedSampleResult` ``(times, mask, marks)`` —
        unpacks like the tuple it replaces. ``marks`` has shape
        ``(max_events,)`` for scalar marks or
        ``(max_events, mark_dim)`` for vector marks; padding
        positions are filled with zeros. Any operator-specific
        sample kwargs (e.g. ``max_candidates`` for a thinning-based
        ground sampler) are forwarded to the ground.
    """
    times, mask, marks = self._sample_joint(key, max_events, **ground_kwargs)
    return MarkedSampleResult(times, mask, marks)

ground_log_prob(event_times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Log-likelihood of the ground events only.

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

def ground_log_prob(
    self,
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Log-likelihood of the ground events only."""
    return self._call_ground_log_prob(event_times, mask)

marks_log_prob(event_times: Float[Array, ...], marks: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Log-likelihood of the marks only, conditional on observed times.

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

def marks_log_prob(
    self,
    event_times: Float[Array, ...],
    marks: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Log-likelihood of the marks only, conditional on observed times."""
    return marks_log_prob(
        event_times,
        marks,
        mask,
        self.mark_distribution_fn,
        history_at_each_event=self.history_at_each_event,
    )

ThinningProcess(base: eqx.Module, retention_fn: Callable[..., Array], observation_window: ArrayLike | None = None, n_integration_points: int = 100)

Bases: Module

Thinned TPP: base ground process + retention callable.

Parameters:

Name	Type	Description	Default
base	Module	Any temporal operator exposing log_prob(times, mask) / sample(key, max_events). The operator must also expose a conditional_intensity_fn way for the retention compensator to query :math:\lambda_\text{base}(t|H) — see :meth:_base_intensity_fn.	required
retention_fn	Callable[..., Array]	Callable (t, history, proposed_mark=None) -> p. Returns the retention probability in [0, 1]. history is the observed (retained) history up to t.	required
observation_window	ArrayLike | None	Window length. Defaults to base.observation_window when None.	None
n_integration_points	int	Trapezoid grid size for the retention compensator.	100

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

def __init__(
    self,
    base: eqx.Module,
    retention_fn: Callable[..., Array],
    observation_window: ArrayLike | None = None,
    n_integration_points: int = 100,
) -> None:
    self.base = base
    self.retention_fn = retention_fn
    if observation_window is None:
        observation_window = base.observation_window
    self.observation_window = jnp.asarray(observation_window)
    self.n_integration_points = n_integration_points

log_prob(event_times: Float[Array, ...], mask: Bool[Array, ...], marks: Float[Array, ...] | None = None, *, mark_sample_key: PRNGKeyArray | None = None, n_mark_samples: int = 8) -> Float[Array, ...]

Log-likelihood of the observed (retained) events.

See the module docstring for the decomposition. Supports a marked base via the optional marks argument — the retention callable is then invoked with (t, history, mark) so e.g. magnitude-gated detectors work out of the box.

For mark-dependent retention against a marked base, pass mark_sample_key so the compensator integral Monte-Carlo- averages retention over the base's mark distribution at each quadrature node. Without a key, retention_fn is called with proposed_mark=None in the compensator path (correct only for mark-independent retention).

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

def log_prob(
    self,
    event_times: Float[Array, ...],
    mask: Bool[Array, ...],
    marks: Float[Array, ...] | None = None,
    *,
    mark_sample_key: PRNGKeyArray | None = None,
    n_mark_samples: int = 8,
) -> Float[Array, ...]:
    """Log-likelihood of the observed (retained) events.

    See the module docstring for the decomposition. Supports a
    marked base via the optional ``marks`` argument — the retention
    callable is then invoked with ``(t, history, mark)`` so e.g.
    magnitude-gated detectors work out of the box.

    For **mark-dependent** retention against a marked base, pass
    ``mark_sample_key`` so the compensator integral Monte-Carlo-
    averages retention over the base's mark distribution at each
    quadrature node. Without a key, ``retention_fn`` is called with
    ``proposed_mark=None`` in the compensator path (correct only
    for mark-independent retention).
    """
    base_logp = self._call_base_log_prob(event_times, mask, marks=marks)

    # Promote scalar marks to 2-D inside EventHistory so that
    # history-dependent retention callables can uniformly look at
    # ``history.marks`` without branching on dimensionality.
    history_marks = None
    if marks is not None:
        history_marks = (
            jnp.expand_dims(marks, axis=-1) if marks.ndim == mask.ndim else marks
        )
    history = EventHistory(times=event_times, mask=mask, marks=history_marks)
    retention_term = thinning_retention_log_prob(
        self.retention_fn, event_times, mask, history, marks=marks
    )

    intensity_fn = self._base_intensity_fn(event_times, mask)

    # If the user wants MC-marginalised retention and the base is
    # a marked process with a ``mark_distribution_fn``, build the
    # sampler by delegating to it.
    mark_sampler_fn = None
    if mark_sample_key is not None:
        base_mark_fn = getattr(self.base, "mark_distribution_fn", None)
        if base_mark_fn is not None:

            def mark_sampler_fn(
                t: Array,
                prefix: EventHistory,
                k: PRNGKeyArray,
            ) -> Array:
                d = base_mark_fn(t, prefix)
                return d.sample(k)

    correction = retention_compensator(
        self.retention_fn,
        intensity_fn,
        history,
        self.observation_window,
        n_points=self.n_integration_points,
        mark_sampler_fn=mark_sampler_fn,
        mark_sample_key=mark_sample_key,
        n_mark_samples=n_mark_samples,
    )
    # The base log_prob subtracted the FULL compensator ∫λ; the thinned
    # process owes only ∫pλ, so the difference ∫(1-p)λ is added back.
    # (Subtracting it here — the old sign — made every fit maximize the
    # wrong objective, off by exactly 2∫(1-p)λ.)
    return base_logp + retention_term + correction

sample(key: PRNGKeyArray, max_events: int, max_candidates: int | None = None, **base_kwargs) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]] | tuple[Float[Array, ...], Bool[Array, ...], Float[Array, ...]]

Sample by first drawing from the base and then Bernoulli-thinning.

Two-pass draw: (1) draw a full latent realisation from the base; (2) for each latent event evaluate the retention probability given the retained history so far and accept with that probability. This correctly propagates history- and mark-dependent retention. max_candidates is forwarded to the base's own sampler for families (Hawkes, IPP) that use thinning internally.

Returns:

Type	Description
tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]] | tuple[Float[Array, ...], Bool[Array, ...], Float[Array, ...]]	(times, mask, n_retained) for an unmarked base, or
tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]] | tuple[Float[Array, ...], Bool[Array, ...], Float[Array, ...]]	(times, mask, retained_marks) when the base is a marked
tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]] | tuple[Float[Array, ...], Bool[Array, ...], Float[Array, ...]]	process. The third element disambiguates the two cases.

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

def sample(
    self,
    key: PRNGKeyArray,
    max_events: int,
    max_candidates: int | None = None,
    **base_kwargs,
) -> (
    tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]
    | tuple[Float[Array, ...], Bool[Array, ...], Float[Array, ...]]
):
    """Sample by first drawing from the base and then Bernoulli-thinning.

    Two-pass draw: (1) draw a full latent realisation from the
    base; (2) for each latent event evaluate the retention
    probability given the *retained* history so far and accept
    with that probability. This correctly propagates history- and
    mark-dependent retention. ``max_candidates`` is forwarded to
    the base's own sampler for families (Hawkes, IPP) that use
    thinning internally.

    Returns:
        ``(times, mask, n_retained)`` for an unmarked base, or
        ``(times, mask, retained_marks)`` when the base is a marked
        process. The third element disambiguates the two cases.
    """
    key_base, key_thin = jax.random.split(key)

    # Base sampler signatures are not uniform; the shared helper
    # dispatches on the base's actual parameter names (see
    # :func:`~xtremax.point_processes.operators._base.call_sampler`).
    base_result = call_sampler(
        self.base,
        key_base,
        max_events,
        max_candidates=max_candidates,
        **base_kwargs,
    )

    # Disambiguate the base result by TYPE first — package samplers
    # return :class:`SampleResult` (unmarked) or
    # :class:`MarkedSampleResult` (marked), so the dispatch is
    # unambiguous even for batched samplers. Bare tuples from
    # user-supplied bases fall back to the legacy shape/dtype
    # heuristic (0-D integer third element ⇒ event count).
    latent_marks = None
    if isinstance(base_result, MarkedSampleResult):
        latent_times, latent_mask, latent_marks = base_result
    elif isinstance(base_result, SampleResult):
        latent_times, latent_mask, _ = base_result
    elif len(base_result) == 2:
        latent_times, latent_mask = base_result
    elif len(base_result) == 3:
        latent_times, latent_mask, third = base_result
        third_arr = jnp.asarray(third)
        if third_arr.ndim == 0 and jnp.issubdtype(third_arr.dtype, jnp.integer):
            latent_marks = None
        else:
            latent_marks = third
    else:
        latent_times, latent_mask = base_result[:2]
        latent_marks = base_result[-1]

    keys = jax.random.split(key_thin, latent_times.shape[-1])

    if latent_marks is None:
        # Plain times + mask scan; retention gets no proposed_mark.
        def step_unmarked(
            carry: EventHistory,
            inp: tuple[Array, Array, PRNGKeyArray],
        ) -> tuple[EventHistory, None]:
            t_i, valid_i, key_i = inp
            p = self.retention_fn(t_i, carry, None)
            u = jax.random.uniform(key_i, dtype=t_i.dtype)
            accepted = valid_i & (u < jnp.clip(p, 0.0, 1.0))
            new_history = carry.append(t_i, accepted=accepted)
            return new_history, None

        initial_history = EventHistory.empty(
            max_events=max_events, mark_dim=None, dtype=latent_times.dtype
        )
        final_history, _ = jax.lax.scan(
            step_unmarked,
            initial_history,
            (
                jnp.moveaxis(latent_times, -1, 0),
                jnp.moveaxis(latent_mask, -1, 0),
                keys,
            ),
        )
    else:
        # Marked base: thread both the proposed mark into the
        # retention evaluation and the retained marks into the
        # carry's history buffer.
        latent_marks_2d = (
            latent_marks if latent_marks.ndim == 2 else latent_marks[..., None]
        )
        mark_dim = latent_marks_2d.shape[-1]

        def step_marked(
            carry: EventHistory,
            inp: tuple[Array, Array, PRNGKeyArray, Array],
        ) -> tuple[EventHistory, None]:
            t_i, valid_i, key_i, mark_row = inp
            # Pass the user a mark in their original shape (scalar
            # if marks were originally 1-D).
            proposed_mark = mark_row[0] if latent_marks.ndim == 1 else mark_row
            p = self.retention_fn(t_i, carry, proposed_mark)
            u = jax.random.uniform(key_i, dtype=t_i.dtype)
            accepted = valid_i & (u < jnp.clip(p, 0.0, 1.0))
            new_history = carry.append(t_i, mark=mark_row, accepted=accepted)
            return new_history, None

        initial_history = EventHistory.empty(
            max_events=max_events,
            mark_dim=mark_dim,
            dtype=latent_times.dtype,
            # Marks keep their own dtype — discrete (integer) marks
            # were previously cast to the float time dtype.
            mark_dtype=latent_marks_2d.dtype,
        )
        final_history, _ = jax.lax.scan(
            step_marked,
            initial_history,
            (
                jnp.moveaxis(latent_times, -1, 0),
                jnp.moveaxis(latent_mask, -1, 0),
                keys,
                jnp.moveaxis(latent_marks_2d, -2, 0),
            ),
        )

    times = jnp.where(
        final_history.mask, final_history.times, self.observation_window
    )
    if latent_marks is None:
        n_retained = jnp.sum(final_history.mask).astype(jnp.int32)
        return SampleResult(times, final_history.mask, n_retained)
    # Marked base: return retained marks in the user's original
    # shape (scalar marks flatten the trailing ``1`` dimension).
    retained_marks = final_history.marks
    if latent_marks.ndim == 1:
        retained_marks = retained_marks[..., 0]
    return MarkedSampleResult(times, final_history.mask, retained_marks)

Spatial processes

Patterns on a rectangular domain with no time dimension. Diagnostics here are function-based rather than methods — see csr_ripleys_k and friends in the primitives.

Equinox-based temporal and spatial point process operators.

HomogeneousSpatialPP(rate: ArrayLike, domain: RectangularDomain)

Bases: Module

Homogeneous spatial Poisson process on a rectangular domain.

Parameters:

Name	Type	Description	Default
rate	ArrayLike	Intensity λ > 0 (points per unit volume). Stored as a JAX array.	required
domain	RectangularDomain	Rectangular domain D ⊂ ℝᵈ.	required

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

def __init__(self, rate: ArrayLike, domain: RectangularDomain) -> None:
    self.rate = jnp.asarray(rate)
    self.domain = domain

n_dims: int property

Spatial dimension d (forwarded from the domain).

log_prob(locations: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Joint log-likelihood of an observed CSR pattern.

Uses the Janossy form n log λ − λ|D| (the -log n! and -n log |D| normalisers are deliberately dropped so the result is directly comparable to the IPP/Hawkes log-probs; see :func:~xtremax.point_processes.primitives.hpp_spatial.hpp_spatial_log_prob). locations is consumed only through its event count via mask — under CSR the location density is constant.

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

def log_prob(
    self,
    locations: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Joint log-likelihood of an observed CSR pattern.

    Uses the Janossy form ``n log λ − λ|D|`` (the ``-log n!`` and
    ``-n log |D|`` normalisers are deliberately dropped so the
    result is directly comparable to the IPP/Hawkes log-probs; see
    :func:`~xtremax.point_processes.primitives.hpp_spatial.hpp_spatial_log_prob`).
    ``locations`` is consumed only through its event count via
    ``mask`` — under CSR the location density is constant.
    """
    del locations  # CSR locations only enter via the count
    n_events = jnp.sum(mask, axis=-1)
    return hpp_spatial_log_prob(n_events, self.rate, self.domain.volume())

count_log_prob(n_events: Int[Array, ...]) -> Float[Array, ...]

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

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

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

sample(key: PRNGKeyArray, max_events: int) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Draw a padded (locations, mask, n_events) realisation.

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

def sample(
    self,
    key: PRNGKeyArray,
    max_events: int,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Draw a padded ``(locations, mask, n_events)`` realisation."""
    return hpp_spatial_sample(key, self.rate, self.domain, max_events)

intensity(locations: Float[Array, ...] | None = None) -> Float[Array, ...]

Pointwise intensity (constant). locations is optional.

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

def intensity(
    self,
    locations: Float[Array, ...] | None = None,
) -> Float[Array, ...]:
    """Pointwise intensity (constant). ``locations`` is optional."""
    if locations is None:
        return self.rate
    return hpp_spatial_intensity(locations, self.rate)

ripleys_k(r: Float[Array, ...]) -> Float[Array, ...]

Theoretical Ripley's K under CSR: :math:V_d r^d.

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

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

l_function(r: Float[Array, ...]) -> Float[Array, ...]

Besag's L-function (linearised K).

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

def l_function(self, r: Float[Array, ...]) -> Float[Array, ...]:
    """Besag's L-function (linearised K)."""
    return csr_l_function(r, self.n_dims)

pair_correlation(r: Float[Array, ...]) -> Float[Array, ...]

Pair correlation function under CSR (constant 1).

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

def pair_correlation(self, r: Float[Array, ...]) -> Float[Array, ...]:
    """Pair correlation function under CSR (constant 1)."""
    return csr_pair_correlation(r, self.n_dims)

predict_count(subdomain: RectangularDomain | None = None) -> tuple[Float[Array, ...], Float[Array, ...]]

Mean and variance of the count over a sub-domain (defaults to D).

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

def predict_count(
    self,
    subdomain: RectangularDomain | None = None,
) -> tuple[Float[Array, ...], Float[Array, ...]]:
    """Mean and variance of the count over a sub-domain (defaults to ``D``)."""
    vol = self.domain.volume() if subdomain is None else subdomain.volume()
    return hpp_spatial_predict_count(self.rate, vol)

predict_nearest_neighbor_distance() -> Float[Array, ...]

Expected nearest-neighbour distance under CSR.

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

def predict_nearest_neighbor_distance(self) -> Float[Array, ...]:
    """Expected nearest-neighbour distance under CSR."""
    return hpp_spatial_nearest_neighbor_distance(self.rate, self.n_dims)

InhomogeneousSpatialPP(log_intensity_fn: Callable[[Array], Array], domain: RectangularDomain, integrated_intensity: ArrayLike | None = None, lambda_max: ArrayLike | None = None, n_integration_points: int = 2048, integration_method: SpatialMethod = 'qmc')

Bases: Module

Inhomogeneous spatial Poisson process on a rectangular domain.

Parameters:

Name	Type	Description	Default
log_intensity_fn	Callable[[Array], Array]	Callable mapping (..., n, d) locations to (..., n) log-intensities. If this is itself an eqx.Module (e.g. a neural intensity), its parameters are real PyTree leaves and gradients flow through :meth:log_prob.	required
domain	RectangularDomain	Rectangular domain D ⊂ ℝᵈ.	required
integrated_intensity	ArrayLike | None	Optional fixed value of Λ(D). Pin only when the intensity is static; during training, stale values silently bias the likelihood. None triggers quadrature on every call.	None
lambda_max	ArrayLike | None	Optional upper bound on λ(s) over D, required by :meth:sample. Must be a true upper bound for thinning to be exact — passing too small a value silently biases the sampled process low in peaks. If None, the operator queries log_intensity_fn.max_intensity() when the intensity module exposes it and otherwise raises at sample time. There is no automatic fallback (e.g. 2 Λ / |D|) because no estimator can guarantee an upper bound for an arbitrary user-supplied intensity.	None
n_integration_points	int	Static node count for quadrature (trapezoid total nodes or QMC samples).	2048
integration_method	SpatialMethod	"qmc" (default) or "trapezoid". QMC is preferred for d ≥ 3; trapezoid is exact for multilinear log-intensities.	'qmc'

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

def __init__(
    self,
    log_intensity_fn: Callable[[Array], Array],
    domain: RectangularDomain,
    integrated_intensity: ArrayLike | None = None,
    lambda_max: ArrayLike | None = None,
    n_integration_points: int = 2048,
    integration_method: SpatialMethod = "qmc",
) -> None:
    self.log_intensity_fn = log_intensity_fn
    self.domain = domain
    self.integrated_intensity = (
        None if integrated_intensity is None else jnp.asarray(integrated_intensity)
    )
    self.lambda_max = None if lambda_max is None else jnp.asarray(lambda_max)
    self.n_integration_points = n_integration_points
    self.integration_method = integration_method

n_dims: int property

Spatial dimension d.

effective_integrated_intensity(subdomain: RectangularDomain | None = None) -> Float[Array, ...]

Compute :math:\Lambda(A) for A = subdomain (default D).

Precedence: (1) pinned :attr:integrated_intensity for the default domain only; (2) quadrature against the live log_intensity_fn otherwise. Sub-domain queries always re-integrate so they stay consistent with the live intensity.

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

def effective_integrated_intensity(
    self,
    subdomain: RectangularDomain | None = None,
) -> Float[Array, ...]:
    """Compute :math:`\\Lambda(A)` for ``A = subdomain`` (default ``D``).

    Precedence: (1) pinned :attr:`integrated_intensity` for the
    default domain only; (2) quadrature against the live
    ``log_intensity_fn`` otherwise. Sub-domain queries always
    re-integrate so they stay consistent with the live intensity.
    """
    target = self.domain if subdomain is None else subdomain
    if subdomain is None and self.integrated_intensity is not None:
        return self.integrated_intensity
    return integrate_log_intensity_spatial(
        self.log_intensity_fn,
        target,
        n_points=self.n_integration_points,
        method=self.integration_method,
    )

effective_lambda_max() -> Float[Array, ...]

Return the thinning bound used by :meth:sample.

Prefers the pinned :attr:lambda_max; otherwise queries the intensity module's .max_intensity() if it has one (e.g. a :class:~xtremax.point_processes.operators.temporal.PiecewiseConstantLogIntensity-style helper that tracks log_rates and reports exp(max(log_rates)) live). Raises when neither is available — Lewis–Shedler thinning requires a true upper bound, and no quadrature-derived estimate (e.g. 2 Λ / |D|, twice the mean intensity) can guarantee that for an arbitrary intensity surface.

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

def effective_lambda_max(self) -> Float[Array, ...]:
    """Return the thinning bound used by :meth:`sample`.

    Prefers the pinned :attr:`lambda_max`; otherwise queries the
    intensity module's ``.max_intensity()`` if it has one (e.g. a
    :class:`~xtremax.point_processes.operators.temporal.PiecewiseConstantLogIntensity`-style
    helper that tracks ``log_rates`` and reports
    ``exp(max(log_rates))`` live). Raises when neither is
    available — Lewis–Shedler thinning requires a *true* upper
    bound, and no quadrature-derived estimate (e.g. ``2 Λ / |D|``,
    twice the mean intensity) can guarantee that for an arbitrary
    intensity surface.
    """
    if self.lambda_max is not None:
        return self.lambda_max
    max_intensity = getattr(self.log_intensity_fn, "max_intensity", None)
    if max_intensity is not None:
        return max_intensity()
    raise ValueError(
        "Cannot sample via thinning: no `lambda_max` pinned on the "
        "operator and `log_intensity_fn` has no `.max_intensity()` "
        "method. Pass an upper bound at construction (it must be a "
        "true upper bound on `λ(s)` over the domain — too small a "
        "value silently biases the sampler low in intensity peaks)."
    )

log_prob(locations: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Log-likelihood :math:\sum \log \lambda(s_i) - \Lambda(D).

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

def log_prob(
    self,
    locations: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Log-likelihood :math:`\\sum \\log \\lambda(s_i) - \\Lambda(D)`."""
    Lambda_D = self.effective_integrated_intensity()
    return ipp_spatial_log_prob(locations, mask, self.log_intensity_fn, Lambda_D)

sample(key: PRNGKeyArray, max_candidates: int) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Lewis–Shedler thinning sampler with a static buffer.

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

def sample(
    self,
    key: PRNGKeyArray,
    max_candidates: int,
) -> tuple[Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Lewis–Shedler thinning sampler with a static buffer."""
    lam_max = self.effective_lambda_max()
    return ipp_spatial_sample_thinning(
        key, self.log_intensity_fn, self.domain, lam_max, max_candidates
    )

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

λ(s) = exp(log_intensity_fn(s)) evaluated at locations.

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

def intensity(self, locations: Float[Array, ...]) -> Float[Array, ...]:
    """``λ(s) = exp(log_intensity_fn(s))`` evaluated at ``locations``."""
    return ipp_spatial_intensity(locations, self.log_intensity_fn)

predict_count(subdomain: RectangularDomain | None = None) -> Float[Array, ...]

Expected count over a sub-domain (defaults to the full D).

Delegates to :meth:effective_integrated_intensity so a pinned integrated_intensity takes precedence for the full-domain query — matching both :meth:log_prob and the spatiotemporal operator, which already honoured the pin.

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

def predict_count(
    self,
    subdomain: RectangularDomain | None = None,
) -> Float[Array, ...]:
    """Expected count over a sub-domain (defaults to the full ``D``).

    Delegates to :meth:`effective_integrated_intensity` so a pinned
    ``integrated_intensity`` takes precedence for the full-domain
    query — matching both :meth:`log_prob` and the spatiotemporal
    operator, which already honoured the pin.
    """
    return self.effective_integrated_intensity(subdomain)

MarkedSpatialPP(ground: object, mark_distribution_fn: Callable[[Array], dist.Distribution], mark_dim: int | None = None)

Bases: SeparableMarkedPP

Separable marked spatial PP: ground process × mark distribution.

Parameters:

Name	Type	Description	Default
ground	object	Spatial PP operator exposing log_prob(locations, mask) and sample(key, max_events) -> (locations, mask, n).	required
mark_distribution_fn	Callable[[Array], Distribution]	Callable s -> Distribution. Called once per real event during log_prob and once per slot during sample (padding slot draws are masked out). Distributions with event_shape == () give scalar marks; non-trivial event_shape gives vector marks.	required
mark_dim	int | None	Static dimensionality of each mark; None for scalar marks. Used only by callers that need the mark shape up front (e.g. to preallocate buffers); the operator itself reads the shape from mark_distribution_fn.	None

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

def __init__(
    self,
    ground: object,
    mark_distribution_fn: Callable[[Array], dist.Distribution],
    mark_dim: int | None = None,
) -> None:
    self.ground = ground
    self.mark_distribution_fn = mark_distribution_fn
    self.mark_dim = mark_dim

domain: RectangularDomain property

Forward to the ground operator's domain.

n_dims: int property

Spatial dimension d.

log_prob(locations: Float[Array, ...], mask: Bool[Array, ...], marks: Float[Array, ...]) -> Float[Array, ...]

Joint log-likelihood: log L_ground(s) + Σ log f(mᵢ | sᵢ).

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

def log_prob(
    self,
    locations: Float[Array, ...],
    mask: Bool[Array, ...],
    marks: Float[Array, ...],
) -> Float[Array, ...]:
    """Joint log-likelihood: ``log L_ground(s) + Σ log f(mᵢ | sᵢ)``."""
    ground_lp = self.ground.log_prob(locations, mask)
    mark_lp = spatial_marks_log_prob(
        locations, marks, mask, self.mark_distribution_fn
    )
    return ground_lp + mark_lp

sample(key: PRNGKeyArray, max_events: int, **ground_kwargs) -> MarkedSpatialSampleResult

Sample ground locations, then draw marks at each location.

Returns :class:MarkedSpatialSampleResult (locations, mask, marks) — unpacks like the strict three-tuple it replaces. The ground's event count is dropped; mask already encodes which slots are real.

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

def sample(
    self,
    key: PRNGKeyArray,
    max_events: int,
    **ground_kwargs,
) -> MarkedSpatialSampleResult:
    """Sample ground locations, then draw marks at each location.

    Returns :class:`MarkedSpatialSampleResult`
    ``(locations, mask, marks)`` — unpacks like the strict
    three-tuple it replaces. The ground's event count is dropped;
    ``mask`` already encodes which slots are real.
    """
    locations, mask, marks = self._sample_joint(key, max_events, **ground_kwargs)
    return MarkedSpatialSampleResult(locations, mask, marks)

ground_log_prob(locations: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Log-likelihood of the ground locations only.

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

def ground_log_prob(
    self,
    locations: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Log-likelihood of the ground locations only."""
    return self.ground.log_prob(locations, mask)

marks_log_prob(locations: Float[Array, ...], marks: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Log-likelihood of the marks only, conditional on locations.

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

def marks_log_prob(
    self,
    locations: Float[Array, ...],
    marks: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Log-likelihood of the marks only, conditional on locations."""
    return spatial_marks_log_prob(locations, marks, mask, self.mark_distribution_fn)

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

Forward to the ground operator's intensity.

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

def intensity(self, locations: Float[Array, ...]) -> Float[Array, ...]:
    """Forward to the ground operator's ``intensity``."""
    # Some operators (HPP) accept ``locations`` optionally; the
    # spatial PP operators all accept it as a positional argument.
    return self.ground.intensity(locations)

Spatiotemporal processes

Events carrying both a location and a time, on \(D \times [t_0, t_1)\) — including the ETAS-style Hawkes process for spatially spreading contagion.

Equinox-based temporal and spatial point process operators.

HomogeneousSpatioTemporalPP(rate: ArrayLike, spatial: RectangularDomain, temporal: TemporalDomain)

Bases: Module

Homogeneous Poisson process on D × [t0, t1).

Parameters:

Name	Type	Description	Default
rate	ArrayLike	Intensity λ > 0 (points per unit space-time volume).	required
spatial	RectangularDomain	Spatial domain.	required
temporal	TemporalDomain	Temporal interval.	required

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

def __init__(
    self,
    rate: ArrayLike,
    spatial: RectangularDomain,
    temporal: TemporalDomain,
) -> None:
    self.rate = jnp.asarray(rate)
    self.spatial = spatial
    self.temporal = temporal

n_dims: int property

Spatial dimension d (forwarded from the spatial domain).

log_prob(locations: Float[Array, ...], times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

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

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

def log_prob(
    self,
    locations: Float[Array, ...],
    times: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Janossy log-likelihood ``n log λ − λ |D| (t1 − t0)``."""
    del locations, times  # the homogeneous likelihood depends only on the count
    n_events = jnp.sum(mask, axis=-1)
    return hpp_spatiotemporal_log_prob(
        n_events, self.rate, self.spatial.volume(), self.temporal.duration
    )

count_log_prob(n_events: Int[Array, ...]) -> Float[Array, ...]

Marginal Poisson-count log-pmf with mean λ |D| T.

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

def count_log_prob(self, n_events: Int[Array, ...]) -> Float[Array, ...]:
    """Marginal Poisson-count log-pmf with mean ``λ |D| T``."""
    return hpp_spatiotemporal_count_log_prob(
        n_events, self.rate, self.spatial.volume(), self.temporal.duration
    )

sample(key: PRNGKeyArray, max_events: int) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Draw (locations, times, mask, n_events); times sorted ascending.

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

def sample(
    self,
    key: PRNGKeyArray,
    max_events: int,
) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Draw ``(locations, times, mask, n_events)``; times sorted ascending."""
    return hpp_spatiotemporal_sample(
        key, self.rate, self.spatial, self.temporal, max_events
    )

intensity(locations: Float[Array, ...] | None = None, times: Float[Array, ...] | None = None) -> Float[Array, ...]

Pointwise intensity (constant).

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

def intensity(
    self,
    locations: Float[Array, ...] | None = None,
    times: Float[Array, ...] | None = None,
) -> Float[Array, ...]:
    """Pointwise intensity (constant)."""
    if locations is None or times is None:
        return self.rate
    return hpp_spatiotemporal_intensity(locations, times, self.rate)

predict_count(sub_spatial: RectangularDomain | None = None, sub_temporal: TemporalDomain | None = None) -> tuple[Float[Array, ...], Float[Array, ...]]

Mean and variance of the count over a sub-slab (defaults to full slab).

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

def predict_count(
    self,
    sub_spatial: RectangularDomain | None = None,
    sub_temporal: TemporalDomain | None = None,
) -> tuple[Float[Array, ...], Float[Array, ...]]:
    """Mean and variance of the count over a sub-slab (defaults to full slab)."""
    # `is None` (not `or`): a falsy-but-valid subdomain must not be
    # silently replaced by the full domain.
    vol = (self.spatial if sub_spatial is None else sub_spatial).volume()
    dur = (self.temporal if sub_temporal is None else sub_temporal).duration
    return hpp_spatiotemporal_predict_count(self.rate, vol, dur)

InhomogeneousSpatioTemporalPP(log_intensity_fn: Callable[[Array, Array], Array], spatial: RectangularDomain, temporal: TemporalDomain, integrated_intensity: ArrayLike | None = None, lambda_max: ArrayLike | None = None, n_integration_points: int = 4096, integration_method: SpatiotemporalMethod = 'qmc')

Bases: Module

Inhomogeneous Poisson process on D × [t0, t1).

Parameters:

Name	Type	Description	Default
log_intensity_fn	Callable[[Array, Array], Array]	Callable (s, t) → log λ(s, t) accepting s: (N, d) and t: (N,). If this is itself an eqx.Module, its parameters are real PyTree leaves and gradients flow through :meth:log_prob.	required
spatial	RectangularDomain	Spatial domain.	required
temporal	TemporalDomain	Temporal interval.	required
integrated_intensity	ArrayLike | None	Optional fixed value of :math:\Lambda. Pin only when the intensity is static; during training, stale values silently bias the likelihood.	None
lambda_max	ArrayLike | None	Optional upper bound on λ(s, t) over the slab, required by :meth:sample. Must be a true upper bound.	None
n_integration_points	int	Static node count for quadrature.	4096
integration_method	SpatiotemporalMethod	"qmc" (default) or "trapezoid".	'qmc'

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

def __init__(
    self,
    log_intensity_fn: Callable[[Array, Array], Array],
    spatial: RectangularDomain,
    temporal: TemporalDomain,
    integrated_intensity: ArrayLike | None = None,
    lambda_max: ArrayLike | None = None,
    n_integration_points: int = 4096,
    integration_method: SpatiotemporalMethod = "qmc",
) -> None:
    self.log_intensity_fn = log_intensity_fn
    self.spatial = spatial
    self.temporal = temporal
    self.integrated_intensity = (
        None if integrated_intensity is None else jnp.asarray(integrated_intensity)
    )
    self.lambda_max = None if lambda_max is None else jnp.asarray(lambda_max)
    self.n_integration_points = n_integration_points
    self.integration_method = integration_method

effective_integrated_intensity(sub_spatial: RectangularDomain | None = None, sub_temporal: TemporalDomain | None = None) -> Float[Array, ...]

Compute :math:\Lambda over a sub-slab (defaults to full slab).

Pinned :attr:integrated_intensity is honoured only when both sub_spatial and sub_temporal are None (i.e. the caller is asking about the originally-pinned slab); subdomain queries always re-integrate against the live intensity.

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

def effective_integrated_intensity(
    self,
    sub_spatial: RectangularDomain | None = None,
    sub_temporal: TemporalDomain | None = None,
) -> Float[Array, ...]:
    """Compute :math:`\\Lambda` over a sub-slab (defaults to full slab).

    Pinned :attr:`integrated_intensity` is honoured **only** when
    both ``sub_spatial`` and ``sub_temporal`` are ``None`` (i.e. the
    caller is asking about the originally-pinned slab); subdomain
    queries always re-integrate against the live intensity.
    """
    sp = self.spatial if sub_spatial is None else sub_spatial
    tp = self.temporal if sub_temporal is None else sub_temporal
    if (
        sub_spatial is None
        and sub_temporal is None
        and self.integrated_intensity is not None
    ):
        return self.integrated_intensity
    return integrate_log_intensity_spatiotemporal(
        self.log_intensity_fn,
        sp,
        tp,
        n_points=self.n_integration_points,
        method=self.integration_method,
    )

effective_lambda_max() -> Float[Array, ...]

Return the thinning bound used by :meth:sample.

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

def effective_lambda_max(self) -> Float[Array, ...]:
    """Return the thinning bound used by :meth:`sample`."""
    if self.lambda_max is not None:
        return self.lambda_max
    max_intensity = getattr(self.log_intensity_fn, "max_intensity", None)
    if max_intensity is not None:
        return max_intensity()
    raise ValueError(
        "Cannot sample via thinning: no `lambda_max` pinned on the "
        "operator and `log_intensity_fn` has no `.max_intensity()` "
        "method. Pass an upper bound at construction (it must be a "
        "true upper bound on `λ(s, t)` over the slab — too small a "
        "value silently biases the sampler low in intensity peaks)."
    )

log_prob(locations: Float[Array, ...], times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Janossy log-likelihood ∑_i log λ(s_i, t_i) − Λ.

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

def log_prob(
    self,
    locations: Float[Array, ...],
    times: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Janossy log-likelihood ``∑_i log λ(s_i, t_i) − Λ``."""
    Lambda = self.effective_integrated_intensity()
    return ipp_spatiotemporal_log_prob(
        locations, times, mask, self.log_intensity_fn, Lambda
    )

sample(key: PRNGKeyArray, max_candidates: int) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Lewis–Shedler thinning sampler on the full slab.

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

def sample(
    self,
    key: PRNGKeyArray,
    max_candidates: int,
) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]:
    """Lewis–Shedler thinning sampler on the full slab."""
    lam_max = self.effective_lambda_max()
    return ipp_spatiotemporal_sample_thinning(
        key,
        self.log_intensity_fn,
        self.spatial,
        self.temporal,
        lam_max,
        max_candidates,
    )

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

λ(s, t) = exp(log_intensity_fn(s, t)) evaluated on a batch.

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

def intensity(
    self,
    locations: Float[Array, ...],
    times: Float[Array, ...],
) -> Float[Array, ...]:
    """``λ(s, t) = exp(log_intensity_fn(s, t))`` evaluated on a batch."""
    return ipp_spatiotemporal_intensity(locations, times, self.log_intensity_fn)

predict_count(sub_spatial: RectangularDomain | None = None, sub_temporal: TemporalDomain | None = None) -> tuple[Float[Array, ...], Float[Array, ...]]

Mean and variance of the count over a sub-slab.

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

def predict_count(
    self,
    sub_spatial: RectangularDomain | None = None,
    sub_temporal: TemporalDomain | None = None,
) -> tuple[Float[Array, ...], Float[Array, ...]]:
    """Mean and variance of the count over a sub-slab."""
    Lambda = self.effective_integrated_intensity(sub_spatial, sub_temporal)
    return ipp_spatiotemporal_predict_count(Lambda)

marginal_spatial_intensity(locations: Float[Array, ...], n_time_points: int | None = None) -> Float[Array, ...]

λ_S(s) = ∫_{t0}^{t1} λ(s, t) dt via trapezoid in time.

The default grid size derives from :attr:n_integration_points (floored at the historical default of 100), so raising the operator's quadrature budget also refines the marginals.

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

def marginal_spatial_intensity(
    self,
    locations: Float[Array, ...],
    n_time_points: int | None = None,
) -> Float[Array, ...]:
    """``λ_S(s) = ∫_{t0}^{t1} λ(s, t) dt`` via trapezoid in time.

    The default grid size derives from :attr:`n_integration_points`
    (floored at the historical default of 100), so raising the
    operator's quadrature budget also refines the marginals.
    """
    if n_time_points is None:
        n_time_points = max(100, round(self.n_integration_points**0.5))
    return ipp_spatiotemporal_marginal_spatial_intensity(
        locations, self.log_intensity_fn, self.temporal, n_time_points=n_time_points
    )

marginal_temporal_intensity(times: Float[Array, ...], n_spatial_points: int | None = None) -> Float[Array, ...]

λ_T(t) = ∫_D λ(s, t) ds via Halton QMC in space.

The default node count derives from :attr:n_integration_points (floored at the historical default of 256).

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

def marginal_temporal_intensity(
    self,
    times: Float[Array, ...],
    n_spatial_points: int | None = None,
) -> Float[Array, ...]:
    """``λ_T(t) = ∫_D λ(s, t) ds`` via Halton QMC in space.

    The default node count derives from :attr:`n_integration_points`
    (floored at the historical default of 256).
    """
    if n_spatial_points is None:
        n_spatial_points = max(256, self.n_integration_points // 16)
    return ipp_spatiotemporal_marginal_temporal_intensity(
        times,
        self.log_intensity_fn,
        self.spatial,
        n_spatial_points=n_spatial_points,
    )

intensity_surface_at_time(t: Float[Array, ...], grid_size: int = 50) -> tuple[Float[Array, ...], Float[Array, ...]]

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

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

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

pearson_residuals(locations: Float[Array, ...], times: Float[Array, ...], mask: Bool[Array, ...], n_spatial_bins: int = 5, n_temporal_bins: int = 5, n_integration_points: int = 16) -> Float[Array, ...]

Cell-wise Pearson residuals on a space-time grid.

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

def pearson_residuals(
    self,
    locations: Float[Array, ...],
    times: Float[Array, ...],
    mask: Bool[Array, ...],
    n_spatial_bins: int = 5,
    n_temporal_bins: int = 5,
    n_integration_points: int = 16,
) -> Float[Array, ...]:
    """Cell-wise Pearson residuals on a space-time grid."""
    return ipp_spatiotemporal_pearson_residuals(
        locations,
        times,
        mask,
        self.log_intensity_fn,
        self.spatial,
        self.temporal,
        n_spatial_bins=n_spatial_bins,
        n_temporal_bins=n_temporal_bins,
        n_integration_points=n_integration_points,
    )

chi_square_gof(locations: Float[Array, ...], times: Float[Array, ...], mask: Bool[Array, ...], 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.

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

def chi_square_gof(
    self,
    locations: Float[Array, ...],
    times: Float[Array, ...],
    mask: Bool[Array, ...],
    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."""
    return ipp_spatiotemporal_chi_square_gof(
        locations,
        times,
        mask,
        self.log_intensity_fn,
        self.spatial,
        self.temporal,
        n_spatial_bins=n_spatial_bins,
        n_temporal_bins=n_temporal_bins,
        n_integration_points=n_integration_points,
    )

MarkedSpatioTemporalPP(ground: object, mark_distribution_fn: Callable[[Array, Array], dist.Distribution], mark_dim: int | None = None)

Bases: SeparableMarkedPP

Separable marked spatiotemporal PP: ground × mark distribution.

Parameters:

Name	Type	Description	Default
ground	object	Spatiotemporal PP operator exposing log_prob(locations, times, mask) and sample(key, max_events) -> (locs, times, mask, n).	required
mark_distribution_fn	Callable[[Array, Array], Distribution]	Callable (s, t) -> Distribution. Used in both log_prob (per real event) and sample (per slot — padding slot draws are masked out by :func:sample_spatiotemporal_marks).	required
mark_dim	int | None	Static dimensionality of the marks; None for scalar marks. Same role as in the marked spatial operator.	None

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

def __init__(
    self,
    ground: object,
    mark_distribution_fn: Callable[[Array, Array], dist.Distribution],
    mark_dim: int | None = None,
) -> None:
    self.ground = ground
    self.mark_distribution_fn = mark_distribution_fn
    self.mark_dim = mark_dim

spatial: RectangularDomain property

Forward to the ground operator's spatial domain.

temporal: TemporalDomain property

Forward to the ground operator's temporal interval.

log_prob(locations: Float[Array, ...], times: Float[Array, ...], mask: Bool[Array, ...], marks: Float[Array, ...]) -> Float[Array, ...]

Joint log-likelihood: log L_ground + Σ log f(mᵢ | sᵢ, tᵢ).

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

def log_prob(
    self,
    locations: Float[Array, ...],
    times: Float[Array, ...],
    mask: Bool[Array, ...],
    marks: Float[Array, ...],
) -> Float[Array, ...]:
    """Joint log-likelihood: ``log L_ground + Σ log f(mᵢ | sᵢ, tᵢ)``."""
    ground_lp = self.ground.log_prob(locations, times, mask)
    mark_lp = spatiotemporal_marks_log_prob(
        locations, times, marks, mask, self.mark_distribution_fn
    )
    return ground_lp + mark_lp

sample(key: PRNGKeyArray, max_events: int, **ground_kwargs) -> MarkedSpatiotemporalSampleResult

Sample ground (locations, times, mask), then marks at each event.

Returns :class:MarkedSpatiotemporalSampleResult — unpacks like the (locations, times, mask, marks) tuple it replaces.

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

def sample(
    self,
    key: PRNGKeyArray,
    max_events: int,
    **ground_kwargs,
) -> MarkedSpatiotemporalSampleResult:
    """Sample ground (locations, times, mask), then marks at each event.

    Returns :class:`MarkedSpatiotemporalSampleResult` — unpacks like
    the ``(locations, times, mask, marks)`` tuple it replaces.
    """
    locations, times, mask, marks = self._sample_joint(
        key, max_events, **ground_kwargs
    )
    return MarkedSpatiotemporalSampleResult(locations, times, mask, marks)

ground_log_prob(locations: Float[Array, ...], times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Log-likelihood of the ground process only.

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

def ground_log_prob(
    self,
    locations: Float[Array, ...],
    times: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Log-likelihood of the ground process only."""
    return self.ground.log_prob(locations, times, mask)

marks_log_prob(locations: Float[Array, ...], times: Float[Array, ...], marks: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Log-likelihood of the marks only, conditional on (s, t).

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

def marks_log_prob(
    self,
    locations: Float[Array, ...],
    times: Float[Array, ...],
    marks: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Log-likelihood of the marks only, conditional on (s, t)."""
    return spatiotemporal_marks_log_prob(
        locations, times, marks, mask, self.mark_distribution_fn
    )

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

Forward to the ground operator's intensity.

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

def intensity(
    self,
    locations: Float[Array, ...],
    times: Float[Array, ...],
) -> Float[Array, ...]:
    """Forward to the ground operator's ``intensity``."""
    return self.ground.intensity(locations, times)

SpatioTemporalHawkes(mu: ArrayLike, alpha: ArrayLike, beta: ArrayLike, sigma: ArrayLike, spatial: RectangularDomain, temporal: TemporalDomain, boundary_correction: bool = True)

Bases: Module

ETAS-style separable exp-Gaussian Hawkes process.

Parameters:

Name	Type	Description	Default
mu	ArrayLike	Baseline rate μ > 0 (constant in space and time).	required
alpha	ArrayLike	Excitation amplitude (α < 1 for sub-criticality).	required
beta	ArrayLike	Temporal decay rate β > 0.	required
sigma	ArrayLike	Spatial bandwidth σ > 0.	required
spatial	RectangularDomain	Spatial domain.	required
temporal	TemporalDomain	Temporal interval.	required
boundary_correction	bool	When True (default) the compensator uses the axis-wise erf-CDF to compute the in-domain mass of each event's spatial kernel. When False it assumes the whole spatial kernel mass falls in D (cheap, exact when σ is small relative to the box and events are interior).	True

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

def __init__(
    self,
    mu: ArrayLike,
    alpha: ArrayLike,
    beta: ArrayLike,
    sigma: ArrayLike,
    spatial: RectangularDomain,
    temporal: TemporalDomain,
    boundary_correction: bool = True,
) -> None:
    self.mu = jnp.asarray(mu)
    self.alpha = jnp.asarray(alpha)
    self.beta = jnp.asarray(beta)
    self.sigma = jnp.asarray(sigma)
    self.spatial = spatial
    self.temporal = temporal
    self.boundary_correction = boundary_correction

intensity(s: Float[Array, ...], t: Float[Array, ...], history_locations: Float[Array, ...], history_times: Float[Array, ...], history_mask: Bool[Array, ...]) -> Float[Array, ...]

Conditional intensity :math:\lambda^*(s, t \mid \mathcal H_t).

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

def intensity(
    self,
    s: Float[Array, ...],
    t: Float[Array, ...],
    history_locations: Float[Array, ...],
    history_times: Float[Array, ...],
    history_mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Conditional intensity :math:`\\lambda^*(s, t \\mid \\mathcal H_t)`."""
    return stpp_hawkes_intensity(
        s,
        t,
        history_locations,
        history_times,
        history_mask,
        self.mu,
        self.alpha,
        self.beta,
        self.sigma,
    )

compensator(locations: Float[Array, ...], times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

Closed-form compensator :math:\Lambda over the slab.

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

def compensator(
    self,
    locations: Float[Array, ...],
    times: Float[Array, ...],
    mask: Bool[Array, ...],
) -> Float[Array, ...]:
    """Closed-form compensator :math:`\\Lambda` over the slab."""
    return stpp_hawkes_compensator(
        locations,
        times,
        mask,
        self.mu,
        self.alpha,
        self.beta,
        self.sigma,
        self.spatial,
        self.temporal,
        boundary_correction=self.boundary_correction,
    )

log_prob(locations: Float[Array, ...], times: Float[Array, ...], mask: Bool[Array, ...]) -> Float[Array, ...]

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

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

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

sample(key: PRNGKeyArray, max_events: int) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Int[Array, ...]]

Ogata-style spatiotemporal thinning sample.

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

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

Building blocks

Components the processes above are assembled from: excitation kernels for Hawkes, a piecewise-constant log-intensity for non-parametric IPP fits, and the diagnostics bundle.

Equinox-based temporal and spatial point process operators.

ExponentialKernel

Bases: Module

Exponential excitation kernel :math:\phi(\tau) = \alpha e^{-\beta \tau}.

Useful as a GeneralHawkesProcess kernel when you want to share the general-kernel plumbing (e.g. with a learnable multi-peak super-kernel made by summing modules) but still get the exact integral for the exponential component.

PiecewiseConstantLogIntensity(bin_edges: ArrayLike, rates: ArrayLike | None = None, log_rates: ArrayLike | None = None)

Bases: Module

Piecewise-constant log-intensity as an eqx.Module.

Storing bin_edges and log_rates as PyTree leaves (rather than captured in a closure) lets them be inspected, vmapped over, and optimised via eqx.filter_grad when used inside an :class:InhomogeneousPoissonProcess. Parameterising the log rates means the rate is always positive without a further constraint — useful for unconstrained gradient-based fitting.

Source code in src/xtremax/point_processes/operators/temporal.py

def __init__(
    self,
    bin_edges: ArrayLike,
    rates: ArrayLike | None = None,
    log_rates: ArrayLike | None = None,
) -> None:
    if (rates is None) == (log_rates is None):
        raise ValueError("Provide exactly one of `rates` or `log_rates`.")
    self.bin_edges = jnp.asarray(bin_edges)
    self.log_rates = (
        jnp.log(jnp.asarray(rates)) if log_rates is None else jnp.asarray(log_rates)
    )

integrate(a: ArrayLike, b: ArrayLike) -> Array

Closed-form :math:\int_a^b \exp(\log \lambda(s))\, ds.

Implemented exactly as a sum over bin overlaps — unlike the operator's default quadrature, this scales linearly in the number of bins rather than the number of quadrature nodes and has no truncation error, so eqx.filter_grad through log_rates is exact. The :class:InhomogeneousPoissonProcess operator prefers this method over quadrature whenever the intensity module exposes it and the caller has not pinned a fixed integrated_intensity.

Source code in src/xtremax/point_processes/operators/temporal.py

def integrate(
    self,
    a: ArrayLike,
    b: ArrayLike,
) -> Array:
    """Closed-form :math:`\\int_a^b \\exp(\\log \\lambda(s))\\, ds`.

    Implemented exactly as a sum over bin overlaps — unlike the
    operator's default quadrature, this scales linearly in the
    number of bins rather than the number of quadrature nodes and
    has no truncation error, so ``eqx.filter_grad`` through
    ``log_rates`` is exact. The :class:`InhomogeneousPoissonProcess`
    operator prefers this method over quadrature whenever the
    intensity module exposes it and the caller has not pinned a
    fixed ``integrated_intensity``.
    """
    a = jnp.asarray(a)
    b = jnp.asarray(b)
    rates = jnp.exp(self.log_rates)
    widths = jnp.diff(self.bin_edges)
    left_edges = self.bin_edges[:-1]
    right_edges = self.bin_edges[1:]
    # Per bin, integrate rate × (min(b, right) − max(a, left))₊ over
    # all bins; vmap-friendly and differentiable through log_rates.
    # The limits get an explicit trailing bin axis and the sum runs
    # over that axis only, so batched limits return a batch of
    # integrals — the previous all-axes sum silently collapsed
    # batched limits into a single scalar.
    overlap = jnp.clip(
        jnp.minimum(b[..., None], right_edges)
        - jnp.maximum(a[..., None], left_edges),
        min=0.0,
        max=widths,
    )
    return jnp.sum(rates * overlap, axis=-1)

max_intensity() -> Array

Live λ_max = max(exp(log_rates)) from the current leaves.

Reading this from log_rates every call means thinning samplers stay safe even after log_rates has been updated by an optimiser.

Source code in src/xtremax/point_processes/operators/temporal.py

def max_intensity(self) -> Array:
    """Live ``λ_max = max(exp(log_rates))`` from the current leaves.

    Reading this from ``log_rates`` every call means thinning
    samplers stay safe even after ``log_rates`` has been updated by
    an optimiser.
    """
    return jnp.exp(jnp.max(self.log_rates))

GoodnessOfFit

Bases: NamedTuple

Bundle of diagnostics returned by goodness_of_fit.

Attributes:

Name	Type	Description
residuals	Float[Array, ...]	Time-rescaled inter-event residuals :math:\tau_i.
mask	Bool[Array, ...]	Real-event mask aligned with residuals.
ks_statistic	Float[Array, ...]	Kolmogorov–Smirnov statistic versus Exp(1).
theoretical_quantiles	Float[Array, ...]	QQ-plot theoretical quantiles.
empirical_quantiles	Float[Array, ...]	QQ-plot empirical quantiles.