WIP: Reconstruct marginalised parameters - updated

Supersedes !317 (merged). Applies a temporary fix for the time marginalized time samples, fixing #323 (closed). Use this branch to fix #324 (closed)

@colm.talbot @shanika.galaudage

bilby/gw/conversion.py

@@ -658,11 +658,10 @@ def _generate_all_cbc_parameters(sample, defaults, base_conversion,
     output_sample, _ = base_conversion(output_sample)
     output_sample = generate_mass_parameters(output_sample)
     output_sample = generate_spin_parameters(output_sample)
+
     if likelihood is not None:
-        if likelihood.distance_marginalization:
-            output_sample = \
-                generate_distance_samples_from_marginalized_likelihood(
-                    output_sample, likelihood)
+        generate_posterior_samples_from_marginalized_likelihood(
+            samples=output_sample, likelihood=likelihood)
     output_sample = generate_source_frame_parameters(output_sample)
     compute_snrs(output_sample, likelihood)
     return output_sample
@@ -691,6 +690,7 @@ def generate_all_bbh_parameters(sample, likelihood=None, priors=None):
         sample, defaults=waveform_defaults,
         base_conversion=convert_to_lal_binary_black_hole_parameters,
         likelihood=likelihood, priors=priors)
+
     return output_sample
 
 
@@ -981,42 +981,63 @@ def compute_snrs(sample, likelihood):
         logger.debug('Not computing SNRs.')
 
 
-def generate_distance_samples_from_marginalized_likelihood(samples, likelihood):
+def generate_posterior_samples_from_marginalized_likelihood(
+        samples, likelihood):
     """
     Reconstruct the distance posterior from a run which used a likelihood which
-    explicitly marginalised over distance.
+    explicitly marginalised over time/distance/phase.
 
     See Eq. (C29-C32) of https://arxiv.org/abs/1809.02293
 
     Parameters
     ----------
     samples: DataFrame
-        Posterior from run with distance marginalisation turned on.
+        Posterior from run with a marginalised likelihood.
     likelihood: bilby.gw.likelihood.GravitationalWaveTransient
         Likelihood used during sampling.
 
     Return
     ------
     sample: DataFrame
-        Returns the posterior with distance samples.
+        Returns the posterior with new samples.
     """
-    if not likelihood.distance_marginalization:
+    if likelihood.time_marginalization:
+        samples.pop('geocent_time')
+        logger.warning('Cannot recover time posterior, distance and phase also will not be recovered')
+
         return samples
-    if likelihood.phase_marginalization or likelihood.time_marginalization:
-        logger.warning('Cannot currently reconstruct distance posterior '
-                       'when other marginalizations are turned on.')
+
+    else:
+        if not any([likelihood.phase_marginalization,
+                    likelihood.distance_marginalization]):
+            return samples
+        else:
+            logger.info('Reconstructing marginalised parameters.')
+        if isinstance(samples, dict):
+            pass
+        elif isinstance(samples, DataFrame):
+            for ii in range(len(samples)):
+                sample = dict(samples.iloc[ii]).copy()
+                signal_polarizations = \
+                    likelihood.waveform_generator.frequency_domain_strain(
+                        sample)
+                if likelihood.distance_marginalization:
+                    sample = generate_distance_sample_from_marginalized_likelihood(
+                        sample=sample, likelihood=likelihood,
+                        signal_polarizations=signal_polarizations)
+                if likelihood.phase_marginalization:
+                    sample = generate_phase_sample_from_marginalized_likelihood(
+                        sample=sample, likelihood=likelihood,
+                        signal_polarizations=signal_polarizations)
+                new_distance_samples.append(sample['luminosity_distance'])
+                new_phase_samples.append(sample['phase'])
+            samples['luminosity_distance'] = new_distance_samples
+            samples['phase'] = new_phase_samples
         return samples
