gh-136: Implementation of variable depth

docs/examples.rst

@@ -46,3 +46,4 @@ More advanced examples doing multiple things at the same time.
    examples/2-advanced/cosmic_shear.ipynb
    examples/2-advanced/stage_4_galaxies.ipynb
    examples/2-advanced/legacy-mode.ipynb
+   examples/2-advanced/variable-depth.ipynb

docs/user/bibliography.rst

@@ -29,6 +29,9 @@ Bibliography
 .. [Johnston19] Johnston, H., Georgiou, C., Joachimi, B., et al., 2019,
     A&A, 624, A30. doi:10.1051/0004-6361/201834714
 
+.. [Joachimi21] Joachimi B., Lin, C.-A., et al., 2021, A&A, 646, A129.
+    doi:10.1051/0004-6361/202038831
+
 .. [Lawson95] Lawson, C. L. and Hanson, R. J. (1995), Solving Least Squares
     Problems. doi: 10.1137/1.9781611971217
 
@@ -43,3 +46,6 @@ Bibliography
 
 .. [Xavier16] Xavier H. S., et al., 2016, MNRAS, 459, 3693.
     doi:10.1093/mnras/stw874
+
+.. [Wietesheim-Kramsta24] von Wietersheim-Kramsta M., Lin, K., et al.,
+    2024, A&A, 695, A223. doi:10.1051/0004-6361/202450487

glass/observations.py

@@ -22,6 +22,8 @@
 ----------
 
 .. autofunction:: vmap_galactic_ecliptic
