Point Processes — Distributions

NumPyro Distribution wrappers around the point-process operators, so a point process can be used directly as a likelihood site in an MCMC or SVI model:

numpyro.sample("events", HomogeneousPoissonProcess(rate, T), obs=times)

Each class mirrors the operator of the same name — the wrapper adds NumPyro's Distribution interface (log_prob, sample, batch/event shapes) and nothing else.

Temporal processes

NumPyro Distribution wrappers for temporal and spatial point processes.

HomogeneousPoissonProcess(rate: ArrayLike, observation_window: ArrayLike, max_events: int = 512, *, validate_args: bool | None = None)

Bases: Distribution

NumPyro wrapper for a homogeneous temporal Poisson process.

Parameters:

Name	Type	Description	Default
rate	ArrayLike	Intensity λ > 0.	required
observation_window	ArrayLike	Window length T > 0.	required
max_events	int	Static buffer size for sample.	512
validate_args	bool | None	Forwarded to numpyro.distributions.Distribution.	None

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

def __init__(
    self,
    rate: ArrayLike,
    observation_window: ArrayLike,
    max_events: int = 512,
    *,
    validate_args: bool | None = None,
) -> None:
    self.rate = jnp.asarray(rate)
    self.observation_window = jnp.asarray(observation_window)
    self._max_events = int(max_events)
    self._op = _HppOp(self.rate, self.observation_window)
    batch_shape = jnp.broadcast_shapes(
        self.rate.shape, self.observation_window.shape
    )
    super().__init__(batch_shape=batch_shape, validate_args=validate_args)

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

Return a (times, mask) PyTree.

This matches the format :meth:log_prob consumes, so dist.log_prob(dist.sample(key)) is a valid round-trip. If you need the uncapped Poisson count, call the operator layer directly or use mask.sum(axis=-1).

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

def sample(
    self,
    key: PRNGKeyArray,
    sample_shape: tuple[int, ...] = (),
) -> tuple[Float[Array, ...], Bool[Array, ...]]:
    """Return a ``(times, mask)`` PyTree.

    This matches the format :meth:`log_prob` consumes, so
    ``dist.log_prob(dist.sample(key))`` is a valid round-trip.
    If you need the uncapped Poisson count, call the operator
    layer directly or use ``mask.sum(axis=-1)``.
    """
    check_prng_key(key)
    times, mask, _ = self._op.sample(key, self._max_events, sample_shape)
    return times, mask

log_prob(value: tuple[Float[Array, ...], Bool[Array, ...]] | tuple[Float[Array, ...], Int[Array, ...]]) -> Float[Array, ...]

Log-likelihood. value may be (times, mask) or (times, n_events).

The event times of an HPP drop out of the likelihood, so only the count is used.

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

def log_prob(
    self,
    value: tuple[Float[Array, ...], Bool[Array, ...]]
    | tuple[Float[Array, ...], Int[Array, ...]],
) -> Float[Array, ...]:
    """Log-likelihood. ``value`` may be ``(times, mask)`` or
    ``(times, n_events)``.

    The event *times* of an HPP drop out of the likelihood, so only
    the count is used.
    """
    _, counts_or_mask = value
    counts_or_mask = jnp.asarray(counts_or_mask)
    if counts_or_mask.dtype == jnp.bool_:
        n_events = jnp.sum(counts_or_mask, axis=-1)
    else:
        n_events = counts_or_mask
    return self._op.log_prob(n_events)

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

Bases: Distribution

NumPyro wrapper for an inhomogeneous temporal Poisson process.

Parameters:

Name	Type	Description	Default
log_intensity_fn	Callable[[Array], Array]	Callable returning log λ(t).	required
observation_window	ArrayLike	Window length.	required
integrated_intensity	ArrayLike | None	:math:\Lambda(T). If None, computed by quadrature at every log_prob call.	None
lambda_max	ArrayLike | None	Upper bound on :math:\lambda(t). Required for sample.	None
max_candidates	int	Static buffer for thinning.	1024
n_integration_points	int	Quadrature nodes.	100
validate_args	bool | None	Forwarded to numpyro.distributions.Distribution.	None

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

