add quantiles in mock_data and in _integ_pzfuncs

clmm/redshift/tools.py

@@ -6,7 +6,7 @@
 import qp
 
 
-def _integ_pzfuncs(pzpdf, pzbins, zmin=0.0, zmax=5, kernel=None, ngrid=1000):
+def _integ_pzfuncs(pzpdf, pzbins, zmin=0.0, zmax=5, kernel=None, ngrid=1000, quantiles=False):
     r"""
     Integrates the product of a photo-z pdf with a given kernel.
     This function was created to allow for data with different photo-z binnings.
@@ -26,36 +26,99 @@ def _integ_pzfuncs(pzpdf, pzbins, zmin=0.0, zmax=5, kernel=None, ngrid=1000):
         Default: kernel(z)=1
     ngrid : int, optional
         Number of points for the interpolation of the redshift pdf.
+    quantiles : bool
+        PDF actually contains quantile information:
+            pzbins (array) - Quantiles on which all PDFs are tabulated.
+            pzpdf (list of arrays) - Location of quantiles for the source galaxies.
 
     Returns
     -------
     numpy.ndarray
         Kernel integrated with the pdf of each galaxy.
     """
+    if quantiles:
+        pzpdf_type = "quantiles"
+        data = {"quantiles": pzbins, "locations": pzpdf}
+    if len(np.array(pzbins).shape) > 1:
+        pzpdf_type = "individual_bins"
+        data = {"pzbins": pzbins, "pzpdf": pzpdf}
+    else:
+        pzpdf_type = "shared_bins"
+        data = {"pzbins": pzbins, "pzpdf": pzpdf}
+    return _integ_pzfuncs_core(pzpdf_type, data, zmin=zmin, zmax=zmax, kernel=kernel, ngrid=ngrid)
 
-    if kernel is None:
 
-        def kernel(z):
-            # pylint: disable=unused-argument
-            return 1.0
+def _integ_pzfuncs_core(pzpdf_type, pzdata, zmin=0.0, zmax=5, kernel=None, ngrid=1000):
+    r"""
+    Integrates the product of a photo-z pdf with a given kernel.
+    This function was created to allow for data with different photo-z binnings.
 
-    z_grid = np.linspace(zmin, zmax, ngrid)
+    Parameters
+    ----------
+    pzpdf_type: str, None
+        Type of photo-z pdf to information provided:
+            `'shared_bins'` - single binning for all galaxies
+            `'individual_bins'` - individual binning for each galaxy
+            `'quantiles'` - quantiles of PDF
+    pzdata : dict
+        Data with Pz PDF, content will dependend on pzpdf_type.
+        If pzpdf_type is `'shared_bins'`, it must contain:
+            pzbins (array) - Redshift axis on which all PDFs are tabulated.
+            pzpdf (list of arrays) - Photometric PDF of the source galaxies.
+        If pzpdf_type is `'individual_bins', it must contain:
+            pzbins (list of arrays) - Redshift axis on which each PDF is tabulated.
+            pzpdf (list of arrays) - Photometric PDF of the source galaxies.
+        If pzpdf_type is `'quantiles', it must contain:
+            quantiles (array) - Quantiles on which all PDFs are tabulated.
+            locations (list of arrays) - Location of quantiles for the source galaxies.
+    zmin : float, optional
+        Minimum redshift for integration. Default: 0
+    zmax : float, optional
+        Maximum redshift for integration. Default: 5
+    kernel : function, optional
+        Function to be integrated with the pdf, must be f(z_array) format.
+        Default: kernel(z)=1
+    ngrid : int, optional
+        Number of points for the interpolation of the redshift pdf.
 
-    if hasattr(pzbins[0], "__len__"):
+    Returns
+    -------
+    numpy.ndarray
+        Kernel integrated with the pdf of each galaxy.
+    """
+
+    if pzpdf_type == "individual_bins":
         # Each galaxy as a diiferent zbin
         _qp_type = qp.interp_irregular
-    else:
+        _data = {
+            "xvals": np.array(pzdata["pzbins"]),
+            "yvals": np.array(pzdata["pzpdf"]),
+        }
+    elif pzpdf_type == "shared_bins":
         _qp_type = qp.interp
+        _data = {
+            "xvals": np.array(pzdata["pzbins"]),
+            "yvals": np.array(pzdata["pzpdf"]),
+        }
+    elif pzpdf_type == "quantiles":
+        _qp_type = qp.quant
+        _data = {
+            "quants": np.array(pzdata["quantiles"]),
+            "locs": pzdata["locations"],
+        }
+    else:
+        raise ValueError(f"pzpdf_type (={pzpdf_type}) value not valid.")
+
+    qp_ensemble = qp.Ensemble(_qp_type, data=_data)
+
+    if kernel is None:
 
-    qp_ensamble = qp.Ensemble(
-        _qp_type,
-        data={
-            "xvals": np.array(pzbins),
-            "yvals": np.array(pzpdf),
-        },
-    )
-    pz_matrix = qp_ensamble.pdf(z_grid)
+        def kernel(z):
+            # pylint: disable=unused-argument
+            return 1.0
 
+    z_grid = np.linspace(zmin, zmax, ngrid)
+    pz_matrix = qp_ensemble.pdf(z_grid)
     kernel_matrix = kernel(z_grid)
 
     return simpson(pz_matrix * kernel_matrix, x=z_grid, axis=1)

clmm/support/mock_data.py

@@ -582,7 +582,10 @@ def _compute_photoz_pdfs(galaxy_catalog, photoz_sigma_unscaled, pz_bins=101):
             gaussian(row["pzbins"], row["z"], row["pzsigma"]) for row in galaxy_catalog
         ]
     elif galaxy_catalog.pzpdf_info["type"] == "quantiles":
-        raise NotImplementedError("PDF storing in quantiles not implemented.")
+        galaxy_catalog.pzpdf_info["quantiles"] = (0.005, 0.025, 0.16, 0.5, 0.84, 0.975, 0.995)
+        galaxy_catalog["pzpdf"] = (
+            galaxy_catalog["z"][:, None] + (np.arange(7) - 3) * galaxy_catalog["pzsigma"][:, None]
+        )
     else:
         raise ValueError(
             "Value of pzpdf_info['type'] " f"(={galaxy_catalog.pzpdf_info['type']}) " "not valid."

tests/test_mockdata.py

@@ -64,6 +64,7 @@ def test_mock_data():
         photoz_sigma_unscaled=0.1,
         pzpdf_type="xxx",
     )
+    """
     assert_raises(
         NotImplementedError,
         mock.generate_galaxy_catalog,
@@ -76,6 +77,7 @@ def test_mock_data():
         photoz_sigma_unscaled=0.1,
         pzpdf_type="quantiles",
     )
+    """
 
     # Test pdz with bad arguments
     assert_raises(