From 2f5d86150560ff7605683cd2f327db4f205329f5 Mon Sep 17 00:00:00 2001
From: Nicolas Tessore <n.tessore@ucl.ac.uk>
Date: Thu, 20 Jul 2023 13:39:20 +0100
Subject: [PATCH] ENH: specify bounds for Gaussian photometric redshifts (#115)

Adds `lower=` and `upper=` keyword parameters to the `gaussian_phz()`
function to sample Gaussian photometric redshifts within the given
limits.

This currently uses a trivial rejection sampler from the full normal
distribution, so it can easily become very inefficient if the bounds
don't contain significant probability mass. Needs to be improved for
serious use.

Closes: #114
Changed: The `gaussian_phz()` function now accepts bounds using `lower=`
  and `upper=` keyword parameters.
---
 glass/galaxies.py           | 31 ++++++++++++++++++++++++++-----
 glass/test/test_galaxies.py | 11 +++++++++++
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/glass/galaxies.py b/glass/galaxies.py
index 70fa58c7..124e3bf0 100644
--- a/glass/galaxies.py
+++ b/glass/galaxies.py
@@ -145,7 +145,9 @@ def galaxy_shear(lon: np.ndarray, lat: np.ndarray, eps: np.ndarray,
     return g
 
 
-def gaussian_phz(z: ArrayLike, sigma_0: float | ArrayLike,
+def gaussian_phz(z: ArrayLike, sigma_0: float | ArrayLike, *,
+                 lower: ArrayLike | None = None,
+                 upper: ArrayLike | None = None,
                  rng: np.random.Generator | None = None) -> np.ndarray:
     r'''Photometric redshifts assuming a Gaussian error.
 
@@ -159,6 +161,8 @@ def gaussian_phz(z: ArrayLike, sigma_0: float | ArrayLike,
         True redshifts.
     sigma_0 : float or array_like
         Redshift error in the tomographic binning at zero redshift.
+    lower, upper : float or array_like, optional
+        Bounds for the returned photometric redshifts.
     rng : :class:`~numpy.random.Generator`, optional
         Random number generator.  If not given, a default RNG is used.
 
@@ -168,6 +172,12 @@ def gaussian_phz(z: ArrayLike, sigma_0: float | ArrayLike,
         Photometric redshifts assuming Gaussian errors, of the same
         shape as *z*.
 
+    Warnings
+    --------
+    The *lower* and *upper* bounds are implemented using plain rejection
+    sampling from the non-truncated normal distribution.  If bounds are
+    used, they should always contain significant probability mass.
+
     See Also
     --------
     glass.observations.tomo_nz_gausserr :
@@ -194,14 +204,25 @@ def gaussian_phz(z: ArrayLike, sigma_0: float | ArrayLike,
 
     zphot = rng.normal(z, sigma)
 
+    print(zphot)
+
+    if lower is None:
+        lower = 0.
+    if upper is None:
+        upper = np.inf
+
+    if not np.all(lower < upper):
+        raise ValueError("requires lower < upper")
+
     if not dims:
-        while zphot < 0:
+        while zphot < lower or zphot > upper:
             zphot = rng.normal(z, sigma)
     else:
         z = np.broadcast_to(z, dims)
-        trunc = np.where(zphot < 0)[0]
+        trunc = np.where((zphot < lower) | (zphot > upper))[0]
         while trunc.size:
-            zphot[trunc] = rng.normal(z[trunc], sigma[trunc])
-            trunc = trunc[zphot[trunc] < 0]
+            znew = rng.normal(z[trunc], sigma[trunc])
+            zphot[trunc] = znew
+            trunc = trunc[(znew < lower) | (znew > upper)]
 
     return zphot
diff --git a/glass/test/test_galaxies.py b/glass/test/test_galaxies.py
index f4f31b52..40999517 100644
--- a/glass/test/test_galaxies.py
+++ b/glass/test/test_galaxies.py
@@ -97,6 +97,17 @@ def test_gaussian_phz():
     assert phz.shape == (100,)
     assert np.all(phz >= 0)
 
+    # case: upper and lower bound
+
+    z = 1.
+    sigma_0 = np.ones(100)
+
+    phz = gaussian_phz(z, sigma_0, lower=0.5, upper=1.5)
+
+    assert phz.shape == (100,)
+    assert np.all(phz >= 0.5)
+    assert np.all(phz <= 1.5)
+
     # test interface
 
     # case: scalar redshift, scalar sigma_0