Point processes

The intensity and probability-of-detection registries as xtremax.point_processes processes: an inhomogeneous Poisson ground process per intensity kernel, a POD curve as a thinning retention, and the detected-event process ThinningProcess(MarkedTemporalPointProcess(...)). xtremax supplies the likelihoods and samplers; plumax supplies the methane-domain kernels. Needs the population extra. Tier V.

processes

Registry intensities and POD curves as xtremax point processes (implementation in plumax.population._src.processes).

LogIntensity

LogIntensity(intensity: Callable[[Array], Array], origin: ArrayLike = 0.0, n_integration_points: int = NODES_PER_DAY + 1)

Bases: Module

A registry intensity as an xtremax log-intensity on window time.

log λ̃(s) = log λ(t₀ + s), with s ∈ [0, T] the time since the window start t₀ in days, so the seasonal and diurnal phases of the kernel stay tied to calendar days.

Parameters:

Name	Type	Description	Default
intensity	Callable[[Array], Array]	A kernel from :mod:plumax.population.intensity.	required
origin	ArrayLike	Window start t₀ [days].	0.0
n_integration_points	int	Simpson nodes for the compensator when the kernel has no closed form.	NODES_PER_DAY + 1

Examples:

>>> import jax.numpy as jnp
>>> from plumax.population.intensity import ConstantIntensity
>>> log_intensity = LogIntensity(ConstantIntensity(lambda_0=2.0), origin=0.0)
>>> bool(jnp.allclose(log_intensity(jnp.zeros(3)), jnp.log(2.0)))
True
>>> float(log_intensity.integrate(0.0, 3.0))
6.0

Source code in src/plumax/population/_src/processes.py

def __init__(
    self,
    intensity: Callable[[Array], Array],
    origin: ArrayLike = 0.0,
    n_integration_points: int = NODES_PER_DAY + 1,
) -> None:
    self.intensity = intensity
    self.origin = jnp.asarray(origin, dtype=float)
    self.n_integration_points = n_integration_points

integrate

integrate(start: ArrayLike, stop: ArrayLike) -> Float[Array, '*time']

Compensator Λ̃(b) − Λ̃(a) = ∫ₐᵇ λ(t₀ + s) ds.

The kernel's own integrate when it has one (exact), else composite Simpson quadrature through xtremax.point_processes.

Parameters:

Name	Type	Description	Default
start	ArrayLike	Lower limit a [days since t₀].	required
stop	ArrayLike	Upper limit b [days since t₀].	required

Returns:

Type	Description
Float[Array, '*time']	Expected event count on [a, b].

Raises:

Type	Description
ImportError	Without the population extra, for kernels without a closed form.

Source code in src/plumax/population/_src/processes.py

def integrate(self, start: ArrayLike, stop: ArrayLike) -> Float[Array, "*time"]:
    """Compensator Λ̃(b) − Λ̃(a) = ∫ₐᵇ λ(t₀ + s) ds.

    The kernel's own ``integrate`` when it has one (exact), else composite
    Simpson quadrature through ``xtremax.point_processes``.

    Args:
        start: Lower limit a [days since t₀].
        stop: Upper limit b [days since t₀].

    Returns:
        Expected event count on ``[a, b]``.

    Raises:
        ImportError: Without the ``population`` extra, for kernels without
            a closed form.
    """
    closed_form = getattr(self.intensity, "integrate", None)
    if closed_form is not None:
        return closed_form(self.origin + start, self.origin + stop)
    point_processes = require_xtremax("LogIntensity.integrate")
    return point_processes.integrate_log_intensity(
        self,
        jnp.asarray(start, dtype=float),
        jnp.asarray(stop, dtype=float),
        n_points=self.n_integration_points,
        method="simpson",
    )

max_intensity

max_intensity() -> Float[Array, '']

Thinning envelope λ_max ≥ λ(t) for all t, from the kernel.

Source code in src/plumax/population/_src/processes.py

def max_intensity(self) -> Float[Array, ""]:
    """Thinning envelope λ_max ≥ λ(t) for all t, from the kernel."""
    return self.intensity.max_intensity()

PODRetention

PODRetention(pod: Callable[..., Array], covariates: Mapping[str, ArrayLike] | None = None, mark_scale: float = 1.0)

Bases: Module

A probability-of-detection curve as an xtremax retention callable.

p(t, H, Q) = P_d(c · Q; covariates), with c converting the mark units to the POD's flux units (the registry PODs take kg/h, so c = 3600 for marks in kg/s). Time and history are ignored: detection depends on the mark alone.