def __init__(
    self,
    log_intensity_fn: Callable[[Array], Array],
    observation_window: ArrayLike,
    integrated_intensity: ArrayLike | None = None,
    lambda_max: ArrayLike | None = None,
    max_candidates: int = 1024,
    n_integration_points: int = 100,
    *,
    validate_args: bool | None = None,
) -> None:
    self.observation_window = jnp.asarray(observation_window)
    self._max_candidates = int(max_candidates)
    self._op = _IppOp(
        log_intensity_fn,
        self.observation_window,
        integrated_intensity=integrated_intensity,
        lambda_max=lambda_max,
        n_integration_points=n_integration_points,
    )
    batch_shape = self.observation_window.shape
    super().__init__(batch_shape=batch_shape, validate_args=validate_args)

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

Return a (times, mask) PyTree (matching :meth:log_prob).

sample_shape must be empty for now — batched IPP sampling over an upper-bound λ_max requires a shared buffer size, which we leave to the operator layer where the user can vmap explicitly. Drops the uncapped candidate count returned by the operator; that is available via mask.sum() for accepted events, or from the operator API if the raw count matters.

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

def sample(
    self,
    key: PRNGKeyArray,
    sample_shape: tuple[int, ...] = (),
) -> tuple[Float[Array, ...], Bool[Array, ...]]:
    """Return a ``(times, mask)`` PyTree (matching :meth:`log_prob`).

    ``sample_shape`` must be empty for now — batched IPP sampling
    over an upper-bound ``λ_max`` requires a shared buffer size,
    which we leave to the operator layer where the user can
    ``vmap`` explicitly. Drops the uncapped candidate count
    returned by the operator; that is available via
    ``mask.sum()`` for accepted events, or from the operator API
    if the raw count matters.
    """
    check_prng_key(key)
    if sample_shape != ():
        raise NotImplementedError(
            "IPP distribution sample_shape=() only; use vmap over "
            "PRNG keys at the operator layer for batched draws."
        )
    times, mask, _ = self._op.sample(key, self._max_candidates)
    return times, mask

log_prob(value: tuple[Float[Array, ...], Bool[Array, ...]]) -> Float[Array, ...]

Log-likelihood value = (event_times, mask).

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

def log_prob(
    self,
    value: tuple[Float[Array, ...], Bool[Array, ...]],
) -> Float[Array, ...]:
    """Log-likelihood ``value = (event_times, mask)``."""
    event_times, mask = value
    return self._op.log_prob(event_times, mask)

RenewalProcess(inter_event_dist: dist.Distribution, observation_window: ArrayLike, max_events: int = 1024, *, validate_args: bool | None = None)

Bases: Distribution

NumPyro wrapper for a renewal temporal point process.

Parameters:

Name	Type	Description	Default
inter_event_dist	Distribution	NumPyro Distribution for inter-event times.	required
observation_window	ArrayLike	Window length T > 0.	required
max_events	int	Static buffer size for sample.	1024
validate_args	bool | None	Forwarded to the Distribution base.	None

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

def __init__(
    self,
    inter_event_dist: dist.Distribution,
    observation_window: ArrayLike,
    max_events: int = 1024,
    *,
    validate_args: bool | None = None,
) -> None:
    self.observation_window = jnp.asarray(observation_window)
    self._max_events = int(max_events)
    self._op = _RenewalOp(inter_event_dist, self.observation_window)
    # A batched inter-event distribution (e.g. vmapped-over rates)
    # adds batch dims to the per-sequence log-likelihood; the
    # window length can also be batched independently. Broadcast
    # both into the NumPyro batch_shape so plated models don't see
    # a phantom "scalar" shape that disagrees with the actual
    # output shape of ``log_prob``.
    batch_shape = jnp.broadcast_shapes(
        self.observation_window.shape, inter_event_dist.batch_shape
    )
    super().__init__(batch_shape=batch_shape, validate_args=validate_args)

ExponentialHawkes(mu: ArrayLike, alpha: ArrayLike, beta: ArrayLike, observation_window: ArrayLike, max_events: int = 1024, *, validate_args: bool | None = None)

Bases: Distribution

NumPyro wrapper for an exponential-kernel Hawkes process.

Parameters:

Name	Type	Description	Default
mu	ArrayLike	Baseline rate.	required
alpha	ArrayLike	Excitation amplitude.	required
beta	ArrayLike	Excitation decay.	required
observation_window	ArrayLike	Window length.	required
max_events	int	Static buffer for sampling.	1024

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

