Add SchechterMag distribution for astrophysical luminosity functions

[arham766]

Adds a SchechterMag distribution — the single Schechter luminosity function (Schechter 1976) in absolute-magnitude space, truncated to [low, high] — as discussed in #2065.

Design

• Magnitude-space density f(m | alpha, m_star) ∝ x^(alpha+1) exp(-x) with x = 10^(-0.4 (m - m_star)); equivalently a doubly-truncated Gamma under the change of variables, which is what makes the density normalizable for faint-end slopes alpha <= -1 (the empirically common regime).
• The normalization needs the upper incomplete gamma at negative shape, which jax.scipy.special.gammaincc doesn't support — implemented via the standard downward recurrence from Gamma(a+3, x), with the removable singularities at a ∈ {0, -1, -2} patched through the exponential integral. Valid for alpha > -4, which covers any physically sensible slope.
• sample via inverse-CDF (64-iteration bisection, custom_jvp with implicit differentiation of cdf(icdf(q)) = q), so samples are reparametrized and usable in gradient-based inference.
• phi_star intentionally omitted: the amplitude doesn't affect the normalized density (docstring points users to a Poisson term on counts for number-density inference).

Tests (in the generic harness + dedicated):

• 4 entries in CONTINUOUS → full generic battery: 92 passed, 40 skipped (skips are the standard pytree/expand exclusions applied to all distributions).
• log_prob and cdf validated against scipy.integrate.quad references across alpha ∈ {-2.5, -2.0, -1.25, -1.0, -0.5, 0.7}.
• Continuity checks at the recurrence poles (alpha ∈ {-1, -2, -3}).
• KS test on 10k samples against the model CDF.

make lint and make format --check are clean.

Happy to follow up with the linear-luminosity-space counterpart and/or a double-Schechter if there's interest — kept this PR to the magnitude form since that's what survey likelihoods typically parametrize.

Closes #2065

[juanitorduz]

@alserene, would you mind helping us review this one? Especially in math, because I do not know much about astrophysics XD

[juanitorduz]

These are some comments from a Claude-assisted review

1. Gradients are silently wrong at exactly alpha ∈ {-1, -2, -3} (medium-high).
The docstring says the derivative at the poles is "not exact," but it's worse than that — the pole patch is constant in a, so the entire normalization contribution to ∂log_prob/∂alpha is dropped. Measured at alpha = -1: autodiff gives -0.4605 while the true gradient is 1.683 (verified by finite differences; just off the pole at alpha = -1.000001, autodiff correctly gives 1.6835). Since alpha = -1 is probably the single most common faint-end slope in the literature, users will initialize HMC/SVI exactly there and get a badly biased first gradient.

Suggested fixes, in order of preference:

• (a) Wrap _upper_incomplete_gamma in a custom_jvp that patches the a-tangent at the poles via a central difference with δ ≈ 1e-3 (the function is smooth in a, so this is accurate).
• (b) Failing that, upgrade the docstring warning from "not exact" to something explicit like: "the gradient w.r.t. alpha is incorrect at exactly alpha ∈ {-1, -2, -3}; avoid initializing at these values."

2. _support belongs in pytree_data_fields, not pytree_aux_fields (medium).
Uniform — the direct precedent for a dependent interval support — puts _support in data fields. This PR puts it in aux, which makes the concrete bound arrays part of the static treedef. Demonstrated consequence: jax.jit over the distribution-as-pytree recompiles for every distinct low/high value (cache size 2 after two calls with different bounds; Uniform compiles once), and traced bounds risk leaked tracers in the treedef. The generic tests don't catch this because they only exercise concrete bounds.

One-line fix: move _support to pytree_data_fields, matching Uniform.

3. Sampling is ~60x slower than it needs to be (medium).
100k samples take ~6.2s on CPU vs ~0.10s for Gamma. The bisection body calls _schechter_mag_cdf, which evaluates _upper_incomplete_gamma three times per iteration — and two of those (the normalization and Γ(a, x_at_low)) are loop-invariant. XLA can't hoist them out of a fori_loop, so they're recomputed 64 times. Hoisting gives an easy ~3x; scaling the iteration count to dtype (float32 needs ~30, not 64) or adding a few Newton steps after coarse bisection would help further. Not blocking, but sample sits inside every predictive call.

