Add Schechter distributions for astrophysical luminosity functions

[alserene]

Feature Summary

Introduce SchechterMag and DoubleSchechterMag distribution classes into NumPyro.

• The Schechter function is the standard parametric form used to describe galaxy luminosity and stellar mass functions.
• Both single and double forms are required in astronomy, with parameters (α, M*, φ*) (single) or two such sets (double).
• Key feature: implementation preserves JAX auto-differentiability, allowing the No-U-Turn Sampler (NUTS) to be applied in hierarchical Bayesian models.
• Our prototype implementations already exist in a standalone package (numpyro-schechter), but native NumPyro support would benefit a broader community.

See github and PyPI for the existing implementation.

Why is this needed?

• Astronomy and cosmology rely heavily on Bayesian hierarchical modeling to study galaxy populations. Current analyses are limited by the absence of a Schechter distribution in PyMC/NumPyro.
• Existing Python packages (e.g., scipy, sympy) can evaluate the Schechter function and integrals, but they are not JAX-compatible and cannot be differentiated automatically.
• A JAX-native Schechter implementation would:
  • Enable efficient gradient-based MCMC (NUTS) for models with hundreds of parameters.
  • Open NumPyro to a large class of astrophysical inference problems (luminosity/mass functions, environmental dependencies, multi-wavelength studies).
  • Provide a “plug-and-play” likelihood component for researchers beyond astronomy who may encounter similar truncated power-law distributions with exponential cutoffs.

Our group has already implemented these distributions externally and confirmed they work with NumPyro’s inference algorithms. We would like to contribute these back to NumPyro, with tests and documentation, so that they are maintained and available to the wider community.

[juanitorduz]

This sounds amazing!

[fehiepsi]

This is great, thanks @alserene and the team!

[alserene]

Thanks 😊

Where in the repo should these go? The contrib folder?

What else is needed from me other than following the contributing guidelines? Are there any design docs you think are relevant to this?

Cheers!

[juanitorduz]

Here @fehiepsi will have additional suggestions (as a maintainer). I, as a contributor, suggest:

• Adding the distributions maybe in https://github.com/pyro-ppl/numpyro/blob/master/numpyro/distributions/continuous.py (e.g. if they are continuous)
• The math_utils, I think you could add in a module in numpyro/contrib/schechter

Here is an example PR to add a new contribution module #1794. You can see the core code, doc strings, tests, and examples. I hope this helps :)

I have not seen any design docs for NumPyro (but maybe they exist).

My suggestion is to open a PR (maybe mark it as draft), and collect feedback from the core developers and the community :)