Importance correction

The importance-corrected TMTPP mark likelihood: per-event evidence in the POINT, GAUSSIAN and FULL regimes, which divides each event's per-event prior out of its posterior; the ESS and Pareto-k̂ diagnostics; log-space quadrature for the mark integrals and the thinned rate; and the NumPyro models and fit that use them (mark_iw_model, tmtpp_iw_model, fit_size_distribution_iw, in plumax.population.fitting). Needs the population extra. Tier V.

importance — per-event evidence and ESS

Importance-corrected per-event evidence (implementation in plumax.population._src.importance).

EventEvidence

Bases: Module

Nodes and log-weights of one event's (or a stack of events') mark integral.

The mark term is log Σₛ exp(ℓₛ + g(qₛ)) for g = log P_d + log f(· | θ).

Attributes:

Name	Type	Description
log_weights	Float[Array, '*event member']	ℓₛ, including −log π_i(qₛ).
q_samples	Float[Array, '*event member']	The nodes qₛ [kg/s].

log_terms

log_terms(log_integrand: Callable[[Array], Array]) -> Float[Array, '*event member']

ℓₛ + g(qₛ), with nodes at Q ≤ 0 excluded (−∞).

g is evaluated on a positive stand-in at excluded nodes, so a POD or mark that is undefined there cannot turn the sum into NaN.

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

def log_terms(
    self, log_integrand: Callable[[Array], Array]
) -> Float[Array, "*event member"]:
    """ℓₛ + g(qₛ), with nodes at Q ≤ 0 excluded (−∞).

    g is evaluated on a positive stand-in at excluded nodes, so a POD or
    mark that is undefined there cannot turn the sum into NaN.
    """
    positive = self.q_samples > 0.0
    safe_nodes = jnp.where(positive, self.q_samples, 1.0)
    terms = self.log_weights + log_integrand(safe_nodes)
    return jnp.where(positive, terms, -jnp.inf)

log_mark_term

log_mark_term(log_integrand: Callable[[Array], Array]) -> Float[Array, '*event']

log Σₛ exp(ℓₛ + g(qₛ)), per event.

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

def log_mark_term(
    self, log_integrand: Callable[[Array], Array]
) -> Float[Array, "*event"]:
    """log Σₛ exp(ℓₛ + g(qₛ)), per event."""
    return logsumexp(self.log_terms(log_integrand), axis=-1)

effective_sample_size

effective_sample_size(log_integrand: Callable[[Array], Array] | None = None) -> Float[Array, '*event']

ESS = (Σₛ wₛ)² / Σₛ wₛ² of wₛ = exp(ℓₛ + g(qₛ)), per event.

Without log_integrand the weights are the prior correction alone.

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

def effective_sample_size(
    self, log_integrand: Callable[[Array], Array] | None = None
) -> Float[Array, "*event"]:
    """ESS = (Σₛ wₛ)² / Σₛ wₛ² of wₛ = exp(ℓₛ + g(qₛ)), per event.

    Without ``log_integrand`` the weights are the prior correction alone.
    """
    if log_integrand is None:
        return effective_sample_size(self.log_weights)
    return effective_sample_size(self.log_terms(log_integrand))

Regime

Bases: StrEnum

How an event's per-event posterior enters the mark integral.

catalog_evidence

catalog_evidence(catalog: EmissionCatalog, *, regime: Regime | str = Regime.FULL, n_nodes: int = 32, measure: Measure = 'lognormal') -> EventEvidence

:func:event_evidence for every event, stacked along a leading axis.

Raises:

Type	Description
ValueError	For an empty catalog, or FULL draws of unequal length.

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

def catalog_evidence(
    catalog: EmissionCatalog,
    *,
    regime: Regime | str = Regime.FULL,
    n_nodes: int = 32,
    measure: Measure = "lognormal",
) -> EventEvidence:
    """:func:`event_evidence` for every event, stacked along a leading axis.

    Raises:
        ValueError: For an empty catalog, or FULL draws of unequal length.
    """
    if catalog.n_events == 0:
        raise ValueError("catalog_evidence: the catalog is empty.")
    per_event = [
        event_evidence(event, regime=regime, n_nodes=n_nodes, measure=measure)
        for event in catalog.events
    ]
    lengths = {evidence.q_samples.shape[0] for evidence in per_event}
    if len(lengths) != 1:
        raise ValueError(
            f"catalog_evidence: events carry different numbers of draws {lengths}."
        )
    return EventEvidence(
        log_weights=jnp.stack([evidence.log_weights for evidence in per_event]),
        q_samples=jnp.stack([evidence.q_samples for evidence in per_event]),
    )

effective_sample_size

effective_sample_size(log_weights: Float[Array, '*event member']) -> Array

ESS = (Σₛ wₛ)² / Σₛ wₛ² from log-weights, along the last axis.

Examples:

>>> import jax.numpy as jnp
>>> bool(jnp.allclose(effective_sample_size(jnp.zeros(8)), 8.0))
True

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

def effective_sample_size(log_weights: Float[Array, "*event member"]) -> Array:
    """ESS = (Σₛ wₛ)² / Σₛ wₛ² from log-weights, along the last axis.

    Examples:
        >>> import jax.numpy as jnp
        >>> bool(jnp.allclose(effective_sample_size(jnp.zeros(8)), 8.0))
        True
    """
    # 2 log Σ w − log Σ w²
    log_ess = 2.0 * logsumexp(log_weights, axis=-1) - logsumexp(
        2.0 * log_weights, axis=-1
    )
    return jnp.exp(log_ess)

ess_report

ess_report(catalog: EmissionCatalog, *, log_integrand: Callable[[Array], Array], regime: Regime | str = Regime.FULL, measure: Measure = 'lognormal', warn: bool = True) -> np.ndarray

Per-event ESS of the importance weights at one θ, with a collapse warning.

wₛ = P_d(qₛ) f(qₛ | θ) / π_i(qₛ) for g = log P_d + log f(· | θ). Warns (does not fail) when an event's ESS < 0.1 S or, in the FULL regime, its Pareto k̂ > 0.7.

Parameters:

