Resolve #430 (Add normalisation flag to constrained prior)

This MR resolves issue #430 (closed) to properly normalize ConstrainedPriorDict.prob() when sampling the prior. This MR allows for users to input a normalisation flag when creating an instance of a Constraint prior and when calling ConstrainedPriorDict.prob() if there are normalisation constants given will return the normalisation*prob and otherwise use the ratio of total samples to kept samples as this normalisation constant. Pass normalisation=1 when creating a Constraint Prior for behaviour before this MR.

bilby/core/prior/dict.py

@@ -3,7 +3,7 @@ from io import open as ioopen
 import json
 import numpy as np
 import os
-from functools import lru_cache
+from methodtools import lru_cache
 
 from future.utils import iteritems
 from matplotlib.cbook import flatten
@@ -381,15 +381,20 @@ class PriorDict(dict):
 
     @lru_cache()
     def normalize_constraint_factor(self, keys):
-        min_accept = 50
-        sampling_chunk = 250
+        min_accept = 1000
+        sampling_chunk = 5000
         samples = self.sample_subset(keys=keys, size=sampling_chunk)
-        keep = np.array(self.evaluate_constraints(samples))
+        keep = np.atleast_1d(self.evaluate_constraints(samples))
+        if len(keep) == 1:
+            return 1
+        all_samples = {key: np.array([]) for key in keys}
+        _first_key = list(all_samples.keys())[0]
         while np.count_nonzero(keep) < min_accept:
-            new_samples = self.sample_subset(keys=keys, size=sampling_chunk)
+            samples = self.sample_subset(keys=keys, size=sampling_chunk)
             for key in samples:
-                samples[key] = np.concatenate(samples[key], new_samples[key])
-            keep = np.array(self.evaluate_constraints(samples))
+                all_samples[key] = np.hstack(
+                    [all_samples[key], samples[key].flatten()])
+            keep = np.array(self.evaluate_constraints(all_samples), dtype=bool)
         return len(keep) / np.count_nonzero(keep)
 
     def prob(self, sample, **kwargs):
@@ -410,14 +415,8 @@ class PriorDict(dict):
         prob = np.product([self[key].prob(sample[key])
                            for key in sample], **kwargs)
 
-        ratio = 1
-        outsample = self.conversion_function(sample)
-        # Check if there is a constraint in sample/outsample
-        if (np.any(isinstance([self[key] for key in sample.keys()], Constraint)) or
-                np.any(isinstance([self[key] for key in outsample.keys()], Constraint))):
-            # If constraint exists in keys, caclulate the cached normalization constant
-            ratio = self.normalize_constraint_factor(sample.keys())
-
+        ratio = self.normalize_constraint_factor(tuple(sample.keys()))
+        #print("ratio = {}".format(ratio))
         if np.all(prob == 0.):
             return prob
         else:
@@ -451,14 +450,7 @@ class PriorDict(dict):
         ln_prob = np.sum([self[key].ln_prob(sample[key])
                           for key in sample], axis=axis)
 
-        ratio = 1
-        outsample = self.conversion_function(sample)
-        # Check if there is a constraint in sample/outsample
-        if (np.any(isinstance([self[key] for key in sample.keys()], Constraint)) or
-                np.any(isinstance([self[key] for key in outsample.keys()], Constraint))):
-            # If constraint exists in keys, caclulate the cached normalization constant
-            ratio = self.normalize_constraint_factor(sample.keys())
-
+        ratio = self.normalize_constraint_factor(tuple(sample.keys()))
         if np.all(np.isinf(ln_prob)):
             return ln_prob
         else:
@@ -520,6 +512,15 @@ class PriorDict(dict):
         """
         return self.__class__(dictionary=dict(self))
 
+    def __key(self):
+        return tuple((k, self[k]) for k in sorted(self))
+
+    def __hash__(self):
+        return hash(self.__key())
+
+    def __eq__(self, other):
+        return self.__key() == other.__key()
+
 
 class PriorSet(PriorDict):