def __init__(
    self,
    mu: ArrayLike,
    alpha: ArrayLike,
    beta: ArrayLike,
    observation_window: ArrayLike,
    max_events: int = 1024,
    *,
    validate_args: bool | None = None,
) -> None:
    self.mu = jnp.asarray(mu)
    self.alpha = jnp.asarray(alpha)
    self.beta = jnp.asarray(beta)
    self.observation_window = jnp.asarray(observation_window)
    self._max_events = int(max_events)
    self._op = _ExpOp(self.mu, self.alpha, self.beta, self.observation_window)
    batch_shape = jnp.broadcast_shapes(
        self.mu.shape,
        self.alpha.shape,
        self.beta.shape,
        self.observation_window.shape,
    )
    super().__init__(batch_shape=batch_shape, validate_args=validate_args)

GeneralHawkesProcess(mu: ArrayLike, kernel: eqx.Module, observation_window: ArrayLike, max_events: int = 1024, n_integration_points: int = 100, max_kernel_value: ArrayLike | None = None, *, validate_args: bool | None = None)

Bases: Distribution

NumPyro wrapper for a general-kernel Hawkes process.

See :class:~xtremax.point_processes.operators.GeneralHawkesProcess for the kernel protocol and sampling details.

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

def __init__(
    self,
    mu: ArrayLike,
    kernel: eqx.Module,
    observation_window: ArrayLike,
    max_events: int = 1024,
    n_integration_points: int = 100,
    max_kernel_value: ArrayLike | None = None,
    *,
    validate_args: bool | None = None,
) -> None:
    self.mu = jnp.asarray(mu)
    self.observation_window = jnp.asarray(observation_window)
    self._max_events = int(max_events)
    self._op = _GeneralOp(
        self.mu,
        kernel,
        self.observation_window,
        n_integration_points=n_integration_points,
        max_kernel_value=max_kernel_value,
    )
    batch_shape = jnp.broadcast_shapes(self.mu.shape, self.observation_window.shape)
    super().__init__(batch_shape=batch_shape, validate_args=validate_args)

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

Bases: Distribution

NumPyro wrapper for a separable marked temporal point process.

Samples and log-probs consume / return the triple (times, mask, marks).

Parameters:

Name	Type	Description	Default
ground	Module	Ground-process operator.	required
mark_distribution_fn	Callable[[Array, EventHistory], Distribution]	(t, history) -> Distribution callable.	required
mark_dim	int | None	Dimensionality of each mark (None for scalar).	None
history_at_each_event	bool	Whether marks see per-event histories.	True
max_events	int	Static buffer size.	1024

Source code in src/xtremax/point_processes/distributions/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,
    max_events: int = 1024,
    *,
    validate_args: bool | None = None,
) -> None:
    self._op = _MarkedOp(
        ground,
        mark_distribution_fn,
        mark_dim=mark_dim,
        history_at_each_event=history_at_each_event,
    )
    self.observation_window = jnp.asarray(self._op.observation_window)
    self._max_events = int(max_events)
    super().__init__(
        batch_shape=self.observation_window.shape, validate_args=validate_args
    )

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

Bases: Distribution

NumPyro wrapper for a thinned TPP.

.. warning:: For history-dependent bases (Hawkes, renewal) log_prob is a pseudo-likelihood — the base intensity is conditioned on retained events only, while sampling thins a full latent realisation. Exact for Poisson-family bases; see the operator module docstring (:mod:xtremax.point_processes.operators.thinning).

Parameters:

Name	Type	Description	Default
base	Module	Base temporal operator.	required
retention_fn	Callable[..., Array]	(t, history, proposed_mark=None) -> p callable.	required
observation_window	ArrayLike | None	Window length. Defaults to the base's.	None
max_events	int	Static buffer size.	1024
n_integration_points	int	Trapezoid grid size for the retention compensator.	100

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

def __init__(
    self,
    base: eqx.Module,
    retention_fn: Callable[..., Array],
    observation_window: ArrayLike | None = None,
    max_events: int = 1024,
    n_integration_points: int = 100,
    *,
    validate_args: bool | None = None,
) -> None:
    self._op = _ThinningOp(
        base,
        retention_fn,
        observation_window=observation_window,
        n_integration_points=n_integration_points,
    )
    self.observation_window = jnp.asarray(self._op.observation_window)
    self._max_events = int(max_events)
    super().__init__(
        batch_shape=self.observation_window.shape, validate_args=validate_args
    )

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

