Resolve "Introduce conditional prior sets"

Closes #270 (closed)

PP test show that the new capabilities do not interfere with the results we are getting so far. This test is done without any conditional priors, it just shows that everything still works as expected.

https://ldas-jobs.ligo.caltech.edu/~moritz.huebner/conditional_prior_pp/outdir_pp_test_4s/pp.png

bilby/core/prior.py

@@ -301,6 +301,7 @@ class CorrelatedPriorDict(PriorDict):
                     cvars = self._get_correlated_variables(key)
                     samples[key] = self[key].sample(size=size, **cvars)
                     correlated_keys.remove(key)
+                    sampled_keys.append(key)
             if not correlated_keys:
                 break
             if i == 999:
@@ -310,13 +311,13 @@ class CorrelatedPriorDict(PriorDict):
 
     def _get_correlated_variables(self, key):
         correlated_variables = dict()
-        for cv in self[key].CORRELATED_VARIABLES:
-            correlated_variables[cv] = self[cv].least_recently_sampled
+        for k in self[key].correlated_variables:
+            correlated_variables[k] = self[k].least_recently_sampled
         return correlated_variables
 
     def _check_correlations_resolved(self, key, sampled_keys):
         correlations_resolved = True
-        for k in self[key].CORRELATED_VARIABLES:
+        for k in self[key].correlated_variables:
             if k not in sampled_keys:
                 correlations_resolved = False
         return correlations_resolved
@@ -447,7 +448,7 @@ class Prior(object):
         self.minimum = minimum
         self.maximum = maximum
         self.least_recently_sampled = None
