Skip to content

Point Processes — Overview

Shared infrastructure for the point-process stack. The three tiers built on top of this live on their own pages — primitives (pure functions), operators (equinox.Module objects), and distributions (NumPyro likelihoods).

Everything here is fixed-size by design: JAX needs static shapes, so event sequences are stored in padded buffers with a companion boolean mask rather than as ragged arrays. That mask threads through every likelihood and sampler in the package.

Domains

The observation window. Its measure appears directly in every likelihood — the compensator term is an integral over exactly this region — so getting it right matters as much as the intensity itself.

Point processes for xtremax.

Three layers, mirroring the rest of the package:

  • :mod:~xtremax.point_processes.primitives — pure functions (no equinox). These take a NumPyro Distribution as the inter-event / mark law where relevant; they never depend on the NumPyro Distribution base machinery (sample / log_prob / cdf is used directly).
  • :mod:~xtremax.point_processes.operators — equinox.Module classes that bundle intensity specs with numerical defaults. The primary user-facing API.
  • :mod:~xtremax.point_processes.distributions — thin numpyro.distributions.Distribution wrappers around the operators so these processes can appear inside NumPyro models.

Families currently available:

  • Homogeneous Poisson (HomogeneousPoissonProcess).
  • Inhomogeneous Poisson (InhomogeneousPoissonProcess).
  • Renewal (RenewalProcess) — any NumPyro inter-event distribution.
  • Self-exciting Hawkes (ExponentialHawkes, GeneralHawkesProcess).
  • Marked (MarkedTemporalPointProcess) — ground × mark-distribution.
  • Thinning (ThinningProcess) — base TPP × retention callable.

Padded-buffer invariant

Every padded event sequence in this package satisfies: the mask is a contiguous prefix (mask == [True]*n + [False]*(max-n)) and padding times sit at the right edge of the window (T; spatial padding rows at domain.lo). All samplers emit sequences in this form — thinning-based samplers compact accepted events before returning — and consumers with sequential structure (time_rescaling_residuals, exp_hawkes_log_prob via Ozaki gaps, renewal_log_prob via gap differences) assume it. Hand-constructed inputs with hole masks are supported only where a function explicitly documents mask-robust behaviour.

The shared :class:~xtremax.point_processes._history.EventHistory pytree is the lingua franca for user-supplied retention and mark callables: every family threads history through in the same shape so user code need not know the underlying family.

Sample-result convention

Every sampler returns a NamedTuple from :mod:~xtremax.point_processes._results — :class:SampleResult (times, mask, n_events) for unmarked temporal processes, :class:MarkedSampleResult (times, mask, marks) for marked ones, and the spatial/spatiotemporal analogues. These unpack positionally exactly like the tuples they replaced, are valid JAX pytrees, and let wrapper operators dispatch on the result type instead of shape/dtype heuristics. The n_events semantics are family-specific and documented per sampler (HPP: uncapped event count; thinning IPPs: uncapped candidate count; Hawkes: proposals consumed; renewal: retained events). sample_shape batching exists only on the temporal HPP; batch every other sampler with jax.vmap / equinox.filter_vmap over keys.

Operator architecture

Cross-family operator surface lives once in :mod:~xtremax.point_processes.operators._base: the time-rescaling diagnostics trio (residuals / goodness_of_fit / compensator_curve) in GoodnessOfFitMixin (families supply only _compensator_fn), the live-λ_max accessor in LiveIntensityMixin, and the separable marked-process plumbing in SeparableMarkedPP.

TemporalDomain(t0: ArrayLike, t1: ArrayLike)

Bases: Module

Half-open temporal interval [t0, t1) ⊂ ℝ.

Parallel to :class:RectangularDomain but specialised to one causally-ordered axis. Used by spatiotemporal point processes so the time domain stays explicit ([t0, t1) instead of an implicit [0, T]) and so callers can shift the origin without re-deriving quadrature offsets.

Parameters:

Name Type Description Default
t0 ArrayLike

Inclusive start time.

 required
t1 ArrayLike

Exclusive end time. Must satisfy t1 > t0.

 required
Source code in src/xtremax/point_processes/_domain.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
def __init__(self, t0: ArrayLike, t1: ArrayLike) -> None:
    t0_arr = jnp.asarray(t0)
    t1_arr = jnp.asarray(t1)
    if t0_arr.shape != t1_arr.shape:
        raise ValueError(
            f"TemporalDomain: `t0` shape {t0_arr.shape} does not match "
            f"`t1` shape {t1_arr.shape}."
        )
    try:
        if not bool(jnp.all(t1_arr > t0_arr)):
            raise ValueError(
                f"TemporalDomain requires `t1 > t0`; got "
                f"t0={t0_arr.tolist()}, t1={t1_arr.tolist()}."
            )
    except jax.errors.ConcretizationTypeError as exc:  # pragma: no cover
        del exc
    self.t0 = t0_arr
    self.t1 = t1_arr

duration: Float[Array, ...] property

Length of the interval t1 - t0.

from_duration(duration: ArrayLike) -> TemporalDomain classmethod

Interval anchored at zero: t0 = 0, t1 = duration.

Mirrors the source-snippet convention where users pass an observation_window scalar directly.

Source code in src/xtremax/point_processes/_domain.py 
163
164
165
166
167
168
169
170
171
@classmethod
def from_duration(cls, duration: ArrayLike) -> TemporalDomain:
    """Interval anchored at zero: ``t0 = 0``, ``t1 = duration``.

    Mirrors the source-snippet convention where users pass an
    ``observation_window`` scalar directly.
    """
    dur = jnp.asarray(duration)
    return cls(jnp.zeros_like(dur), dur)

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

Lebesgue measure of the interval — alias for :attr:duration.

Provided for API symmetry with :class:RectangularDomain so joint-domain code can call spatial.volume() * temporal.volume() without special-casing the temporal axis.