Name	Type	Description	Default
catalog	EmissionCatalog	The events.	required
log_integrand	Callable[[Array], Array]	g, e.g. lambda q: mark.log_prob(q, **theta).	required
regime	Regime | str	How the events enter.	FULL
measure	Measure	The GAUSSIAN-regime summary (see :func:event_evidence).	'lognormal'
warn	bool	Emit the collapse warning.	True

Returns:

Type	Description
ndarray	ESS per event.

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

def ess_report(
    catalog: EmissionCatalog,
    *,
    log_integrand: Callable[[Array], Array],
    regime: Regime | str = Regime.FULL,
    measure: Measure = "lognormal",
    warn: bool = True,
) -> np.ndarray:
    """Per-event ESS of the importance weights at one θ, with a collapse warning.

    wₛ = P_d(qₛ) f(qₛ | θ) / π_i(qₛ) for g = log P_d + log f(· | θ). Warns (does
    not fail) when an event's ESS < 0.1 S or, in the FULL regime, its Pareto
    k̂ > 0.7.

    Args:
        catalog: The events.
        log_integrand: g, e.g. ``lambda q: mark.log_prob(q, **theta)``.
        regime: How the events enter.
        measure: The GAUSSIAN-regime summary (see :func:`event_evidence`).
        warn: Emit the collapse warning.

    Returns:
        ESS per event.
    """
    evidence = catalog_evidence(catalog, regime=regime, measure=measure)
    log_w = np.asarray(evidence.log_terms(log_integrand))
    ess = np.asarray(effective_sample_size(jnp.asarray(log_w)))
    if warn:
        pareto_k = None
        if Regime(regime) is Regime.FULL:
            pareto_k = np.array([pareto_smoothed_log_weights(row)[1] for row in log_w])
        warn_on_collapse(ess, log_w.shape[-1], pareto_k)
    return ess

event_evidence

event_evidence(event: EmissionEvent, *, regime: Regime | str = Regime.FULL, n_nodes: int = 32, measure: Measure = 'lognormal') -> EventEvidence

Nodes and log-weights of one event's importance-corrected mark integral.

In the GAUSSIAN regime the posterior is replaced by its two-moment summary, a lognormal (the 06a default) or a normal in Q (measure="normal", the family of the Gaussian Tier II–IV posteriors). The prior is divided out exactly, so the summary's tails must not be heavier than the prior's: when the summary's likelihood p̃/π rises at an outermost Hermite node more than TAIL_TOLERANCE nats above its interior maximum, the approximation is rejected, since the quadrature would then be driven by its extreme nodes. Nodes at Q ≤ 0 carry no weight (every registry mark lives on Q > 0).

Parameters:

Name	Type	Description	Default
event	EmissionEvent	The event, with prior_logpdf (and q_samples for FULL).	required
regime	Regime | str	POINT, GAUSSIAN or FULL.	FULL
n_nodes	int	Gauss–Hermite nodes for the GAUSSIAN regime.	32
measure	Measure	"lognormal" or "normal", the GAUSSIAN-regime summary.	'lognormal'

Returns:

Name	Type	Description
The	EventEvidence	class:EventEvidence.

Raises:

Type	Description
PayloadError	Without the per-event prior, without draws for FULL, or when the GAUSSIAN summary's tails outweigh the prior's.
ValueError	For a non-positive emission rate outside FULL, or an unknown measure.

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

def event_evidence(
    event: EmissionEvent,
    *,
    regime: Regime | str = Regime.FULL,
    n_nodes: int = 32,
    measure: Measure = "lognormal",
) -> EventEvidence:
    """Nodes and log-weights of one event's importance-corrected mark integral.

    In the GAUSSIAN regime the posterior is replaced by its two-moment summary,
    a lognormal (the 06a default) or a normal in Q (``measure="normal"``, the
    family of the Gaussian Tier II–IV posteriors). The prior is divided out
    exactly, so the summary's tails must not be heavier than the prior's: when
    the summary's likelihood p̃/π rises at an outermost Hermite node more than
    ``TAIL_TOLERANCE`` nats above its interior maximum, the approximation is
    rejected, since the quadrature would then be driven by its extreme nodes.
    Nodes at Q ≤ 0 carry no weight (every registry mark lives on Q > 0).

    Args:
        event: The event, with ``prior_logpdf`` (and ``q_samples`` for FULL).
        regime: POINT, GAUSSIAN or FULL.
        n_nodes: Gauss–Hermite nodes for the GAUSSIAN regime.
        measure: ``"lognormal"`` or ``"normal"``, the GAUSSIAN-regime summary.

    Returns:
        The :class:`EventEvidence`.

    Raises:
        PayloadError: Without the per-event prior, without draws for FULL, or
            when the GAUSSIAN summary's tails outweigh the prior's.
        ValueError: For a non-positive emission rate outside FULL, or an
            unknown measure.
    """
    regime = Regime(regime)
    prior_logpdf = _prior_logpdf(event)
    if regime is Regime.FULL:
        if event.q_samples is None:
            raise PayloadError(
                "event_evidence: the FULL regime needs posterior draws; pass a key "
                "to event_from_posterior."
            )
        nodes = jnp.ravel(jnp.asarray(event.q_samples, dtype=float))
        # ℓₛ = −log π(qₛ) − log S
        log_weights = -prior_logpdf(nodes) - math.log(nodes.shape[0])
        return EventEvidence(
            log_weights=_positive_only(log_weights, nodes), q_samples=nodes
        )
    if not event.emission_rate > 0.0:
        raise ValueError(
            f"event_evidence: the {regime.value} regime needs a positive "
            f"emission rate, got {event.emission_rate!r}."
        )
    if regime is Regime.POINT:
        nodes = jnp.asarray([event.emission_rate], dtype=float)
        return EventEvidence(log_weights=-prior_logpdf(nodes), q_samples=nodes)
    hermite, weights = np.polynomial.hermite.hermgauss(n_nodes)
    if measure == "lognormal":
        # lognormal summary moment-matched to (mean, sd), as catalog.log_moments
        summary = LognormalEmission.from_moments(
            event.emission_rate, event.emission_std
        )
        # qₖ = exp(μ + √2 σ xₖ)
        nodes = jnp.exp(summary.loc + math.sqrt(2.0) * summary.scale * hermite)
    elif measure == "normal":
        summary = GaussianEmission(event.emission_rate, event.emission_std)
        # qₖ = m + √2 s xₖ
        nodes = summary.mean + math.sqrt(2.0) * summary.std * jnp.asarray(hermite)
    else:
        raise ValueError(f"event_evidence: unknown measure {measure!r}.")
    log_prior = prior_logpdf(nodes)
    # ℓₖ = log(wₖ/√π) − log π(qₖ)
    quadrature = jnp.log(jnp.asarray(weights)) - 0.5 * math.log(math.pi)
    log_weights = _positive_only(quadrature - log_prior, nodes)
    # the summary's likelihood log p̃(q) − log π(q) must not climb towards
    # the outermost nodes, or they would dominate the corrected integral
    log_likelihood = np.asarray(
        _positive_only(summary.log_prob(nodes) - log_prior, nodes)
    )
    edge = max(log_likelihood[0], log_likelihood[-1])
    interior = np.max(log_likelihood[1:-1])
    if edge > interior + TAIL_TOLERANCE:
        raise PayloadError(
            f"event_evidence: the {measure} summary has heavier tails than the "
            "per-event prior, so dividing the prior out is dominated by the "
            "outermost Hermite node; use the FULL regime or the measure that "
            "matches the posterior's family."
        )
    return EventEvidence(log_weights=log_weights, q_samples=nodes)