Parameters:

Name	Type	Description	Default
pod	Callable[..., Array]	A model from :mod:plumax.population.pod.	required
covariates	Mapping[str, ArrayLike] | None	Fixed keyword covariates of the POD beyond the flux (wind speed U, pixel size p, …), by its argument names.	None
mark_scale	float	The factor c from mark units to the POD's flux units.	1.0

Examples:

>>> import jax.numpy as jnp
>>> from plumax.population.pod import LogLogisticPOD
>>> retention = PODRetention(LogLogisticPOD(beta_0=0.0, beta_1=2.0))
>>> float(retention(0.0, None, jnp.asarray(1.0)))
0.5

Source code in src/plumax/population/_src/processes.py

def __init__(
    self,
    pod: Callable[..., Array],
    covariates: Mapping[str, ArrayLike] | None = None,
    mark_scale: float = 1.0,
) -> None:
    self.pod = pod
    self.covariates = {
        name: jnp.asarray(value) for name, value in (covariates or {}).items()
    }
    self.mark_scale = mark_scale

detection_process

detection_process(ground: Any, mark_distribution: Any, retention: Callable[..., Array], *, n_integration_points: int | None = None) -> Any

The detected-event process: POD-thinned, marked ground process.

λ_det(t, Q) = p(t, H, Q) · λ(t) · f(Q), as ThinningProcess(MarkedTemporalPointProcess(ground, f), p).

Parameters:

Name	Type	Description	Default
ground	Any	The ground process, e.g. from :func:make_process.	required
mark_distribution	Any	The mark law f(Q), a NumPyro distribution.	required
retention	Callable[..., Array]	The retention p, e.g. a :class:PODRetention.	required
n_integration_points	int | None	Trapezoid nodes of the retention compensator. Default: the ground process's own.	None

Returns:

