Repository: skypyproject/skypy
Pull request #243 (2020-08-03 12:03:14)

Scaled Schechter samples
## Description
Optionally rescale the outputs of the Schechter sampler by a multiplicative scale factor. If the scale factor is an array a sample is generated for each element of that array.

## Checklist
- [x] Follow the [Contributor Guidelines](https://github.com/skypyproject/skypy/blob/master/CONTRIBUTING.md)
- [x] Write unit tests
- [x] Write documentation strings
- [x] Assign someone from your working team to review this pull request
- [x] Assign someone from the infrastructure team to review this pull request



Patch:
diff --git a/skypy/utils/random.py b/skypy/utils/random.py
index 387ffb1..a855116 100644
--- a/skypy/utils/random.py
+++ b/skypy/utils/random.py
@@ -18,7 +18,7 @@ import numpy as np
 import skypy.utils.special as special
 
 
-def schechter(alpha, x_min, x_max, resolution=100, size=None):
+def schechter(alpha, x_min, x_max, resolution=100, size=None, scale=1.):
     """Sample from the Schechter function.
 
     Parameters
@@ -30,7 +30,11 @@ def schechter(alpha, x_min, x_max, resolution=100, size=None):
     resolution : int
         Resolution of the inverse transform sampling spline. Default is 100.
     size: int, optional
-        Output shape of samples. Default is None.
+        Output shape of samples. If size is None and scale is a scalar, a
+        single sample is returned. If size is None and scale is an array, an
+        array of samples is returned with the same shape as scale.
+    scale: array-like, optional
+        Scale factor for the returned samples. Default is 1.
 
     Returns
     -------
@@ -50,6 +54,9 @@ def schechter(alpha, x_min, x_max, resolution=100, size=None):
     .. [1] https://en.wikipedia.org/wiki/Luminosity_function_(astronomy)
     """
 
+    if size is None:
+        size = np.broadcast(x_min, x_max, scale).shape or None
+
     x = np.logspace(np.log10(np.min(x_min)), np.log10(np.max(x_max)),
                     resolution)
     cdf = _schechter_cdf(x, np.min(x_min), np.max(x_max), alpha)
@@ -58,7 +65,7 @@ def schechter(alpha, x_min, x_max, resolution=100, size=None):
     u = np.random.uniform(t_lower, t_upper, size=size)
     x_sample = np.interp(u, cdf, x)
 
-    return x_sample
+    return x_sample * scale
 
 
 def _schechter_cdf(x, x_min, x_max, alpha):


Test patch:
diff --git a/skypy/utils/tests/test_random.py b/skypy/utils/tests/test_random.py
index 3109eb6..6ae276b 100644
--- a/skypy/utils/tests/test_random.py
+++ b/skypy/utils/tests/test_random.py
@@ -34,3 +34,19 @@ def test_schechter():
     # Test the distribution of galaxy properties follows the right distribution
     p_value = scipy.stats.kstest(sample, calc_cdf)[1]
     assert p_value > 0.01
+
+    # Test output shape when scale is a scalar
+    scale = 5
+    samples = random.schechter(alpha, x_min, x_max, scale=scale)
+    assert np.shape(samples) == np.shape(scale)
+
+    # Test output shape when scale is an array
+    scale = np.random.uniform(size=10)
+    samples = random.schechter(alpha, x_min, x_max, scale=scale)
+    assert np.shape(samples) == np.shape(scale)
+
+    # When alpha = 0 check the samples match an exponential distribution
+    alpha, scale, size = 0, 5, 1000
+    samples = random.schechter(alpha, x_min, x_max, size=size, scale=scale)
+    p_value = scipy.stats.kstest(samples, 'expon', args=(0, scale))[1]
+    assert p_value > 0.01