numpyro_schechter package

Submodules

numpyro_schechter.distribution module

class numpyro_schechter.distribution.SchechterMag(alpha, M_star, logphi, mag_obs, include_poisson_term=False, volume=1.0, alpha_domain_depth=3, validate_args=None)

Bases: Distribution

NumPyro-compatible distribution based on the single-component Schechter luminosity function in absolute magnitude space.

This distribution models galaxy number densities using the Schechter parameterisation in magnitude space:

φ(M) ∝ 10**[0.4(α + 1)(M* − M)] * exp[−10**(0.4(M* − M))]

Parameters:

• alpha (float) – Faint-end slope of the Schechter function.

• M_star (float) – Characteristic magnitude (M*).

• logphi (float) – Logarithm (base 10) of the normalisation φ* (number density per unit volume per magnitude). The internal phi_star is computed as 10**logphi.

• mag_obs (array-like) – Observed magnitudes used for computing the likelihood.

• include_poisson_term (bool, optional) – If True, includes a Poisson likelihood term on the total observed counts to constrain φ*.

• volume (float, optional) – Survey volume, used to scale expected counts for the Poisson term. Volume is in the same units as phi_star (e.g., if phi_star is per Mpc³ per mag, volume should be in Mpc³).

• alpha_domain_depth (int, optional) – Depth parameter for recurrence-based incomplete gamma calculation (default=3).

• Constraints

• -----------

• default (- alpha must be real and non-integer. By) – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• (−3 (the valid domain for alpha + 1 is) – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• 3) – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• corresponding – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• 0. (- mag_obs must contain values such that 10**(0.4(M* − M)) >)

Notes

• Normalisation is handled via a custom, recurrence-based computation of the upper incomplete gamma function, allowing support for automatic differentiation in JAX/NumPyro. To ensure compatibility with NUTS/HMC (which uses JAX’s reverse-mode autodiff), only a fixed set of alpha_domain_depth values are supported. See SchechterMag.supported_depths() for available options (e.g., 3, 5, 10, 15).

• If include_poisson_term=True, the log-likelihood includes:

  log L = Σ log(pdf_shape) + Poisson(N_obs | N_exp)

  where N_exp = volume * normalisation integral over observed magnitudes.

support = Real()

property has_rsample: bool

static supported_depths()

Returns a list of supported values for alpha_domain_depth, corresponding to increasing valid alpha ranges. Larger alpha_domain_depth will see reduced performance due to the corresponding increase in recursions.

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

numpy.ndarray

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

class numpyro_schechter.distribution.DoubleSchechterMag(alpha1, M_star1, logphi1, alpha2, M_star2, logphi2, mag_obs, include_poisson_term=False, volume=1.0, alpha_domain_depth=3, validate_args=None)

Bases: Distribution

NumPyro-compatible distribution based on the double-component Schechter luminosity function in absolute magnitude space.

This distribution models galaxy number densities using the sum of two Schechter functions:

φ(M) = φ1(M) + φ2(M)

with each component of the form:

φ_i(M) ∝ 10**[0.4(α_i + 1)(M*_i − M)] * exp[−10**(0.4(M*_i − M))]

Parameters:

• alpha1 (float) – Faint-end slopes of the two Schechter components.

• alpha2 (float) – Faint-end slopes of the two Schechter components.

• M_star1 (float) – Characteristic magnitudes (M*) of the first and second components.

• M_star2 (float) – Characteristic magnitudes (M*) of the first and second components.

• logphi1 (float) – Logarithm (base 10) of the normalisation φ*_i (number density per unit volume per magnitude) for the first and second Schechter components, respectively. Internally converted as phi_star_i = 10**logphi_i.

• logphi2 (float) – Logarithm (base 10) of the normalisation φ*_i (number density per unit volume per magnitude) for the first and second Schechter components, respectively. Internally converted as phi_star_i = 10**logphi_i.

• mag_obs (array-like) – Observed magnitudes used for computing the likelihood.

• include_poisson_term (bool, optional) – If True, includes a Poisson likelihood term on the total observed counts to constrain φ*_i.

• volume (float, optional) – Survey volume, used to scale expected counts for the Poisson term. Volume is in the same units as phi_star_i (e.g., if phi_star_i is per Mpc³ per mag, volume should be in Mpc³).

