diff --git a/docs/source/distributions.rst b/docs/source/distributions.rst index 1165c30ef..645ac8df1 100644 --- a/docs/source/distributions.rst +++ b/docs/source/distributions.rst @@ -412,6 +412,14 @@ RelaxedBernoulliLogits :show-inheritance: :member-order: bysource +SchechterMag +^^^^^^^^^^^^ +.. autoclass:: numpyro.distributions.continuous.SchechterMag + :members: + :undoc-members: + :show-inheritance: + :member-order: bysource + SoftLaplace ^^^^^^^^^^^ .. autoclass:: numpyro.distributions.continuous.SoftLaplace diff --git a/numpyro/distributions/__init__.py b/numpyro/distributions/__init__.py index 2255e2335..4abb277f0 100644 --- a/numpyro/distributions/__init__.py +++ b/numpyro/distributions/__init__.py @@ -59,6 +59,7 @@ Pareto, RelaxedBernoulli, RelaxedBernoulliLogits, + SchechterMag, SoftLaplace, StudentT, Uniform, @@ -226,6 +227,7 @@ "LeftCensoredDistribution", "RightCensoredDistribution", "IntervalCensoredDistribution", + "SchechterMag", "SineBivariateVonMises", "SineSkewed", "SoftLaplace", diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py index 57a3d2baa..6be38542b 100644 --- a/numpyro/distributions/continuous.py +++ b/numpyro/distributions/continuous.py @@ -44,6 +44,7 @@ expi, expit, gammainc, + gammaincc, gammaln, logit, multigammaln, @@ -4516,6 +4517,294 @@ def variance(self) -> ArrayLike: ) +def _upper_incomplete_gamma_impl(a: ArrayLike, x: ArrayLike) -> ArrayLike: + a_round = jnp.round(a) + is_pole = (a == a_round) & (a_round <= 0.0) + # keep the recurrence free of divisions by zero so that `jnp.where` does not + # propagate NaNs through the gradients of the branch that is not taken + a_safe = jnp.where(is_pole, 0.5, a) + log_x = jnp.log(x) + g = jnp.exp(gammaln(a_safe + 3.0)) * gammaincc(a_safe + 3.0, x) + for k in (2.0, 1.0, 0.0): + s = a_safe + k + g = (g - jnp.exp(s * log_x - x)) / s + g_zero = -expi(-x) + g_minus_one = jnp.exp(-log_x - x) - g_zero + g_minus_two = 0.5 * (jnp.exp(-2.0 * log_x - x) - g_minus_one) + g_pole = jnp.where( + a_round == 0.0, g_zero, jnp.where(a_round == -1.0, g_minus_one, g_minus_two) + ) + return jnp.where(is_pole, g_pole, g) + + +@jax.custom_jvp +def _upper_incomplete_gamma(a: ArrayLike, x: ArrayLike) -> ArrayLike: + r"""Upper incomplete gamma function :math:`\Gamma(a, x)` for :math:`x > 0`, + extended to negative shape values :math:`a > -3` via the downward recurrence + + .. math:: + \Gamma(a, x) = \frac{\Gamma(a + 1, x) - x^a e^{-x}}{a}. + + The removable singularities of the recurrence at :math:`a \in \{0, -1, -2\}` + are patched with :math:`\Gamma(0, x) = -\mathrm{Ei}(-x)` and its recurrences. + The patch is constant in ``a``, so the ``a``-tangent at those isolated + points is supplied by the custom JVP below (central difference across the + removable singularity); the ``x``-tangent uses the analytic derivative + :math:`\partial\Gamma(a, x)/\partial x = -x^{a-1} e^{-x}`. + """ + return _upper_incomplete_gamma_impl(a, x) + + +@_upper_incomplete_gamma.defjvp +def _upper_incomplete_gamma_jvp(primals, tangents): + a, x = primals + a_dot, x_dot = tangents + primal = _upper_incomplete_gamma_impl(a, x) + a_round = jnp.round(a) + is_pole = (a == a_round) & (a_round <= 0.0) + # Exact a-derivative through the recurrence away from the poles; the pole + # patch is constant in `a`, so autodiff there would silently drop the + # normalization gradient. The function is smooth in `a` across the + # removable singularities, so a central difference is accurate there. + a_safe = jnp.where(is_pole, 0.5, a) + _, da_recurrence = jax.jvp( + lambda aa: _upper_incomplete_gamma_impl(aa, x), (a_safe,), (jnp.ones_like(a),) + ) + delta = 1e-3 + da_pole = ( + _upper_incomplete_gamma_impl(a + delta, x) + - _upper_incomplete_gamma_impl(a - delta, x) + ) / (2.0 * delta) + da = jnp.where(is_pole, da_pole, da_recurrence) + dx = -jnp.exp((a - 1.0) * jnp.log(x) - x) + return primal, da * a_dot + dx * x_dot + + +def _schechter_mag_parts( + alpha: ArrayLike, m_star: ArrayLike, low: ArrayLike, high: ArrayLike +) -> tuple[ArrayLike, ArrayLike, ArrayLike]: + """Concentration, Gamma(concentration, x(low)), and normalization constant.""" + c = 0.4 * np.log(10.0) + concentration = alpha + 1.0 + x_at_high = jnp.exp(c * (m_star - high)) + x_at_low = jnp.exp(c * (m_star - low)) + gamma_at_low = _upper_incomplete_gamma(concentration, x_at_low) + norm = _upper_incomplete_gamma(concentration, x_at_high) - gamma_at_low + return concentration, gamma_at_low, norm + + +def _schechter_mag_log_prob( + value: ArrayLike, + alpha: ArrayLike, + m_star: ArrayLike, + low: ArrayLike, + high: ArrayLike, +) -> ArrayLike: + c = 0.4 * np.log(10.0) + _, _, norm = _schechter_mag_parts(alpha, m_star, low, high) + # `norm` underflows to zero when the whole support is far brighter than + # m_star (x >> 1 across the interval); return -inf rather than +inf there. + norm_ok = norm > 0.0 + log_norm = jnp.log(jnp.where(norm_ok, norm, 1.0)) + log_prob = ( + np.log(c) + + (alpha + 1.0) * c * (m_star - value) + - jnp.exp(c * (m_star - value)) + - log_norm + ) + return jnp.where(norm_ok, log_prob, -jnp.inf) + + +def _schechter_mag_cdf( + value: ArrayLike, + alpha: ArrayLike, + m_star: ArrayLike, + low: ArrayLike, + high: ArrayLike, +) -> ArrayLike: + c = 0.4 * np.log(10.0) + concentration, gamma_at_low, norm = _schechter_mag_parts(alpha, m_star, low, high) + x = jnp.exp(c * (m_star - value)) + # guard the normalization-underflow regime (see _schechter_mag_log_prob): + # 0/0 would produce an accidental NaN here; make it a deliberate one + norm_ok = norm > 0.0 + safe_norm = jnp.where(norm_ok, norm, 1.0) + cdf = jnp.clip( + (_upper_incomplete_gamma(concentration, x) - gamma_at_low) / safe_norm, 0.0, 1.0 + ) + return jnp.where(norm_ok, cdf, jnp.nan) + + +@jax.custom_jvp +def _schechter_mag_icdf( + q: ArrayLike, + alpha: ArrayLike, + m_star: ArrayLike, + low: ArrayLike, + high: ArrayLike, +) -> ArrayLike: + shape = lax.broadcast_shapes( + jnp.shape(q), + jnp.shape(alpha), + jnp.shape(m_star), + jnp.shape(low), + jnp.shape(high), + ) + lower = jnp.broadcast_to(low * jnp.ones_like(q), shape) + upper = jnp.broadcast_to(high * jnp.ones_like(q), shape) + c = 0.4 * np.log(10.0) + # hoist the loop-invariant incomplete-gamma evaluations out of the + # bisection: only Gamma(concentration, x(mid)) changes per iteration + concentration, gamma_at_low, norm = _schechter_mag_parts(alpha, m_star, low, high) + + def body_fn(_, carry): + lower, upper = carry + mid = 0.5 * (lower + upper) + x_mid = jnp.exp(c * (m_star - mid)) + cdf_mid = jnp.clip( + (_upper_incomplete_gamma(concentration, x_mid) - gamma_at_low) / norm, + 0.0, + 1.0, + ) + go_right = cdf_mid < q + return jnp.where(go_right, mid, lower), jnp.where(go_right, upper, mid) + + # bisection halves the interval once per iteration, so ~32 iterations + # already exceed float32 resolution; float64 needs the full 64 + num_iters = 64 if jnp.result_type(float) == jnp.float64 else 32 + lower, upper = lax.fori_loop(0, num_iters, body_fn, (lower, upper)) + # in the normalization-underflow regime the cdf is NaN and the bisection + # would silently collapse to `low`; return NaN loudly instead + return jnp.where(norm > 0.0, 0.5 * (lower + upper), jnp.nan) + + +@_schechter_mag_icdf.defjvp +def _schechter_mag_icdf_jvp(primals, tangents): + q, alpha, m_star, low, high = primals + q_dot, alpha_dot, m_star_dot, low_dot, high_dot = tangents + value = _schechter_mag_icdf(q, alpha, m_star, low, high) + # implicit differentiation of cdf(value, alpha, m_star, low, high) = q + density = jnp.exp(_schechter_mag_log_prob(value, alpha, m_star, low, high)) + _, cdf_dot = jax.jvp( + lambda *params: _schechter_mag_cdf(value, *params), + (alpha, m_star, low, high), + (alpha_dot, m_star_dot, low_dot, high_dot), + ) + return value, (q_dot - cdf_dot) / density + + +class SchechterMag(Distribution): + r"""Schechter distribution over absolute magnitudes, i.e. the single Schechter + luminosity function [1] expressed in absolute-magnitude space and truncated to + the interval :math:`[\mathrm{low}, \mathrm{high}]`. Its probability density + function is given by, + + .. math:: + f(m\mid \alpha, m^*) = \frac{0.4 \ln(10)}{Z}\, + x(m)^{\alpha + 1} e^{-x(m)}, + \qquad x(m) = 10^{-0.4 (m - m^*)}, + + for :math:`\mathrm{low} \le m \le \mathrm{high}`, where :math:`\alpha` is the + faint-end slope, :math:`m^*` is the characteristic magnitude, and the + normalization constant is expressed in terms of the upper incomplete gamma + function as + + .. math:: + Z = \Gamma\left(\alpha + 1, x(\mathrm{high})\right) + - \Gamma\left(\alpha + 1, x(\mathrm{low})\right). + + Under the change of variables :math:`x = 10^{-0.4 (m - m^*)}`, this is the + distribution of :math:`m^* - 2.5 \log_{10} X` for a Gamma variate + :math:`X` with concentration :math:`\alpha + 1` doubly truncated to + :math:`[x(\mathrm{high}), x(\mathrm{low})]`. Thanks to the truncation, the + density remains normalizable for faint-end slopes :math:`\alpha \le -1`, for + which :math:`\Gamma(\alpha + 1, x)` is evaluated with a recurrence-based + extension of the upper incomplete gamma function to negative shape values. + + .. note:: In single precision, the normalization underflows when the whole + interval sits several magnitudes bright of :math:`m^*` (roughly + :math:`m^* - \mathrm{high} \gtrsim 5`), in which case ``log_prob`` + returns ``-inf`` (``cdf`` and ``sample`` return NaN); enable double + precision for such bright-truncated configurations. Precision also degrades within a machine-epsilon-sized + neighborhood of :math:`\alpha \in \{-1, -2, -3\}` (exact values at the + removable singularities are handled). Second-order differentiation with + respect to ``alpha`` is unavailable: JAX has no derivative rule for + ``igamma_grad_a``, a pre-existing limitation shared by other + ``gammaincc``-based code. + + The number-density amplitude :math:`\phi^*` of the luminosity function does + not affect the normalized density. To also constrain :math:`\phi^*`, add a + Poisson likelihood term on the total number of observed objects. + + :param alpha: faint-end slope :math:`\alpha \in (-4, \infty)`. + :param m_star: characteristic magnitude :math:`m^*`. + :param low: lower (bright-end) bound of the magnitude interval. + :param high: upper (faint-end) bound of the magnitude interval. + + **References:** + + 1. Schechter, P. (1976). An analytic expression for the luminosity function + for galaxies. *The Astrophysical Journal*, 203, 297-306. + """ + + arg_constraints = { + "alpha": constraints.greater_than(-4.0), + "m_star": constraints.real, + "low": constraints.dependent(is_discrete=False, event_dim=0), + "high": constraints.dependent(is_discrete=False, event_dim=0), + } + reparametrized_params = ["alpha", "m_star", "low", "high"] + pytree_data_fields = ("alpha", "m_star", "low", "high", "_support") + + def __init__( + self, + alpha: ArrayLike, + m_star: ArrayLike, + low: ArrayLike, + high: ArrayLike, + *, + validate_args: Optional[bool] = None, + ) -> None: + self.alpha, self.m_star, self.low, self.high = promote_shapes( + alpha, m_star, low, high + ) + self._support = constraints.interval(low, high) + batch_shape = lax.broadcast_shapes( + jnp.shape(alpha), jnp.shape(m_star), jnp.shape(low), jnp.shape(high) + ) + super().__init__(batch_shape=batch_shape, validate_args=validate_args) + + @constraints.dependent_property(is_discrete=False, event_dim=0) + def support(self) -> constraints.Constraint: + return self._support + + def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike: + assert is_prng_key(key) + u = random.uniform(key, shape=sample_shape + self.batch_shape) + return self.icdf(u) + + @property + def normalization(self) -> ArrayLike: + r"""The normalization constant :math:`Z = \Gamma(lpha + 1, x(\mathrm{high})) + - \Gamma(lpha + 1, x(\mathrm{low}))`, i.e. the integral of the unnormalized + magnitude-space Schechter density over the support. Useful for building + number-density likelihoods that infer the amplitude :math:`\phi^*`. + """ + return _schechter_mag_parts(self.alpha, self.m_star, self.low, self.high)[2] + + @validate_sample + def log_prob(self, value: ArrayLike) -> ArrayLike: + return _schechter_mag_log_prob( + value, self.alpha, self.m_star, self.low, self.high + ) + + def cdf(self, value: ArrayLike) -> ArrayLike: + return _schechter_mag_cdf(value, self.alpha, self.m_star, self.low, self.high) + + def icdf(self, q: ArrayLike) -> ArrayLike: + return _schechter_mag_icdf(q, self.alpha, self.m_star, self.low, self.high) + + class HurdleGamma(HurdleProbs): r"""A hurdle Gamma distribution: a two-part model in which a structural zero occurs with probability :math:`g` and, conditional on a positive outcome, diff --git a/test/test_distributions.py b/test/test_distributions.py index 9604ddde8..6ad019a2c 100644 --- a/test/test_distributions.py +++ b/test/test_distributions.py @@ -14,6 +14,7 @@ from numpy.testing import assert_allclose, assert_array_equal import pytest import scipy +from scipy.integrate import quad from scipy.sparse import csr_matrix import scipy.stats as osp @@ -1150,6 +1151,17 @@ def get_sp_dist(jax_dist): T(dist.Dagum, 2.0, np.array([1.0, 2.0, 10.0]), 4.0), T(dist.Dagum, 2.0, 3.0, np.array([0.5, 2.0, 1.0])), T(dist.Dagum, np.array([5.0, 2.0, 10.0]), 3.0, 5.0), + T(dist.SchechterMag, -1.25, -20.5, -24.0, -16.0), + T(dist.SchechterMag, 0.7, 10.0, 6.0, 14.0), + T(dist.SchechterMag, -2.4, -20.0, -25.0, -17.0), + T(dist.SchechterMag, np.array([-1.3, -0.5]), -20.5, -24.0, -16.0), + T( + dist.SchechterMag, + -1.25, + -20.5, + np.array([-24.0, -25.0]), + np.array([-16.0, -15.0]), + ), T(dist.HurdleGamma, 0.35, 2.0, 1.0), T( dist.HurdleGamma, @@ -2117,6 +2129,104 @@ def test_gamma_poisson_log_prob(shape): assert_allclose(actual, expected, rtol=0.05) +def _schechter_mag_unnormalized_pdf(m, alpha, m_star): + x = 10.0 ** (0.4 * (m_star - m)) + return 0.4 * np.log(10.0) * x ** (alpha + 1.0) * np.exp(-x) + + +@pytest.mark.parametrize("alpha", [-2.5, -2.0, -1.25, -1.0, -0.5, 0.7]) +def test_schechter_mag_log_prob_and_cdf(alpha): + m_star, low, high = -20.5, -24.0, -16.0 + d = dist.SchechterMag(alpha, m_star, low, high) + normalization, _ = quad( + _schechter_mag_unnormalized_pdf, low, high, args=(alpha, m_star) + ) + values = np.linspace(low + 0.5, high - 0.5, 9) + expected_log_prob = np.log( + _schechter_mag_unnormalized_pdf(values, alpha, m_star) / normalization + ) + expected_cdf = np.array( + [ + quad(_schechter_mag_unnormalized_pdf, low, v, args=(alpha, m_star))[0] + / normalization + for v in values + ] + ) + tol = 1e-4 if jnp.result_type(float) == jnp.float32 else 1e-8 + assert_allclose(d.log_prob(values), expected_log_prob, rtol=tol, atol=tol) + assert_allclose(d.cdf(values), expected_cdf, rtol=tol, atol=tol) + + +@pytest.mark.parametrize("alpha", [-1.0, -2.0, -3.0]) +def test_schechter_mag_alpha_at_integer_pole(alpha): + # the normalization at alpha + 1 in {0, -1, -2} hits the removable + # singularities of the incomplete gamma recurrence, which are patched via + # the exponential integral; check finiteness and continuity in alpha there + m_star, low, high = -20.5, -24.0, -16.0 + values = np.linspace(-23.0, -17.0, 5) + log_prob = dist.SchechterMag(alpha, m_star, low, high).log_prob(values) + assert np.isfinite(log_prob).all() + for eps in (-1e-4, 1e-4): + log_prob_nearby = dist.SchechterMag(alpha + eps, m_star, low, high).log_prob( + values + ) + assert_allclose(log_prob, log_prob_nearby, rtol=1e-3, atol=1e-3) + + +def test_schechter_mag_normalization_matches_quadrature(): + alpha, m_star, low, high = -1.25, -20.5, -24.0, -16.0 + expected, _ = quad(_schechter_mag_unnormalized_pdf, low, high, args=(alpha, m_star)) + actual = dist.SchechterMag(alpha, m_star, low, high).normalization + tol = 1e-4 if jnp.result_type(float) == jnp.float32 else 1e-8 + assert_allclose(actual, expected, rtol=tol) + + +@pytest.mark.parametrize("alpha", [-1.0, -2.0, -3.0]) +def test_schechter_mag_alpha_pole_gradient(alpha): + # the incomplete-gamma pole patch is constant in alpha, so without the + # custom JVP the normalization contribution to the gradient is dropped + m_star, low, high = -20.5, -24.0, -16.0 + values = np.linspace(-23.0, -17.0, 5) + + def log_prob_sum(a): + return dist.SchechterMag(a, m_star, low, high).log_prob(values).sum() + + actual = jax.grad(log_prob_sum)(alpha) + eps = 1e-3 if jnp.result_type(float) == jnp.float32 else 1e-5 + expected = (log_prob_sum(alpha + eps) - log_prob_sum(alpha - eps)) / (2 * eps) + assert_allclose(actual, expected, rtol=5e-3) + + +def test_schechter_mag_pytree_structure_is_bound_independent(): + # _support must be a data field (as in Uniform); keeping the bounds in the + # static treedef would force a jit recompile for every distinct interval + d1 = dist.SchechterMag(-1.25, -20.5, -24.0, -16.0) + d2 = dist.SchechterMag(-1.25, -20.5, -25.0, -15.0) + assert jax.tree_util.tree_structure(d1) == jax.tree_util.tree_structure(d2) + + +def test_schechter_mag_bright_truncated_underflow_is_neg_inf(): + # the normalization underflows when the support is far bright of m_star; + # log_prob must return -inf (impossible), not +inf, and cdf/sample must + # return NaN loudly rather than 0/0 garbage or a silent collapse to `low` + d = dist.SchechterMag(-1.25, -20.5, -32.0, -30.0) + log_prob = d.log_prob(jnp.asarray(-31.0)) + assert np.isneginf(np.asarray(log_prob)) + assert np.isnan(np.asarray(d.cdf(jnp.asarray(-31.0)))) + samples = d.sample(random.key(0), (3,)) + assert np.isnan(np.asarray(samples)).all() + + +@pytest.mark.parametrize("alpha", [-1.7, -1.0, 0.3]) +def test_schechter_mag_sample_ks(alpha): + m_star, low, high = -21.0, -25.5, -17.0 + d = dist.SchechterMag(alpha, m_star, low, high) + samples = d.sample(random.key(1), (10000,)) + assert (samples >= low).all() and (samples <= high).all() + _, p_value = osp.kstest(np.asarray(samples), lambda m: np.asarray(d.cdf(m))) + assert p_value > 0.01 + + @pytest.mark.parametrize("conc", [15.0, 20.0, 30.0]) def test_inverse_wishart_variance(conc): """Test InverseWishart variance formula against Monte Carlo samples.