Source code in src/xtremax/point_processes/_domain.py 
178
179
180
181
182
183
184
185
def volume(self) -> Float[Array, ...]:
    """Lebesgue measure of the interval — alias for :attr:`duration`.

    Provided for API symmetry with :class:`RectangularDomain` so
    joint-domain code can call ``spatial.volume() * temporal.volume()``
    without special-casing the temporal axis.
    """
    return self.duration

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

Boolean True where t ∈ [t0, t1).

Source code in src/xtremax/point_processes/_domain.py 
187
188
189
def contains(self, t: Float[Array, ...]) -> Array:
    """Boolean ``True`` where ``t ∈ [t0, t1)``."""
    return (t >= self.t0) & (t < self.t1)

sample_uniform(key: PRNGKeyArray, shape: tuple[int, ...] = ()) -> Float[Array, ...]

Draw uniform times in the interval.

Output shape is exactly shape; no trailing dimension axis, because time is scalar.

Source code in src/xtremax/point_processes/_domain.py 
191
192
193
194
195
196
197
198
199
200
201
202
def sample_uniform(
    self,
    key: PRNGKeyArray,
    shape: tuple[int, ...] = (),
) -> Float[Array, ...]:
    """Draw uniform times in the interval.

    Output shape is exactly ``shape``; no trailing dimension axis,
    because time is scalar.
    """
    u = random.uniform(key, shape=tuple(shape))
    return self.t0 + u * self.duration

RectangularDomain(lo: ArrayLike, hi: ArrayLike)

Bases: Module

Axis-aligned rectangle [lo, hi] ⊂ ℝᵈ.

Both lo and hi are real PyTree leaves so the domain can be vmapped over (e.g. for batches with different bounding boxes) and so that, in principle, the bounds could be optimised through gradient descent.

Parameters:

Name Type Description Default
lo ArrayLike

Lower corner (d,). Defaults to the origin when only hi is provided via :meth:from_size.

 required
hi ArrayLike

Upper corner (d,). Must be strictly above lo elementwise.

 required
Notes

The convention is half-open [lo, hi) for sampling — the upper edge has measure zero so any choice is fine, but uniform draws on [lo, hi) are the natural pairing with lo + U·side.

Source code in src/xtremax/point_processes/_domain.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def __init__(self, lo: ArrayLike, hi: ArrayLike) -> None:
    lo_arr = jnp.asarray(lo)
    hi_arr = jnp.asarray(hi)
    # Validate eagerly at construction. Shape and ordering errors
    # are user-facing input bugs that should surface here rather
    # than producing silently negative volumes / out-of-box draws
    # downstream. Using ``bool(...)`` forces concretisation, which
    # is fine because RectangularDomain is invariably built from
    # static (non-traced) bounds; tracing through __init__ would
    # be a misuse.
    if lo_arr.shape != hi_arr.shape:
        raise ValueError(
            f"RectangularDomain: `lo` shape {lo_arr.shape} does not "
            f"match `hi` shape {hi_arr.shape}."
        )
    try:
        if not bool(jnp.all(hi_arr > lo_arr)):
            raise ValueError(
                f"RectangularDomain requires `hi > lo` elementwise; "
                f"got lo={lo_arr.tolist()}, hi={hi_arr.tolist()}."
            )
    except jax.errors.ConcretizationTypeError as exc:  # pragma: no cover
        # Building a domain inside a trace (e.g. inside ``jit``) is
        # not the supported path; skip the eager check rather than
        # erroring out of an otherwise valid jit-compiled function.
        del exc
    self.lo = lo_arr
    self.hi = hi_arr

n_dims: int property

Spatial dimension d.

Static at PyTree-construction time so callers can use it as a Python int (e.g. for kernel-grid construction).

side_lengths: Float[Array, ...] property

Per-axis box widths hi - lo with shape (d,).

from_size(size: ArrayLike) -> RectangularDomain classmethod

Box anchored at the origin: lo = 0, hi = size.

Mirrors the source-snippet convention where users pass a domain_size array directly. Always returns a 1-D hi array even for scalar input.

Source code in src/xtremax/point_processes/_domain.py 
80
81
82
83
84
85
86
87
88
89
@classmethod
def from_size(cls, size: ArrayLike) -> RectangularDomain:
    """Box anchored at the origin: ``lo = 0``, ``hi = size``.

    Mirrors the source-snippet convention where users pass a
    ``domain_size`` array directly. Always returns a 1-D ``hi``
    array even for scalar input.
    """
    size_arr = jnp.atleast_1d(jnp.asarray(size))
    return cls(jnp.zeros_like(size_arr), size_arr)

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

Lebesgue measure :math:|D| = \prod_i (hi_i - lo_i).

Source code in src/xtremax/point_processes/_domain.py 
105
106
107
def volume(self) -> Float[Array, ...]:
    """Lebesgue measure :math:`|D| = \\prod_i (hi_i - lo_i)`."""
    return jnp.prod(self.side_lengths, axis=-1)

contains(x: Float[Array, ...]) -> Array

Boolean True where every coordinate lies in the half-open box.

Source code in src/xtremax/point_processes/_domain.py 
109
110
111
def contains(self, x: Float[Array, ...]) -> Array:
    """Boolean ``True`` where every coordinate lies in the half-open box."""
    return jnp.all((x >= self.lo) & (x < self.hi), axis=-1)

sample_uniform(key: PRNGKeyArray, shape: tuple[int, ...] = ()) -> Float[Array, ...]

Draw uniform points in the box.

The output shape is shape + (d,) — shape is the leading sample/batch axes and d is the spatial dimension.

Source code in src/xtremax/point_processes/_domain.py 
113
114
115
116
117
118
119
120
121
122
123
124
def sample_uniform(
    self,
    key: PRNGKeyArray,
    shape: tuple[int, ...] = (),
) -> Float[Array, ...]:
    """Draw uniform points in the box.

    The output shape is ``shape + (d,)`` — ``shape`` is the leading
    sample/batch axes and ``d`` is the spatial dimension.
    """
    u = random.uniform(key, shape=(*tuple(shape), self.n_dims))
    return self.lo + u * self.side_lengths

Event history

The conditioning buffer for self-exciting processes: which events have already happened, and (optionally) their marks. Hawkes intensities read from this.