warn_on_collapse

warn_on_collapse(ess: ndarray, n_draws: int, pareto_k: ndarray | None = None) -> None

Warn when an event's ESS < 0.1 S or its Pareto k̂ > 0.7.

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

def warn_on_collapse(
    ess: np.ndarray, n_draws: int, pareto_k: np.ndarray | None = None
) -> None:
    """Warn when an event's ESS < 0.1 S or its Pareto k̂ > 0.7."""
    low = np.flatnonzero(ess < ESS_WARNING_FRACTION * n_draws)
    heavy = np.array([], dtype=int)
    if pareto_k is not None:
        heavy = np.flatnonzero(np.asarray(pareto_k) > PARETO_K_WARNING)
    if low.size or heavy.size:
        warnings.warn(
            "importance weights collapsed: ESS < "
            f"{ESS_WARNING_FRACTION:g}·S for events {low.tolist()}, Pareto k̂ > "
            f"{PARETO_K_WARNING:g} for events {heavy.tolist()}; the per-event "
            "posteriors sit far from the population mark distribution there.",
            RuntimeWarning,
            stacklevel=3,
        )

integrate — log-space quadrature

Log-space integration for the Tier V mark likelihood (implementation in plumax.population._src.integrate).

log_gauss_hermite

log_gauss_hermite(log_f: Callable[[Array], Array], mu: ArrayLike, sigma: ArrayLike, *, n: int = 32) -> Float[Array, '']

log ∫ f(q) LN(q; μ, σ²) dq by Gauss–Hermite quadrature in log space.

With q = exp(μ + √2 σ x), the integral is π^(−½) Σₖ wₖ f(qₖ) for the Hermite nodes xₖ and weights wₖ, evaluated as a log-sum-exp so that f(qₖ) never leaves log space.

Parameters:

Name	Type	Description	Default
log_f	Callable[[Array], Array]	log f, on an array of rates.	required
mu	ArrayLike	μ, the mean of log Q.	required
sigma	ArrayLike	σ > 0, the standard deviation of log Q.	required
n	int	Number of Hermite nodes (static).	32

Returns:

Type	Description
Float[Array, '']	The log-integral.

Examples:

>>> import jax.numpy as jnp
>>> # E[Q] = exp(μ + σ²/2) for a lognormal Q
>>> value = log_gauss_hermite(jnp.log, 0.2, 0.5)
>>> bool(jnp.allclose(value, 0.2 + 0.5 * 0.5**2))
True

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

def log_gauss_hermite(
    log_f: Callable[[Array], Array],
    mu: ArrayLike,
    sigma: ArrayLike,
    *,
    n: int = 32,
) -> Float[Array, ""]:
    """log ∫ f(q) LN(q; μ, σ²) dq by Gauss–Hermite quadrature in log space.

    With q = exp(μ + √2 σ x), the integral is π^(−½) Σₖ wₖ f(qₖ) for the
    Hermite nodes xₖ and weights wₖ, evaluated as a log-sum-exp so that
    f(qₖ) never leaves log space.

    Args:
        log_f: log f, on an array of rates.
        mu: μ, the mean of log Q.
        sigma: σ > 0, the standard deviation of log Q.
        n: Number of Hermite nodes (static).

    Returns:
        The log-integral.

    Examples:
        >>> import jax.numpy as jnp
        >>> # E[Q] = exp(μ + σ²/2) for a lognormal Q
        >>> value = log_gauss_hermite(jnp.log, 0.2, 0.5)
        >>> bool(jnp.allclose(value, 0.2 + 0.5 * 0.5**2))
        True
    """
    nodes, weights = np.polynomial.hermite.hermgauss(n)
    mu = jnp.asarray(mu, dtype=float)
    sigma = jnp.asarray(sigma, dtype=float)
    # qₖ = exp(μ + √2 σ xₖ)
    rates = jnp.exp(mu + math.sqrt(2.0) * sigma * jnp.asarray(nodes))
    # log(wₖ / √π) + log f(qₖ)
    terms = jnp.log(jnp.asarray(weights)) - 0.5 * math.log(math.pi) + log_f(rates)
    return logsumexp(terms)

log_simpson

log_simpson(log_f: Callable[[Array], Array], q_min: float, q_max: float, *, n: int = 2001) -> Float[Array, '']

log ∫ f(q) dq over [q_min, q_max] by composite Simpson in u = log q.

∫ f(q) dq = ∫ f(eᵘ) eᵘ du, with Simpson weights on a uniform u grid, as a log-sum-exp. A uniform grid in log q resolves a lognormal body and a power-law tail alike, which a uniform grid in q cannot.

Parameters:

Name	Type	Description	Default
log_f	Callable[[Array], Array]	log f, on an array of rates.	required
q_min	float	Lower limit, > 0.	required
q_max	float	Upper limit.	required
n	int	Number of nodes (static; made odd).	2001

Returns:

Type	Description
Float[Array, '']	The log-integral.

Examples:

>>> import jax.numpy as jnp
>>> # ∫₁^∞ 3 q⁻⁴ dq = 1, truncated at 1e6
>>> value = log_simpson(lambda q: jnp.log(3.0) - 4.0 * jnp.log(q), 1.0, 1e6)
>>> bool(jnp.abs(value) < 1e-8)
True

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

def log_simpson(
    log_f: Callable[[Array], Array],
    q_min: float,
    q_max: float,
    *,
    n: int = 2001,
) -> Float[Array, ""]:
    """log ∫ f(q) dq over [q_min, q_max] by composite Simpson in u = log q.

    ∫ f(q) dq = ∫ f(eᵘ) eᵘ du, with Simpson weights on a uniform u grid, as a
    log-sum-exp. A uniform grid in log q resolves a lognormal body and a
    power-law tail alike, which a uniform grid in q cannot.

    Args:
        log_f: log f, on an array of rates.
        q_min: Lower limit, > 0.
        q_max: Upper limit.
        n: Number of nodes (static; made odd).

    Returns:
        The log-integral.

    Examples:
        >>> import jax.numpy as jnp
        >>> # ∫₁^∞ 3 q⁻⁴ dq = 1, truncated at 1e6
        >>> value = log_simpson(lambda q: jnp.log(3.0) - 4.0 * jnp.log(q), 1.0, 1e6)
        >>> bool(jnp.abs(value) < 1e-8)
        True
    """
    nodes = n if n % 2 == 1 else n + 1
    grid = np.linspace(math.log(q_min), math.log(q_max), nodes)
    step = grid[1] - grid[0]
    simpson = np.ones(nodes)
    simpson[1:-1:2] = 4.0
    simpson[2:-1:2] = 2.0
    log_weights = jnp.asarray(np.log(simpson * step / 3.0))
    log_grid = jnp.asarray(grid)
    # log[wₖ f(e^uₖ) e^uₖ]
    terms = log_weights + log_f(jnp.exp(log_grid)) + log_grid
    return logsumexp(terms)

pareto_smoothed_log_weights

pareto_smoothed_log_weights(log_weights: ArrayLike) -> tuple[np.ndarray, float]

Pareto-smoothed importance log-weights and the tail shape k̂.

The largest M = ⌈min(S/5, 3√S)⌉ weights are replaced by the expected order statistics of a generalised Pareto fitted to their exceedances, then the weights are normalised. k̂ > 0.7 flags weights too heavy-tailed for the importance estimate to be trusted. A host-side diagnostic (NumPy, not traced): the fit sorts the weights.

Parameters:

Name	Type	Description	Default
log_weights	ArrayLike	Unnormalised log importance weights.	required

Returns:

Type	Description
ndarray	(smoothed, k̂): normalised smoothed log-weights (log-sum-exp 0) and
float	the fitted shape (NaN when fewer than five weights exceed the cutoff,
tuple[ndarray, float]	e.g. equal weights, which need no smoothing).

Raises:

Type	Description
ImportError	Without the population extra.

Examples:

>>> import numpy as np
>>> log_weights = np.random.default_rng(0).normal(0.0, 0.3, 1000)
>>> smoothed, k_hat = pareto_smoothed_log_weights(log_weights)
>>> bool(np.isclose(np.logaddexp.reduce(smoothed), 0.0)), k_hat < 0.7
(True, True)

References

Vehtari, Simpson, Gelman, Yao & Gabry, Pareto smoothed importance sampling, arXiv:1507.02646. Zhang & Stephens (2009), Technometrics.

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

def pareto_smoothed_log_weights(log_weights: ArrayLike) -> tuple[np.ndarray, float]:
    """Pareto-smoothed importance log-weights and the tail shape k̂.

    The largest M = ⌈min(S/5, 3√S)⌉ weights are replaced by the expected
    order statistics of a generalised Pareto fitted to their exceedances, then
    the weights are normalised. k̂ > 0.7 flags weights too heavy-tailed for
    the importance estimate to be trusted. A host-side diagnostic (NumPy, not
    traced): the fit sorts the weights.

    Args:
        log_weights: Unnormalised log importance weights.

    Returns:
        ``(smoothed, k̂)``: normalised smoothed log-weights (log-sum-exp 0) and
        the fitted shape (NaN when fewer than five weights exceed the cutoff,
        e.g. equal weights, which need no smoothing).

    Raises:
        ImportError: Without the ``population`` extra.

    Examples:
        >>> import numpy as np
        >>> log_weights = np.random.default_rng(0).normal(0.0, 0.3, 1000)
        >>> smoothed, k_hat = pareto_smoothed_log_weights(log_weights)
        >>> bool(np.isclose(np.logaddexp.reduce(smoothed), 0.0)), k_hat < 0.7
        (True, True)

    References:
        Vehtari, Simpson, Gelman, Yao & Gabry, Pareto smoothed importance
        sampling, arXiv:1507.02646. Zhang & Stephens (2009), Technometrics.
    """
    require_xtremax("pareto_smoothed_log_weights")
    import xtremax.primitives as extremes

    log_w = np.asarray(log_weights, dtype=float).ravel()
    log_w = log_w - log_w.max()
    size = log_w.size
    tail_length = math.ceil(min(0.2 * size, 3.0 * math.sqrt(size)))
    if tail_length >= size:
        # too few weights to separate a tail: normalise, no diagnostic
        return log_w - np.logaddexp.reduce(log_w), math.nan
    order = np.argsort(log_w)
    cutoff = max(log_w[order[-tail_length - 1]], math.log(np.finfo(float).tiny))
    tail = order[log_w[order] > cutoff]
    k_hat = math.nan
    if tail.size > 4:
        exceedances = np.exp(log_w[tail]) - math.exp(cutoff)
        k_hat, scale = _generalized_pareto_fit(exceedances)
        if math.isfinite(k_hat):
            # expected order statistics, F⁻¹((j − ½)/M)
            levels = (np.arange(tail.size) + 0.5) / tail.size
            quantiles = np.asarray(extremes.gpd_icdf(levels, scale, k_hat))
            # the largest weight stays the largest raw weight, at most 1
            log_w[tail] = np.minimum(np.log(quantiles + math.exp(cutoff)), 0.0)
    smoothed = log_w - np.logaddexp.reduce(log_w)
    return smoothed, float(k_hat)