Type	Description
Any	xtremax.point_processes.operators.ThinningProcess. Its
Any	log_prob(times, mask, marks, mark_sample_key=key) averages the
Any	retention over f by Monte Carlo inside the compensator (xtremax's
Any	n_mark_samples, see jejjohnson/xtremax#103).

Raises:

Type	Description
ImportError	Without the population extra.

Source code in src/plumax/population/_src/processes.py

def detection_process(
    ground: Any,
    mark_distribution: Any,
    retention: Callable[..., Array],
    *,
    n_integration_points: int | None = None,
) -> Any:
    """The detected-event process: POD-thinned, marked ground process.

    λ_det(t, Q) = p(t, H, Q) · λ(t) · f(Q), as
    ``ThinningProcess(MarkedTemporalPointProcess(ground, f), p)``.

    Args:
        ground: The ground process, e.g. from :func:`make_process`.
        mark_distribution: The mark law f(Q), a NumPyro distribution.
        retention: The retention p, e.g. a :class:`PODRetention`.
        n_integration_points: Trapezoid nodes of the retention compensator.
            Default: the ground process's own.

    Returns:
        ``xtremax.point_processes.operators.ThinningProcess``. Its
        ``log_prob(times, mask, marks, mark_sample_key=key)`` averages the
        retention over f by Monte Carlo inside the compensator (xtremax's
        ``n_mark_samples``, see jejjohnson/xtremax#103).

    Raises:
        ImportError: Without the ``population`` extra.
    """
    point_processes = require_xtremax("detection_process")
    operators = point_processes.operators
    marked = operators.MarkedTemporalPointProcess(
        ground, point_processes.constant_mark_distribution(mark_distribution)
    )
    if n_integration_points is None:
        n_integration_points = getattr(ground, "n_integration_points", 100)
    return operators.ThinningProcess(
        marked, retention, n_integration_points=n_integration_points
    )

expected_events_by_tod_bin

expected_events_by_tod_bin(intensity: Callable[[Array], Array], start: float, stop: float, *, n_bins: int, tod_origin: float = 0.0, nodes_per_day: int = NODES_PER_DAY) -> Float[Array, ' bin']

Expected events in each hour-of-day bin over a time interval.

Λ_b = Σ_c ∫_c λ(t) dt, over the cells c of [start, stop] that fall in bin b.

Parameters:

Name	Type	Description	Default
intensity	Callable[[Array], Array]	The rate λ(t) [events per day].	required
start	float	Start of the interval [days].	required
stop	float	End of the interval [days].	required
n_bins	int	Bins per day.	required
tod_origin	float	Local hour of day at t = 0 [hours].	0.0
nodes_per_day	int	Simpson nodes per day, without an exact integrate.	NODES_PER_DAY

Returns:

Type	Description
Float[Array, ' bin']	Λ_b for each bin.

Source code in src/plumax/population/_src/processes.py

def expected_events_by_tod_bin(
    intensity: Callable[[Array], Array],
    start: float,
    stop: float,
    *,
    n_bins: int,
    tod_origin: float = 0.0,
    nodes_per_day: int = NODES_PER_DAY,
) -> Float[Array, " bin"]:
    """Expected events in each hour-of-day bin over a time interval.

    Λ_b = Σ_c ∫_c λ(t) dt, over the cells c of [start, stop] that fall in bin b.

    Args:
        intensity: The rate λ(t) [events per day].
        start: Start of the interval [days].
        stop: End of the interval [days].
        n_bins: Bins per day.
        tod_origin: Local hour of day at t = 0 [hours].
        nodes_per_day: Simpson nodes per day, without an exact ``integrate``.

    Returns:
        Λ_b for each bin.
    """
    # host arithmetic: the edges depend on the interval, never on a traced parameter
    bins_per_hour = n_bins / 24.0
    shift = bins_per_hour * tod_origin
    first = math.floor(n_bins * float(start) + shift)
    last = math.ceil(n_bins * float(stop) + shift)
    indices = np.arange(first, last + 1)
    edges = np.clip((indices - shift) / n_bins, float(start), float(stop))
    cell_bins = jnp.asarray(indices[:-1] % n_bins)
    log_intensity = LogIntensity(
        intensity,
        n_integration_points=_quadrature_points(1.0 / n_bins, nodes_per_day),
    )
    # ∫_c λ(t) dt
    per_cell = jax.vmap(log_intensity.integrate)(
        jnp.asarray(edges[:-1]), jnp.asarray(edges[1:])
    )
    # Λ_b
    return jax.ops.segment_sum(per_cell, cell_bins, num_segments=n_bins)

make_process

make_process(intensity: str | Module, window: tuple[float, float], *, nodes_per_day: int = NODES_PER_DAY, **params: Any) -> Any

An xtremax inhomogeneous Poisson process for a registry intensity.

N(t) on s ∈ [0, T], T = t₁ − t₀, with log-intensity log λ(t₀ + s).

Parameters:

Name	Type	Description	Default
intensity	str | Module	A registry name (see :data:~plumax.population.intensity.INTENSITY_REGISTRY) or a kernel instance.	required
window	tuple[float, float]	(t₀, t₁), the observation window [days].	required
nodes_per_day	int	Quadrature density for kernels without a closed-form compensator.	NODES_PER_DAY
**params	Any	Kernel parameters, when intensity is a name.	{}

Returns:

Type	Description
Any	xtremax.point_processes.operators.InhomogeneousPoissonProcess over
Any	class:LogIntensity; its log_prob(times, mask) takes times since
Any	t₀ (see :func:padded_event_buffer).

Raises:

Type	Description
KeyError	For an unknown registry name.
ValueError	For an empty window, parameters passed with an instance, or the Weibull renewal hazard (a renewal process, not a Poisson intensity: use xtremax.point_processes.operators.RenewalProcess).
ImportError	Without the population extra.

Examples:

>>> import jax.numpy as jnp
>>> process = make_process("constant", (0.0, 10.0), lambda_0=0.5)
>>> times = jnp.array([1.0, 4.0, 10.0])
>>> mask = jnp.array([True, True, False])
>>> bool(jnp.allclose(process.log_prob(times, mask), 2 * jnp.log(0.5) - 5.0))
True

Source code in src/plumax/population/_src/processes.py

def make_process(
    intensity: str | eqx.Module,
    window: tuple[float, float],
    *,
    nodes_per_day: int = NODES_PER_DAY,
    **params: Any,
) -> Any:
    """An xtremax inhomogeneous Poisson process for a registry intensity.

    N(t) on s ∈ [0, T], T = t₁ − t₀, with log-intensity log λ(t₀ + s).

    Args:
        intensity: A registry name (see
            :data:`~plumax.population.intensity.INTENSITY_REGISTRY`) or a
            kernel instance.
        window: ``(t₀, t₁)``, the observation window [days].
        nodes_per_day: Quadrature density for kernels without a closed-form
            compensator.
        **params: Kernel parameters, when ``intensity`` is a name.

    Returns:
        ``xtremax.point_processes.operators.InhomogeneousPoissonProcess`` over
        :class:`LogIntensity`; its ``log_prob(times, mask)`` takes times since
        t₀ (see :func:`padded_event_buffer`).

    Raises:
        KeyError: For an unknown registry name.
        ValueError: For an empty window, parameters passed with an instance,
            or the Weibull renewal hazard (a renewal process, not a Poisson
            intensity: use ``xtremax.point_processes.operators.RenewalProcess``).
        ImportError: Without the ``population`` extra.

    Examples:
        >>> import jax.numpy as jnp
        >>> process = make_process("constant", (0.0, 10.0), lambda_0=0.5)
        >>> times = jnp.array([1.0, 4.0, 10.0])
        >>> mask = jnp.array([True, True, False])
        >>> bool(jnp.allclose(process.log_prob(times, mask), 2 * jnp.log(0.5) - 5.0))
        True
    """
    point_processes = require_xtremax("make_process")
    kernel = poisson_kernel(intensity, params, caller="make_process")
    start, stop = (float(edge) for edge in window)
    if not stop > start:
        raise ValueError(f"make_process: window must have t₁ > t₀, got {window}.")
    window_length = stop - start
    n_points = _quadrature_points(window_length, nodes_per_day)
    log_intensity = LogIntensity(kernel, origin=start, n_integration_points=n_points)
    return point_processes.operators.InhomogeneousPoissonProcess(
        log_intensity,
        observation_window=window_length,
        n_integration_points=n_points,
    )

padded_event_buffer

padded_event_buffer(times: ArrayLike | Sequence[float], window: tuple[float, float], *, max_events: int | None = None) -> tuple[Float[Array, ' event'], Bool[Array, ' event']]

Detection times as xtremax's padded buffer on window time.

Sorts s_i = t_i − t₀, pads to max_events with T = t₁ − t₀ and returns the mask, which is a contiguous prefix of True.

Parameters:

Name	Type	Description	Default
times	ArrayLike | Sequence[float]	Detection times t_i [days].	required
window	tuple[float, float]	(t₀, t₁) [days]; every time must lie inside it.	required
max_events	int | None	Buffer length. Default: the number of events.	None

Returns:

Type	Description
tuple[Float[Array, ' event'], Bool[Array, ' event']]	(event_times, mask) for xtremax log_prob calls.

Raises:

Type	Description
ValueError	If a time lies outside the window or the buffer is too short.

Examples:

>>> event_times, mask = padded_event_buffer([3.0, 1.0], (0.0, 5.0), max_events=3)
>>> event_times.tolist(), mask.tolist()
([1.0, 3.0, 5.0], [True, True, False])

Source code in src/plumax/population/_src/processes.py

def padded_event_buffer(
    times: ArrayLike | Sequence[float],
    window: tuple[float, float],
    *,
    max_events: int | None = None,
) -> tuple[Float[Array, " event"], Bool[Array, " event"]]:
    """Detection times as xtremax's padded buffer on window time.

    Sorts s_i = t_i − t₀, pads to ``max_events`` with T = t₁ − t₀ and returns
    the mask, which is a contiguous prefix of ``True``.

    Args:
        times: Detection times t_i [days].
        window: ``(t₀, t₁)`` [days]; every time must lie inside it.
        max_events: Buffer length. Default: the number of events.

    Returns:
        ``(event_times, mask)`` for xtremax ``log_prob`` calls.

    Raises:
        ValueError: If a time lies outside the window or the buffer is too
            short.

    Examples:
        >>> event_times, mask = padded_event_buffer([3.0, 1.0], (0.0, 5.0), max_events=3)
        >>> event_times.tolist(), mask.tolist()
        ([1.0, 3.0, 5.0], [True, True, False])
    """
    start, stop = (float(edge) for edge in window)
    host_times = np.sort(np.asarray(times, dtype=float).ravel())
    if host_times.size and (host_times[0] < start or host_times[-1] > stop):
        raise ValueError(
            f"padded_event_buffer: every time must lie in the window {window}."
        )
    n_events = host_times.size
    length = n_events if max_events is None else max_events
    if length < n_events:
        raise ValueError(
            f"padded_event_buffer: max_events={length} < {n_events} events."
        )
    window_time = np.full(length, stop - start)
    window_time[:n_events] = host_times - start
    mask = np.arange(length) < n_events
    return jnp.asarray(window_time), jnp.asarray(mask)