Merge branch 'tify-up-of-inject-signal' of git.ligo.org:Monash/tupak into tify-up-of-inject-signal

examples/injection_examples/basic_tutorial.py

@@ -41,17 +41,13 @@ waveform_generator = tupak.WaveformGenerator(time_duration=time_duration,
                                              frequency_domain_source_model=tupak.gw.source.lal_binary_black_hole,
                                              parameters=injection_parameters,
                                              waveform_arguments=waveform_arguments)
+hf_signal = waveform_generator.frequency_domain_strain()
 
 # Set up interferometers.  In this case we'll use three interferometers (LIGO-Hanford (H1), LIGO-Livingston (L1),
 # and Virgo (V1)).  These default to their design sensitivity
-IFOs = []
-for name in ['H1', 'L1', 'V1']:
-    IFOs.append(
-        tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(
-            name, waveform_generator=waveform_generator,
-            injection_parameters=injection_parameters,
-            time_duration=time_duration,
-            sampling_frequency=sampling_frequency, outdir=outdir))
+IFOs = [tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(
+    name, injection_polarizations=hf_signal, injection_parameters=injection_parameters, time_duration=time_duration,
+    sampling_frequency=sampling_frequency, outdir=outdir) for name in ['H1', 'L1']]
 
 # Set up prior, which is a dictionary
 # By default we will sample all terms in the signal models.  However, this will take a long time for the calculation,

examples/injection_examples/basic_tutorial_dist_time_phase_marg.py

@@ -39,11 +39,12 @@ waveform_generator = tupak.WaveformGenerator(time_duration=time_duration,
                                              frequency_domain_source_model=tupak.gw.source.lal_binary_black_hole,
                                              parameters=injection_parameters,
                                              waveform_arguments=waveform_arguments)
+hf_signal = waveform_generator.frequency_domain_strain()
 
 # Set up interferometers.  In this case we'll use three interferometers (LIGO-Hanford (H1), LIGO-Livingston (L1),
 # and Virgo (V1)).  These default to their design sensitivity
 IFOs = [tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(
-    name, waveform_generator=waveform_generator, injection_parameters=injection_parameters, time_duration=time_duration,
+    name, injection_polarizations=hf_signal, injection_parameters=injection_parameters, time_duration=time_duration,
     sampling_frequency=sampling_frequency, outdir=outdir) for name in ['H1', 'L1']]
 
 # Set up prior, which is a dictionary

examples/injection_examples/basic_tutorial_time_phase_marg.py

@@ -38,11 +38,12 @@ waveform_generator = tupak.WaveformGenerator(time_duration=time_duration,
                                              frequency_domain_source_model=tupak.gw.source.lal_binary_black_hole,
                                              parameters=injection_parameters,
                                              waveform_arguments=waveform_arguments)
+hf_signal = waveform_generator.frequency_domain_strain()
 
 # Set up interferometers.  In this case we'll use three interferometers (LIGO-Hanford (H1), LIGO-Livingston (L1),
 # and Virgo (V1)).  These default to their design sensitivity
 IFOs = [tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(
-    name, waveform_generator=waveform_generator, injection_parameters=injection_parameters, time_duration=time_duration,
+    name, injection_polarizations=hf_signal, injection_parameters=injection_parameters, time_duration=time_duration,
     sampling_frequency=sampling_frequency, outdir=outdir) for name in ['H1', 'L1']]
 
 # Set up prior, which is a dictionary

examples/injection_examples/change_sampled_parameters.py

@@ -32,16 +32,12 @@ waveform_generator = tupak.gw.waveform_generator.WaveformGenerator(
     parameter_conversion=tupak.gw.conversion.convert_to_lal_binary_black_hole_parameters,
     non_standard_sampling_parameter_keys=['chirp_mass', 'mass_ratio'],
     parameters=injection_parameters, waveform_arguments=waveform_arguments)
+hf_signal = waveform_generator.frequency_domain_strain()
 
 # Set up interferometers.