Importance-corrected models

Importance-corrected population fits: the mark-only fit and the full TMTPP.

Both score detection i by the importance-corrected mark term of :mod:~plumax.population.importance,

mᵢ(θ) = log Σₛ exp(ℓᵢₛ + log P_d(qᵢₛ) + log f(qᵢₛ | θ)),

with the thinned rate p̄(θ) = ∫ P_d(q) f(q | θ) dq in log space (:func:~plumax.population.integrate.log_simpson; p̄ = 1 without a POD).

• :func:mark_iw_model: the marks given the detections, Σᵢ mᵢ(θ) − N log p̄(θ), the λ-free part of the likelihood.
• :func:tmtpp_iw_model: the full TMTPP, Σᵢ [log λ(tᵢ) + mᵢ(θ)] − Λ(T) p̄(θ), with the ground intensity from the intensity registry on xtremax (:func:~plumax.population.processes.make_process).
• :func:fit_size_distribution_iw: NUTS on :func:mark_iw_model, with the per-event ESS and Pareto k̂ at the posterior-mean θ and a collapse warning.
• :class:TMTPPPosterior: the draws of :func:tmtpp_iw_model, as the intensity kernels and mark parameters the persistency and total-emission products take.

The POD is a fixed callable log_pod, or the union of calibrated instruments pods (:mod:~plumax.population.pod), whose (Q₅₀, k) are then sampled too.

MarkPosterior dataclass

MarkPosterior(mark: str, regime: Regime, samples: dict[str, ndarray], ess: ndarray, pareto_k: ndarray)

Posterior over a mark family's parameters, with importance diagnostics.

Attributes:

Name	Type	Description
mark	str	The registry name of the family.
regime	Regime	How the events entered.
samples	dict[str, ndarray]	Posterior draws of θ by parameter name.
ess	ndarray	Per-event ESS of the importance weights at the posterior-mean θ.
pareto_k	ndarray	Per-event Pareto k̂ (FULL regime; NaN otherwise).

posterior_mean

posterior_mean() -> dict[str, np.floating | np.ndarray]

θ̄, the posterior mean of every parameter.

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

def posterior_mean(self) -> dict[str, np.floating | np.ndarray]:
    """θ̄, the posterior mean of every parameter."""
    return {name: np.mean(draws, axis=0) for name, draws in self.samples.items()}

TMTPPPosterior dataclass

TMTPPPosterior(intensity: str, mark: Mark, samples: dict[str, Array])

Posterior draws of the full TMTPP, by the site names of its model.

Attributes:

Name	Type	Description
intensity	str	The registry name of the ground intensity.
mark	Mark	The mark family.
samples	dict[str, Array]	Draws from :func:tmtpp_iw_model by site name, as MCMC.get_samples() returns them: the kernel's sample_priors sites, and the mark's under mark..

num_samples property

num_samples: int

The number of posterior draws.

kernel

kernel(draw: dict[str, Array]) -> Any

The intensity kernel of one draw.

A kernel's site names need not match its fields (lambda_bg is the field lambda_background), so the draw goes through sample_priors.

Parameters:

Name	Type	Description	Default
draw	dict[str, Array]	One draw of :attr:samples.	required

Returns:

Type	Description
Any	The kernel at that draw.

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

def kernel(self, draw: dict[str, Array]) -> Any:
    """The intensity kernel of one draw.

    A kernel's site names need not match its fields (``lambda_bg`` is the
    field ``lambda_background``), so the draw goes through ``sample_priors``.

    Args:
        draw: One draw of :attr:`samples`.

    Returns:
        The kernel at that draw.
    """
    from numpyro import handlers

    from plumax.population.intensity import INTENSITY_REGISTRY

    prior = INTENSITY_REGISTRY[self.intensity].sample_priors
    return handlers.substitute(prior, data=draw)()

theta

theta(draw: dict[str, Array]) -> dict[str, Array]

The mark parameters θ of one draw.

Parameters:

Name	Type	Description	Default
draw	dict[str, Array]	One draw of :attr:samples.	required

Returns:

Type	Description
dict[str, Array]	θ by parameter name, as the mark's methods take it.

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

def theta(self, draw: dict[str, Array]) -> dict[str, Array]:
    """The mark parameters θ of one draw.

    Args:
        draw: One draw of :attr:`samples`.

    Returns:
        θ by parameter name, as the mark's methods take it.
    """
    prefix = "mark."
    return {
        name[len(prefix) :]: value
        for name, value in draw.items()
        if name.startswith(prefix) and not name.startswith(f"{prefix}raw_")
    }

log_thinned_rate

log_thinned_rate(mark: Mark, theta: dict[str, Any], log_pod: Callable[[Array], Array] | None, support: tuple[float, float], *, n_nodes: int = 2001) -> Float[Array, '']

log p̄(θ) = log ∫ P_d(q) f(q | θ) dq, 0 without a POD.

Log-space Simpson on [q_lo, q_hi], plus the mass above q_hi at the POD's value there, P_d(q_hi) P(Q > q_hi | θ), so a heavy tail beyond the grid is not dropped. The mass below q_lo is left out: a detection curve is nondecreasing in Q, and q_lo sits far below every detection.

Parameters:

Name	Type	Description	Default
mark	Mark	The mark family.	required
theta	dict[str, Any]	Its parameters.	required
log_pod	Callable[[Array], Array] | None	log P_d, or None for certain detection.	required
support	tuple[float, float]	(q_lo, q_hi) of the log-space Simpson rule [kg/s].	required
n_nodes	int	Simpson nodes.	2001

Returns:

Type	Description
Float[Array, '']	log p̄.

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