Point processes for xtremax.

Three layers, mirroring the rest of the package:

  • :mod:~xtremax.point_processes.primitives — pure functions (no equinox). These take a NumPyro Distribution as the inter-event / mark law where relevant; they never depend on the NumPyro Distribution base machinery (sample / log_prob / cdf is used directly).
  • :mod:~xtremax.point_processes.operators — equinox.Module classes that bundle intensity specs with numerical defaults. The primary user-facing API.
  • :mod:~xtremax.point_processes.distributions — thin numpyro.distributions.Distribution wrappers around the operators so these processes can appear inside NumPyro models.

Families currently available:

  • Homogeneous Poisson (HomogeneousPoissonProcess).
  • Inhomogeneous Poisson (InhomogeneousPoissonProcess).
  • Renewal (RenewalProcess) — any NumPyro inter-event distribution.
  • Self-exciting Hawkes (ExponentialHawkes, GeneralHawkesProcess).
  • Marked (MarkedTemporalPointProcess) — ground × mark-distribution.
  • Thinning (ThinningProcess) — base TPP × retention callable.

Padded-buffer invariant

Every padded event sequence in this package satisfies: the mask is a contiguous prefix (mask == [True]*n + [False]*(max-n)) and padding times sit at the right edge of the window (T; spatial padding rows at domain.lo). All samplers emit sequences in this form — thinning-based samplers compact accepted events before returning — and consumers with sequential structure (time_rescaling_residuals, exp_hawkes_log_prob via Ozaki gaps, renewal_log_prob via gap differences) assume it. Hand-constructed inputs with hole masks are supported only where a function explicitly documents mask-robust behaviour.

The shared :class:~xtremax.point_processes._history.EventHistory pytree is the lingua franca for user-supplied retention and mark callables: every family threads history through in the same shape so user code need not know the underlying family.

Sample-result convention

Every sampler returns a NamedTuple from :mod:~xtremax.point_processes._results — :class:SampleResult (times, mask, n_events) for unmarked temporal processes, :class:MarkedSampleResult (times, mask, marks) for marked ones, and the spatial/spatiotemporal analogues. These unpack positionally exactly like the tuples they replaced, are valid JAX pytrees, and let wrapper operators dispatch on the result type instead of shape/dtype heuristics. The n_events semantics are family-specific and documented per sampler (HPP: uncapped event count; thinning IPPs: uncapped candidate count; Hawkes: proposals consumed; renewal: retained events). sample_shape batching exists only on the temporal HPP; batch every other sampler with jax.vmap / equinox.filter_vmap over keys.

Operator architecture

Cross-family operator surface lives once in :mod:~xtremax.point_processes.operators._base: the time-rescaling diagnostics trio (residuals / goodness_of_fit / compensator_curve) in GoodnessOfFitMixin (families supply only _compensator_fn), the live-λ_max accessor in LiveIntensityMixin, and the separable marked-process plumbing in SeparableMarkedPP.

EventHistory

Bases: Module

Fixed-size buffer of events + optional marks.

Attributes:

Name Type Description
times Float[Array, ...]

Sorted event times with padding positions (typically) at the right edge of the observation window.
mask Bool[Array, ...]

True where times is a real event, False at padding positions.
marks Float[Array, ...] | None

Optional marks aligned with times. None for unmarked processes. For marked processes this is a 2-D (max_events, mark_dim) array with padding positions filled with zeros.

empty(max_events: int, mark_dim: int | None = None, dtype: jnp.dtype = jnp.float32, mark_dtype: jnp.dtype | None = None) -> EventHistory classmethod

Buffer with no events yet: all times 0, mask all False.

mark_dtype sets the marks-buffer dtype independently of the times dtype (defaults to dtype) — discrete marks would otherwise be silently cast to float.

Source code in src/xtremax/point_processes/_history.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
@classmethod
def empty(
    cls,
    max_events: int,
    mark_dim: int | None = None,
    dtype: jnp.dtype = jnp.float32,
    mark_dtype: jnp.dtype | None = None,
) -> EventHistory:
    """Buffer with no events yet: all times ``0``, mask all ``False``.

    ``mark_dtype`` sets the marks-buffer dtype independently of the
    times dtype (defaults to ``dtype``) — discrete marks would
    otherwise be silently cast to float.
    """
    times = jnp.zeros(max_events, dtype=dtype)
    mask = jnp.zeros(max_events, dtype=jnp.bool_)
    marks = (
        None
        if mark_dim is None
        else jnp.zeros(
            (max_events, mark_dim),
            dtype=dtype if mark_dtype is None else mark_dtype,
        )
    )
    return cls(times=times, mask=mask, marks=marks)

n_events() -> Int[Array, ...]

Number of real events currently in the buffer.

Source code in src/xtremax/point_processes/_history.py 
68
69
70
def n_events(self) -> Int[Array, ...]:
    """Number of real events currently in the buffer."""
    return jnp.sum(self.mask).astype(jnp.int32)

last_time(default: float | Float[Array, ...] = 0.0) -> Float[Array, ...]

Most recent real event time, or default if the history is empty.

Uses the mask directly rather than indexing by count, so it is safe under jit without triggering a dynamic gather.

Source code in src/xtremax/point_processes/_history.py 
72
73
74
75
76
77
78
79
80
81
82
83
84
def last_time(
    self,
    default: float | Float[Array, ...] = 0.0,
) -> Float[Array, ...]:
    """Most recent real event time, or ``default`` if the history is empty.

    Uses the mask directly rather than indexing by count, so it is
    safe under ``jit`` without triggering a dynamic gather.
    """
    default_arr = jnp.asarray(default, dtype=self.times.dtype)
    safe_times = jnp.where(self.mask, self.times, -jnp.inf)
    last = jnp.max(safe_times)
    return jnp.where(jnp.any(self.mask), last, default_arr)

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

t - last_time with default returned for empty history.