-IFOs = []
-for name in ['H1', 'L1', 'V1']:
-    IFOs.append(
-        tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(
-            name, waveform_generator=waveform_generator,
-            injection_parameters=injection_parameters,
-            time_duration=time_duration,
-            sampling_frequency=sampling_frequency, outdir=outdir))
+IFOs = [tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(
+    name, injection_polarizations=hf_signal, injection_parameters=injection_parameters, time_duration=time_duration,
+    sampling_frequency=sampling_frequency, outdir=outdir) for name in ['H1', 'L1', 'V1']]
 
 # Set up prior
 priors = tupak.gw.prior.BBHPriorSet()

examples/injection_examples/create_your_own_source_model.py

@@ -29,10 +29,11 @@ waveform_generator = tupak.gw.waveform_generator.WaveformGenerator(time_duration
                                                                    sampling_frequency=sampling_frequency,
                                                                    frequency_domain_source_model=sine_gaussian,
                                                                    parameters=injection_parameters)
+hf_signal = waveform_generator.frequency_domain_strain()
 
 # Set up interferometers.
 IFOs = [tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(
-    name, waveform_generator=waveform_generator,
+    name, injection_polarizations=hf_signal,
     injection_parameters=injection_parameters, time_duration=time_duration,
     sampling_frequency=sampling_frequency, outdir=outdir)
     for name in ['H1', 'L1', 'V1']]

examples/injection_examples/create_your_own_time_domain_source_model.py

@@ -36,6 +36,7 @@ waveform = tupak.gw.waveform_generator.WaveformGenerator(time_duration=time_dura
                                                          time_domain_source_model=time_domain_damped_sinusoid,
                                                          parameters=injection_parameters)
 
+hf_signal = waveform.frequency_domain_strain()
 #note we could plot the time domain signal with the following code
 # import matplotlib.pyplot as plt
 # plt.plot(waveform.time_array, waveform.time_domain_strain()['plus'])
@@ -46,14 +47,12 @@ waveform = tupak.gw.waveform_generator.WaveformGenerator(time_duration=time_dura
 
 
 # inject the signal into three interferometers
-IFOs = []
-for name in ['H1', 'L1', 'V1']:
-    IFOs.append(
-        tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(
-            name, waveform_generator=waveform,
-            injection_parameters=injection_parameters,
-            time_duration=time_duration,
-            sampling_frequency=sampling_frequency, outdir=outdir))
+IFOs = [tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(
+        name, injection_polarizations=hf_signal,
+        injection_parameters=injection_parameters, time_duration=time_duration,
+        sampling_frequency=sampling_frequency, outdir=outdir)
+        for name in ['H1', 'L1']]
+
 
 #  create the priors
 prior = injection_parameters.copy()

examples/injection_examples/how_to_specify_the_prior.py

@@ -28,10 +28,11 @@ waveform_generator = tupak.WaveformGenerator(time_duration=time_duration,
                                              frequency_domain_source_model=tupak.gw.source.lal_binary_black_hole,
                                              parameters=injection_parameters,
                                              waveform_arguments=waveform_arguments)
+hf_signal = waveform_generator.frequency_domain_strain()
 
 # Set up interferometers.
 IFOs = [tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(
-    name, waveform_generator=waveform_generator, injection_parameters=injection_parameters, time_duration=time_duration,
+    name, injection_polarizations=hf_signal, injection_parameters=injection_parameters, time_duration=time_duration,
     sampling_frequency=sampling_frequency, outdir=outdir) for name in ['H1', 'L1', 'V1']]
 
 # Set up prior

examples/injection_examples/marginalized_likelihood.py

@@ -26,16 +26,12 @@ waveform_generator = tupak.WaveformGenerator(
     time_duration=time_duration, sampling_frequency=sampling_frequency,
     frequency_domain_source_model=tupak.gw.source.lal_binary_black_hole, parameters=injection_parameters,
     waveform_arguments=waveform_arguments)
+hf_signal = waveform_generator.frequency_domain_strain()
 
 # Set up interferometers.
-IFOs = []
-for name in ['H1', 'L1', 'V1']:
-    IFOs.append(
-        tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(
-            name, waveform_generator=waveform_generator,
-            injection_parameters=injection_parameters,
-            time_duration=time_duration,
-            sampling_frequency=sampling_frequency, outdir=outdir))
+IFOs = [tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(
+    name, injection_polarizations=hf_signal, injection_parameters=injection_parameters, time_duration=time_duration,
+    sampling_frequency=sampling_frequency, outdir=outdir) for name in ['H1', 'L1', 'V1']]
 
 # Set up prior
 priors = tupak.gw.prior.BBHPriorSet()

examples/injection_examples/sine_gaussian_example.py

@@ -29,11 +29,12 @@ waveform_generator = tupak.gw.waveform_generator.WaveformGenerator(time_duration
                                                                    sampling_frequency=sampling_frequency,
                                                                    frequency_domain_source_model=tupak.gw.source.sinegaussian,
                                                                    parameters=injection_parameters)
+hf_signal = waveform_generator.frequency_domain_strain()
 
 # Set up interferometers.  In this case we'll use three interferometers (LIGO-Hanford (H1), LIGO-Livingston (L1),
 # and Virgo (V1)).  These default to their design sensitivity
 IFOs = [tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(
-    name, waveform_generator=waveform_generator, injection_parameters=injection_parameters, time_duration=time_duration,
+    name, injection_polarizations=hf_signal, injection_parameters=injection_parameters, time_duration=time_duration,
     sampling_frequency=sampling_frequency, outdir=outdir) for name in ['H1', 'L1', 'V1']]
 
 # Set up prior, which is a dictionary

test/detector_tests.py

@@ -247,9 +247,9 @@ class TestDetector(unittest.TestCase):
             parameters=dict(ra=0, dec=0, geocent_time=0, psi=0))
         self.assertTrue(np.array_equal(response, (plus+cross)*self.ifo.frequency_mask*np.exp(-0j)))
 