+.. class:: AngularVariableDepthMask
+.. class:: AngularLosVariableDepthMask
 
 """  # noqa: D400
 
@@ -37,6 +39,8 @@
 import glass.arraytools
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from numpy.typing import NDArray
 
 
@@ -328,3 +332,206 @@ def tomo_nz_gausserr(
     binned_nz *= nz
 
     return binned_nz  # type: ignore[no-any-return]
+
+
+class AngularVariableDepthMask:
+    """Variable depth mask for tomographic bins.
+
+    This class allows to create a mask with a different variable
+    depth mask in the angular direction for each tomographic bin.
+
+    Parameters
+    ----------
+    vardepth_map
+        Map of variable which traces the depth per tomographic bin.
+    n_bins
+        Number of tomographic bins.
+    zbins
+        Shell redshift limits.
+
+    """
+
+    def __init__(
+        self,
+        vardepth_map: NDArray[np.float64],
+        n_bins: int,
+        zbins: list[tuple[float, float]],
+    ) -> None:
+        self.vardepth_map = vardepth_map
+        self.n_bins = n_bins
+        self.zbins = zbins
+
+    def check_index(self, index: tuple[int, int]) -> None:
+        """Check the index for validity."""
+        if not isinstance(index, tuple):
+            raise TypeError("Index must be an tuple of two integers")
+
+        if index[0] >= self.n_bins:
+            raise ValueError("Leading index cannot exceed number of tomographic bins")
+
+        if index[1] >= len(self.zbins):
+            raise ValueError("Trailing index cannot exceed number of shells")
+
+    def __getitem__(self, index: tuple[int, int]) -> NDArray[np.float64]:
+        """Get the mask for the given index.
+
+        Parameters
+        ----------
+        index
+            Indices of the tomographic bin and shell pair.
+
+        Returns
+        -------
+            Mask for the given index.
+
+        Raises
+        ------
+        ValueError
+            If the index is invalid.
+
+        """
+        self.check_index(index)
+
+        return self.vardepth_map[index[0]]  # type: ignore[no-any-return]
+
+
+class AngularLosVariableDepthMask(AngularVariableDepthMask):
+    """Variable depth mask for tomographic bins.
+
+    This class allows to create a mask with a different variable
+    depth mask in both the angular and the line-of-sight directions
+    for each tomographic bin.
+
+    Parameters
+    ----------
+    vardepth_map
+        Map of variable which traces the depth per tomographic bin.
+        If vardepth_tomo_functions is not provided, the values are
+        treated like a map of galaxy count ratios.
+    n_bins
+        Number of tomographic bins.
+    zbins
+        Shell redshift limits.
+    ztomo
+        Tomographic redshift bin limits.
+    dndz
+        Redshift distributions per tomographic bin.
+    z
+        Redshift domain of dndz.
+    dndz_vardepth
+        Redshift distribution affected by variable depth
+        (n_bins x len(vardepth_values) x len(z)).
+    vardepth_values
+        Variable depth tracer domain/values of dndz_vardepth.
+    vardepth_los_tracer
+        Map of the variable depth tracer for line-of-sight direction. If
+        provided, it is assumed to cover the same domain as vardepth_values.
+    vardepth_tomo_functions
+        List of functions which map the input vardepth_map to the ratio
+        between the galaxy count due to the variable depth and the galaxy
+        count without variable depth (for each tomographic bin). If
+        provided, it is assumed that there is one vardepth_map which
+        traces the variable depth for all tomographic bins.
+
+    """
+
+    def __init__(  # noqa: PLR0913
+        self,
+        vardepth_map: NDArray[np.float64],
+        n_bins: int,
+        zbins: list[tuple[float, float]],
+        ztomo: list[tuple[float, float]],
+        dndz: NDArray[np.float64],
+        z: NDArray[np.float64],
+        dndz_vardepth: NDArray[np.float64],
+        vardepth_values: NDArray[np.float64],
+        vardepth_los_tracer: NDArray[np.float64] | None = None,
+        vardepth_tomo_functions: list[
+            Callable[[NDArray[np.float64]], NDArray[np.float64]]
+        ]
+        | None = None,
+    ) -> None:
+        super().__init__(vardepth_map, n_bins, zbins)
+        self.ztomo = ztomo
+        self.dndz = dndz
+        self.z = z
+        self.dndz_vardepth = dndz_vardepth
+        self.vardepth_values = vardepth_values
+        self.vardepth_los_tracer = vardepth_los_tracer
+        self.vardepth_tomo_functions = vardepth_tomo_functions
+
+        if vardepth_tomo_functions is not None:
+            self.vardepth_map = np.atleast_2d(vardepth_map).reshape(1, -1)
+
+    def get_los_fraction(self, index: tuple[int, int]) -> NDArray[np.float64]:
+        """Gets the fraction of galaxies affected by variable depth in the
+        line-of-sight direction for the given tomographic bin and shell.
+
+        Parameters
+        ----------
+        index
+            Indices of the tomographic bin and shell pair.
+
+        Returns
+        -------
+            Fraction of galaxies affected by variable depth in the
+            line-of-sight direction.
+
+        """
+        is_in_shell = (self.zbins[index[1]][0] < self.z) & (
+            self.z <= self.zbins[index[1]][1]
+        )
+
+        n_gal_in_tomo_vardepth = np.trapezoid(
+            self.dndz_vardepth[index[0]][:, is_in_shell], self.z[is_in_shell]
+        )
+        n_gal_in_tomo = np.trapezoid(
+            self.dndz[index[0]][is_in_shell], self.z[is_in_shell]
+        )
+        return np.divide(  # type: ignore[no-any-return]
+            n_gal_in_tomo_vardepth,
+            n_gal_in_tomo,
+            out=np.ones_like(n_gal_in_tomo_vardepth),
+            where=n_gal_in_tomo != 0,
+        )
+
+    def __getitem__(self, index: tuple[int, int]) -> NDArray[np.float64]:
+        """Get the mask for the given index.
+
+        Parameters
+        ----------
+        index
+            Indices of the tomographic bin and shell pair.
+
+        Returns
+        -------
+            Mask for the given index.
+
+        Raises
+        ------
+        ValueError
+            If the index is invalid.
+
+        """
+        self.check_index(index)
+
+        if self.vardepth_tomo_functions is None:
+            angular_vardepth_map = angular_tracer_map = self.vardepth_map[index[0]]
+        else:
+            angular_vardepth_map = self.vardepth_tomo_functions[index[0]](
+                self.vardepth_map[0]
+            )
+            angular_tracer_map = self.vardepth_map[0]
+
+        los_fraction_vardepth = self.get_los_fraction(index)
+
+        if self.vardepth_los_tracer is None:
+            los_vardepth_map = np.interp(
+                angular_tracer_map, self.vardepth_values, los_fraction_vardepth
+            )
+        else:
+            los_vardepth_map = np.interp(
+                self.vardepth_los_tracer, self.vardepth_values, los_fraction_vardepth
+            )
+
+        return np.multiply(angular_vardepth_map, los_vardepth_map)  # type: ignore[no-any-return]

tests/test_observations.py

@@ -1,8 +1,10 @@
 import healpix
 import numpy as np
 import pytest
+from numpy.typing import NDArray
 
 import glass
+import glass.observations
 
 
 def test_vmap_galactic_ecliptic() -> None:
@@ -136,3 +138,86 @@ def test_tomo_nz_gausserr() -> None:
     # check the shape of the output
 
     np.testing.assert_array_equal(binned_nz.shape, (len(zbins), len(z)))
+
+
+@pytest.mark.parametrize(
+    ("vardepth_map", "n_bins", "zbins", "index", "expected_mask"),
+    [
+        (
+            np.array([[1.0, 2.0], [3.0, 4.0]]),
+            2,
+            [(0.0, 0.5), (0.5, 1.0)],
+            (0, 0),
+            np.array([1.0, 2.0]),
+        ),
+        (
+            np.array([[1.0, 2.0], [3.0, 4.0]]),
+            2,
+            [(0.0, 0.5), (0.5, 1.0)],
+            (1, 1),
+            np.array([3.0, 4.0]),
+        ),
+    ],
+    ids=["test_valid_index_1", "test_valid_index_2"],
+)
+def test_getitem_happy_path(
+    vardepth_map: NDArray[np.float64],
+    n_bins: int,
+    zbins: list[tuple[float, float]],
+    index: tuple[int, int],
+    expected_mask: NDArray[np.float64],
+) -> None:
+    # Arrange
+    mask = glass.observations.AngularVariableDepthMask(vardepth_map, n_bins, zbins)
+
+    # Act
+    result = mask[index]
+
+    # Assert
+    np.testing.assert_array_equal(result, expected_mask)
+
+
+@pytest.mark.parametrize(
+    ("vardepth_map", "n_bins", "zbins", "index", "expected_error"),
+    [
+        (
+            np.array([[1.0, 2.0], [3.0, 4.0]]),
+            2,
+            [(0.0, 0.5), (0.5, 1.0)],
+            (2, 0),
+            ValueError,
+        ),
+        (
+            np.array([[1.0, 2.0], [3.0, 4.0]]),
+            2,
+            [(0.0, 0.5), (0.5, 1.0)],
+            (0, 2),
+            ValueError,
+        ),
+        (
+            np.array([[1.0, 2.0], [3.0, 4.0]]),
+            2,
+            [(0.0, 0.5), (0.5, 1.0)],
+            0,
+            TypeError,
+        ),
+    ],
+    ids=[
+        "test_invalid_index_1",
+        "test_invalid_index_2",
+        "test_invalid_index_type",
+    ],
+)
+def test_getitem_error_cases(
+    vardepth_map: NDArray[np.float64],
+    n_bins: int,
+    zbins: list[tuple[float, float]],
+    index: tuple[int, int],
+    expected_error: type[BaseException],
+) -> None:
+    # Arrange
+    mask = glass.observations.AngularVariableDepthMask(vardepth_map, n_bins, zbins)
+
+    # Act & Assert
+    with pytest.raises(expected_error):
+        mask[index]