Source code in src/xtremax/point_processes/_history.py 
86
87
88
89
90
91
92
def time_since_last(
    self,
    t: Float[Array, ...],
    default: float | Float[Array, ...] = 0.0,
) -> Float[Array, ...]:
    """``t - last_time`` with ``default`` returned for empty history."""
    return jnp.asarray(t) - self.last_time(default=default)

append(t: Float[Array, ...], mark: Float[Array, ...] | None = None, *, accepted: Bool[Array, ...] | bool = True) -> EventHistory

Return a new history with t appended if accepted.

Writes at position n_events(); a False accepted or an overrun past the buffer is a no-op so the routine stays shape- and jit-stable. Callers that care about overruns can compare n_events() before/after.

Source code in src/xtremax/point_processes/_history.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def append(
    self,
    t: Float[Array, ...],
    mark: Float[Array, ...] | None = None,
    *,
    accepted: Bool[Array, ...] | bool = True,
) -> EventHistory:
    """Return a new history with ``t`` appended if ``accepted``.

    Writes at position ``n_events()``; a ``False`` ``accepted`` or an
    overrun past the buffer is a no-op so the routine stays shape-
    and jit-stable. Callers that care about overruns can compare
    ``n_events()`` before/after.
    """
    accepted_arr = jnp.asarray(accepted, dtype=jnp.bool_)
    max_events = self.times.shape[-1]
    idx = self.n_events()
    can_write = accepted_arr & (idx < max_events)
    safe_idx = jnp.clip(idx, 0, max_events - 1)

    current_time = self.times[safe_idx]
    current_mask = self.mask[safe_idx]
    new_time_at_idx = jnp.where(can_write, t, current_time)
    new_mask_at_idx = jnp.where(can_write, jnp.bool_(True), current_mask)
    new_times = self.times.at[safe_idx].set(new_time_at_idx)
    new_mask = self.mask.at[safe_idx].set(new_mask_at_idx)

    if self.marks is None:
        return EventHistory(times=new_times, mask=new_mask, marks=None)

    if mark is None:
        raise ValueError(
            "EventHistory has a marks buffer but append() was called "
            "without a mark; supply mark=<value> or use an unmarked history."
        )
    current_mark_row = self.marks[safe_idx]
    new_mark_row = jnp.where(can_write, jnp.asarray(mark), current_mark_row)
    new_marks = self.marks.at[safe_idx].set(new_mark_row)
    return EventHistory(times=new_times, mask=new_mask, marks=new_marks)

Sample results

What the samplers return: padded event arrays plus the mask that says which entries are real. One type per event geometry (temporal / spatial / spatiotemporal), each with a marked counterpart that carries the mark array alongside.

Point processes for xtremax.

Three layers, mirroring the rest of the package:

  • :mod:~xtremax.point_processes.primitives — pure functions (no equinox). These take a NumPyro Distribution as the inter-event / mark law where relevant; they never depend on the NumPyro Distribution base machinery (sample / log_prob / cdf is used directly).
  • :mod:~xtremax.point_processes.operators — equinox.Module classes that bundle intensity specs with numerical defaults. The primary user-facing API.
  • :mod:~xtremax.point_processes.distributions — thin numpyro.distributions.Distribution wrappers around the operators so these processes can appear inside NumPyro models.

Families currently available:

  • Homogeneous Poisson (HomogeneousPoissonProcess).
  • Inhomogeneous Poisson (InhomogeneousPoissonProcess).
  • Renewal (RenewalProcess) — any NumPyro inter-event distribution.
  • Self-exciting Hawkes (ExponentialHawkes, GeneralHawkesProcess).
  • Marked (MarkedTemporalPointProcess) — ground × mark-distribution.
  • Thinning (ThinningProcess) — base TPP × retention callable.

Padded-buffer invariant

Every padded event sequence in this package satisfies: the mask is a contiguous prefix (mask == [True]*n + [False]*(max-n)) and padding times sit at the right edge of the window (T; spatial padding rows at domain.lo). All samplers emit sequences in this form — thinning-based samplers compact accepted events before returning — and consumers with sequential structure (time_rescaling_residuals, exp_hawkes_log_prob via Ozaki gaps, renewal_log_prob via gap differences) assume it. Hand-constructed inputs with hole masks are supported only where a function explicitly documents mask-robust behaviour.

The shared :class:~xtremax.point_processes._history.EventHistory pytree is the lingua franca for user-supplied retention and mark callables: every family threads history through in the same shape so user code need not know the underlying family.

Sample-result convention

Every sampler returns a NamedTuple from :mod:~xtremax.point_processes._results — :class:SampleResult (times, mask, n_events) for unmarked temporal processes, :class:MarkedSampleResult (times, mask, marks) for marked ones, and the spatial/spatiotemporal analogues. These unpack positionally exactly like the tuples they replaced, are valid JAX pytrees, and let wrapper operators dispatch on the result type instead of shape/dtype heuristics. The n_events semantics are family-specific and documented per sampler (HPP: uncapped event count; thinning IPPs: uncapped candidate count; Hawkes: proposals consumed; renewal: retained events). sample_shape batching exists only on the temporal HPP; batch every other sampler with jax.vmap / equinox.filter_vmap over keys.

Operator architecture

Cross-family operator surface lives once in :mod:~xtremax.point_processes.operators._base: the time-rescaling diagnostics trio (residuals / goodness_of_fit / compensator_curve) in GoodnessOfFitMixin (families supply only _compensator_fn), the live-λ_max accessor in LiveIntensityMixin, and the separable marked-process plumbing in SeparableMarkedPP.

SampleResult

Bases: NamedTuple

Unmarked temporal sample: (times, mask, n_events).

SpatialSampleResult

Bases: NamedTuple

Unmarked spatial sample: (locations, mask, n_events).

SpatiotemporalSampleResult

Bases: NamedTuple

Unmarked spatiotemporal sample: (locations, times, mask, n_events).

MarkedSampleResult

Bases: NamedTuple

Marked temporal sample: (times, mask, marks).

Three fields (not four) so it unpacks identically to the tuple it replaces; the retained-event count is mask.sum().

MarkedSpatialSampleResult