-    #def test_inject_signal_no_waveform_generator(self):
-    #    with self.assertRaises(ValueError):
-    #        self.ifo.inject_signal(waveform_generator=None, parameters=None)
+    def test_inject_signal_no_waveform_polarizations(self):
+        with self.assertRaises(ValueError):
+            self.ifo.inject_signal(injection_polarizations=None, parameters=None)
 
     def test_unit_vector_along_arm_default(self):
         with self.assertRaises(ValueError):

test/make_standard_data.py

@@ -36,8 +36,9 @@ waveform_generator = WaveformGenerator(time_duration=time_duration, sampling_fre
                                        frequency_domain_source_model=tupak.gw.source.lal_binary_black_hole,
                                        parameters=simulation_parameters)
 
+signal = waveform_generator.frequency_domain_strain()
 
-IFO = tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(name='H1', waveform_generator=waveform_generator,
+IFO = tupak.gw.detector.get_interferometer_with_fake_noise_and_injection(name='H1', injection_polarizations=signal,
                                                                          injection_parameters=simulation_parameters,
                                                                          time_duration=time_duration, plot=False,
                                                                          sampling_frequency=sampling_frequency)

tupak/gw/detector.py

@@ -732,21 +732,43 @@ class Interferometer(object):
 
         return signal_ifo
 
-    def inject_signal(self, waveform_generator, parameters):
+    def inject_signal(self,  parameters=None, injection_polarizations=None,
+                      waveform_generator=None):
         """ Inject a signal into noise
 
         Parameters
         ----------
-        waveform_generator: tupak.gw.waveform_generator
-            A WaveformGenerator instance using the source model to inject
         parameters: dict
-            parameters describing position and time of arrival of the signal
+            Parameters of the injection.
+        injection_polarizations: dict
+           Polarizations of waveform to inject, output of
+           `waveform_generator.frequency_domain_strain()`. If
+           `waveform_generator` is also given, the injection_polarizations will
+           be calculated directly and this argument can be ignored.
+        waveform_generator: tupak.gw.waveform_generator
+            A WaveformGenerator instance using the source model to inject. If
+            `injection_polarizations` is given, this will be ignored.
+
+        Note: if your signal takes a substantial amount of time to generate, or
+        you experience buggy behaviour. It is preferable to provide the
+        injection_polarizations directly.
+
+        Returns
+        -------
+        injection_polarizations: dict
+
         """
 
-        waveform_generator.parameters = parameters
-        waveform_polarizations = waveform_generator.frequency_domain_strain()
+        if injection_polarizations is None:
+            if waveform_generator is not None:
+                waveform_generator.parameters = parameters
+                injection_polarizations = waveform_generator.frequency_domain_strain()
+            else:
+                raise ValueError(
+                    "inject_signal needs one of waveform_generator or "
+                    "injection_polarizations.")
 
-        if waveform_polarizations is None:
+        if injection_polarizations is None:
             raise ValueError(
                 'Trying to inject signal which is None. The most likely cause'
                 ' is that waveform_generator.frequency_domain_strain returned'
@@ -757,7 +779,7 @@ class Interferometer(object):
                 'Injecting signal outside segment, start_time={}, merger time={}.'
                 .format(self.strain_data.start_time, parameters['geocent_time']))
 
-        signal_ifo = self.get_detector_response(waveform_polarizations, parameters)
+        signal_ifo = self.get_detector_response(injection_polarizations, parameters)
         if np.shape(self.frequency_domain_strain).__eq__(np.shape(signal_ifo)):
             self.strain_data.add_to_frequency_domain_strain(signal_ifo)
         else:
@@ -780,6 +802,8 @@ class Interferometer(object):
         for key in parameters:
             logging.info('  {} = {}'.format(key, parameters[key]))
 
+        return injection_polarizations
+
     def unit_vector_along_arm(self, arm):
         """
         Calculate the unit vector pointing along the specified arm in cartesian Earth-based coordinates.
@@ -1296,8 +1320,8 @@ def get_interferometer_with_open_data(
 
 
 def get_interferometer_with_fake_noise_and_injection(
-        name, injection_parameters, waveform_generator=None,
-        waveform_polarizations=None, sampling_frequency=4096, time_duration=4,
+        name, injection_parameters, injection_polarizations=None,
+        waveform_generator=None, sampling_frequency=4096, time_duration=4,
         start_time=None, outdir='outdir', label=None, plot=True, save=True,
         zero_noise=False):
     """
@@ -1310,8 +1334,14 @@ def get_interferometer_with_fake_noise_and_injection(
         Detector name, e.g., 'H1'.
     injection_parameters: dict
         injection parameters, needed for sky position and timing
+    injection_polarizations: dict
+       Polarizations of waveform to inject, output of
+       `waveform_generator.frequency_domain_strain()`. If
+       `waveform_generator` is also given, the injection_polarizations will
+       be calculated directly and this argument can be ignored.
     waveform_generator: tupak.gw.waveform_generator
-        A WaveformGenerator instance using the source model to inject
+        A WaveformGenerator instance using the source model to inject. If
+        `injection_polarizations` is given, this will be ignored.
     sampling_frequency: float
         sampling frequency for data, should match injection signal
     time_duration: float
@@ -1350,11 +1380,12 @@ def get_interferometer_with_fake_noise_and_injection(
         interferometer.set_strain_data_from_power_spectral_density(
             sampling_frequency=sampling_frequency, duration=time_duration,
             start_time=start_time)
-    interferometer.inject_signal(
-        waveform_generator=waveform_generator,
-        parameters=injection_parameters)
 
-    injection_polarizations = waveform_generator.frequency_domain_strain()
+    injection_polarizations = interferometer.inject_signal(
+        parameters=injection_parameters,
+        injection_polarizations=injection_polarizations,
+        waveform_generator=waveform_generator)
+
     signal = interferometer.get_detector_response(
         injection_polarizations, injection_parameters)