Return (times, mask) for an unmarked base, or (times, mask, marks) when the base is a marked process.

The triple matches what :meth:log_prob accepts so that d.log_prob(d.sample(key)) round-trips for both cases.

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

def sample(
    self,
    key: PRNGKeyArray,
    sample_shape: tuple[int, ...] = (),
) -> (
    tuple[Float[Array, ...], Bool[Array, ...]]
    | tuple[Float[Array, ...], Bool[Array, ...], Float[Array, ...]]
):
    """Return ``(times, mask)`` for an unmarked base, or
    ``(times, mask, marks)`` when the base is a marked process.

    The triple matches what :meth:`log_prob` accepts so that
    ``d.log_prob(d.sample(key))`` round-trips for both cases.
    """
    check_prng_key(key)
    if sample_shape != ():
        raise NotImplementedError(
            "ThinningProcess distribution supports sample_shape=() only."
        )
    times, mask, third = self._op.sample(key, self._max_events)
    # Marked base: ``third`` is a mark array. Otherwise a scalar
    # integer count that we drop (users who want the count can
    # read ``mask.sum()``).
    third_arr = jnp.asarray(third)
    if third_arr.ndim == 0 and jnp.issubdtype(third_arr.dtype, jnp.integer):
        return times, mask
    return times, mask, third

Spatial processes

NumPyro Distribution wrappers for temporal and spatial point processes.

HomogeneousSpatialPP(rate: ArrayLike, domain: RectangularDomain, max_events: int = 512, *, validate_args: bool | None = None)

Bases: Distribution

NumPyro wrapper for a homogeneous spatial Poisson process.

Parameters:

Name	Type	Description	Default
rate	ArrayLike	Intensity λ > 0.	required
domain	RectangularDomain	Rectangular domain.	required
max_events	int	Static buffer size for sample.	512
validate_args	bool | None	Forwarded to numpyro.distributions.Distribution.	None

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

def __init__(
    self,
    rate: ArrayLike,
    domain: RectangularDomain,
    max_events: int = 512,
    *,
    validate_args: bool | None = None,
) -> None:
    self.rate = jnp.asarray(rate)
    if self.rate.shape != ():
        raise ValueError(
            "HomogeneousSpatialPP distribution requires a scalar `rate`; "
            f"got shape {self.rate.shape}. Spatial sampling produces "
            "variable-length point patterns and is not batch-safe at the "
            "distribution layer — vmap over `rate` at the operator layer "
            "(`xtremax.point_processes.operators.HomogeneousSpatialPP`) "
            "if you need batched draws."
        )
    self.domain = domain
    self._max_events = int(max_events)
    self._op = _HppSpatialOp(self.rate, self.domain)
    # Empty batch_shape — `rate` is required scalar so the
    # distribution carries no batched axes.
    super().__init__(batch_shape=(), validate_args=validate_args)

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

Return (locations, mask).

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

def sample(
    self,
    key: PRNGKeyArray,
    sample_shape: tuple[int, ...] = (),
) -> tuple[Float[Array, ...], Bool[Array, ...]]:
    """Return ``(locations, mask)``."""
    check_prng_key(key)
    if sample_shape != ():
        raise NotImplementedError(
            "HomogeneousSpatialPP distribution supports sample_shape=() "
            "only; vmap over PRNG keys at the operator layer for batched draws."
        )
    locations, mask, _ = self._op.sample(key, self._max_events)
    return locations, mask

log_prob(value: tuple[Float[Array, ...], Bool[Array, ...]] | tuple[Float[Array, ...], Int[Array, ...]]) -> Float[Array, ...]

Joint log-likelihood. value may be (locations, mask) or (locations, n_events).

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