Bases: NamedTuple

Marked spatial sample: (locations, mask, marks).

MarkedSpatiotemporalSampleResult

Bases: NamedTuple

Marked spatiotemporal sample: (locations, times, mask, marks).

Goodness of fit

The bundle returned by the temporal operators' goodness_of_fit method — time-rescaled residuals plus the KS statistic and QQ quantiles computed from them.

Point processes for xtremax.

Three layers, mirroring the rest of the package:

  • :mod:~xtremax.point_processes.primitives — pure functions (no equinox). These take a NumPyro Distribution as the inter-event / mark law where relevant; they never depend on the NumPyro Distribution base machinery (sample / log_prob / cdf is used directly).
  • :mod:~xtremax.point_processes.operators — equinox.Module classes that bundle intensity specs with numerical defaults. The primary user-facing API.
  • :mod:~xtremax.point_processes.distributions — thin numpyro.distributions.Distribution wrappers around the operators so these processes can appear inside NumPyro models.

Families currently available:

  • Homogeneous Poisson (HomogeneousPoissonProcess).
  • Inhomogeneous Poisson (InhomogeneousPoissonProcess).
  • Renewal (RenewalProcess) — any NumPyro inter-event distribution.
  • Self-exciting Hawkes (ExponentialHawkes, GeneralHawkesProcess).
  • Marked (MarkedTemporalPointProcess) — ground × mark-distribution.
  • Thinning (ThinningProcess) — base TPP × retention callable.

Padded-buffer invariant

Every padded event sequence in this package satisfies: the mask is a contiguous prefix (mask == [True]*n + [False]*(max-n)) and padding times sit at the right edge of the window (T; spatial padding rows at domain.lo). All samplers emit sequences in this form — thinning-based samplers compact accepted events before returning — and consumers with sequential structure (time_rescaling_residuals, exp_hawkes_log_prob via Ozaki gaps, renewal_log_prob via gap differences) assume it. Hand-constructed inputs with hole masks are supported only where a function explicitly documents mask-robust behaviour.

The shared :class:~xtremax.point_processes._history.EventHistory pytree is the lingua franca for user-supplied retention and mark callables: every family threads history through in the same shape so user code need not know the underlying family.

Sample-result convention

Every sampler returns a NamedTuple from :mod:~xtremax.point_processes._results — :class:SampleResult (times, mask, n_events) for unmarked temporal processes, :class:MarkedSampleResult (times, mask, marks) for marked ones, and the spatial/spatiotemporal analogues. These unpack positionally exactly like the tuples they replaced, are valid JAX pytrees, and let wrapper operators dispatch on the result type instead of shape/dtype heuristics. The n_events semantics are family-specific and documented per sampler (HPP: uncapped event count; thinning IPPs: uncapped candidate count; Hawkes: proposals consumed; renewal: retained events). sample_shape batching exists only on the temporal HPP; batch every other sampler with jax.vmap / equinox.filter_vmap over keys.

Operator architecture

Cross-family operator surface lives once in :mod:~xtremax.point_processes.operators._base: the time-rescaling diagnostics trio (residuals / goodness_of_fit / compensator_curve) in GoodnessOfFitMixin (families supply only _compensator_fn), the live-λ_max accessor in LiveIntensityMixin, and the separable marked-process plumbing in SeparableMarkedPP.

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.

Marks & retention builders

Small adapters that lift a simple callable into the full (coords, history) -> Distribution signature the marked and thinned processes expect — so a constant mark distribution or a time-varying retention probability doesn't require writing the whole closure by hand.

Point processes for xtremax.

Three layers, mirroring the rest of the package:

  • :mod:~xtremax.point_processes.primitives — pure functions (no equinox). These take a NumPyro Distribution as the inter-event / mark law where relevant; they never depend on the NumPyro Distribution base machinery (sample / log_prob / cdf is used directly).
  • :mod:~xtremax.point_processes.operators — equinox.Module classes that bundle intensity specs with numerical defaults. The primary user-facing API.
  • :mod:~xtremax.point_processes.distributions — thin numpyro.distributions.Distribution wrappers around the operators so these processes can appear inside NumPyro models.

Families currently available:

  • Homogeneous Poisson (HomogeneousPoissonProcess).
  • Inhomogeneous Poisson (InhomogeneousPoissonProcess).
  • Renewal (RenewalProcess) — any NumPyro inter-event distribution.
  • Self-exciting Hawkes (ExponentialHawkes, GeneralHawkesProcess).
  • Marked (MarkedTemporalPointProcess) — ground × mark-distribution.
  • Thinning (ThinningProcess) — base TPP × retention callable.

Padded-buffer invariant

Every padded event sequence in this package satisfies: the mask is a contiguous prefix (mask == [True]*n + [False]*(max-n)) and padding times sit at the right edge of the window (T; spatial padding rows at domain.lo). All samplers emit sequences in this form — thinning-based samplers compact accepted events before returning — and consumers with sequential structure (time_rescaling_residuals, exp_hawkes_log_prob via Ozaki gaps, renewal_log_prob via gap differences) assume it. Hand-constructed inputs with hole masks are supported only where a function explicitly documents mask-robust behaviour.

The shared :class:~xtremax.point_processes._history.EventHistory pytree is the lingua franca for user-supplied retention and mark callables: every family threads history through in the same shape so user code need not know the underlying family.

Sample-result convention

Every sampler returns a NamedTuple from :mod:~xtremax.point_processes._results — :class:SampleResult (times, mask, n_events) for unmarked temporal processes, :class:MarkedSampleResult (times, mask, marks) for marked ones, and the spatial/spatiotemporal analogues. These unpack positionally exactly like the tuples they replaced, are valid JAX pytrees, and let wrapper operators dispatch on the result type instead of shape/dtype heuristics. The n_events semantics are family-specific and documented per sampler (HPP: uncapped event count; thinning IPPs: uncapped candidate count; Hawkes: proposals consumed; renewal: retained events). sample_shape batching exists only on the temporal HPP; batch every other sampler with jax.vmap / equinox.filter_vmap over keys.

