DEV: make sure all priors return float when needed (#979)

* DEV: make sure all priors return float when needed

* FEAT: add fermi-dirac CDF

* TST: add testing of fermi dirac and symmetric log uniform priors

* FMT: remove extraneous whitespace

* BUG: revert bad changes to equal comparisons

* BUG: Fix syntax error in return statement

* BUG: Fix list return

* BUG: fix array type output for rescale

* BUG: Fix missing parentheses in squeeze method call

* BUG: fix flattening logic

* BUG: stop using arr[np.where(cond)]

* Address comments

* BUG: fix scipy integrate imports

* BUG: Remove unused import of interp1d

Removed unused import of interp1d from scipy.interpolate.

Author: ColmTalbot

bilby/core/prior/analytical.py

@@ -169,8 +169,10 @@ def cdf(self, val):
             _cdf = (np.log(val / self.minimum) /
                     np.log(self.maximum / self.minimum))
         else:
-            _cdf = np.atleast_1d(val ** (self.alpha + 1) - self.minimum ** (self.alpha + 1)) / \
-                (self.maximum ** (self.alpha + 1) - self.minimum ** (self.alpha + 1))
+            _cdf = (
+                (val ** (self.alpha + 1) - self.minimum ** (self.alpha + 1))
+                / (self.maximum ** (self.alpha + 1) - self.minimum ** (self.alpha + 1))
+            )
         _cdf = np.minimum(_cdf, 1)
         _cdf = np.maximum(_cdf, 0)
         return _cdf
@@ -367,16 +369,16 @@ def ln_prob(self, val):
         return np.nan_to_num(- np.log(2 * np.abs(val)) - np.log(np.log(self.maximum / self.minimum)))
 
     def cdf(self, val):
-        val = np.atleast_1d(val)
         norm = 0.5 / np.log(self.maximum / self.minimum)
-        cdf = np.zeros((len(val)))
-        lower_indices = np.where(np.logical_and(-self.maximum <= val, val <= -self.minimum))[0]
-        upper_indices = np.where(np.logical_and(self.minimum <= val, val <= self.maximum))[0]
-        cdf[lower_indices] = -norm * np.log(-val[lower_indices] / self.maximum)
-        cdf[np.where(np.logical_and(-self.minimum < val, val < self.minimum))] = 0.5
-        cdf[upper_indices] = 0.5 + norm * np.log(val[upper_indices] / self.minimum)
-        cdf[np.where(self.maximum < val)] = 1
-        return cdf
+        _cdf = (
+            -norm * np.log(abs(val) / self.maximum)
+            * (val <= -self.minimum) * (val >= -self.maximum)
+            + (0.5 + norm * np.log(abs(val) / self.minimum))
+            * (val >= self.minimum) * (val <= self.maximum)
+            + 0.5 * (val > -self.minimum) * (val < self.minimum)
+            + 1 * (val > self.maximum)
+        )
+        return _cdf
 
 
 class Cosine(Prior):
@@ -426,10 +428,12 @@ def prob(self, val):
         return np.cos(val) / 2 * self.is_in_prior_range(val)
 
     def cdf(self, val):
-        _cdf = np.atleast_1d((np.sin(val) - np.sin(self.minimum)) /
-                             (np.sin(self.maximum) - np.sin(self.minimum)))
-        _cdf[val > self.maximum] = 1
-        _cdf[val < self.minimum] = 0
+        _cdf = (
+            (np.sin(val) - np.sin(self.minimum))
+            / (np.sin(self.maximum) - np.sin(self.minimum))
+            * (val >= self.minimum) * (val <= self.maximum)
+            + 1 * (val > self.maximum)
+        )
         return _cdf
 
 
@@ -480,10 +484,12 @@ def prob(self, val):
         return np.sin(val) / 2 * self.is_in_prior_range(val)
 
     def cdf(self, val):
-        _cdf = np.atleast_1d((np.cos(val) - np.cos(self.minimum)) /
-                             (np.cos(self.maximum) - np.cos(self.minimum)))
-        _cdf[val > self.maximum] = 1
-        _cdf[val < self.minimum] = 0
+        _cdf = (
+            (np.cos(val) - np.cos(self.minimum))
+            / (np.cos(self.maximum) - np.cos(self.minimum))
+            * (val >= self.minimum) * (val <= self.maximum)
+            + 1 * (val > self.maximum)
+        )
         return _cdf
 
 
@@ -625,11 +631,13 @@ def prob(self, val):
             / self.sigma / self.normalisation * self.is_in_prior_range(val)
 
     def cdf(self, val):
-        val = np.atleast_1d(val)
-        _cdf = (erf((val - self.mu) / 2 ** 0.5 / self.sigma) - erf(
-            (self.minimum - self.mu) / 2 ** 0.5 / self.sigma)) / 2 / self.normalisation
-        _cdf[val > self.maximum] = 1
-        _cdf[val < self.minimum] = 0
+        _cdf = (
+            (
+                erf((val - self.mu) / 2 ** 0.5 / self.sigma)
+                - erf((self.minimum - self.mu) / 2 ** 0.5 / self.sigma)
+            ) / 2 / self.normalisation * (val >= self.minimum) * (val <= self.maximum)
+            + 1 * (val > self.maximum)
+        )
         return _cdf
 
 
@@ -1367,6 +1375,8 @@ def __init__(self, sigma, mu=None, r=None, name=None, latex_label=None,
             raise ValueError("For the Fermi-Dirac prior the values of sigma and r "
                              "must be positive.")
 
+        self.expr = np.exp(self.r)
+
     def rescale(self, val):
         """
         'Rescale' a sample from the unit line element to the appropriate Fermi-Dirac prior.
@@ -1384,21 +1394,8 @@ def rescale(self, val):
         .. [1] M. Pitkin, M. Isi, J. Veitch & G. Woan, `arXiv:1705.08978v1
            <https:arxiv.org/abs/1705.08978v1>`_, 2017.
         """
-        inv = (-np.exp(-1. * self.r) + (1. + np.exp(self.r)) ** -val +
-               np.exp(-1. * self.r) * (1. + np.exp(self.r)) ** -val)
-
-        # if val is 1 this will cause inv to be negative (due to numerical
-        # issues), so return np.inf
-        if isinstance(val, (float, int)):
-            if inv < 0:
-                return np.inf
-            else:
-                return -self.sigma * np.log(inv)
-        else:
-            idx = inv >= 0.
-            tmpinv = np.inf * np.ones(len(np.atleast_1d(val)))
-            tmpinv[idx] = -self.sigma * np.log(inv[idx])
-            return tmpinv
+        inv = -1 / self.expr + (1 + self.expr)**-val + (1 + self.expr)**-val / self.expr
+        return -self.sigma * np.log(np.maximum(inv, 0))
 
     def prob(self, val):
         """Return the prior probability of val.
@@ -1411,7 +1408,11 @@ def prob(self, val):
         =======
         float: Prior probability of val
         """
-        return np.exp(self.ln_prob(val))
+        return (
+            (np.exp((val - self.mu) / self.sigma) + 1)**-1
+            / (self.sigma * np.log1p(self.expr))
+            * (val >= self.minimum)
+        )
 
     def ln_prob(self, val):
         """Return the log prior probability of val.
@@ -1424,19 +1425,34 @@ def ln_prob(self, val):
         =======
         Union[float, array_like]: Log prior probability of val
         """
+        return np.log(self.prob(val))
 
-        norm = -np.log(self.sigma * np.log(1. + np.exp(self.r)))
-        if isinstance(val, (float, int)):
-            if val < self.minimum:
-                return -np.inf
-            else:
-                return norm - np.logaddexp((val / self.sigma) - self.r, 0.)
-        else:
-            val = np.atleast_1d(val)
-            lnp = -np.inf * np.ones(len(val))
-            idx = val >= self.minimum
-            lnp[idx] = norm - np.logaddexp((val[idx] / self.sigma) - self.r, 0.)
-            return lnp
+    def cdf(self, val):
+        """
+        Evaluate the CDF of the Fermi-Dirac distribution using a slightly
+        modified form of Equation 23 of [1]_.
+
+        Parameters
+        ==========
+        val: Union[float, int, array_like]
+            The value(s) to evaluate the CDF at
+
+        Returns
+        =======
+        Union[float, array_like]:
+            The CDF value(s)
+
+        References
+        ==========
+
+        .. [1] M. Pitkin, M. Isi, J. Veitch & G. Woan, `arXiv:1705.08978v1
+           <https:arxiv.org/abs/1705.08978v1>`_, 2017.
+        """
+        result = (
+            (np.logaddexp(0, -self.r) - np.logaddexp(-val / self.sigma, -self.r))
+            / np.logaddexp(0, self.r)
+        )
+        return np.clip(result, 0, 1)
 
 
 class DiscreteValues(Prior):

bilby/core/prior/base.py

@@ -5,14 +5,14 @@
 
 import numpy as np
 import scipy.stats
-from scipy.interpolate import interp1d
 
 from ..utils import (
     infer_args_from_method,
     BilbyJsonEncoder,
     decode_bilby_json,
     logger,
     get_dict_with_properties,
+    WrappedInterp1d as interp1d,
 )
 
 
@@ -178,7 +178,10 @@ def cdf(self, val):
         cdf = cumulative_trapezoid(pdf, x, initial=0)
         interp = interp1d(x, cdf, assume_sorted=True, bounds_error=False,
                           fill_value=(0, 1))
-        return interp(val)
+        output = interp(val)
+        if isinstance(val, (int, float)):
+            output = float(output)
+        return output
 
     def ln_prob(self, val):
         """Return the prior ln probability of val, this should be overwritten

bilby/core/prior/dict.py

@@ -504,7 +504,9 @@ def check_efficiency(n_tested, n_valid):
     def normalize_constraint_factor(
         self, keys, min_accept=10000, sampling_chunk=50000, nrepeats=10
     ):
-        if keys in self._cached_normalizations.keys():
+        if len(self.constraint_keys) == 0:
+            return 1
+        elif keys in self._cached_normalizations.keys():
             return self._cached_normalizations[keys]
         else:
             factor_estimates = [
@@ -566,8 +568,10 @@ def check_prob(self, sample, prob):
                     return 0.0
             else:
                 constrained_prob = np.zeros_like(prob)
-                keep = np.array(self.evaluate_constraints(sample), dtype=bool)
-                constrained_prob[keep] = prob[keep] * ratio
+                in_bounds = np.isfinite(prob)
+                subsample = {key: sample[key][in_bounds] for key in sample}
+                keep = np.array(self.evaluate_constraints(subsample), dtype=bool)
+                constrained_prob[in_bounds] = prob[in_bounds] * keep * ratio
                 return constrained_prob
 
     def ln_prob(self, sample, axis=None, normalized=True):
@@ -608,8 +612,10 @@ def check_ln_prob(self, sample, ln_prob, normalized=True):
                     return -np.inf
             else:
                 constrained_ln_prob = -np.inf * np.ones_like(ln_prob)
-                keep = np.array(self.evaluate_constraints(sample), dtype=bool)
-                constrained_ln_prob[keep] = ln_prob[keep] + np.log(ratio)
+                in_bounds = np.isfinite(ln_prob)
+                subsample = {key: sample[key][in_bounds] for key in sample}
+                keep = np.log(np.array(self.evaluate_constraints(subsample), dtype=bool))
+                constrained_ln_prob[in_bounds] = ln_prob[in_bounds] + keep + np.log(ratio)
                 return constrained_ln_prob
 
     def cdf(self, sample):
@@ -643,12 +649,9 @@ def rescale(self, keys, theta):
         =======
         list: List of floats containing the rescaled sample
         """
-        theta = list(theta)
-        samples = []
-        for key, units in zip(keys, theta):
-            samps = self[key].rescale(units)
-            samples += list(np.asarray(samps).flatten())
-        return samples
+        return list(
+            [self[key].rescale(sample) for key, sample in zip(keys, theta)]
+        )
 
     def test_redundancy(self, key, disable_logging=False):
         """Empty redundancy test, should be overwritten in subclasses"""
@@ -670,9 +673,7 @@ def test_has_redundant_keys(self):
             del temp[key]
             if temp.test_redundancy(key, disable_logging=True):
                 logger.warning(
-                    "{} is a redundant key in this {}.".format(
-                        key, self.__class__.__name__
-                    )
+                    f"{key} is a redundant key in this {self.__class__.__name__}."
                 )
                 redundant = True
         return redundant
@@ -880,17 +881,43 @@ def rescale(self, keys, theta):
         self._check_resolved()
         self._update_rescale_keys(keys)
         result = dict()
+        joint = dict()
         for key, index in zip(
             self.sorted_keys_without_fixed_parameters, self._rescale_indexes
         ):
             result[key] = self[key].rescale(
                 theta[index], **self.get_required_variables(key)
             )
             self[key].least_recently_sampled = result[key]
-        samples = []
-        for key in keys:
-            samples += list(np.asarray(result[key]).flatten())
-        return samples
+            if isinstance(self[key], JointPrior) and self[key].dist.distname not in joint:
+                joint[self[key].dist.distname] = [key]
+            elif isinstance(self[key], JointPrior):
+                joint[self[key].dist.distname].append(key)
+        for names in joint.values():
+            # this is needed to unpack how joint prior rescaling works
+            # as an example of a joint prior over {a, b, c, d} we might
+            # get the following based on the order within the joint prior
+            # {a: [], b: [], c: [1, 2, 3, 4], d: []}
+            # -> [1, 2, 3, 4]
+            # -> {a: 1, b: 2, c: 3, d: 4}
+            values = list()
+            for key in names:
+                values = np.concatenate([values, result[key]])
+            for key, value in zip(names, values):
+                result[key] = value
+
+        def safe_flatten(value):
+            """
+            this is gross but can be removed whenever we switch to returning
+            arrays, flatten converts 0-d arrays to 1-d so has to be special
+            cased
+            """
+            if isinstance(value, (float, int)):
+                return value
+            else:
+                return result[key].flatten()
+
+        return [safe_flatten(result[key]) for key in keys]
 
     def _update_rescale_keys(self, keys):
         if not keys == self._least_recently_rescaled_keys:

bilby/core/prior/interpolated.py

@@ -1,9 +1,8 @@
 import numpy as np
 from scipy.integrate import trapezoid
-from scipy.interpolate import interp1d
 
 from .base import Prior
-from ..utils import logger
+from ..utils import logger, WrappedInterp1d as interp1d
 
 
 class Interped(Prior):
@@ -87,10 +86,7 @@ def rescale(self, val):
 
         This maps to the inverse CDF. This is done using interpolation.
         """
-        rescaled = self.inverse_cumulative_distribution(val)
-        if rescaled.shape == ():
-            rescaled = float(rescaled)
-        return rescaled
+        return self.inverse_cumulative_distribution(val)
 
     @property
     def minimum(self):

bilby/core/prior/slabspike.py

@@ -88,11 +88,11 @@ def rescale(self, val):
         original_is_number = isinstance(val, Number)
         val = np.atleast_1d(val)
 
-        lower_indices = np.where(val < self.inverse_cdf_below_spike)[0]
-        intermediate_indices = np.where(np.logical_and(
+        lower_indices = val < self.inverse_cdf_below_spike
+        intermediate_indices = np.logical_and(
             self.inverse_cdf_below_spike <= val,
-            val <= self.inverse_cdf_below_spike + self.spike_height))[0]
-        higher_indices = np.where(val > self.inverse_cdf_below_spike + self.spike_height)[0]
+            val <= (self.inverse_cdf_below_spike + self.spike_height))
+        higher_indices = val > (self.inverse_cdf_below_spike + self.spike_height)
 
         res = np.zeros(len(val))
         res[lower_indices] = self._contracted_rescale(val[lower_indices])
@@ -137,7 +137,7 @@ def prob(self, val):
         original_is_number = isinstance(val, Number)
         res = self.slab.prob(val) * self.slab_fraction
         res = np.atleast_1d(res)
-        res[np.where(val == self.spike_location)] = np.inf
+        res[val == self.spike_location] = np.inf
         if original_is_number:
             try:
                 res = res[0]
@@ -161,7 +161,7 @@ def ln_prob(self, val):
         original_is_number = isinstance(val, Number)
         res = self.slab.ln_prob(val) + np.log(self.slab_fraction)
         res = np.atleast_1d(res)
-        res[np.where(val == self.spike_location)] = np.inf
+        res[val == self.spike_location] = np.inf
         if original_is_number:
             try:
                 res = res[0]
@@ -185,7 +185,5 @@ def cdf(self, val):
 
         """
         res = self.slab.cdf(val) * self.slab_fraction
-        res = np.atleast_1d(res)
-        indices_above_spike = np.where(val > self.spike_location)[0]
-        res[indices_above_spike] += self.spike_height
+        res += self.spike_height * (val > self.spike_location)
         return res