def log_prob(
    self,
    value: tuple[Float[Array, ...], Bool[Array, ...]]
    | tuple[Float[Array, ...], Int[Array, ...]],
) -> Float[Array, ...]:
    """Joint log-likelihood. ``value`` may be ``(locations, mask)`` or
    ``(locations, n_events)``.
    """
    locations, counts_or_mask = value
    del locations  # the HPP log-likelihood depends only on the count
    counts_or_mask = jnp.asarray(counts_or_mask)
    if counts_or_mask.dtype == jnp.bool_:
        n_events = jnp.sum(counts_or_mask, axis=-1)
    else:
        # Use the raw count directly. Fabricating a mask from
        # ``ranks < count`` would silently clip ``count > max_events``
        # to ``max_events`` and return the wrong likelihood for any
        # caller scoring observed counts above the buffer size
        # (Codex P1 in PR #12 review).
        n_events = counts_or_mask
    return hpp_spatial_log_prob(n_events, self._op.rate, self._op.domain.volume())

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

Bases: Distribution

NumPyro wrapper for an inhomogeneous spatial Poisson process.

Parameters:

Name	Type	Description	Default
log_intensity_fn	Callable[[Array], Array]	Log-intensity callable (..., n, d) -> (..., n).	required
domain	RectangularDomain	Rectangular domain.	required
integrated_intensity	ArrayLike | None	Optional pinned Λ(D); None triggers quadrature.	None
lambda_max	ArrayLike | None	Optional thinning bound. None defers to the intensity module's .max_intensity() when available and otherwise raises at sample time — there is no automatic fallback, since no derived estimate can guarantee a true upper bound for an arbitrary intensity surface.	None
max_candidates	int	Static buffer for thinning.	1024
n_integration_points	int	Quadrature node count.	2048
integration_method	SpatialMethod	"qmc" (default) or "trapezoid".	'qmc'
validate_args	bool | None	Forwarded to numpyro.distributions.Distribution.	None

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

def __init__(
    self,
    log_intensity_fn: Callable[[Array], Array],
    domain: RectangularDomain,
    integrated_intensity: ArrayLike | None = None,
    lambda_max: ArrayLike | None = None,
    max_candidates: int = 1024,
    n_integration_points: int = 2048,
    integration_method: SpatialMethod = "qmc",
    *,
    validate_args: bool | None = None,
) -> None:
    self.domain = domain
    self._max_candidates = int(max_candidates)
    self._op = _IppSpatialOp(
        log_intensity_fn,
        domain,
        integrated_intensity=integrated_intensity,
        lambda_max=lambda_max,
        n_integration_points=n_integration_points,
        integration_method=integration_method,
    )
    # No batched parameters here — the domain is a structural
    # element. Use an empty batch shape; vmap externally for
    # multiple intensities.
    super().__init__(batch_shape=(), validate_args=validate_args)

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

Return (locations, mask).

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

def sample(
    self,
    key: PRNGKeyArray,
    sample_shape: tuple[int, ...] = (),
) -> tuple[Float[Array, ...], Bool[Array, ...]]:
    """Return ``(locations, mask)``."""
    check_prng_key(key)
    if sample_shape != ():
        raise NotImplementedError(
            "InhomogeneousSpatialPP distribution supports sample_shape=() "
            "only; vmap over PRNG keys at the operator layer for batched draws."
        )
    locations, mask, _ = self._op.sample(key, self._max_candidates)
    return locations, mask

log_prob(value: tuple[Float[Array, ...], Bool[Array, ...]]) -> Float[Array, ...]

Log-likelihood Σ log λ(sᵢ) − Λ(D).

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

def log_prob(
    self,
    value: tuple[Float[Array, ...], Bool[Array, ...]],
) -> Float[Array, ...]:
    """Log-likelihood ``Σ log λ(sᵢ) − Λ(D)``."""
    locations, mask = value
    return self._op.log_prob(locations, mask)

MarkedSpatialPP(ground: eqx.Module, mark_distribution_fn: Callable[[Array], dist.Distribution], max_events: int = 1024, mark_dim: int | None = None, *, validate_args: bool | None = None)

Bases: Distribution

NumPyro wrapper for a separable marked spatial Poisson process.

Parameters:

Name	Type	Description	Default
ground	Module	Spatial PP operator (HomogeneousSpatialPP or InhomogeneousSpatialPP) — the operator-layer object, not its distribution wrapper. The buffer size lives on the wrapper, not the operator, so we accept it explicitly.	required
mark_distribution_fn	Callable[[Array], Distribution]	Callable s -> Distribution for marks.	required
max_events	int	Static buffer size for sample — forwarded to the ground operator's sampler. Choose comfortably above the expected count.	1024
mark_dim	int | None	Static mark dimensionality (None for scalar).	None
validate_args	bool | None	Forwarded to numpyro.distributions.Distribution.	None

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

