ENH: Added 'source' flag to generate_component_masses and generate_mass_parameters

AUTHORS.md

@@ -59,6 +59,7 @@ Paul Easter
 Paul Lasky
 Philip Relton
 Rhys Green
+R. Udall
 Rico Lo
 Roberto Cotesta
 Rory Smith

bilby/gw/conversion.py

@@ -1001,7 +1001,7 @@ def fill_from_fixed_priors(sample, priors):
     return output_sample
 
 
-def generate_component_masses(sample, require_add=False):
+def generate_component_masses(sample, require_add=False, source=False):
     """"
     Add the component masses to the dataframe/dictionary
     We add:
@@ -1022,6 +1022,9 @@ def generate_component_masses(sample, require_add=False):
         component with overall mass scaling (i.e.
         chirp_mass, mass_1, mass_2, total_mass) and
         then any other mass parameter.
+    source : bool, default False
+        If True, then perform the conversions for source mass parameters
+        i.e. mass_1_source instead of mass_1
 
     Returns
     dict : the updated dictionary
@@ -1032,12 +1035,24 @@ def generate_component_masses(sample, require_add=False):
         else:
             return sample
     output_sample = sample.copy()
-    if "mass_1" in sample.keys():
-        if "mass_2" in sample.keys():
+
+    if source:
+        mass_1_key = "mass_1_source"
+        mass_2_key = "mass_2_source"
+        total_mass_key = "total_mass_source"
+        chirp_mass_key = "chirp_mass_source"
+    else:
+        mass_1_key = "mass_1"
+        mass_2_key = "mass_2"
+        total_mass_key = "total_mass"
+        chirp_mass_key = "chirp_mass"
+
+    if mass_1_key in sample.keys():
+        if mass_2_key in sample.keys():
             return output_sample
-        if "total_mass" in sample.keys():
-            output_sample["mass_2"] = output_sample["total_mass"] - (
-                output_sample["mass_1"]
+        if total_mass_key in sample.keys():
+            output_sample[mass_2_key] = output_sample[total_mass_key] - (
+                output_sample[mass_1_key]
             )
             return output_sample
 
@@ -1048,26 +1063,26 @@ def generate_component_masses(sample, require_add=False):
                 symmetric_mass_ratio_to_mass_ratio(
                     output_sample["symmetric_mass_ratio"])
             )
-        elif "chirp_mass" in sample.keys():
+        elif chirp_mass_key in sample.keys():
             output_sample["mass_ratio"] = (
                 mass_1_and_chirp_mass_to_mass_ratio(
-                    mass_1=output_sample["mass_1"],
-                    chirp_mass=output_sample["chirp_mass"])
+                    mass_1=output_sample[mass_1_key],
+                    chirp_mass=output_sample[chirp_mass_key])
             )
         else:
             return check_and_return_quietly(require_add, sample)
 
-        output_sample["mass_2"] = (
-            output_sample["mass_ratio"] * output_sample["mass_1"]
+        output_sample[mass_2_key] = (
+            output_sample["mass_ratio"] * output_sample[mass_1_key]
         )
 
         return output_sample
 
-    elif "mass_2" in sample.keys():
+    elif mass_2_key in sample.keys():
         # mass_1 is not in the dict
-        if "total_mass" in sample.keys():
-            output_sample["mass_1"] = (
-                output_sample["total_mass"] - output_sample["mass_2"]
+        if total_mass_key in sample.keys():
+            output_sample[mass_1_key] = (
+                output_sample[total_mass_key] - output_sample[mass_2_key]
             )
             return output_sample
         elif "mass_ratio" in sample.keys():
@@ -1077,23 +1092,23 @@ def generate_component_masses(sample, require_add=False):
                 symmetric_mass_ratio_to_mass_ratio(
                     output_sample["symmetric_mass_ratio"])
             )
-        elif "chirp_mass" in sample.keys():
+        elif chirp_mass_key in sample.keys():
             output_sample["mass_ratio"] = (
                 mass_2_and_chirp_mass_to_mass_ratio(
-                    mass_2=output_sample["mass_2"],
-                    chirp_mass=output_sample["chirp_mass"])
+                    mass_2=output_sample[mass_2_key],
+                    chirp_mass=output_sample[chirp_mass_key])
             )
         else:
             check_and_return_quietly(require_add, sample)
 
-        output_sample["mass_1"] = 1 / output_sample["mass_ratio"] * (
-            output_sample["mass_2"]
+        output_sample[mass_1_key] = 1 / output_sample["mass_ratio"] * (
+            output_sample[mass_2_key]
         )
 
         return output_sample
 
     # Only if neither mass_1 or mass_2 is in the input sample
-    if "total_mass" in sample.keys():
+    if total_mass_key in sample.keys():
         if "mass_ratio" in sample.keys():
             pass  # We have everything we need already
         elif "symmetric_mass_ratio" in sample.keys():
@@ -1101,11 +1116,11 @@ def generate_component_masses(sample, require_add=False):
                 symmetric_mass_ratio_to_mass_ratio(
                     output_sample["symmetric_mass_ratio"])
             )
-        elif "chirp_mass" in sample.keys():
+        elif chirp_mass_key in sample.keys():
             output_sample["symmetric_mass_ratio"] = (
                 chirp_mass_and_total_mass_to_symmetric_mass_ratio(
-                    chirp_mass=output_sample["chirp_mass"],
-                    total_mass=output_sample["total_mass"])
+                    chirp_mass=output_sample[chirp_mass_key],
+                    total_mass=output_sample[total_mass_key])
             )
             output_sample["mass_ratio"] = (
                 symmetric_mass_ratio_to_mass_ratio(
@@ -1114,7 +1129,7 @@ def generate_component_masses(sample, require_add=False):
         else:
             return check_and_return_quietly(require_add, sample)
 
-    elif "chirp_mass" in sample.keys():
+    elif chirp_mass_key in sample.keys():
         if "mass_ratio" in sample.keys():
             pass
         elif "symmetric_mass_ratio" in sample.keys():
@@ -1125,27 +1140,28 @@ def generate_component_masses(sample, require_add=False):
         else:
             return check_and_return_quietly(require_add, sample)
 
-        output_sample["total_mass"] = (
+        output_sample[total_mass_key] = (
             chirp_mass_and_mass_ratio_to_total_mass(
-                chirp_mass=output_sample["chirp_mass"],
+                chirp_mass=output_sample[chirp_mass_key],
                 mass_ratio=output_sample["mass_ratio"])
         )
 
     # We haven't matched any of the criteria
-    if "total_mass" not in output_sample.keys() or (
+    if total_mass_key not in output_sample.keys() or (
             "mass_ratio" not in output_sample.keys()):
         return check_and_return_quietly(require_add, sample)
     mass_1, mass_2 = (
         total_mass_and_mass_ratio_to_component_masses(
-            total_mass=output_sample["total_mass"],
+            total_mass=output_sample[total_mass_key],
             mass_ratio=output_sample["mass_ratio"])
     )
-    output_sample["mass_1"] = mass_1
-    output_sample["mass_2"] = mass_2
+    output_sample[mass_1_key] = mass_1
+    output_sample[mass_2_key] = mass_2
+
     return output_sample
 
 
-def generate_mass_parameters(sample):
+def generate_mass_parameters(sample, source=False):
     """
     Add the known mass parameters to the data frame/dictionary.  We do
     not recompute keys already present in the dictionary