Operator architecture

Cross-family operator surface lives once in :mod:~xtremax.point_processes.operators._base: the time-rescaling diagnostics trio (residuals / goodness_of_fit / compensator_curve) in GoodnessOfFitMixin (families supply only _compensator_fn), the live-λ_max accessor in LiveIntensityMixin, and the separable marked-process plumbing in SeparableMarkedPP.

constant_mark_distribution(d: dist.Distribution) -> Callable[[Float[Array, ...], EventHistory], dist.Distribution]

Mark distribution that ignores time and history and returns d.

Source code in src/xtremax/point_processes/_adapters.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
def constant_mark_distribution(
    d: dist.Distribution,
) -> Callable[[Float[Array, ...], EventHistory], dist.Distribution]:
    """Mark distribution that ignores time and history and returns ``d``."""

    def _fn(
        t: Float[Array, ...],
        history: EventHistory,
    ) -> dist.Distribution:
        del t, history
        return d

    return _fn

time_varying_marks(fn: Callable[[Float[Array, ...]], dist.Distribution]) -> Callable[[Float[Array, ...], EventHistory], dist.Distribution]

Wrap a time → Distribution callable into the full signature.

Source code in src/xtremax/point_processes/_adapters.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
def time_varying_marks(
    fn: Callable[[Float[Array, ...]], dist.Distribution],
) -> Callable[[Float[Array, ...], EventHistory], dist.Distribution]:
    """Wrap a time → Distribution callable into the full signature."""

    def _fn(
        t: Float[Array, ...],
        history: EventHistory,
    ) -> dist.Distribution:
        del history
        return fn(t)

    return _fn

constant_retention(p: float | Float[Array, ...]) -> Callable[[Float[Array, ...], EventHistory, Float[Array, ...] | None], Array]

Retention that always returns the same probability p.

Source code in src/xtremax/point_processes/_adapters.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def constant_retention(
    p: float | Float[Array, ...],
) -> Callable[[Float[Array, ...], EventHistory, Float[Array, ...] | None], Array]:
    """Retention that always returns the same probability ``p``."""
    p_arr = jnp.asarray(p)

    def _fn(
        t: Float[Array, ...],
        history: EventHistory,
        proposed_mark: Float[Array, ...] | None = None,
    ) -> Float[Array, ...]:
        del t, history, proposed_mark
        return p_arr

    return _fn

time_varying_retention(fn: Callable[[Float[Array, ...]], Float[Array, ...]]) -> Callable[[Float[Array, ...], EventHistory, Float[Array, ...] | None], Array]

Wrap a pure function of time into the full retention signature.

Source code in src/xtremax/point_processes/_adapters.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def time_varying_retention(
    fn: Callable[[Float[Array, ...]], Float[Array, ...]],
) -> Callable[[Float[Array, ...], EventHistory, Float[Array, ...] | None], Array]:
    """Wrap a pure function of time into the full retention signature."""

    def _fn(
        t: Float[Array, ...],
        history: EventHistory,
        proposed_mark: Float[Array, ...] | None = None,
    ) -> Float[Array, ...]:
        del history, proposed_mark
        return fn(t)

    return _fn

Log-intensity integration

Quadrature for the compensator \(\Lambda = \int \exp(\log \lambda)\) over each domain type. Intensities are parameterised in log space for positivity, so these integrate the exponential rather than the callable directly.

Point processes for xtremax.

Three layers, mirroring the rest of the package:

  • :mod:~xtremax.point_processes.primitives — pure functions (no equinox). These take a NumPyro Distribution as the inter-event / mark law where relevant; they never depend on the NumPyro Distribution base machinery (sample / log_prob / cdf is used directly).
  • :mod:~xtremax.point_processes.operators — equinox.Module classes that bundle intensity specs with numerical defaults. The primary user-facing API.
  • :mod:~xtremax.point_processes.distributions — thin numpyro.distributions.Distribution wrappers around the operators so these processes can appear inside NumPyro models.

Families currently available:

  • Homogeneous Poisson (HomogeneousPoissonProcess).
  • Inhomogeneous Poisson (InhomogeneousPoissonProcess).
  • Renewal (RenewalProcess) — any NumPyro inter-event distribution.
  • Self-exciting Hawkes (ExponentialHawkes, GeneralHawkesProcess).
  • Marked (MarkedTemporalPointProcess) — ground × mark-distribution.
  • Thinning (ThinningProcess) — base TPP × retention callable.

Padded-buffer invariant

Every padded event sequence in this package satisfies: the mask is a contiguous prefix (mask == [True]*n + [False]*(max-n)) and padding times sit at the right edge of the window (T; spatial padding rows at domain.lo). All samplers emit sequences in this form — thinning-based samplers compact accepted events before returning — and consumers with sequential structure (time_rescaling_residuals, exp_hawkes_log_prob via Ozaki gaps, renewal_log_prob via gap differences) assume it. Hand-constructed inputs with hole masks are supported only where a function explicitly documents mask-robust behaviour.

The shared :class:~xtremax.point_processes._history.EventHistory pytree is the lingua franca for user-supplied retention and mark callables: every family threads history through in the same shape so user code need not know the underlying family.

Sample-result convention

Every sampler returns a NamedTuple from :mod:~xtremax.point_processes._results — :class:SampleResult (times, mask, n_events) for unmarked temporal processes, :class:MarkedSampleResult (times, mask, marks) for marked ones, and the spatial/spatiotemporal analogues. These unpack positionally exactly like the tuples they replaced, are valid JAX pytrees, and let wrapper operators dispatch on the result type instead of shape/dtype heuristics. The n_events semantics are family-specific and documented per sampler (HPP: uncapped event count; thinning IPPs: uncapped candidate count; Hawkes: proposals consumed; renewal: retained events). sample_shape batching exists only on the temporal HPP; batch every other sampler with jax.vmap / equinox.filter_vmap over keys.

Operator architecture