4. float32 failure mode returns +inf, not -inf (medium-low).
When the support sits ≳5 mag bright of m_star, exp(-x) underflows in float32, Z → 0, and log_prob returns +inf — which produces NaN gradients rather than a recoverable -inf. In float64 the same configuration is exact to machine precision (linear-space Γ survives to x ≈ 700). Since many users run float32 by default, either compute the normalization in log space or add a docstring note recommending enable_x64 for bright-truncated configurations.

5. Minor items.

• Near-pole precision degrades like eps_machine / |alpha - pole| — negligible in float64 (~1e-4 relative error only within |Δalpha| < 1e-10), but in float32 the fuzzy zone is wider; a one-line note suffices.
• No low < high check under validate_args (consistent with Uniform, so arguably fine).
• super(SchechterMag, self).__init__ should be plain super().__init__ per the style elsewhere in the file.
• The CONTINUOUS entries batch only alpha; adding one entry with batched low/high would exercise the neighborhood of issue 2.
• mean/variance properties are absent — fine to defer, though both are expressible through the truncated-gamma representation.

[arham766]

Thanks @juanitorduz — I validated each point locally before changing anything. Results:

1. Pole gradients — confirmed, fixed. Reproduced (worse than your numbers with a batch of values: autodiff gave the same -2.30 at all three poles vs true 8.42 / 14.03 / 15.99). Fixed with your option (a): _upper_incomplete_gamma now has a custom_jvp — exact recurrence tangent away from the poles, central difference (δ=1e-3) across the removable singularities, analytic x-tangent. Autodiff now matches finite differences at alpha ∈ {-1, -2, -3} to <0.5%, with a regression test per pole (fails on the previous commit).

2. _support pytree placement — confirmed, fixed. Moved to pytree_data_fields, matching Uniform; added a test asserting the treedef is identical across different bounds (fails on the previous commit).

3. Sampling speed — confirmed (67x vs Gamma on my machine), but the suggested fix doesn't deliver. I implemented the hoist and measured back-to-back under identical conditions: 5.79s → 5.64s per 100k samples. XLA was evidently already hoisting the loop-invariant evaluations out of the fori_loop; the real cost is the genuinely per-iteration Γ(α+1, x_mid) over 64 bisection steps. I kept the hoisted structure (clearer, and guarantees the optimization) and scaled iterations to dtype (32 for float32, which does halve its cost). A real speedup needs bisection→Newton refinement — happy to do that as a follow-up if you want it, but I'd rather not grow this PR.

4. float32 underflow — confirmed (+inf log_prob, nan grad), fixed. log_prob now returns -inf when the normalization underflows, with a docstring note recommending x64 for bright-truncated configurations, and a regression test.

5. super().__init__ fixed, batched low/high harness entry added, near-pole precision note added. mean/variance via the truncated-gamma representation I'd also prefer as a follow-up.

Full battery after the changes: 117 passed / 50 skipped; lint/format clean.

[juanitorduz]

Thanks @arham766 ! Here is another review round (pair review with Claude)

• float32 underflow — fixed, with one leftover inconsistency. log_prob now returns −inf instead of +inf when the normalization underflows (verified on the new test case, which cleverly triggers in both precisions), and the docstring note is accurate. However, cdf in the same regime returns NaN (0/0, unguarded) and sample silently returns the low bound (the NaN comparison makes go_right always False, so bisection collapses to the bright edge). Since log_prob was given a guard, cdf deserves the same one-line treatment; the current behavior means a user in the underflow regime gets a loud −inf from inference but plausible-looking garbage from prior predictive sampling. Minor, but cheap to fix while the file is open.

• Minor items — mostly done. super().init fixed, batched-bounds test added, precision docstring note added. The low < high validation and mean/variance properties weren't done, but those were flagged as optional and I'd stand by that
  .
  -One non-issue worth knowing about: Second-order differentiation (grad(grad(log_prob)) w.r.t. α) raises NotImplementedError: Differentiation rule for 'igamma_grad_a' not implemented. I checked the old commit — identical failure, so this is a pre-existing JAX limitation, not something the custom JVP introduced. Anything needing Hessians w.r.t. α (Laplace approximations, some mass-matrix adaptation schemes) will hit it, but that's true of other gammaincc-based code in the ecosystem too. At most a docstring footnote; arguably not this PR's problem at all.