def log_thinned_rate(
    mark: Mark,
    theta: dict[str, Any],
    log_pod: Callable[[Array], Array] | None,
    support: tuple[float, float],
    *,
    n_nodes: int = 2001,
) -> Float[Array, ""]:
    """log p̄(θ) = log ∫ P_d(q) f(q | θ) dq, 0 without a POD.

    Log-space Simpson on [q_lo, q_hi], plus the mass above q_hi at the POD's
    value there, P_d(q_hi) P(Q > q_hi | θ), so a heavy tail beyond the grid is
    not dropped. The mass below q_lo is left out: a detection curve is
    nondecreasing in Q, and q_lo sits far below every detection.

    Args:
        mark: The mark family.
        theta: Its parameters.
        log_pod: log P_d, or ``None`` for certain detection.
        support: ``(q_lo, q_hi)`` of the log-space Simpson rule [kg/s].
        n_nodes: Simpson nodes.

    Returns:
        log p̄.
    """
    if log_pod is None:
        # a normalised f integrates to one
        return jnp.zeros(())
    integrand = _log_integrand(mark, theta, log_pod)
    inside = log_simpson(integrand, support[0], support[1], n=n_nodes)
    # P_d(q_hi) P(Q > q_hi), the tail beyond the last node
    upper = jnp.asarray(support[1], dtype=float)
    beyond = log_pod(upper) + mark.log_survival(upper, **theta)
    return jnp.logaddexp(inside, beyond)

default_support

default_support(evidence: EventEvidence) -> tuple[float, float]

SUPPORT_DECADES decades either side of the evidence nodes.

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

def default_support(evidence: EventEvidence) -> tuple[float, float]:
    """``SUPPORT_DECADES`` decades either side of the evidence nodes."""
    nodes = np.asarray(evidence.q_samples)
    positive = nodes[nodes > 0.0]
    lower = float(positive.min()) * 10.0 ** (-SUPPORT_DECADES)
    upper = float(positive.max()) * 10.0**SUPPORT_DECADES
    return lower, upper

mark_iw_model

mark_iw_model(evidence: EventEvidence, *, mark: Mark, log_pod: Callable[[Array], Array] | None = None, support: tuple[float, float] | None = None) -> None

NumPyro model of the marks given the detections, importance-corrected.

log L(θ) = Σᵢ mᵢ(θ) − N log p̄(θ), θ under mark.numpyro_prior("mark").

Parameters:

Name	Type	Description	Default
evidence	EventEvidence	Stacked per-event evidence (:func:catalog_evidence).	required
mark	Mark	The mark family.	required
log_pod	Callable[[Array], Array] | None	log P_d on rates in the marks' units, or None.	None
support	tuple[float, float] | None	Rate range of the p̄ integral; default from the evidence.	None

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

def mark_iw_model(
    evidence: EventEvidence,
    *,
    mark: Mark,
    log_pod: Callable[[Array], Array] | None = None,
    support: tuple[float, float] | None = None,
) -> None:
    """NumPyro model of the marks given the detections, importance-corrected.

    log L(θ) = Σᵢ mᵢ(θ) − N log p̄(θ), θ under ``mark.numpyro_prior("mark")``.

    Args:
        evidence: Stacked per-event evidence (:func:`catalog_evidence`).
        mark: The mark family.
        log_pod: log P_d on rates in the marks' units, or ``None``.
        support: Rate range of the p̄ integral; default from the evidence.
    """
    import numpyro

    theta = mark.numpyro_prior("mark")
    integrand = _log_integrand(mark, theta, log_pod)
    # mᵢ(θ) = log Σₛ exp(ℓᵢₛ + g(qᵢₛ))
    mark_terms = evidence.log_mark_term(integrand)
    support = support or default_support(evidence)
    log_rate = log_thinned_rate(mark, theta, log_pod, support)
    n_events = mark_terms.shape[0]
    # Σᵢ mᵢ − N log p̄
    numpyro.factor("marks", jnp.sum(mark_terms) - n_events * log_rate)

tmtpp_iw_model

tmtpp_iw_model(event_times: Float[Array, ' event'], mask: Bool[Array, ' event'], evidence: EventEvidence, *, mark: Mark, intensity: str, window: tuple[float, float], log_pod: Callable[[Array], Array] | None = None, support: tuple[float, float] | None = None, pods: Sequence[InstrumentPOD] | None = None, coverage: Bool[Array, 'event instrument'] | None = None, coverage_intervals: Sequence[tuple[float, float, Bool[Array, ' instrument']]] | None = None, tod_origin: float = 0.0, n_nodes: int = 2001) -> None

NumPyro model of the full importance-corrected TMTPP likelihood.

log L = Σᵢ [log λ(tᵢ) + mᵢ(θ)] − Λ(T) p̄(θ), with λ a registry intensity under its sample_priors on the xtremax Poisson process of the window.

With pods, P_d is the union P_d^∪ under :func:calibration_prior. With coverage, event i is thinned by the instruments that overpassed it and the compensator follows who was looking when, as 06a § Non-detection events asks: Σ_w Λ_w p̄_w(θ), over intervals w in which the covering set is constant. A per-event mask alone does not fix that integral, so the two come together. An instrument with tod_bins shifts its Q₅₀ by the hour of each event, and the compensator sums over the hours too: Σ_w Σ_b Λ_w,b p̄_w,b(θ).

Parameters:

Name	Type	Description	Default
event_times	Float[Array, ' event']	Detection times since t₀ in xtremax's padded buffer (:func:~plumax.population.processes.padded_event_buffer) [days].	required
mask	Bool[Array, ' event']	Its mask.	required
evidence	EventEvidence	Stacked per-event evidence, one row per detection in the order of the buffer's real entries. It may also carry the buffer's padding, one row per buffer entry: the mask drops those rows, so catalogs of different sizes can share one shape under vmap.	required
mark	Mark	The mark family.	required
intensity	str	A registry name of a Poisson intensity.	required
window	tuple[float, float]	(t₀, t₁) [days].	required
log_pod	Callable[[Array], Array] | None	log P_d on rates in the marks' units, or None.	None
support	tuple[float, float] | None	Rate range of the p̄ integral; default from the evidence.	None
pods	Sequence[InstrumentPOD] | None	Calibrated instruments, instead of log_pod.	None
coverage	Bool[Array, 'event instrument'] | None	Which instruments overpassed each row of evidence. Default: all of them, at all times.	None
coverage_intervals	Sequence[tuple[float, float, Bool[Array, ' instrument']]] | None	(start, stop, covering) in days, disjoint and inside the window. No detection is possible outside them.	None
tod_origin	float	Local hour of day at t = 0 of window [hours], for tod_bins.	0.0
n_nodes	int	Simpson nodes of each p̄ integral.	2001