Cross-family operator surface lives once in :mod:~xtremax.point_processes.operators._base: the time-rescaling diagnostics trio (residuals / goodness_of_fit / compensator_curve) in GoodnessOfFitMixin (families supply only _compensator_fn), the live-λ_max accessor in LiveIntensityMixin, and the separable marked-process plumbing in SeparableMarkedPP.

cumulative_log_intensity(log_intensity_fn: Callable[[Float[Array, ...]], Float[Array, ...]], t: Float[Array, ...], n_points: int = 100, method: Method = 'trapezoid') -> Float[Array, ...]

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

Parameters:

Name Type Description Default
log_intensity_fn Callable[[Float[Array, ...]], Float[Array, ...]]

Log-intensity callable (see :func:integrate_log_intensity).

 required
t Float[Array, ...]

Upper limit(s). Must be >= 0.

 required
n_points int

Quadrature nodes.

 100
method Method

Quadrature rule.

 'trapezoid'

Returns:

Type Description
Float[Array, ...]

Cumulative intensity with the same shape as t.
Source code in src/xtremax/point_processes/_integration.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def cumulative_log_intensity(
    log_intensity_fn: Callable[[Float[Array, ...]], Float[Array, ...]],
    t: Float[Array, ...],
    n_points: int = 100,
    method: Method = "trapezoid",
) -> Float[Array, ...]:
    """Shorthand for :math:`\\Lambda(t) = \\int_0^t \\lambda(s) ds`.

    Args:
        log_intensity_fn: Log-intensity callable (see
            :func:`integrate_log_intensity`).
        t: Upper limit(s). Must be ``>= 0``.
        n_points: Quadrature nodes.
        method: Quadrature rule.

    Returns:
        Cumulative intensity with the same shape as ``t``.
    """
    t = jnp.asarray(t)
    return integrate_log_intensity(
        log_intensity_fn, jnp.zeros_like(t), t, n_points=n_points, method=method
    )

integrate_log_intensity(log_intensity_fn: Callable[[Float[Array, ...]], Float[Array, ...]], t0: Float[Array, ...], t1: Float[Array, ...], n_points: int = 100, method: Method = 'trapezoid') -> Float[Array, ...]

Compute :math:\int_{t_0}^{t_1} \exp(\log\lambda(s))\, ds.

Parameters:

Name Type Description Default
log_intensity_fn Callable[[Float[Array, ...]], Float[Array, ...]]

Callable mapping a 1-D (n_points,) array of times to an array of log-intensities with the same shape. Must support vectorised evaluation along the grid axis.

 required
t0 Float[Array, ...]

Lower limit (scalar or batched). Broadcast against t1.

 required
t1 Float[Array, ...]

Upper limit (scalar or batched). Broadcast against t0.

 required
n_points int

Number of quadrature nodes (static Python int). Must be >= 2 for trapezoid and >= 3 (odd) for Simpson.

 100
method Method

Quadrature rule — "trapezoid" (default) or "simpson".

 'trapezoid'

Returns:

Type Description
Float[Array, ...]

Array with shape broadcast(t0, t1) giving the integral.
Notes

For batched limits, each pair :math:(t_0^{(i)}, t_1^{(i)}) gets its own linspace grid via vmap; the user's log_intensity_fn only ever sees a 1-D grid of times.

Source code in src/xtremax/point_processes/_integration.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def integrate_log_intensity(
    log_intensity_fn: Callable[[Float[Array, ...]], Float[Array, ...]],
    t0: Float[Array, ...],
    t1: Float[Array, ...],
    n_points: int = 100,
    method: Method = "trapezoid",
) -> Float[Array, ...]:
    """Compute :math:`\\int_{t_0}^{t_1} \\exp(\\log\\lambda(s))\\, ds`.

    Args:
        log_intensity_fn: Callable mapping a 1-D ``(n_points,)`` array of
            times to an array of log-intensities with the same shape.
            Must support vectorised evaluation along the grid axis.
        t0: Lower limit (scalar or batched). Broadcast against ``t1``.
        t1: Upper limit (scalar or batched). Broadcast against ``t0``.
        n_points: Number of quadrature nodes (static Python int). Must
            be ``>= 2`` for trapezoid and ``>= 3`` (odd) for Simpson.
        method: Quadrature rule — ``"trapezoid"`` (default) or
            ``"simpson"``.

    Returns:
        Array with shape ``broadcast(t0, t1)`` giving the integral.

    Notes:
        For batched limits, each pair :math:`(t_0^{(i)}, t_1^{(i)})`
        gets its own linspace grid via ``vmap``; the user's
        ``log_intensity_fn`` only ever sees a 1-D grid of times.
    """
    if n_points < 2:
        raise ValueError(f"n_points must be >= 2, got {n_points}")
    if method == "simpson" and n_points % 2 == 0:
        # Composite Simpson's rule requires an odd number of nodes
        # (equivalently, an even number of sub-intervals).
        n_points = n_points + 1

    t0 = jnp.asarray(t0)
    t1 = jnp.asarray(t1)

    def _single(a: Array, b: Array) -> Array:
        grid = jnp.linspace(a, b, n_points)
        intensities = jnp.exp(log_intensity_fn(grid))
        if method == "trapezoid":
            return jnp.trapezoid(intensities, grid)
        # Simpson's composite rule on uniform grid.
        h = (b - a) / (n_points - 1)
        # Weight pattern: 1, 4, 2, 4, 2, ..., 4, 1.
        idx = jnp.arange(n_points)
        weights = jnp.where(
            (idx == 0) | (idx == n_points - 1),
            1.0,
            jnp.where(idx % 2 == 1, 4.0, 2.0),
        )
        return (h / 3.0) * jnp.sum(weights * intensities)

    broadcast_shape = jnp.broadcast_shapes(t0.shape, t1.shape)
    t0_b = jnp.broadcast_to(t0, broadcast_shape)
    t1_b = jnp.broadcast_to(t1, broadcast_shape)

    flat0 = t0_b.reshape(-1)
    flat1 = t1_b.reshape(-1)
    flat_result = jax.vmap(_single)(flat0, flat1)
    return flat_result.reshape(broadcast_shape)

