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

@@ -261,6 +261,134 @@ class PriorSet(PriorDict):
         super(PriorSet, self).__init__(dictionary, filename)
 
 
+class CorrelatedPriorDict(PriorDict):
+
+    def __init__(self, dictionary=None, filename=None):
+        super(CorrelatedPriorDict, self).__init__(dictionary=dictionary, filename=filename)
+        self.correlated_keys = []
+        self.uncorrelated_keys = []
+        self._resolve_correlations()
+
+    def _resolve_correlations(self):
+        self.convert_floats_to_delta_functions()
+        for key in self.keys():
+            if self[key].has_correlated_variables():
+                self.correlated_keys.append(key)
+            else:
+                self.uncorrelated_keys.append(key)
+
+    def sample_subset(self, keys=iter([]), size=None):
+        self.convert_floats_to_delta_functions()
+        correlated_keys = []
+        uncorrelated_keys = []
+        sampled_keys = []
+        for key in keys:
+            if key in self.uncorrelated_keys:
+                uncorrelated_keys.append(key)
+            elif key in self.correlated_keys:
+                correlated_keys.append(key)
+            else:
+                raise KeyError('Invalid key in keys argument')
+
+        samples = dict()
+        for key in uncorrelated_keys:
+            samples[key] = self[key].sample(size=size)
+            sampled_keys.append(key)
+
+        for i in range(0, 1000):
+            for key in correlated_keys:
+                if self._check_correlations_resolved(key, sampled_keys):
+                    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:
+                raise Exception('This set contains unresolvable correlations')
+
+        return samples
+
+    def _get_correlated_variables(self, key):
+        correlated_variables = dict()
+        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:
+            if k not in sampled_keys:
+                correlations_resolved = False
+        return correlations_resolved
+
+    def prob(self, sample, **kwargs):
+        """
+
+        Parameters
+        ----------
+        sample: dict
+            Dictionary of the samples of which we want to have the probability of
+        kwargs:
+            The keyword arguments are passed directly to `np.product`
+
+        Returns
+        -------
+        float: Joint probability of all individual sample probabilities
+
+        """
+        ls = []
+        for key in sample:
+            method_kwargs = infer_args_from_method(self[key].prob)
+            method_kwargs.remove('val')
+            correlated_variables = {key: sample[key] for key in method_kwargs}
+            ls.append(self[key].prob(sample[key], **correlated_variables))
+        return np.product(ls, **kwargs)
+
+    def ln_prob(self, sample):
+        """
+
+        Parameters
+        ----------
+        sample: dict
+            Dictionary of the samples of which we want to have the log probability of
+
+        Returns
+        -------
+        float: Joint log probability of all the individual sample probabilities
+
+        """
+        ls = []
+        for key in sample:
+            method_kwargs = infer_args_from_method(self[key].prob)
+            method_kwargs.remove('val')
+            correlated_variables = {key: sample[key] for key in method_kwargs}
+            ls.append(self[key].ln_prob(sample[key], **correlated_variables))
+        return np.sum(ls)
+
+    def rescale(self, keys, theta):
+        """Rescale samples from unit cube to prior
+
+        Parameters
+        ----------
+        keys: list
+            List of prior keys to be rescaled
+        theta: list
+            List of randomly drawn values on a unit cube associated with the prior keys
+
+        Returns
+        -------
+        list: List of floats containing the rescaled sample
+        """
+        ls = []
+        for key in theta:
+            method_kwargs = infer_args_from_method(self[key].prob)
+            method_kwargs.remove('val')
+            correlated_variables = {key: theta for key in method_kwargs}
+            ls.append(self[key].rescale(sample, correlated_variables) for key, sample in zip(keys, theta))
+        return ls
+
+
 def create_default_prior(name, default_priors_file=None):
     """Make a default prior for a parameter with a known name.
 
@@ -319,6 +447,8 @@ class Prior(object):
         self.unit = unit
         self.minimum = minimum
         self.maximum = maximum
+        self.least_recently_sampled = None
+        self._correlated_variables = []
 
     def __call__(self):
         """Overrides the __call__ special method. Calls the sample method.
@@ -356,7 +486,8 @@ class Prior(object):
         float: A random number between 0 and 1, rescaled to match the distribution of this Prior
 
         """
-        return self.rescale(np.random.uniform(0, 1, size))
+        self.least_recently_sampled = self.rescale(np.random.uniform(0, 1, size))
+        return self.least_recently_sampled
 
     def rescale(self, val):
         """
@@ -530,6 +661,19 @@ class Prior(object):
             label = self.name
         return label
 
+    def has_correlated_variables(self):
+        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):
 
@@ -1709,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)