def __init__(
    self,
    ground: eqx.Module,
    mark_distribution_fn: Callable[[Array], dist.Distribution],
    max_events: int = 1024,
    mark_dim: int | None = None,
    *,
    validate_args: bool | None = None,
) -> None:
    self._op = _MarkedSpatialOp(
        ground=ground,
        mark_distribution_fn=mark_distribution_fn,
        mark_dim=mark_dim,
    )
    self._max_events = int(max_events)
    super().__init__(batch_shape=(), validate_args=validate_args)

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

Return (locations, mask, marks).

Sample shapes other than () are not supported here — vmap over PRNG keys at the operator layer for batched draws.

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

def sample(
    self,
    key: PRNGKeyArray,
    sample_shape: tuple[int, ...] = (),
) -> tuple[Float[Array, ...], Bool[Array, ...], Float[Array, ...]]:
    """Return ``(locations, mask, marks)``.

    Sample shapes other than ``()`` are not supported here — vmap
    over PRNG keys at the operator layer for batched draws.
    """
    check_prng_key(key)
    if sample_shape != ():
        raise NotImplementedError(
            "MarkedSpatialPP distribution supports sample_shape=() only; "
            "vmap over PRNG keys at the operator layer for batched draws."
        )
    return self._op.sample(key, self._max_events)

log_prob(value: tuple[Float[Array, ...], Bool[Array, ...], Float[Array, ...]]) -> Float[Array, ...]

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

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

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

Spatiotemporal processes

NumPyro Distribution wrappers for temporal and spatial point processes.

HomogeneousSpatioTemporalPP(rate: ArrayLike, spatial: RectangularDomain, temporal: TemporalDomain, max_events: int = 512, *, validate_args: bool | None = None)

Bases: Distribution

NumPyro wrapper for a homogeneous spatiotemporal Poisson process.

Parameters:

Name	Type	Description	Default
rate	ArrayLike	Intensity λ > 0 (scalar — vmap at the operator layer for batched rates).	required
spatial	RectangularDomain	Spatial domain.	required
temporal	TemporalDomain	Temporal interval.	required
max_events	int	Static buffer size for sample.	512
validate_args	bool | None	Forwarded to numpyro.	None

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

def __init__(
    self,
    rate: ArrayLike,
    spatial: RectangularDomain,
    temporal: TemporalDomain,
    max_events: int = 512,
    *,
    validate_args: bool | None = None,
) -> None:
    self.rate = jnp.asarray(rate)
    if self.rate.shape != ():
        raise ValueError(
            "HomogeneousSpatioTemporalPP distribution requires a scalar "
            f"`rate`; got shape {self.rate.shape}. Spatiotemporal sampling "
            "produces variable-length point patterns and is not batch-safe "
            "at the distribution layer — vmap over `rate` at the operator "
            "layer "
            "(`xtremax.point_processes.operators.HomogeneousSpatioTemporalPP`) "
            "if you need batched draws."
        )
    self.spatial = spatial
    self.temporal = temporal
    self._max_events = int(max_events)
    self._op = _HppStOp(self.rate, self.spatial, self.temporal)
    super().__init__(batch_shape=(), validate_args=validate_args)

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

Return (locations, times, mask).

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

def sample(
    self,
    key: PRNGKeyArray,
    sample_shape: tuple[int, ...] = (),
) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...]]:
    """Return ``(locations, times, mask)``."""
    check_prng_key(key)
    if sample_shape != ():
        raise NotImplementedError(
            "HomogeneousSpatioTemporalPP supports sample_shape=() only; "
            "vmap over PRNG keys at the operator layer for batched draws."
        )
    locations, times, mask, _ = self._op.sample(key, self._max_events)
    return locations, times, mask

log_prob(value: tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...]] | tuple[Float[Array, ...], Float[Array, ...], Int[Array, ...]]) -> Float[Array, ...]

Joint log-likelihood. value may be (locations, times, mask) or (locations, times, n_events).

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