integrate_log_intensity_spatial(log_intensity_fn: Callable[[Float[Array, ...]], Float[Array, ...]], domain: RectangularDomain, n_points: int = 2048, method: SpatialMethod = 'qmc') -> Float[Array, ...]

Compute :math:\int_D \exp(\log\lambda(s))\, ds.

Parameters:

Name Type Description Default
log_intensity_fn Callable[[Float[Array, ...]], Float[Array, ...]]

Callable mapping (N, d) locations to (N,) log-intensities. Must accept a 2-D leading-batch array — the same convention as the temporal IPP, extended with a trailing dimension axis.

 required
domain RectangularDomain

Spatial domain over which to integrate.

 required
n_points int

Static number of quadrature/QMC nodes. For method="trapezoid" this is the total node count; the per-axis grid is :math:n^{1/d} rounded up. For method="qmc" it is the literal sample count.

 2048
method SpatialMethod

"trapezoid" (closed-form on tensor-product grid, best for d ≤ 2) or "qmc" (Halton sequence, best for d ≥ 3 and the default for any d).

 'qmc'

Returns:

Type Description
Float[Array, ...]

Scalar Λ, differentiable through log_intensity_fn's
Float[Array, ...]

parameters.
Notes

Both methods are deterministic given a fixed n_points — no PRNG key is needed. QMC uses a digit-reversed Halton sequence on the first d primes, which is reproducible and avoids the wrap-around bias of plain MC.

Source code in src/xtremax/point_processes/_integration_spatial.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def integrate_log_intensity_spatial(
    log_intensity_fn: Callable[[Float[Array, ...]], Float[Array, ...]],
    domain: RectangularDomain,
    n_points: int = 2048,
    method: SpatialMethod = "qmc",
) -> Float[Array, ...]:
    """Compute :math:`\\int_D \\exp(\\log\\lambda(s))\\, ds`.

    Args:
        log_intensity_fn: Callable mapping ``(N, d)`` locations to
            ``(N,)`` log-intensities. Must accept a 2-D leading-batch
            array — the same convention as the temporal IPP, extended
            with a trailing dimension axis.
        domain: Spatial domain over which to integrate.
        n_points: Static number of quadrature/QMC nodes. For
            ``method="trapezoid"`` this is the *total* node count; the
            per-axis grid is :math:`n^{1/d}` rounded up. For
            ``method="qmc"`` it is the literal sample count.
        method: ``"trapezoid"`` (closed-form on tensor-product grid,
            best for ``d ≤ 2``) or ``"qmc"`` (Halton sequence, best for
            ``d ≥ 3`` and the default for any ``d``).

    Returns:
        Scalar Λ, differentiable through ``log_intensity_fn``'s
        parameters.

    Notes:
        Both methods are deterministic given a fixed ``n_points`` —
        no PRNG key is needed. QMC uses a digit-reversed Halton
        sequence on the first ``d`` primes, which is reproducible and
        avoids the wrap-around bias of plain MC.
    """
    if method == "trapezoid":
        return _integrate_trapezoid(log_intensity_fn, domain, n_points)
    if method == "qmc":
        return _integrate_qmc(log_intensity_fn, domain, n_points)
    raise ValueError(f"Unknown method {method!r}; expected 'trapezoid' or 'qmc'.")

integrate_log_intensity_spatiotemporal(log_intensity_fn: Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]], spatial: RectangularDomain, temporal: TemporalDomain, n_points: int = 4096, method: SpatiotemporalMethod = 'qmc') -> Float[Array, ...]

Compute :math:\int_{t_0}^{t_1} \int_D \exp(\log\lambda(s, t))\, ds\, dt.

Parameters:

Name Type Description Default
log_intensity_fn Callable[[Float[Array, ...], Float[Array, ...]], Float[Array, ...]]

Callable (s, t) → log λ(s, t) where s has shape (N, d_space) and t has shape (N,).

 required
spatial RectangularDomain

Spatial domain.

 required
temporal TemporalDomain

Temporal interval.

 required
n_points int

Static node budget across all d_space + 1 axes. For trapezoid this is the total cell count (per-axis nodes are :math:n^{1/(d+1)} rounded up).

 4096
method SpatiotemporalMethod

"trapezoid" (closed form, best for d_space ≤ 2) or "qmc" (Halton sequence, default).

 'qmc'

Returns:

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

math:\Lambda, differentiable through
Float[Array, ...]

log_intensity_fn's parameters.
Source code in src/xtremax/point_processes/_integration_spatiotemporal.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def integrate_log_intensity_spatiotemporal(
    log_intensity_fn: Callable[
        [Float[Array, ...], Float[Array, ...]], Float[Array, ...]
    ],
    spatial: RectangularDomain,
    temporal: TemporalDomain,
    n_points: int = 4096,
    method: SpatiotemporalMethod = "qmc",
) -> Float[Array, ...]:
    """Compute :math:`\\int_{t_0}^{t_1} \\int_D \\exp(\\log\\lambda(s, t))\\, ds\\, dt`.

    Args:
        log_intensity_fn: Callable ``(s, t) → log λ(s, t)`` where ``s``
            has shape ``(N, d_space)`` and ``t`` has shape ``(N,)``.
        spatial: Spatial domain.
        temporal: Temporal interval.
        n_points: Static node budget across all ``d_space + 1`` axes.
            For ``trapezoid`` this is the total cell count
            (per-axis nodes are :math:`n^{1/(d+1)}` rounded up).
        method: ``"trapezoid"`` (closed form, best for ``d_space ≤ 2``)
            or ``"qmc"`` (Halton sequence, default).

    Returns:
        Scalar :math:`\\Lambda`, differentiable through
        ``log_intensity_fn``'s parameters.
    """
    if method == "trapezoid":
        return _integrate_trapezoid(log_intensity_fn, spatial, temporal, n_points)
    if method == "qmc":
        return _integrate_qmc(log_intensity_fn, spatial, temporal, n_points)
    raise ValueError(f"Unknown method {method!r}; expected 'trapezoid' or 'qmc'.")