-        self.CORRELATED_VARIABLES = []
+        self._correlated_variables = []
 
     def __call__(self):
         """Overrides the __call__ special method. Calls the sample method.
@@ -661,10 +662,18 @@ class Prior(object):
         return label
 
     def has_correlated_variables(self):
-        if len(self.CORRELATED_VARIABLES) > 0:
+        if len(self.correlated_variables) > 0:
             return True
         return False
 
+    @property
+    def correlated_variables(self):
+        return self._correlated_variables
+
+    @correlated_variables.setter
+    def correlated_variables(self, correlated_variables):
+        self._correlated_variables = correlated_variables
+
 
 class DeltaFunction(Prior):
 
@@ -1844,3 +1853,176 @@ class FromFile(Interped):
             logger.warning("Can't load {}.".format(self.id))
             logger.warning("Format should be:")
             logger.warning(r"x\tp(x)")
+
+
+class CorrelatedGaussian(Prior):
+
+    def __init__(self, mu, sigma, name=None, latex_label=None, unit=None, correlation_func=None):
+        """Gaussian prior with mean mu and width sigma
+
+        Parameters
+        ----------
+        mu: float
+            Mean of the Gaussian prior
+        sigma:
+            Width/Standard deviation of the Gaussian prior
+        name: str
+            See superclass
+        latex_label: str
+            See superclass
+        unit: str
+            See superclass
+        """
+        Prior.__init__(self, name=name, latex_label=latex_label, unit=unit)
+        self.mu = mu
+        self.sigma = sigma
+        if not correlation_func:
+            self.correlation_func = lambda x, **y: x
+        else:
+            self.correlation_func = correlation_func
+
+    @property
+    def correlated_variables(self):
+        from .utils import infer_parameters_from_function
+        return infer_parameters_from_function(self.correlation_func)
+
+    def sample(self, size=None, **cvars):
+        """Draw a sample from the prior
+
+        Parameters
+        ----------
+        size: int or tuple of ints, optional
+            See numpy.random.uniform docs
+
+        Returns
+        -------
+        float: A random number between 0 and 1, rescaled to match the distribution of this Prior
+
+        """
+        self.least_recently_sampled = self.rescale(np.random.uniform(0, 1, size), **cvars)
+        return self.least_recently_sampled
+
+    def mean(self, **correlated_variables):
+        return self.correlation_func(self.mu, **correlated_variables)
+
+    def rescale(self, val, **correlated_variables):
+        """
+        'Rescale' a sample from the unit line element to the appropriate Gaussian prior.
+
+        This maps to the inverse CDF. This has been analytically solved for this case.
+        """
+        Prior.test_valid_for_rescaling(val)
+        return self.mean(**correlated_variables) + \
+            erfinv(2 * val - 1) * 2 ** 0.5 * self.sigma
+
+    def prob(self, val, **correlated_variables):
+        """Return the prior probability of val.
+
+        Parameters
+        ----------
+        val: float
+
+        Returns
+        -------
+        float: Prior probability of val
+        """
+        return np.exp(-(self.mean(**correlated_variables) - val) ** 2 /
+                      (2 * self.sigma ** 2)) / (2 * np.pi) ** 0.5 / self.sigma
+
+    def ln_prob(self, val, **correlated_variables):
+        return -0.5 * ((self.mean(**correlated_variables) - val) ** 2 /
+                       self.sigma ** 2 + np.log(2 * np.pi * self.sigma ** 2))
+
+
+class CorrelatedUniform(Prior):
+
+    def __init__(self, minimum, maximum, name=None, latex_label=None,
+                 unit=None, correlation_func=None):
+        """Uniform prior with bounds
+
+        Parameters
+        ----------
+        minimum: float
+            See superclass
+        maximum: float
+            See superclass
+        name: str
+            See superclass
+        latex_label: str
+            See superclass
+        unit: str
+            See superclass
+        """
+        Prior.__init__(self, name=name, latex_label=latex_label,
+                       minimum=minimum, maximum=maximum, unit=unit)
+        if not correlation_func:
+            self.correlation_func = lambda x, **y: x
+        else:
+            self.correlation_func = correlation_func
+
+    def mean(self, **correlated_variables):
+        mean = (self.maximum + self.minimum)/2
+        return self.correlation_func(mean, **correlated_variables)
+
+    @property
+    def width(self):
+        return self.maximum - self.minimum
+
+    @property
+    def correlated_variables(self):
+        from .utils import infer_parameters_from_function
+        return infer_parameters_from_function(self.correlation_func)
+
+    def sample(self, size=None, **correlated_variables):
+        """Draw a sample from the prior
+
+        Parameters
+        ----------
+        size: int or tuple of ints, optional
+            See numpy.random.uniform docs
+
+        Returns
+        -------
+        float: A random number between 0 and 1, rescaled to match the distribution of this Prior
+
+        """
+        self.least_recently_sampled = self.rescale(np.random.uniform(0, 1, size), **correlated_variables)
+        return self.least_recently_sampled
+
+    def rescale(self, val, **correlated_variables):
+        Prior.test_valid_for_rescaling(val)
+        minimum = self.mean(**correlated_variables) - self.width/2
+        maximum = self.mean(**correlated_variables) + self.width/2
+        return minimum + val * (maximum - minimum)
+
+    def prob(self, val, **correlated_variables):
+        """Return the prior probability of val
+
+        Parameters
+        ----------
+        val: float
+
+        Returns
+        -------
+        float: Prior probability of val
+        """
+        minimum = self.mean(**correlated_variables) - self.width/2
+        maximum = self.mean(**correlated_variables) + self.width/2
+        return scipy.stats.uniform.pdf(val, loc=minimum,
+                                       scale=maximum - minimum)
+
+    def ln_prob(self, val, **correlated_variables):
+        """Return the log prior probability of val
+
+        Parameters
+        ----------
+        val: float
+
+        Returns
+        -------
+        float: log probability of val
+        """
+        minimum = self.mean(**correlated_variables) - self.width/2
+        maximum = self.mean(**correlated_variables) + self.width/2
+        return scipy.stats.uniform.logpdf(val, loc=minimum,
+                                          scale=maximum - minimum)