gh-271: add tests for glass.observations (#436)

glass/observations.py

@@ -89,7 +89,7 @@ def gaussian_nz(
     mean: float | npt.NDArray[np.float64],
     sigma: float | npt.NDArray[np.float64],
     *,
-    norm: npt.NDArray[np.float64] | None = None,
+    norm: float | npt.NDArray[np.float64] | None = None,
 ) -> npt.NDArray[np.float64]:
     """
     Gaussian redshift distribution.
@@ -130,11 +130,11 @@ def gaussian_nz(
 
 def smail_nz(
     z: npt.NDArray[np.float64],
-    z_mode: npt.NDArray[np.float64],
-    alpha: npt.NDArray[np.float64],
-    beta: npt.NDArray[np.float64],
+    z_mode: float | npt.NDArray[np.float64],
+    alpha: float | npt.NDArray[np.float64],
+    beta: float | npt.NDArray[np.float64],
     *,
-    norm: npt.NDArray[np.float64] | None = None,
+    norm: float | npt.NDArray[np.float64] | None = None,
 ) -> npt.NDArray[np.float64]:
     r"""
     Redshift distribution following Smail et al. (1994).

tests/test_observations.py

@@ -0,0 +1,144 @@
+import numpy as np
+import pytest
+
+from glass import (
+    equal_dens_zbins,
+    fixed_zbins,
+    gaussian_nz,
+    smail_nz,
+    tomo_nz_gausserr,
+    vmap_galactic_ecliptic,
+)
+
+
+def test_vmap_galactic_ecliptic() -> None:
+    """Add unit tests for :func:`vmap_galactic_ecliptic`."""
+    n_side = 4
+
+    # check shape
+
+    vmap = vmap_galactic_ecliptic(n_side)
+    np.testing.assert_array_equal(len(vmap), 12 * n_side**2)
+
+    # no rotation
+
+    vmap = vmap_galactic_ecliptic(n_side, galactic=(0, 0), ecliptic=(0, 0))
+    np.testing.assert_array_equal(vmap, np.zeros_like(vmap))
+
+    # check errors raised
+
+    with pytest.raises(TypeError, match="galactic stripe must be a pair of numbers"):
+        vmap_galactic_ecliptic(n_side, galactic=(1,))  # type: ignore[arg-type]
+
+    with pytest.raises(TypeError, match="ecliptic stripe must be a pair of numbers"):
+        vmap_galactic_ecliptic(n_side, ecliptic=(1,))  # type: ignore[arg-type]
+
+    with pytest.raises(TypeError, match="galactic stripe must be a pair of numbers"):
+        vmap_galactic_ecliptic(n_side, galactic=(1, 2, 3))  # type: ignore[arg-type]
+
+    with pytest.raises(TypeError, match="ecliptic stripe must be a pair of numbers"):
+        vmap_galactic_ecliptic(n_side, ecliptic=(1, 2, 3))  # type: ignore[arg-type]
+
+
+def test_gaussian_nz(rng: np.random.Generator) -> None:
+    """Add unit tests for :func:`gaussian_nz`."""
+    mean = 0
+    sigma = 1
+    z = np.linspace(0, 1, 11)
+
+    # check passing in the norm
+
+    nz = gaussian_nz(z, mean, sigma, norm=0)
+    np.testing.assert_array_equal(nz, np.zeros_like(nz))
+
+    # check the value of each entry is close to the norm
+
+    norm = 1
+    nz = gaussian_nz(z, mean, sigma, norm=norm)
+    np.testing.assert_allclose(nz.sum() / nz.shape, norm, rtol=1e-2)
+
+    # check multidimensionality size
+
+    nz = gaussian_nz(
+        z,
+        np.tile(mean, z.shape),
+        np.tile(sigma, z.shape),
+        norm=rng.normal(size=z.shape),
+    )
+    np.testing.assert_array_equal(nz.shape, (len(z), len(z)))
+
+
+def test_smail_nz() -> None:
+    """Add unit tests for :func:`smail_nz`."""
+    alpha = 1
+    beta = 1
+    mode = 1
+    z = np.linspace(0, 1, 11)
+
+    # check passing in the norm
+
+    pz = smail_nz(z, mode, alpha, beta, norm=0)
+    np.testing.assert_array_equal(pz, np.zeros_like(pz))
+
+
+def test_fixed_zbins() -> None:
+    """Add unit tests for :func:`fixed_zbins`."""
+    zmin = 0
+    zmax = 1
+
+    # check nbins input
+
+    nbins = 5
+    expected_zbins = [(0.0, 0.2), (0.2, 0.4), (0.4, 0.6), (0.6, 0.8), (0.8, 1.0)]
+    zbins = fixed_zbins(zmin, zmax, nbins=nbins)
+    np.testing.assert_array_equal(len(zbins), nbins)
+    np.testing.assert_allclose(zbins, expected_zbins, rtol=1e-15)
+
+    # check dz input
+
+    dz = 0.2
+    zbins = fixed_zbins(zmin, zmax, dz=dz)
+    np.testing.assert_array_equal(len(zbins), np.ceil((zmax - zmin) / dz))
+    np.testing.assert_allclose(zbins, expected_zbins, rtol=1e-15)
+
+    # check dz for spacing which results in a max value above zmax
+
+    zbins = fixed_zbins(zmin, zmax, dz=0.3)
+    np.testing.assert_array_less(zmax, zbins[-1][1])
+
+    # check error raised
+
+    with pytest.raises(ValueError, match="exactly one of nbins and dz must be given"):
+        fixed_zbins(zmin, zmax, nbins=nbins, dz=dz)
+
+
+def test_equal_dens_zbins() -> None:
+    """Add unit tests for :func:`equal_dens_zbins`."""
+    z = np.linspace(0, 1, 11)
+    nbins = 5
+
+    # check expected zbins returned
+
+    expected_zbins = [(0.0, 0.2), (0.2, 0.4), (0.4, 0.6), (0.6, 0.8), (0.8, 1.0)]
+    zbins = equal_dens_zbins(z, np.ones_like(z), nbins)
+    np.testing.assert_allclose(zbins, expected_zbins, rtol=1e-15)
+
+    # check output shape
+
+    np.testing.assert_array_equal(len(zbins), nbins)
+
+
+def test_tomo_nz_gausserr() -> None:
+    """Add unit tests for :func:`tomo_nz_gausserr`."""
+    sigma_0 = 0.1
+    z = np.linspace(0, 1, 11)
+    zbins = [(0, 0.2), (0.2, 0.4), (0.4, 0.6), (0.6, 0.8), (0.8, 1.0)]
+
+    # check zeros returned
+
+    binned_nz = tomo_nz_gausserr(z, np.zeros_like(z), sigma_0, zbins)
+    np.testing.assert_array_equal(binned_nz, np.zeros_like(binned_nz))
+
+    # check the shape of the output
+
+    np.testing.assert_array_equal(binned_nz.shape, (len(zbins), len(z)))