Mark registry

The population size distribution f(Q | θ) as one registry shared by the population fit, the missing-mass simulator and the total-emission estimator: a power law, a lognormal, a lognormal body spliced continuously to a Pareto tail, and a mixture of lognormals. Each family has log_prob, sample, the partial first moment truncated_mean that the totals need, and numpyro_prior for its default priors. Tier V.

marks

The mark registry, the population size distribution f(Q | θ) (implementation in plumax.population._src.marks).

LognormalMark

Bases: Module

Lognormal f(q) = LN(q; μ, σ²).

θ = {μ, σ}.

Attributes:

Name	Type	Description
mu_prior_scale	float	Scale of the Normal(0, ·) prior on μ.
sigma_prior_scale	float	Scale of the HalfNormal prior on σ.

Examples:

>>> import jax.numpy as jnp
>>> mark = LognormalMark()
>>> mean = mark.truncated_mean(0.0, jnp.inf, mu=0.0, sigma=1.0)
>>> bool(jnp.allclose(mean, jnp.exp(0.5)))
True

log_survival

log_survival(q: ArrayLike, *, mu: ArrayLike, sigma: ArrayLike) -> Float[Array, '*q']

log P(Q > q) = log Φ(−(log q − μ)/σ), 0 for q ≤ 0.

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

def log_survival(
    self, q: ArrayLike, *, mu: ArrayLike, sigma: ArrayLike
) -> Float[Array, "*q"]:
    """log P(Q > q) = log Φ(−(log q − μ)/σ), 0 for q ≤ 0."""
    q = jnp.asarray(q, dtype=float)
    positive = q > 0.0
    log_q = jnp.log(jnp.where(positive, q, 1.0))
    standardised = (log_q - jnp.asarray(mu)) / jnp.asarray(sigma)
    return jnp.where(positive, norm.logcdf(-standardised), 0.0)

LognormalParetoMark

Bases: Module

Lognormal body spliced to a Pareto tail at Q_break, continuous there.

f(q) = w LN(q; μ, σ²)/F_LN(Q_b) for q < Q_b and f(q) = (1 − w) α Q_bᵅ q^(−α−1) for q ≥ Q_b, with the body weight w = t/(t + b) set by continuity, where b = LN(Q_b; μ, σ²)/F_LN(Q_b) and t = α/Q_b are the two normalised pieces' densities at the break. One free tail index α; E[Q] is finite for α > 1 and Var[Q] for α > 2.

θ = {μ, σ, α, Q_b}.

Attributes:

Name	Type	Description
q_break_loc	float	Centre of the LogNormal(log ·, 0.5²) prior on Q_b, a published detection-floor value for the instrument class [kg/s].
q_break_log_scale	float	Log-scale of that prior.
mu_prior_scale	float	Scale of the Normal(0, ·) prior on μ.
sigma_prior_scale	float	Scale of the HalfNormal prior on σ.
alpha_bounds	tuple[float, float]	(low, high) of the uniform prior on α.

Examples:

>>> import jax.numpy as jnp
>>> mark = LognormalParetoMark(q_break_loc=2.0)
>>> theta = dict(mu=0.0, sigma=1.0, alpha=2.5, q_break=2.0)
>>> below = mark.log_prob(jnp.asarray(2.0 - 1e-9), **theta)
>>> above = mark.log_prob(jnp.asarray(2.0), **theta)
>>> bool(jnp.allclose(below, above))
True

log_survival

log_survival(q: ArrayLike, *, mu: ArrayLike, sigma: ArrayLike, alpha: ArrayLike, q_break: ArrayLike) -> Float[Array, '*q']

log P(Q > q): (1 − w)(Q_b/q)ᵅ above the break, plus the body above.

P(Q > q) = (1 − w) + w [Φ(z_b) − Φ(z)]/Φ(z_b) below Q_b.

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

def log_survival(
    self,
    q: ArrayLike,
    *,
    mu: ArrayLike,
    sigma: ArrayLike,
    alpha: ArrayLike,
    q_break: ArrayLike,
) -> Float[Array, "*q"]:
    """log P(Q > q): (1 − w)(Q_b/q)ᵅ above the break, plus the body above.

    P(Q > q) = (1 − w) + w [Φ(z_b) − Φ(z)]/Φ(z_b) below Q_b.
    """
    q = jnp.asarray(q, dtype=float)
    mu, sigma, alpha, q_break = (
        jnp.asarray(value, dtype=float) for value in (mu, sigma, alpha, q_break)
    )
    log_body_weight, log_tail_weight, log_body_mass = _splice_log_weights(
        mu, sigma, alpha, q_break
    )
    above = q >= q_break
    positive = q > 0.0
    # tail: log(1 − w) + α log(Q_b/q)
    safe_tail_q = jnp.where(above, q, q_break)
    tail = log_tail_weight + alpha * (jnp.log(q_break) - jnp.log(safe_tail_q))
    # body: log[(1 − w) + w (Φ(z_b) − Φ(z))/Φ(z_b)]
    safe_body_q = jnp.where(above | ~positive, 0.5 * q_break, q)
    lower_z = (jnp.log(safe_body_q) - mu) / sigma
    upper_z = (jnp.log(q_break) - mu) / sigma
    log_body_above = _log_normal_interval(lower_z, upper_z) - log_body_mass
    body = jnp.logaddexp(log_tail_weight, log_body_weight + log_body_above)
    return jnp.where(above, tail, jnp.where(positive, body, 0.0))

Mark

Bases: Protocol

A mark family f(Q | θ), with θ passed by keyword.

log_prob

log_prob(q: ArrayLike, **theta: Any) -> Array