def log_prob(
    self,
    value: tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...]]
    | tuple[Float[Array, ...], Float[Array, ...], Int[Array, ...]],
) -> Float[Array, ...]:
    """Joint log-likelihood. ``value`` may be ``(locations, times, mask)``
    or ``(locations, times, n_events)``.
    """
    locations, times, counts_or_mask = value
    del locations, times
    counts_or_mask = jnp.asarray(counts_or_mask)
    if counts_or_mask.dtype == jnp.bool_:
        n_events = jnp.sum(counts_or_mask, axis=-1)
    else:
        n_events = counts_or_mask
    return hpp_spatiotemporal_log_prob(
        n_events,
        self._op.rate,
        self._op.spatial.volume(),
        self._op.temporal.duration,
    )

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

Bases: Distribution

NumPyro wrapper for an inhomogeneous spatiotemporal Poisson process.

Parameters:

Name	Type	Description	Default
log_intensity_fn	Callable[[Array, Array], Array]	Callable (s, t) → log λ(s, t).	required
spatial	RectangularDomain	Spatial domain.	required
temporal	TemporalDomain	Temporal interval.	required
integrated_intensity	ArrayLike | None	Optional pinned :math:\Lambda.	None
lambda_max	ArrayLike | None	Optional thinning bound (must be a true upper bound; the operator raises if neither this nor a log_intensity_fn.max_intensity() is available).	None
max_candidates	int	Static thinning buffer.	1024
n_integration_points	int	Quadrature node count.	4096
integration_method	SpatiotemporalMethod	"qmc" (default) or "trapezoid".	'qmc'
validate_args	bool | None	Forwarded to numpyro.	None

Source code in src/xtremax/point_processes/distributions/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,
    max_candidates: int = 1024,
    n_integration_points: int = 4096,
    integration_method: SpatiotemporalMethod = "qmc",
    *,
    validate_args: bool | None = None,
) -> None:
    self.spatial = spatial
    self.temporal = temporal
    self._max_candidates = int(max_candidates)
    self._op = _IppStOp(
        log_intensity_fn,
        spatial,
        temporal,
        integrated_intensity=integrated_intensity,
        lambda_max=lambda_max,
        n_integration_points=n_integration_points,
        integration_method=integration_method,
    )
    super().__init__(batch_shape=(), validate_args=validate_args)

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

Return (locations, times, mask).

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

def sample(
    self,
    key: PRNGKeyArray,
    sample_shape: tuple[int, ...] = (),
) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...]]:
    """Return ``(locations, times, mask)``."""
    check_prng_key(key)
    if sample_shape != ():
        raise NotImplementedError(
            "InhomogeneousSpatioTemporalPP supports sample_shape=() only; "
            "vmap over PRNG keys at the operator layer for batched draws."
        )
    locations, times, mask, _ = self._op.sample(key, self._max_candidates)
    return locations, times, mask

log_prob(value: tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...]]) -> Float[Array, ...]

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

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

def log_prob(
    self,
    value: tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...]],
) -> Float[Array, ...]:
    """Log-likelihood :math:`\\sum_i \\log \\lambda(s_i, t_i) - \\Lambda`."""
    locations, times, mask = value
    return self._op.log_prob(locations, times, mask)

MarkedSpatioTemporalPP(ground: eqx.Module, mark_distribution_fn: Callable[[Array, Array], dist.Distribution], max_events: int = 1024, mark_dim: int | None = None, *, validate_args: bool | None = None)

Bases: Distribution

NumPyro wrapper for a separable marked spatiotemporal Poisson process.

Parameters:

Name	Type	Description	Default
ground	Module	Spatiotemporal PP operator (HomogeneousSpatioTemporalPP, InhomogeneousSpatioTemporalPP, or SpatioTemporalHawkes).	required
mark_distribution_fn	Callable[[Array, Array], Distribution]	Callable (s, t) -> Distribution.	required
max_events	int	Static buffer size for sample.	1024
mark_dim	int | None	Static mark dimensionality (None for scalar).	None
validate_args	bool | None	Forwarded to numpyro.	None

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

def __init__(
    self,
    ground: eqx.Module,
    mark_distribution_fn: Callable[[Array, Array], dist.Distribution],
    max_events: int = 1024,
    mark_dim: int | None = None,
    *,
    validate_args: bool | None = None,
) -> None:
    self._op = _MarkedStOp(
        ground=ground,
        mark_distribution_fn=mark_distribution_fn,
        mark_dim=mark_dim,
    )
    self._max_events = int(max_events)
    super().__init__(batch_shape=(), validate_args=validate_args)

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

