Repository: glass-dev/glass

Issue:
Shells: weight functions have incompatible signature
### Add your issue here

The shell functions such as `glass.tophat_shells()` have a `weight=` parameter for passing a unary weight function.

The existing weight functions, such as `glass.volume_weight()`, do not have the correct signature to be passed, since they require an additional `cosmo` argument. The entire point of these weight functions is to work in conjunction with the shell definition.

The awkward workaround `weight=functools.partial(glass.volume_weight, cosmo)` should be replaced by something better out of the box. For example, we could replace the weight functions by tiny callable class instances: `glass.tophat_windows(..., weight=glass.VolumeWeight(cosmo))`.

Hints from issue comments:
@paddyroddy: I assume this meant to say `glass.tophat_windows`?@paddyroddy: I've got an idea, will implement

Solution patch:
diff --git a/glass/__init__.py b/glass/__init__.py
index 5fd9b17..8c74db4 100644
--- a/glass/__init__.py
+++ b/glass/__init__.py
@@ -1,8 +1,11 @@
 """GLASS package."""
 
 __all__ = [
+    "DensityWeight",
+    "DistanceWeight",
     "MultiPlaneConvergence",
     "RadialWindow",
+    "VolumeWeight",
     "algorithm",
     "check_posdef_spectra",
     "cls2cov",
@@ -11,10 +14,8 @@ __all__ = [
     "cov_from_spectra",
     "cubic_windows",
     "deflect",
-    "density_weight",
     "discretized_cls",
     "distance_grid",
-    "distance_weight",
     "effective_bias",
     "effective_cls",
     "ellipticity_gaussian",
@@ -64,7 +65,6 @@ __all__ = [
     "triaxial_axis_ratio",
     "uniform_positions",
     "vmap_galactic_ecliptic",
-    "volume_weight",
     "write_catalog",
 ]
 
@@ -137,17 +137,17 @@ from glass.shapes import (
     triaxial_axis_ratio,
 )
 from glass.shells import (
+    DensityWeight,
+    DistanceWeight,
     RadialWindow,
+    VolumeWeight,
     combine,
     cubic_windows,
-    density_weight,
     distance_grid,
-    distance_weight,
     linear_windows,
     partition,
     redshift_grid,
     restrict,
     tophat_windows,
-    volume_weight,
 )
 from glass.user import load_cls, save_cls, write_catalog
diff --git a/glass/shells.py b/glass/shells.py
index cc6d085..26b27bc 100644
--- a/glass/shells.py
+++ b/glass/shells.py
@@ -37,14 +37,15 @@ Redshift grids
 Weight functions
 ----------------
 
-.. autofunction:: distance_weight
-.. autofunction:: volume_weight
-.. autofunction:: density_weight
+.. autoclass:: DistanceWeight
+.. autoclass:: VolumeWeight
+.. autoclass:: DensityWeight
 
 """  # noqa: D400
 
 from __future__ import annotations
 
+import dataclasses
 import itertools
 import warnings
 from typing import TYPE_CHECKING, NamedTuple
@@ -65,70 +66,95 @@ if TYPE_CHECKING:
     WeightFunc = Callable[[ArrayLike1D], NDArray[np.float64]]
 
 
-def distance_weight(
-    z: NDArray[np.float64],
-    cosmo: Cosmology,
-) -> NDArray[np.float64]:
-    """
-    Uniform weight in comoving distance.
+@dataclasses.dataclass
+class DistanceWeight:
+    """Uniform weight in comoving distance.
 
-    Parameters
+    Attributes
     ----------
-    z
-        The redshifts at which to evaluate the weight.
     cosmo
         Cosmology instance.
 
-    Returns
-    -------
-        The weight function evaluated at redshifts *z*.
-
     """
-    return 1 / cosmo.ef(z)  # type: ignore[no-any-return]
 
+    cosmo: Cosmology
 
-def volume_weight(
-    z: NDArray[np.float64],
-    cosmo: Cosmology,
-) -> NDArray[np.float64]:
-    """
-    Uniform weight in comoving volume.
+    def __call__(self, z: NDArray[np.float64]) -> NDArray[np.float64]:
+        """
+        Uniform weight in comoving distance.
 
-    Parameters
+        Parameters
+        ----------
+        z
+            The redshifts at which to evaluate the weight.
+
+        Returns
+        -------
+            The weight function evaluated at redshifts *z*.
+
+        """
+        return 1 / self.cosmo.ef(z)  # type: ignore[no-any-return]
+
+
+@dataclasses.dataclass
+class VolumeWeight:
+    """Uniform weight in comoving volume.
+
+    Attributes
     ----------
-    z
-        The redshifts at which to evaluate the weight.
     cosmo
         Cosmology instance.
 
-    Returns
-    -------
-        The weight function evaluated at redshifts *z*.
-
     """
-    return cosmo.xm(z) ** 2 / cosmo.ef(z)  # type: ignore[no-any-return]
 
+    cosmo: Cosmology
 
-def density_weight(
-    z: NDArray[np.float64],
-    cosmo: Cosmology,
-) -> NDArray[np.float64]:
+    def __call__(self, z: NDArray[np.float64]) -> NDArray[np.float64]:
+        """
+        Uniform weight in comoving distance.
+
+        Parameters
+        ----------
+        z
+            The redshifts at which to evaluate the weight.
+
+        Returns
+        -------
+            The weight function evaluated at redshifts *z*.
+
+        """
+        return self.cosmo.xm(z) ** 2 / self.cosmo.ef(z)  # type: ignore[no-any-return]
+
+
+@dataclasses.dataclass
+class DensityWeight:
     """
     Uniform weight in matter density.
 
-    Parameters
+    Attributes
     ----------
-    z
-        The redshifts at which to evaluate the weight.
     cosmo
         Cosmology instance.
 
-    Returns
-    -------
-        The weight function evaluated at redshifts *z*.
-
     """
-    return cosmo.rho_m_z(z) * cosmo.xm(z) ** 2 / cosmo.ef(z)  # type: ignore[no-any-return]
+
+    cosmo: Cosmology
+
+    def __call__(self, z: NDArray[np.float64]) -> NDArray[np.float64]:
+        """
+        Uniform weight in comoving distance.
+
+        Parameters
+        ----------
+        z
+            The redshifts at which to evaluate the weight.
+
+        Returns
+        -------
+            The weight function evaluated at redshifts *z*.
+
+        """
+        return self.cosmo.rho_m_z(z) * self.cosmo.xm(z) ** 2 / self.cosmo.ef(z)  # type: ignore[no-any-return]
 
 
 class RadialWindow(NamedTuple):


Test patch:
diff --git a/tests/test_shells.py b/tests/test_shells.py
index ab97f83..22b5beb 100644
--- a/tests/test_shells.py
+++ b/tests/test_shells.py
@@ -7,12 +7,12 @@ import glass
 
 
 def test_distance_weight(cosmo: Cosmology) -> None:
-    """Add unit tests for :func:`glass.distance_weight`."""
+    """Add unit tests for :class:`glass.DistanceWeight`."""
     z = np.linspace(0, 1, 6)
 
     # check shape
 
-    w = glass.distance_weight(z, cosmo)
+    w = glass.DistanceWeight(cosmo)(z)
     np.testing.assert_array_equal(w.shape, z.shape)
 
     # check first value is 1
@@ -25,12 +25,12 @@ def test_distance_weight(cosmo: Cosmology) -> None:
 
 
 def test_volume_weight(cosmo: Cosmology) -> None:
-    """Add unit tests for :func:`glass.volume_weight`."""
+    """Add unit tests for :class:`glass.VolumeWeight`."""
     z = np.linspace(0, 1, 6)
 
     # check shape
 
-    w = glass.volume_weight(z, cosmo)
+    w = glass.VolumeWeight(cosmo)(z)
     np.testing.assert_array_equal(w.shape, z.shape)
 
     # check first value is 0
@@ -43,12 +43,12 @@ def test_volume_weight(cosmo: Cosmology) -> None:
 
 
 def test_density_weight(cosmo: Cosmology) -> None:
-    """Add unit tests for :func:`glass.density_weight`."""
+    """Add unit tests for :class:`glass.DensityWeight`."""
     z = np.linspace(0, 1, 6)
 
     # check shape
 
-    w = glass.density_weight(z, cosmo)
+    w = glass.DensityWeight(cosmo)(z)
     np.testing.assert_array_equal(w.shape, z.shape)
 
     # check first value is 0