• alpha_domain_depth (int, optional) – Depth parameter for recurrence-based incomplete gamma calculation (default=3).

• Constraints

• -----------

• default (- alpha1 and alpha2 must be real and non-integer. By) – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• (−3 (the valid domain for alpha_i + 1 is) – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• 3) – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• corresponding – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• 0. (- mag_obs must contain values such that 10**(0.4(M*_i − M)) >)

Notes

• Normalisation is handled via a custom, recurrence-based computation of the upper incomplete gamma function, allowing support for automatic differentiation in JAX/NumPyro. To ensure compatibility with NUTS/HMC (which uses JAX’s reverse-mode autodiff), only a fixed set of alpha_domain_depth values are supported. See SchechterMag.supported_depths() for available options (e.g., 3, 5, 10, 15).

• If include_poisson_term=True, the log-likelihood includes:

  log L = Σ log(pdf_shape) + Poisson(N_obs | N_exp)

  where N_exp = volume * (normalisation1 + normalisation2) over observed magnitudes.

support = Real()

property has_rsample

static supported_depths()

Returns a list of supported values for alpha_domain_depth, corresponding to increasing valid alpha ranges. Larger alpha_domain_depth will see reduced performance due to the corresponding increase in recursions.

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

numpy.ndarray

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

numpyro_schechter.math_utils module

numpyro_schechter.math_utils.schechter_mag(phi_star, M_star, alpha, M)

Schechter density in magnitude space (unnormalised).

numpyro_schechter.math_utils.custom_gammaincc(s: Array | ndarray | bool_ | number | bool | int | float | complex, x: Array | ndarray | bool_ | number | bool | int | float | complex, recur_depth: int = 3) → Array

Computes Γ(s, x) using a recurrence-based method.

Valid for real, non-integer s. The default supported domain for s is approximately (-3, 3), corresponding to recur_depth=3. Increase recur_depth to extend this range. Input x must be positive.

Raises:

ValueError – If recur_depth is not in the supported set.

numpyro_schechter.math_utils.s_positive(s: Array | ndarray | bool_ | number | bool | int | float | complex, x: Array | ndarray | bool_ | number | bool | int | float | complex) → Array

Regularised upper incomplete gamma function * Gamma(s)

numpyro_schechter.math_utils.compute_gamma(s: Array | ndarray | bool_ | number | bool | int | float | complex, x: Array | ndarray | bool_ | number | bool | int | float | complex, recur_depth: int) → Array

Module contents

class numpyro_schechter.SchechterMag(alpha, M_star, logphi, mag_obs, include_poisson_term=False, volume=1.0, alpha_domain_depth=3, validate_args=None)

Bases: Distribution

NumPyro-compatible distribution based on the single-component Schechter luminosity function in absolute magnitude space.

This distribution models galaxy number densities using the Schechter parameterisation in magnitude space:

φ(M) ∝ 10**[0.4(α + 1)(M* − M)] * exp[−10**(0.4(M* − M))]

Parameters:

• alpha (float) – Faint-end slope of the Schechter function.

• M_star (float) – Characteristic magnitude (M*).

• logphi (float) – Logarithm (base 10) of the normalisation φ* (number density per unit volume per magnitude). The internal phi_star is computed as 10**logphi.

• mag_obs (array-like) – Observed magnitudes used for computing the likelihood.

• include_poisson_term (bool, optional) – If True, includes a Poisson likelihood term on the total observed counts to constrain φ*.

• volume (float, optional) – Survey volume, used to scale expected counts for the Poisson term. Volume is in the same units as phi_star (e.g., if phi_star is per Mpc³ per mag, volume should be in Mpc³).

• alpha_domain_depth (int, optional) – Depth parameter for recurrence-based incomplete gamma calculation (default=3).

• Constraints

• -----------

• default (- alpha must be real and non-integer. By) – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• (−3 (the valid domain for alpha + 1 is) – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• 3) – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• corresponding – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• 0. (- mag_obs must contain values such that 10**(0.4(M* − M)) >)

Notes

• Normalisation is handled via a custom, recurrence-based computation of the upper incomplete gamma function, allowing support for automatic differentiation in JAX/NumPyro. To ensure compatibility with NUTS/HMC (which uses JAX’s reverse-mode autodiff), only a fixed set of alpha_domain_depth values are supported. See SchechterMag.supported_depths() for available options (e.g., 3, 5, 10, 15).