Return (locations, times, mask, marks).

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

def sample(
    self,
    key: PRNGKeyArray,
    sample_shape: tuple[int, ...] = (),
) -> tuple[
    Float[Array, ...], Float[Array, ...], Bool[Array, ...], Float[Array, ...]
]:
    """Return ``(locations, times, mask, marks)``."""
    check_prng_key(key)
    if sample_shape != ():
        raise NotImplementedError(
            "MarkedSpatioTemporalPP supports sample_shape=() only; "
            "vmap over PRNG keys at the operator layer for batched draws."
        )
    return self._op.sample(key, self._max_events)

log_prob(value: tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...], Float[Array, ...]]) -> Float[Array, ...]

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

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

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

SpatioTemporalHawkes(mu: ArrayLike, alpha: ArrayLike, beta: ArrayLike, sigma: ArrayLike, spatial: RectangularDomain, temporal: TemporalDomain, max_events: int = 512, boundary_correction: bool = True, *, validate_args: bool | None = None)

Bases: Distribution

NumPyro wrapper for a separable exp-Gaussian spatiotemporal Hawkes process.

Parameters:

Name	Type	Description	Default
mu	ArrayLike	Baseline rate (scalar > 0).	required
alpha	ArrayLike	Excitation amplitude (scalar; alpha < 1 sub-critical).	required
beta	ArrayLike	Temporal decay (scalar > 0).	required
sigma	ArrayLike	Spatial bandwidth (scalar > 0).	required
spatial	RectangularDomain	Spatial domain.	required
temporal	TemporalDomain	Temporal interval.	required
max_events	int	Static buffer for sample.	512
boundary_correction	bool	Whether the compensator uses the axis-wise erf-CDF for in-domain spatial-kernel mass (default True).	True
validate_args	bool | None	Forwarded to numpyro.	None

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

def __init__(
    self,
    mu: ArrayLike,
    alpha: ArrayLike,
    beta: ArrayLike,
    sigma: ArrayLike,
    spatial: RectangularDomain,
    temporal: TemporalDomain,
    max_events: int = 512,
    boundary_correction: bool = True,
    *,
    validate_args: bool | None = None,
) -> None:
    self.mu = jnp.asarray(mu)
    self.alpha = jnp.asarray(alpha)
    self.beta = jnp.asarray(beta)
    self.sigma = jnp.asarray(sigma)
    for name, val in (
        ("mu", self.mu),
        ("alpha", self.alpha),
        ("beta", self.beta),
        ("sigma", self.sigma),
    ):
        if val.shape != ():
            raise ValueError(
                f"SpatioTemporalHawkes distribution requires scalar `{name}`; "
                f"got shape {val.shape}. Vmap over parameters at the "
                "operator layer for batched draws."
            )
    self.spatial = spatial
    self.temporal = temporal
    self._max_events = int(max_events)
    self._op = _StHawkesOp(
        mu=self.mu,
        alpha=self.alpha,
        beta=self.beta,
        sigma=self.sigma,
        spatial=spatial,
        temporal=temporal,
        boundary_correction=boundary_correction,
    )
    super().__init__(batch_shape=(), validate_args=validate_args)

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

Return (locations, times, mask).

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

def sample(
    self,
    key: PRNGKeyArray,
    sample_shape: tuple[int, ...] = (),
) -> tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...]]:
    """Return ``(locations, times, mask)``."""
    check_prng_key(key)
    if sample_shape != ():
        raise NotImplementedError(
            "SpatioTemporalHawkes supports sample_shape=() only; "
            "vmap over PRNG keys at the operator layer for batched draws."
        )
    locations, times, mask, _ = self._op.sample(key, self._max_events)
    return locations, times, mask

log_prob(value: tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...]]) -> Float[Array, ...]

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

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

def log_prob(
    self,
    value: tuple[Float[Array, ...], Float[Array, ...], Bool[Array, ...]],
) -> Float[Array, ...]:
    """Log-likelihood :math:`\\sum_i \\log \\lambda^*(s_i, t_i) - \\Lambda`."""
    locations, times, mask = value
    return self._op.log_prob(locations, times, mask)