log f(q | θ).

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

def log_prob(self, q: ArrayLike, **theta: Any) -> Array:
    """log f(q | θ)."""
    ...

sample

sample(key: Array, n: int, **theta: Any) -> Array

n draws of Q.

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

def sample(self, key: Array, n: int, **theta: Any) -> Array:
    """n draws of Q."""
    ...

truncated_mean

truncated_mean(q_min: ArrayLike, q_max: ArrayLike, **theta: Any) -> Array

∫ q f(q | θ) dq over [q_min, q_max].

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

def truncated_mean(self, q_min: ArrayLike, q_max: ArrayLike, **theta: Any) -> Array:
    """∫ q f(q | θ) dq over [q_min, q_max]."""
    ...

log_survival

log_survival(q: ArrayLike, **theta: Any) -> Array

log P(Q > q | θ).

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

def log_survival(self, q: ArrayLike, **theta: Any) -> Array:
    """log P(Q > q | θ)."""
    ...

numpyro_prior

numpyro_prior(name: str) -> dict[str, Any]

NumPyro sites for θ under the default priors.

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

def numpyro_prior(self, name: str) -> dict[str, Any]:
    """NumPyro sites for θ under the default priors."""
    ...

MixtureLognormalMark

Bases: Module

Mixture of K lognormals f(q) = Σₖ wₖ LN(q; μₖ, σₖ²).

θ = {weights, locs, scales}, each of length K.

Attributes:

Name	Type	Description
n_components	int	K.
concentration	float	Dirichlet concentration of the prior on the weights (the same for every component).
mu_prior_scale	float	Scale of the Normal(0, ·) prior on each μₖ; the sampled μₖ are sorted, which removes label switching.
sigma_prior_scale	float	Scale of the HalfNormal prior on each σₖ.

Examples:

>>> import jax.numpy as jnp
>>> mark = MixtureLognormalMark(n_components=2)
>>> theta = dict(
...     weights=jnp.array([0.5, 0.5]),
...     locs=jnp.array([0.0, 0.0]),
...     scales=jnp.array([1.0, 1.0]),
... )
>>> single = LognormalMark().log_prob(2.0, mu=0.0, sigma=1.0)
>>> bool(jnp.allclose(mark.log_prob(2.0, **theta), single))
True

log_survival

log_survival(q: ArrayLike, *, weights: ArrayLike, locs: ArrayLike, scales: ArrayLike) -> Float[Array, '*q']

log P(Q > q) = log Σₖ wₖ Φ(−(log q − μₖ)/σₖ).

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

def log_survival(
    self,
    q: ArrayLike,
    *,
    weights: ArrayLike,
    locs: ArrayLike,
    scales: ArrayLike,
) -> Float[Array, "*q"]:
    """log P(Q > q) = log Σₖ wₖ Φ(−(log q − μₖ)/σₖ)."""
    q = jnp.asarray(q, dtype=float)
    weights, locs, scales = (
        jnp.asarray(value, dtype=float) for value in (weights, locs, scales)
    )
    positive = q > 0.0
    log_q = einx.rearrange("... -> ... 1", jnp.log(jnp.where(positive, q, 1.0)))
    component = norm.logcdf(-(log_q - locs) / scales)
    mixture = einx.logsumexp("... [k]", component + jnp.log(weights))
    return jnp.where(positive, mixture, 0.0)

PowerLawMark

Bases: Module

Pareto power law f(q) = (α − 1)/q_min · (q/q_min)^(−α) on q ≥ q_min.

θ = {α}, α > 1. E[Q] is finite for α > 2; for α ≤ 2 truncated_mean over an unbounded interval is +∞ (logarithmic in the upper limit at α = 2).

Attributes:

Name	Type	Description
q_min	float	Lower end of the support [kg/s].
alpha_bounds	tuple[float, float]	(low, high) of the uniform prior on α.

Examples:

>>> import jax.numpy as jnp
>>> mark = PowerLawMark(q_min=1.0)
>>> float(jnp.exp(mark.log_prob(jnp.asarray(1.0), alpha=3.0)))
2.0
>>> float(mark.truncated_mean(1.0, jnp.inf, alpha=3.0))
2.0

log_survival

log_survival(q: ArrayLike, *, alpha: ArrayLike) -> Float[Array, '*q']

log P(Q > q) = (α − 1) log(q_min/q) above q_min, 0 below.

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

def log_survival(self, q: ArrayLike, *, alpha: ArrayLike) -> Float[Array, "*q"]:
    """log P(Q > q) = (α − 1) log(q_min/q) above q_min, 0 below."""
    q = jnp.asarray(q, dtype=float)
    alpha = jnp.asarray(alpha, dtype=float)
    ratio = jnp.maximum(q, self.q_min) / self.q_min
    return -(alpha - 1.0) * jnp.log(ratio)

make_mark

make_mark(name: str, **fields: Any) -> Mark

Build a registry mark family by name.

Parameters:

Name	Type	Description	Default
name	str	A key of :data:MARKS.	required
**fields	Any	The family's fields (support bounds, prior settings).	{}

Returns:

Type	Description
Mark	The mark family.

Raises:

Type	Description
KeyError	For an unknown name.

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

def make_mark(name: str, **fields: Any) -> Mark:
    """Build a registry mark family by name.

    Args:
        name: A key of :data:`MARKS`.
        **fields: The family's fields (support bounds, prior settings).

    Returns:
        The mark family.

    Raises:
        KeyError: For an unknown name.
    """
    if name not in MARKS:
        raise KeyError(
            f"make_mark: unknown mark {name!r}; expected one of {sorted(MARKS)}."
        )
    return MARKS[name](**fields)