[arham766]

Thanks — validated and fixed. Confirmed both behaviors first: cdf returned NaN via the unguarded 0/0, and sample returned the low bound for every draw ([-32, -32, ...]) — plausible-looking garbage exactly as you said.

• cdf now carries the same guard as log_prob and returns a deliberate NaN in the underflow regime (with a comment; NaN is the honest answer for an undefined distribution, but now it's on purpose rather than arithmetic accident).
• _schechter_mag_icdf returns NaN there too, so sample fails loudly instead of silently collapsing to the bright edge.
• The existing underflow test now also asserts cdf is NaN and all samples are NaN (it failed before the guards, passes after).
• Added the igamma_grad_a second-order limitation as a docstring footnote per your suggestion — agreed it's pre-existing and not this PR's problem, but cheap to document while the note block exists.

Full schechter_mag selection: 17 passed; lint/format clean.

[juanitorduz]

ok, thanks :)

I would wait for @fehiepsi 's comments before we merge :) (and ideally @alserene :D )

[alserene]

Thanks for your work on this @arham766, it's great to see this happen!

I've compared it against the standalone implementation I originally wrote, and the mathematics looks sound 👍

One suggestion from an astronomy perspective: would it be possible to expose the normalisation integral (or an equivalent helper) through the public API? I agree that the Poisson count term belongs in the user's model rather than in SchechterMag, but exposing the normalisation would make it straightforward to construct the standard number-density likelihood (i.e. infer φ*) without duplicating the internal normalisation code.

Also, my original proposal included both single- and double-Schechter distributions. I'd certainly welcome a follow-up PR adding a DoubleSchechterMag distribution, as that remains a common use case in luminosity and stellar mass function analyses :)

[arham766]

Thanks so much for checking the math against your implementation, @alserene — that's exactly the review this needed!

Both suggestions taken:

• Normalization exposed: SchechterMag now has a public normalization property returning Z = Γ(α+1, x(high)) − Γ(α+1, x(low)) (the integral of the unnormalized magnitude-space density over the support), with a docstring pointing at the φ* number-density use case and a test verifying it against scipy.integrate.quad.
• DoubleSchechterMag: gladly — I'll open a follow-up PR for it once this one lands, keeping the same parametrization conventions so the two compose naturally for LF/SMF analyses.

[entaylor]

Late to the party, but ... a couple of user-level comments ...

On the domain for the power-law index a — for a < -1 the indefinite integral (i.e. the total population number) diverges; for -1 < a < -2 the number diverges but the integrated light does converge. For this reason, and for mine, i would say that a > -2 is a sensible physical limit for astronomical purposes.

That said, the difference between limits of -1, -2, -3, ... is just the number of iterations using the recursion relation. This might be easily accommodated by including a 'num_recursions' keyword argument? Then you can have a limit of, say, -2 by default, but with the facility for users to explore more extreme solutions if they want to (at the cost of extra compute time).

[arham766]

Thanks @entaylor — good points, and worth spelling out the reasoning in the code's terms.

On the domain: agreed that for physical luminosity/mass functions alpha > -2 is where sensible fits live (and untruncated, the number integral diverges below -1). The reason the class allows more is that this distribution is defined on a truncated interval, so it's a proper, normalizable distribution for any alpha the incomplete-gamma evaluation can handle — the alpha > -4 constraint reflects the numerical validity of the fixed three-step recursion, not physics. My leaning is to keep the class mathematically permissive and let users encode the physical prior in their model (a Uniform(-2, ...) prior on alpha, say), since a hard limit in the distribution would turn an extreme-but-testable hypothesis into a validation error. Happy to add a sentence to the docstring pointing out the physically sensible range, which I think addresses the practical concern.

On num_recursions: you're right that the depth is exactly what sets the bound, and since it would be a static Python int it unrolls cleanly at trace time — implementable without touching the JVP logic. My hesitation is API surface for a regime (alpha < -4) that I've not seen used; each extra step also compounds the cancellation error in the downward recursion, so "more depth" isn't free accuracy-wise. If the maintainers would rather have it configurable I'm glad to add it, either here or as a follow-up alongside DoubleSchechterMag. @fehiepsi @juanitorduz — preference?