BUG: schechter_smf callable input and docs (#525)

Author: philipp128

skypy/galaxies/__init__.py

@@ -5,6 +5,7 @@
 
 __all__ = [
     'schechter_lf',
+    'schechter_smf',
 ]
 
 from . import luminosity  # noqa F401,F403

skypy/galaxies/stellar_mass.py

@@ -4,13 +4,16 @@
 import numpy as np
 
 from ..utils.random import schechter
+from ..utils import dependent_argument
 
 
 __all__ = [
     'schechter_smf_mass',
 ]
 
 
+@dependent_argument('m_star', 'redshift')
+@dependent_argument('alpha', 'redshift')
 def schechter_smf_mass(redshift, alpha, m_star, m_min, m_max, size=None,
                        resolution=1000):
     r""" Stellar masses following the Schechter mass function [1]_.
@@ -19,9 +22,10 @@ def schechter_smf_mass(redshift, alpha, m_star, m_min, m_max, size=None,
     ----------
     redshift : array_like
         Galaxy redshifts for which to sample magnitudes.
-    alpha : float
-        The alpha parameter in the Schechter stellar mass function.
-    m_star : (nm,) array-like
+    alpha : float or function
+        The alpha parameter in the Schechter stellar mass function. If function,
+        it must return a scalar value.
+    m_star : (nm,) array-like or function
         Characteristic stellar mass m_*.
     size: int, optional
          Output shape of stellar mass samples. If size is None and m_star

skypy/galaxies/tests/test_stellar_mass.py

@@ -3,6 +3,7 @@
 import scipy.integrate
 from scipy.special import gammaln
 import pytest
+from astropy.modeling.models import Exponential1D
 
 from skypy.galaxies import stellar_mass
 from skypy.utils import special
@@ -25,12 +26,30 @@ def test_stellar_masses():
     with pytest.raises(ValueError):
         stellar_mass.schechter_smf_mass(0., -1.4, np.array([1e10, 2e10]), 1.e7, 1.e13, size=3)
 
-    # Test that an array with the sme shape as m_star is returned if m_star is
+    # Test that an array with the same shape as m_star is returned if m_star is
     # an array and size = None
     m_star = np.array([1e10, 2e10])
     sample = stellar_mass.schechter_smf_mass(0., -1.4, m_star, 1.e7, 1.e13, size=None)
     assert m_star.shape == sample.shape
 
+    # Test m_star can be a function that returns an array of values for each redshift
+    redshift = np.linspace(0, 2, 100)
+    alpha = 0.357
+    m_min = 10 ** 7
+    m_max = 10 ** 14
+    m_star_function = Exponential1D(10**10.626, np.log(10)/0.095)
+    sample = stellar_mass.schechter_smf_mass(redshift, alpha, m_star_function, m_min, m_max)
+    assert sample.shape == redshift.shape
+
+    # Test alpha can be a function returning a scalar value.
+    sample = stellar_mass.schechter_smf_mass(redshift, lambda z: alpha, 10**10.626, m_min, m_max)
+    assert sample.shape == redshift.shape
+
+    # Sampling with an array for alpha is not implemented
+    alpha_array = np.full_like(redshift, alpha)
+    with pytest.raises(NotImplementedError, match='only scalar alpha is supported'):
+        sample = stellar_mass.schechter_smf_mass(redshift, alpha_array, m_star, m_min, m_max)
+
     # Test that sampling corresponds to sampling from the right pdf.
     # For this, we sample an array of luminosities for redshift z = 1.0 and we
     # compare it to the corresponding cdf.