ENH(galaxies): sample redshifts from radial window (#128)

Adds a function `redshifts()` that samples redshifts from a distribution
which follows the profile of a radial window function.

Closes: #127
Added: Function `redshifts()` to sample redshifts following a radial
  window function.

glass/galaxies.py

@@ -12,6 +12,7 @@
 Functions
 ---------
 
+.. autofunction:: redshifts
 .. autofunction:: redshifts_from_nz
 .. autofunction:: galaxy_shear
 .. autofunction:: gaussian_phz
@@ -26,6 +27,34 @@
 from numpy.typing import ArrayLike
 
 from .core.array import broadcast_leading_axes, cumtrapz
+from .shells import RadialWindow
+
+
+def redshifts(n: int | ArrayLike, w: RadialWindow, *,
+              rng: np.random.Generator | None = None
+              ) -> np.ndarray:
+    '''Sample redshifts from a radial window function.
+
+    This function samples *n* redshifts from a distribution that follows
+    the given radial window function *w*.
+
+    Parameters
+    ----------
+    n : int or array_like
+        Number of redshifts to sample.  If an array is given, the
+        results are concatenated.
+    w : :class:`~glass.shells.RadialWindow`
+        Radial window function.
+    rng : :class:`~numpy.random.Generator`, optional
+        Random number generator.  If not given, a default RNG is used.
+
+    Returns
+    -------
+    z : array_like
+        Random redshifts following the radial window function.
+
+    '''
+    return redshifts_from_nz(n, w.za, w.wa, rng=rng)
 
 
 def redshifts_from_nz(count: int | ArrayLike, z: ArrayLike, nz: ArrayLike, *,

glass/test/test_galaxies.py

@@ -1,6 +1,26 @@
 import pytest
 
 
+def test_redshifts():
+    from unittest.mock import Mock
+    import numpy as np
+    from glass.galaxies import redshifts
+
+    # create a mock radial window function
+    w = Mock()
+    w.za = np.linspace(0., 1., 20)
+    w.wa = np.exp(-0.5*(w.za - 0.5)**2/0.1**2)
+
+    # sample redshifts (scalar)
+    z = redshifts(13, w)
+    assert z.shape == (13,)
+    assert z.min() >= 0. and z.max() <= 1.
+
+    # sample redshifts (array)
+    z = redshifts([[1, 2], [3, 4]], w)
+    assert z.shape == (10,)
+
+
 def test_redshifts_from_nz():
     import numpy as np
     from glass.galaxies import redshifts_from_nz