tidy up likelihood

peyote/likelihood.py

@@ -49,30 +49,42 @@ class MarginalizedLikelihood(Likelihood):
         self.phase_marginalization = phase_marginalization
         self.time_marginalization = time_marginalization
         self.prior = prior
+
         if self.distance_marginalization:
-            if 'luminosity_distance' not in self.prior.keys():
-                logging.info('No prior provided for distance, using default prior.')
-                self.prior['luminosity_distance'] = peyote.prior.create_default_prior('luminosity_distance')
-            self.distance_array = np.linspace(self.prior['luminosity_distance'].minimum,
-                                              self.prior['luminosity_distance'].maximum, 1e4)
-            self.delta_distance = self.distance_array[1] - self.distance_array[0]
-            self.distance_prior_array = np.array([self.prior['luminosity_distance'].prob(distance)
-                                                  for distance in self.distance_array])
+            self.distance_array = np.array([])
+            self.delta_distance = 0
+            self.distance_prior_array = np.array([])
+            self.setup_distance_marginalization()
             prior['luminosity_distance'] = 1
+
         if self.phase_marginalization:
-            if 'psi' not in self.prior.keys() or not isinstance(prior['psi'], peyote.prior.Prior):
-                logging.info('No prior provided for polarization, using default prior.')
-                self.prior['psi'] = peyote.prior.create_default_prior('psi')
-            # self.phase_array = np.exp(1j * np.linspace(0, 2 * np.pi, 100))
-            self.bessel_function_interped = interp1d(np.linspace(0, 1e6, 1e5),
-                                                     np.log([i0e(snr) for snr in np.linspace(0, 1e6, 1e5)])
-                                                     + np.linspace(0, 1e6, 1e5),
-                                                     bounds_error=False, fill_value=-np.inf)
-            self.phase_array = np.exp(1j * np.linspace(self.prior['psi'].minimum, self.prior['psi'].maximum, 500))
-            self.delta_phase = self.phase_array[1] - self.phase_array[0]
-            self.phase_prior_array = np.array([self.prior['psi'].prob(phase) for phase in self.phase_array])
+            self.bessel_function_interped = None
+            self.setup_phase_marginalization()
             prior['psi'] = 0
 
+        if self.time_marginalization:
+            logging.warning('Time marginalisation not yet implemented.')
+            self.time_marginalization = False
+
+    def setup_distance_marginalization(self):
+        if 'luminosity_distance' not in self.prior.keys():
+            logging.info('No prior provided for distance, using default prior.')
+            self.prior['luminosity_distance'] = peyote.prior.create_default_prior('luminosity_distance')
+        self.distance_array = np.linspace(self.prior['luminosity_distance'].minimum,
+                                          self.prior['luminosity_distance'].maximum, 1e4)
+        self.delta_distance = self.distance_array[1] - self.distance_array[0]
+        self.distance_prior_array = np.array([self.prior['luminosity_distance'].prob(distance)
+                                              for distance in self.distance_array])
+
+    def setup_phase_marginalization(self):
+        if 'psi' not in self.prior.keys() or not isinstance(prior['psi'], peyote.prior.Prior):
+            logging.info('No prior provided for polarization, using default prior.')
+            self.prior['psi'] = peyote.prior.create_default_prior('psi')
+        self.bessel_function_interped = interp1d(np.linspace(0, 1e6, 1e5),
+                                                 np.log([i0e(snr) for snr in np.linspace(0, 1e6, 1e5)])
+                                                 + np.linspace(0, 1e6, 1e5),
+                                                 bounds_error=False, fill_value=-np.inf)
+
     def noise_log_likelihood(self):
         log_l = 0
         for interferometer in self.interferometers:
@@ -111,9 +123,6 @@ class MarginalizedLikelihood(Likelihood):
 
         if self.phase_marginalization and not self.distance_marginalization and not self.time_marginalization:
             matched_filter_snr_squared = self.bessel_function_interped(abs(matched_filter_snr_squared))
-            # matched_filter_snr_squared = logsumexp([np.real(matched_filter_snr_squared * phase)
-            #                                         for phase in self.phase_array],
-            #                                        b=self.phase_prior_array * self.delta_phase)
 
             log_l = matched_filter_snr_squared - optimal_snr_squared / 2