@@ -1165,27 +1181,39 @@ def generate_mass_parameters(sample):
 
     """
     # Only add the parameters if they're not already present
-    intermediate_sample = generate_component_masses(sample)
+    intermediate_sample = generate_component_masses(sample, source=source)
     output_sample = intermediate_sample.copy()
-    if "chirp_mass" not in output_sample.keys():
-        output_sample['chirp_mass'] = (
-            component_masses_to_chirp_mass(output_sample['mass_1'],
-                                           output_sample['mass_2'])
+
+    if source:
+        mass_1_key = 'mass_1_source'
+        mass_2_key = 'mass_2_source'
+        total_mass_key = 'total_mass_source'
+        chirp_mass_key = 'chirp_mass_source'
+    else:
+        mass_1_key = 'mass_1'
+        mass_2_key = 'mass_2'
+        total_mass_key = 'total_mass'
+        chirp_mass_key = 'chirp_mass'
+
+    if chirp_mass_key not in output_sample.keys():
+        output_sample[chirp_mass_key] = (
+            component_masses_to_chirp_mass(output_sample[mass_1_key],
+                                           output_sample[mass_2_key])
         )
-    if "total_mass" not in output_sample.keys():
-        output_sample['total_mass'] = (
-            component_masses_to_total_mass(output_sample['mass_1'],
-                                           output_sample['mass_2'])
+    if total_mass_key not in output_sample.keys():
+        output_sample[total_mass_key] = (
+            component_masses_to_total_mass(output_sample[mass_1_key],
+                                           output_sample[mass_2_key])
         )
-    if "symmetric_mass_ratio" not in output_sample.keys():
+    if 'symmetric_mass_ratio' not in output_sample.keys():
         output_sample['symmetric_mass_ratio'] = (
-            component_masses_to_symmetric_mass_ratio(output_sample['mass_1'],
-                                                     output_sample['mass_2'])
+            component_masses_to_symmetric_mass_ratio(output_sample[mass_1_key],
+                                                     output_sample[mass_2_key])
         )
-    if "mass_ratio" not in output_sample.keys():
+    if 'mass_ratio' not in output_sample.keys():
         output_sample['mass_ratio'] = (
-            component_masses_to_mass_ratio(output_sample['mass_1'],
-                                           output_sample['mass_2'])
+            component_masses_to_mass_ratio(output_sample[mass_1_key],
+                                           output_sample[mass_2_key])
         )
 
     return output_sample

test/gw/conversion_test.py

@@ -537,17 +537,25 @@ class TestGenerateMassParameters(unittest.TestCase):
                                 'mass_2': 1.0,
                                 'chirp_mass': 1.2167286837864113,
                                 'total_mass': 3.0,
+                                'mass_1_source': 4.0,
+                                'mass_2_source': 2.0,
+                                'chirp_mass_source': 2.433457367572823,
+                                'total_mass_source': 6,
                                 'symmetric_mass_ratio': 0.2222222222222222,
                                 'mass_ratio': 0.5}
 
-    def helper_generation_from_keys(self, keys, expected_values,):
+    def helper_generation_from_keys(self, keys, expected_values, source=False):
         # Explicitly test the helper generate_component_masses
         local_test_vars = \
             {key: expected_values[key] for key in keys}
         local_test_vars_with_component_masses = \
-            conversion.generate_component_masses(local_test_vars)
-        self.assertTrue("mass_1" in local_test_vars_with_component_masses.keys())
-        self.assertTrue("mass_2" in local_test_vars_with_component_masses.keys())
+            conversion.generate_component_masses(local_test_vars, source=source)
+        if source:
+            self.assertTrue("mass_1_source" in local_test_vars_with_component_masses.keys())
+            self.assertTrue("mass_2_source" in local_test_vars_with_component_masses.keys())
+        else:
+            self.assertTrue("mass_1" in local_test_vars_with_component_masses.keys())
+            self.assertTrue("mass_2" in local_test_vars_with_component_masses.keys())
         for key in local_test_vars_with_component_masses.keys():
             self.assertAlmostEqual(
                 local_test_vars_with_component_masses[key],
@@ -555,10 +563,32 @@ class TestGenerateMassParameters(unittest.TestCase):
 
         # Test the function more generally
         local_all_mass_parameters = \
-            conversion.generate_mass_parameters(local_test_vars)
-        self.assertEqual(local_all_mass_parameters.keys(),
-                         self.expected_values.keys())
-        for key in expected_values.keys():
+            conversion.generate_mass_parameters(local_test_vars, source=source)
+        if source:
+            self.assertEqual(
+                set(local_all_mass_parameters.keys()),
+                set(["mass_1_source",
+                     "mass_2_source",
+                     "chirp_mass_source",
+                     "total_mass_source",
+                     "symmetric_mass_ratio",
+                     "mass_ratio",
+                     ]
+                    )
+            )
+        else:
+            self.assertEqual(
+                set(local_all_mass_parameters.keys()),
+                set(["mass_1",
+                     "mass_2",
+                     "chirp_mass",
+                     "total_mass",
+                     "symmetric_mass_ratio",
+                     "mass_ratio",
+                     ]
+                    )
+            )
+        for key in local_all_mass_parameters.keys():
             self.assertAlmostEqual(expected_values[key], local_all_mass_parameters[key])
 
     def test_from_mass_1_and_mass_2(self):
@@ -593,6 +623,38 @@ class TestGenerateMassParameters(unittest.TestCase):
         self.helper_generation_from_keys(["chirp_mass", "mass_2"],
                                          self.expected_values)
 
+    def test_from_mass_1_source_and_mass_2_source(self):
+        self.helper_generation_from_keys(["mass_1_source", "mass_2_source"],
+                                         self.expected_values, source=True)
+
+    def test_from_mass_1_source_and_mass_ratio(self):
+        self.helper_generation_from_keys(["mass_1_source", "mass_ratio"],
+                                         self.expected_values, source=True)
+
+    def test_from_mass_2_source_and_mass_ratio(self):
+        self.helper_generation_from_keys(["mass_2_source", "mass_ratio"],
+                                         self.expected_values, source=True)
+
+    def test_from_mass_1_source_and_total_mass(self):
+        self.helper_generation_from_keys(["mass_2_source", "total_mass_source"],
+                                         self.expected_values, source=True)
+
+    def test_from_chirp_mass_source_and_mass_ratio(self):
+        self.helper_generation_from_keys(["chirp_mass_source", "mass_ratio"],
+                                         self.expected_values, source=True)
+
+    def test_from_chirp_mass_source_and_symmetric_mass_ratio(self):
+        self.helper_generation_from_keys(["chirp_mass_source", "symmetric_mass_ratio"],
+                                         self.expected_values, source=True)
+
+    def test_from_chirp_mass_source_and_symmetric_mass_1(self):
+        self.helper_generation_from_keys(["chirp_mass_source", "mass_1_source"],
+                                         self.expected_values, source=True)
+
+    def test_from_chirp_mass_source_and_symmetric_mass_2(self):
+        self.helper_generation_from_keys(["chirp_mass_source", "mass_2_source"],
+                                         self.expected_values, source=True)
+
 
 if __name__ == "__main__":
     unittest.main()