Raises:

Type	Description
ImportError	Without the population extra.
ValueError	For log_pod with pods, coverage without pods or without its intervals, or overlapping or out-of-window intervals.
EquinoxRuntimeError	If a detection has no covering instrument.

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

def tmtpp_iw_model(
    event_times: Float[Array, " event"],
    mask: Bool[Array, " event"],
    evidence: EventEvidence,
    *,
    mark: Mark,
    intensity: str,
    window: tuple[float, float],
    log_pod: Callable[[Array], Array] | None = None,
    support: tuple[float, float] | None = None,
    pods: Sequence[InstrumentPOD] | None = None,
    coverage: Bool[Array, "event instrument"] | None = None,
    coverage_intervals: Sequence[tuple[float, float, Bool[Array, " instrument"]]]
    | None = None,
    tod_origin: float = 0.0,
    n_nodes: int = 2001,
) -> None:
    """NumPyro model of the full importance-corrected TMTPP likelihood.

    log L = Σᵢ [log λ(tᵢ) + mᵢ(θ)] − Λ(T) p̄(θ), with λ a registry intensity
    under its ``sample_priors`` on the xtremax Poisson process of the window.

    With ``pods``, P_d is the union P_d^∪ under :func:`calibration_prior`. With
    ``coverage``, event i is thinned by the instruments that overpassed it and the
    compensator follows who was looking when, as 06a § Non-detection events asks:
    Σ_w Λ_w p̄_w(θ), over intervals w in which the covering set is constant. A
    per-event mask alone does not fix that integral, so the two come together. An
    instrument with ``tod_bins`` shifts its Q₅₀ by the hour of each event, and the
    compensator sums over the hours too: Σ_w Σ_b Λ_w,b p̄_w,b(θ).

    Args:
        event_times: Detection times since t₀ in xtremax's padded buffer
            (:func:`~plumax.population.processes.padded_event_buffer`) [days].
        mask: Its mask.
        evidence: Stacked per-event evidence, one row per detection in the
            order of the buffer's real entries. It may also carry the buffer's
            padding, one row per buffer entry: the mask drops those rows, so
            catalogs of different sizes can share one shape under ``vmap``.
        mark: The mark family.
        intensity: A registry name of a Poisson intensity.
        window: ``(t₀, t₁)`` [days].
        log_pod: log P_d on rates in the marks' units, or ``None``.
        support: Rate range of the p̄ integral; default from the evidence.
        pods: Calibrated instruments, instead of ``log_pod``.
        coverage: Which instruments overpassed each row of ``evidence``. Default:
            all of them, at all times.
        coverage_intervals: ``(start, stop, covering)`` in days, disjoint and inside
            the window. No detection is possible outside them.
        tod_origin: Local hour of day at t = 0 of ``window`` [hours], for ``tod_bins``.
        n_nodes: Simpson nodes of each p̄ integral.

    Raises:
        ImportError: Without the ``population`` extra.
        ValueError: For ``log_pod`` with ``pods``, ``coverage`` without ``pods`` or
            without its intervals, or overlapping or out-of-window intervals.
        equinox.EquinoxRuntimeError: If a detection has no covering instrument.
    """
    import numpyro

    from plumax.population._src.processes import make_process
    from plumax.population.intensity import INTENSITY_REGISTRY

    intervals = _check_coverage_arguments(
        log_pod, pods, coverage, coverage_intervals, window
    )
    theta = mark.numpyro_prior("mark")
    kernel = INTENSITY_REGISTRY[intensity].sample_priors()
    process = make_process(kernel, window)
    # Σᵢ log λ(tᵢ) over the real events
    log_rates = process.log_intensity_fn(event_times)
    ground = jnp.sum(jnp.where(mask, log_rates, 0.0))

    # --- detection ---
    binned = pods is not None and is_tod_binned(pods)
    calibration = None
    event_log_pod = log_pod
    n_rows = evidence.q_samples.shape[0]
    # mask is a contiguous prefix, so this is every row of an unpadded evidence
    event_mask = mask[:n_rows]
    if pods is not None:
        calibration = calibration_prior(pods)
        event_coverage = None
        if coverage is not None:
            coverage = jnp.asarray(coverage, dtype=bool)
            covered = jnp.any(coverage, axis=-1)
            event_mask = eqx.error_if(
                event_mask,
                jnp.any(event_mask & ~covered),
                "tmtpp_iw_model: a detection has no covering instrument",
            )
            # a padding row needs some instrument for a finite log P_d^∪
            event_coverage = jnp.where(covered[:, None], coverage, True)[:, None, :]
        event_tod_bin = None
        if binned:
            # event_times count from t₀, the bins from t = 0
            event_bins = tod_bin_index(
                window[0] + event_times[:n_rows], tod_origin=tod_origin
            )
            event_tod_bin = event_bins[:, None]
        event_log_pod = log_pod_union(
            pods,
            coverage=event_coverage,
            calibration=calibration,
            tod_bin=event_tod_bin,
        )

    # --- marks ---
    integrand = _log_integrand(mark, theta, event_log_pod)
    mark_terms = evidence.log_mark_term(integrand)
    marks = jnp.sum(jnp.where(event_mask, mark_terms, 0.0))

    # --- compensator ---
    support = support or default_support(evidence)
    if binned:
        # no coverage given: every instrument looks throughout; an empty list stays empty
        everyone = np.ones(len(pods), dtype=bool)
        whole_window = [(float(window[0]), float(window[1]), everyone)]
        thinned = _binned_compensator(
            kernel,
            mark,
            theta,
            pods,
            calibration,
            whole_window if intervals is None else intervals,
            support,
            n_nodes,
            tod_origin,
        )
    elif intervals is None:
        window_log_pod = (
            log_pod if pods is None else log_pod_union(pods, calibration=calibration)
        )
        log_rate = log_thinned_rate(
            mark, theta, window_log_pod, support, n_nodes=n_nodes
        )
        # Λ(T) p̄
        thinned = process.effective_integrated_intensity() * jnp.exp(log_rate)
    else:
        thinned = _covered_compensator(
            kernel, mark, theta, pods, calibration, intervals, support, n_nodes
        )
    # Σᵢ [log λ(tᵢ) + mᵢ] − Λ(T) p̄
    numpyro.factor("tmtpp", ground + marks - thinned)

