Source-frame computation added

CHANGELOG.md

 # Changelog
 
 ## Unreleased 0.8
+- Adding computation of EM-Bright probabilities using source-frame samples.
 - Neutron star compactness can now be computed from equation of
   states available in lalsimulation.
 - Refactored methods `computeDiskMass` and `computeCompactness`.

ligo/em_bright.py

@@ -8,6 +8,9 @@ import h5py
 
 import numpy as np
 import pkg_resources
+from scipy.interpolate import interp1d
+from astropy import cosmology, units as u
+
 
 from .computeDiskMass import compute_isco
 from .computeDiskMass import computeDiskMass
@@ -80,13 +83,62 @@ def source_classification(m1, m2, chi1, chi2, snr,
     return prediction_ns[0], prediction_em[0]
 
 
-def source_classification_pe(posterior_samples_file, hdf5=True, threshold=3.0):
+def get_redshifts(distances, N=10000):
+    """
+    This function computes the redshifts using the Planck15 cosmology.
+    It accepts as input the posterior samples from a PE file and then
+    computes redshift by interpolating the distance-redshift function.
+
+    Parameters
+    ----------
+    distances: float or numpy.ndarray
+              distance(s) in Mpc
+
+    N: int
+      Number of steps for the computation of the interpolation function
+
+    Example
+    -------
+    >>> distances = np.linspace(10, 100, 10)
+    >>> em_bright.get_redshifts(distances)
+    array([0.00225566, 0.00450357, 0.00674384, 0.00897655,
+           0.01120181, 0.0134197 , 0.01563032, 0.01783375
+           0.02003009, 0.02221941])
+    """
+    function = cosmology.Planck15.luminosity_distance
+    min_dist = np.min(distances)
+    max_dist = np.max(distances)
+    z_min = cosmology.z_at_value(func=function, fval=min_dist*u.Mpc)
+    z_max = cosmology.z_at_value(func=function, fval=max_dist*u.Mpc)
+    z_steps = np.linspace(z_min - (0.1*z_min), z_max + (0.1*z_min), N)
+    lum_dists = cosmology.Planck15.luminosity_distance(z_steps)
+    s = interp1d(lum_dists, z_steps)
+    redshifts = s(distances)
+    return redshifts
+
+
+def source_classification_pe(posterior_samples_file, hdf5=True,
+                             threshold=3.0, sourceframe=True):
     """
     Returns the predicted probability of whether a
     compact binary either has an EM counterpart or
     has a NS as one of its components from PE
     results.
 
+    Parameters
+    ----------
+    posterior_samples_file: string
+                           Posterior-samples file name (*.hdf5 or *.dat)
+
+    hdf5: Bool
+         Set False if you want to use a .dat file
+
+    threshold: float
+              Maximum neutron star mass
+
+    sourceframe: Bool
+                Set False if you want to use detector frame
+
     Returns
     -------
     tuple
@@ -96,16 +148,29 @@ def source_classification_pe(posterior_samples_file, hdf5=True, threshold=3.0):
     Examples
     --------
     >>> from ligo import em_bright
-    >>> em_bright.source_classification_pe('posterior_samples.hdf5')
-    (1.0, 1.0)
+    >>> em_bright.source_classification_pe('posterior_V1H1L1_1240327333.3365-0.hdf5')
+    (1.0, 0.9616727412238634)
+    >>> em_bright.source_classification_pe('posterior_samples_online.dat', hdf5=False)
+    (0.0, 0.0)
     """
     if hdf5:
         data = h5py.File(posterior_samples_file)
         engine = list(data['lalinference'].keys())[0]
         samples = data['lalinference'][engine]['posterior_samples'][()]
+        mc_det_frame = samples['mc']
+        lum_dist = samples['dist']
+        redshifts = get_redshifts(lum_dist)
+        if sourceframe:
+            mc = mc_det_frame/(1 + redshifts)
+        else:
+            mc = mc_det_frame
+
     else:
         samples = np.recfromtxt(posterior_samples_file, names=True)
-    mc = samples['mc']
+        if sourceframe:
+            mc = samples['mc_source']
+        else:
+            mc = samples['mc']
     q = samples['q']
     m1 = mc * (1 + q)**(1/5) * (q)**(-3/5)
     m2 = mc * (1 + q)**(1/5) * (q)**(2/5)

ligo/tests/test_em_bright.py

@@ -12,9 +12,9 @@ def test_source_classification_pe():
     f = NamedTemporaryFile()
     filename = f.name
     with h5py.File(f) as tmp_h5:
-        data = np.array([(1.2, 1.0, 0.0, 0.0), (2.0, 0.5, 0.99, 0.99)],
+        data = np.array([(1.2, 1.0, 0.0, 0.0, 100.0), (2.0, 0.5, 0.99, 0.99, 150.0)],
                         dtype=[('mc', '<f8'), ('q', '<f8'),
-                               ('a1', '<f8'), ('a2', '<f8')])
+                               ('a1', '<f8'), ('a2', '<f8'), ('dist', '<f8')])
         tmp_h5.create_dataset(
             'lalinference/lalinference_mcmc/posterior_samples',
             data=data