-    if isinstance(samples, dict):
-        pass
-    elif isinstance(samples, DataFrame):
-        for ii in range(len(samples)):
-            temp = _generate_distance_sample_from_marginalized_likelihood(
-                dict(samples.iloc[ii]), likelihood)
-            samples['luminosity_distance'][ii] = temp['luminosity_distance']
-    return samples
 
 
-def _generate_distance_sample_from_marginalized_likelihood(sample, likelihood):
+def generate_distance_sample_from_marginalized_likelihood(
+        sample, likelihood, signal_polarizations=None):
     """
     Generate a single sample from the posterior distribution for luminosity
     distance when using a likelihood which explicitly marginalises over
@@ -1030,14 +1051,19 @@ def _generate_distance_sample_from_marginalized_likelihood(sample, likelihood):
         The set of parameters used with the marginalised likelihood.
     likelihood: bilby.gw.likelihood.GravitationalWaveTransient
         The likelihood used.
+    signal_polarizations: dict, optional
+        Polarizations modes of the template.
+        Note: These are rescaled in place after the distance sample is
+              generated to allow further parameter reconstruction to occur.
 
     Returns
     -------
     sample: dict
-        Modifed dictionary with the distance sampled from the posterior.
+        Modified dictionary with the distance sampled from the posterior.
     """
-    signal_polarizations = \
-        likelihood.waveform_generator.frequency_domain_strain(sample)
+    if signal_polarizations is None:
+        signal_polarizations = \
+            likelihood.waveform_generator.frequency_domain_strain(sample)
     d_inner_h = 0
     rho_opt_sq = 0
     for ifo in likelihood.interferometers:
@@ -1051,11 +1077,64 @@ def _generate_distance_sample_from_marginalized_likelihood(sample, likelihood):
     rho_opt_sq_dist = (rho_opt_sq * sample['luminosity_distance']**2 /
                        likelihood._distance_array**2)
 
-    distance_log_like = (d_inner_h_dist.real - rho_opt_sq_dist.real / 2)
+    if likelihood.phase_marginalization:
+        distance_log_like = (
+            likelihood._bessel_function_interped(abs(d_inner_h_dist)) -
+            rho_opt_sq_dist.real / 2)
+    else:
+        distance_log_like = (d_inner_h_dist.real - rho_opt_sq_dist.real / 2)
 
-    distance_post = np.exp(distance_log_like - max(distance_log_like)) *\
-        likelihood.distance_prior_array
+    distance_post = (np.exp(distance_log_like - max(distance_log_like)) *
+                     likelihood.distance_prior_array)
 
     sample['luminosity_distance'] = Interped(
         likelihood._distance_array, distance_post).sample()
+
+    for mode in signal_polarizations:
+        signal_polarizations[mode] *= (
+            likelihood._ref_dist / sample['luminosity_distance'])
+    return sample
+
+
+def generate_phase_sample_from_marginalized_likelihood(
+        sample, likelihood, signal_polarizations=None):
+    """
+    Generate a single sample from the posterior distribution for phase when
+    using a likelihood which explicitly marginalises over phase.
+
+    See Eq. (C29-C32) of https://arxiv.org/abs/1809.02293
+
+    Parameters
+    ----------
+    sample: dict
+        The set of parameters used with the marginalised likelihood.
+    likelihood: bilby.gw.likelihood.GravitationalWaveTransient
+        The likelihood used.
+    signal_polarizations: dict, optional
+        Polarizations modes of the template.
+
+    Returns
+    -------
+    sample: dict
+        Modified dictionary with the phase sampled from the posterior.
+
+    Notes
+    -----
+    This is only valid when assumes that mu(phi) \propto exp(-2i phi).
+    """
+    if signal_polarizations is None:
+        signal_polarizations = \
+            likelihood.waveform_generator.frequency_domain_strain(sample)
+    d_inner_h = 0
+    rho_opt_sq = 0
+    for ifo in likelihood.interferometers:
+        signal = ifo.get_detector_response(signal_polarizations, sample)
+        d_inner_h += ifo.inner_product(signal=signal)
+        rho_opt_sq += ifo.optimal_snr_squared(signal=signal)
+
+    phases = np.linspace(0, 2 * np.pi, 101)
+    phasor = np.exp(-2j * phases)
+    phase_log_post = d_inner_h * phasor - rho_opt_sq / 2
+    phase_post = np.exp(phase_log_post.real - max(phase_log_post.real))
+    sample['phase'] = Interped(phases, phase_post).sample()
     return sample