fit_size_distribution_iw

fit_size_distribution_iw(catalog: EmissionCatalog, *, mark: str = 'lognormal', mark_fields: dict[str, Any] | None = None, regime: Regime | str = Regime.FULL, log_pod: Callable[[Array], Array] | None = None, support: tuple[float, float] | None = None, n_nodes: int = 32, measure: Measure = 'lognormal', num_warmup: int = 500, num_samples: int = 500, num_chains: int = 1, seed: int = 0) -> MarkPosterior

Importance-corrected NUTS fit of the population size distribution.

Runs :func:mark_iw_model, then reports the per-event ESS (and, in the FULL regime, the Pareto k̂) of the weights at the posterior-mean θ, warning in every regime when an event collapses (ESS < 0.1 S, or k̂ > 0.7).

Parameters:

Name	Type	Description	Default
catalog	EmissionCatalog	Events with the payload-v2 fields (prior_logpdf, and q_samples for FULL).	required
mark	str	A registry name (:data:~plumax.population.marks.MARKS).	'lognormal'
mark_fields	dict[str, Any] | None	The family's fields (support, prior settings).	None
regime	Regime | str	POINT, GAUSSIAN or FULL.	FULL
log_pod	Callable[[Array], Array] | None	log P_d on rates in the marks' units, or None.	None
support	tuple[float, float] | None	Rate range of the p̄ integral; default from the evidence.	None
n_nodes	int	Gauss–Hermite nodes of the GAUSSIAN regime.	32
measure	Measure	The GAUSSIAN-regime summary, "lognormal" or "normal" (see :func:~plumax.population.importance.event_evidence).	'lognormal'
num_warmup	int	NUTS warm-up steps.	500
num_samples	int	Posterior draws per chain.	500
num_chains	int	NUTS chains.	1
seed	int	PRNG seed.	0

Returns:

Name	Type	Description
The	MarkPosterior	class:MarkPosterior.

Raises:

Type	Description
PayloadError	When the events lack the payload the regime needs.

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

def fit_size_distribution_iw(
    catalog: EmissionCatalog,
    *,
    mark: str = "lognormal",
    mark_fields: dict[str, Any] | None = None,
    regime: Regime | str = Regime.FULL,
    log_pod: Callable[[Array], Array] | None = None,
    support: tuple[float, float] | None = None,
    n_nodes: int = 32,
    measure: Measure = "lognormal",
    num_warmup: int = 500,
    num_samples: int = 500,
    num_chains: int = 1,
    seed: int = 0,
) -> MarkPosterior:
    """Importance-corrected NUTS fit of the population size distribution.

    Runs :func:`mark_iw_model`, then reports the per-event ESS (and, in the
    FULL regime, the Pareto k̂) of the weights at the posterior-mean θ,
    warning in every regime when an event collapses (ESS < 0.1 S, or
    k̂ > 0.7).

    Args:
        catalog: Events with the payload-v2 fields (``prior_logpdf``, and
            ``q_samples`` for FULL).
        mark: A registry name (:data:`~plumax.population.marks.MARKS`).
        mark_fields: The family's fields (support, prior settings).
        regime: POINT, GAUSSIAN or FULL.
        log_pod: log P_d on rates in the marks' units, or ``None``.
        support: Rate range of the p̄ integral; default from the evidence.
        n_nodes: Gauss–Hermite nodes of the GAUSSIAN regime.
        measure: The GAUSSIAN-regime summary, ``"lognormal"`` or ``"normal"``
            (see :func:`~plumax.population.importance.event_evidence`).
        num_warmup: NUTS warm-up steps.
        num_samples: Posterior draws per chain.
        num_chains: NUTS chains.
        seed: PRNG seed.

    Returns:
        The :class:`MarkPosterior`.

    Raises:
        PayloadError: When the events lack the payload the regime needs.
    """
    from numpyro.infer import MCMC, NUTS

    regime = Regime(regime)
    family = MARKS[mark](**(mark_fields or {}))
    evidence = catalog_evidence(
        catalog, regime=regime, n_nodes=n_nodes, measure=measure
    )
    mcmc = MCMC(
        NUTS(mark_iw_model),
        num_warmup=num_warmup,
        num_samples=num_samples,
        num_chains=num_chains,
        progress_bar=False,
    )
    mcmc.run(
        jax.random.PRNGKey(seed),
        evidence,
        mark=family,
        log_pod=log_pod,
        support=support,
    )
    draws = mcmc.get_samples()
    prefix = "mark."
    samples = {
        name[len(prefix) :]: np.asarray(value)
        for name, value in draws.items()
        if name.startswith(prefix) and not name.startswith(f"{prefix}raw_")
    }
    theta = {
        name: jnp.asarray(np.mean(value, axis=0)) for name, value in samples.items()
    }
    log_weights = np.asarray(evidence.log_terms(_log_integrand(family, theta, log_pod)))
    ess = np.asarray(effective_sample_size(jnp.asarray(log_weights)))
    pareto_k = np.full(ess.shape, math.nan)
    if regime is Regime.FULL:
        pareto_k = np.array(
            [pareto_smoothed_log_weights(row)[1] for row in log_weights]
        )
    warn_on_collapse(ess, log_weights.shape[-1], pareto_k)
    return MarkPosterior(
        mark=mark, regime=regime, samples=samples, ess=ess, pareto_k=pareto_k
    )