• If include_poisson_term=True, the log-likelihood includes:

  log L = Σ log(pdf_shape) + Poisson(N_obs | N_exp)

  where N_exp = volume * normalisation integral over observed magnitudes.

support = Real()

property has_rsample: bool

static supported_depths()

Returns a list of supported values for alpha_domain_depth, corresponding to increasing valid alpha ranges. Larger alpha_domain_depth will see reduced performance due to the corresponding increase in recursions.

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

numpy.ndarray

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

class numpyro_schechter.DoubleSchechterMag(alpha1, M_star1, logphi1, alpha2, M_star2, logphi2, mag_obs, include_poisson_term=False, volume=1.0, alpha_domain_depth=3, validate_args=None)

Bases: Distribution

NumPyro-compatible distribution based on the double-component Schechter luminosity function in absolute magnitude space.

This distribution models galaxy number densities using the sum of two Schechter functions:

φ(M) = φ1(M) + φ2(M)

with each component of the form:

φ_i(M) ∝ 10**[0.4(α_i + 1)(M*_i − M)] * exp[−10**(0.4(M*_i − M))]

Parameters:

• alpha1 (float) – Faint-end slopes of the two Schechter components.

• alpha2 (float) – Faint-end slopes of the two Schechter components.

• M_star1 (float) – Characteristic magnitudes (M*) of the first and second components.

• M_star2 (float) – Characteristic magnitudes (M*) of the first and second components.

• logphi1 (float) – Logarithm (base 10) of the normalisation φ*_i (number density per unit volume per magnitude) for the first and second Schechter components, respectively. Internally converted as phi_star_i = 10**logphi_i.

• logphi2 (float) – Logarithm (base 10) of the normalisation φ*_i (number density per unit volume per magnitude) for the first and second Schechter components, respectively. Internally converted as phi_star_i = 10**logphi_i.

• mag_obs (array-like) – Observed magnitudes used for computing the likelihood.

• include_poisson_term (bool, optional) – If True, includes a Poisson likelihood term on the total observed counts to constrain φ*_i.

• volume (float, optional) – Survey volume, used to scale expected counts for the Poisson term. Volume is in the same units as phi_star_i (e.g., if phi_star_i is per Mpc³ per mag, volume should be in Mpc³).

• alpha_domain_depth (int, optional) – Depth parameter for recurrence-based incomplete gamma calculation (default=3).

• Constraints

• -----------

• default (- alpha1 and alpha2 must be real and non-integer. By) – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• (−3 (the valid domain for alpha_i + 1 is) – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• 3) – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• corresponding – to alpha_domain_depth=3. To support broader domains, increase alpha_domain_depth (see note).

• 0. (- mag_obs must contain values such that 10**(0.4(M*_i − M)) >)

Notes

• Normalisation is handled via a custom, recurrence-based computation of the upper incomplete gamma function, allowing support for automatic differentiation in JAX/NumPyro. To ensure compatibility with NUTS/HMC (which uses JAX’s reverse-mode autodiff), only a fixed set of alpha_domain_depth values are supported. See SchechterMag.supported_depths() for available options (e.g., 3, 5, 10, 15).

• If include_poisson_term=True, the log-likelihood includes:

  log L = Σ log(pdf_shape) + Poisson(N_obs | N_exp)

  where N_exp = volume * (normalisation1 + normalisation2) over observed magnitudes.

support = Real()

property has_rsample

static supported_depths()

Returns a list of supported values for alpha_domain_depth, corresponding to increasing valid alpha ranges. Larger alpha_domain_depth will see reduced performance due to the corresponding increase in recursions.

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

numpy.ndarray

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

numpyro_schechter.custom_gammaincc(s: Array | ndarray | bool_ | number | bool | int | float | complex, x: Array | ndarray | bool_ | number | bool | int | float | complex, recur_depth: int = 3) → Array

Computes Γ(s, x) using a recurrence-based method.

Valid for real, non-integer s. The default supported domain for s is approximately (-3, 3), corresponding to recur_depth=3. Increase recur_depth to extend this range. Input x must be positive.

Raises:

ValueError – If recur_depth is not in the supported set.