Simplify wg redundant code

There is a lot of duplicate code in frequency_domain_strain and time_domain_strain that I eliminated. This also reduces the size of the methods and updates the test suite for the module. It now also has an identity expression as a default for parameter_conversion.

test/waveform_generator_tests.py

@@ -6,21 +6,13 @@ import mock
 from mock import MagicMock
 
 
-def gaussian_frequency_domain_strain(frequency_array, amplitude, mu, sigma, ra, dec, geocent_time, psi, **kwargs):
-    ht = {'plus': amplitude * np.exp(-(mu - frequency_array) ** 2 / sigma ** 2 / 2),
-          'cross': amplitude * np.exp(-(mu - frequency_array) ** 2 / sigma ** 2 / 2)}
-    return ht
-
-
-def gaussian_frequency_domain_strain_2(frequency_array, a, m, s, ra, dec, geocent_time, psi, **kwargs):
-    ht = {'plus': a * np.exp(-(m - frequency_array) ** 2 / s ** 2 / 2),
-          'cross': a * np.exp(-(m - frequency_array) ** 2 / s ** 2 / 2)}
-    return ht
+def dummy_func_array_return_value(frequency_array, amplitude, mu, sigma, ra, dec, geocent_time, psi, **kwargs):
+    return amplitude + mu + frequency_array + sigma + ra + dec + geocent_time + psi
 
 
-def gaussian_time_domain_strain_2(time_array, a, m, s, ra, dec, geocent_time, psi, **kwargs):
-    ht = {'plus': a * np.exp(-(m - time_array) ** 2 / s ** 2 / 2),
-          'cross': a * np.exp(-(m - time_array) ** 2 / s ** 2 / 2)}
+def dummy_func_dict_return_value(frequency_array, amplitude, mu, sigma, ra, dec, geocent_time, psi, **kwargs):
+    ht = {'plus': amplitude + mu + frequency_array + sigma + ra + dec + geocent_time + psi,
+          'cross': amplitude + mu + frequency_array + sigma + ra + dec + geocent_time + psi}
     return ht
 
 
@@ -29,7 +21,7 @@ class TestWaveformGeneratorInstantiationWithoutOptionalParameters(unittest.TestC
     def setUp(self):
         self.waveform_generator = \
             tupak.gw.waveform_generator.WaveformGenerator(1, 4096,
-                                                          frequency_domain_source_model=gaussian_frequency_domain_strain)
+                                                          frequency_domain_source_model=dummy_func_dict_return_value)
         self.simulation_parameters = dict(amplitude=1e-21, mu=100, sigma=1,
                                           ra=1.375,
                                           dec=-1.2108,
@@ -47,7 +39,7 @@ class TestWaveformGeneratorInstantiationWithoutOptionalParameters(unittest.TestC
         self.assertEqual(self.waveform_generator.sampling_frequency, 4096)
 
     def test_source_model(self):
-        self.assertEqual(self.waveform_generator.frequency_domain_source_model, gaussian_frequency_domain_strain)
+        self.assertEqual(self.waveform_generator.frequency_domain_source_model, dummy_func_dict_return_value)
 
     def test_frequency_array_type(self):
         self.assertIsInstance(self.waveform_generator.frequency_array, np.ndarray)
@@ -64,7 +56,7 @@ class TestWaveformArgumentsSetting(unittest.TestCase):
     def setUp(self):
         self.waveform_generator = \
             tupak.gw.waveform_generator.WaveformGenerator(1, 4096,
-                                                          frequency_domain_source_model=gaussian_frequency_domain_strain,
+                                                          frequency_domain_source_model=dummy_func_dict_return_value,
                                                           waveform_arguments=dict(test='test', arguments='arguments'))
 
     def tearDown(self):
@@ -80,7 +72,7 @@ class TestSetters(unittest.TestCase):
     def setUp(self):
         self.waveform_generator = \
             tupak.gw.waveform_generator.WaveformGenerator(1, 4096,
-                                                          frequency_domain_source_model=gaussian_frequency_domain_strain)
+                                                          frequency_domain_source_model=dummy_func_dict_return_value)
         self.simulation_parameters = dict(amplitude=1e-21, mu=100, sigma=1,
                                           ra=1.375,
                                           dec=-1.2108,
@@ -112,12 +104,12 @@ class TestSetters(unittest.TestCase):
         self.assertTrue(np.array_equal(new_time_array, self.waveform_generator.time_array))
 
     def test_parameters_set_from_frequency_domain_source_model(self):
-        self.waveform_generator.frequency_domain_source_model = gaussian_frequency_domain_strain_2
+        self.waveform_generator.frequency_domain_source_model = dummy_func_dict_return_value
         self.assertListEqual(sorted(list(self.waveform_generator.parameters.keys())),
                              sorted(list(self.simulation_parameters.keys())))
 
     def test_parameters_set_from_time_domain_source_model(self):
-        self.waveform_generator.time_domain_source_model = gaussian_time_domain_strain_2
+        self.waveform_generator.time_domain_source_model = dummy_func_dict_return_value
         self.assertListEqual(sorted(list(self.waveform_generator.parameters.keys())),
                              sorted(list(self.simulation_parameters.keys())))
 
@@ -126,9 +118,9 @@ class TestFrequencyDomainStrainMethod(unittest.TestCase):
 
     def setUp(self):
         self.waveform_generator = \
-            tupak.gw.waveform_generator.WaveformGenerator(1, 4096,
-                                                          frequency_domain_source_model=gaussian_frequency_domain_strain)
-        self.simulation_parameters = dict(amplitude=1e-21, mu=100, sigma=1,
+            tupak.gw.waveform_generator.WaveformGenerator(duration=1, sampling_frequency=4096,
+                                                          frequency_domain_source_model=dummy_func_dict_return_value)
+        self.simulation_parameters = dict(amplitude=1e-2, mu=100, sigma=1,
                                           ra=1.375,
                                           dec=-1.2108,
                                           geocent_time=1126259642.413,
@@ -144,30 +136,47 @@ class TestFrequencyDomainStrainMethod(unittest.TestCase):
             self.waveform_generator.frequency_domain_strain()
 
     def test_frequency_domain_source_model_call(self):
-        self.waveform_generator.frequency_domain_source_model = MagicMock(return_value=3)
-        self.assertEqual(3, self.waveform_generator.frequency_domain_strain())
+        self.waveform_generator.parameters = self.simulation_parameters
+        expected = self.waveform_generator.frequency_domain_source_model(self.waveform_generator.frequency_array,
+                                                                         self.simulation_parameters['amplitude'],
+                                                                         self.simulation_parameters['mu'],
+                                                                         self.simulation_parameters['sigma'],
+                                                                         self.simulation_parameters['ra'],
+                                                                         self.simulation_parameters['dec'],
+                                                                         self.simulation_parameters['geocent_time'],
+                                                                         self.simulation_parameters['psi'])
+        actual = self.waveform_generator.frequency_domain_strain()
+        self.assertTrue(np.array_equal(expected['plus'], actual['plus']))
+        self.assertTrue(np.array_equal(expected['cross'], actual['cross']))
 
     def test_time_domain_source_model_call_with_ndarray(self):
         self.waveform_generator.frequency_domain_source_model = None
-        self.waveform_generator.time_domain_source_model = MagicMock(return_value=np.array([1, 2, 3]))
+        self.waveform_generator.time_domain_source_model = dummy_func_array_return_value
+        self.waveform_generator.parameters = self.simulation_parameters
 
         def side_effect(value, value2):
             return value
 
         with mock.patch('tupak.core.utils.nfft') as m:
             m.side_effect = side_effect
-            self.assertTrue(np.array_equal(np.array([1, 2, 3]), self.waveform_generator.frequency_domain_strain()))
+            expected = self.waveform_generator.time_domain_strain()
+            actual = self.waveform_generator.frequency_domain_strain()
+            self.assertTrue(np.array_equal(expected, actual))
 
     def test_time_domain_source_model_call_with_dict(self):
         self.waveform_generator.frequency_domain_source_model = None
-        self.waveform_generator.time_domain_source_model = MagicMock(return_value=dict(plus=1, cross=2))
+        self.waveform_generator.time_domain_source_model = dummy_func_dict_return_value
+        self.waveform_generator.parameters = self.simulation_parameters
 
         def side_effect(value, value2):
-            return value, value2
+            return value, self.waveform_generator.frequency_array
 
         with mock.patch('tupak.core.utils.nfft') as m:
             m.side_effect = side_effect
-            self.assertDictEqual(dict(plus=1, cross=2), self.waveform_generator.frequency_domain_strain())
+            expected = self.waveform_generator.time_domain_strain()
+            actual = self.waveform_generator.frequency_domain_strain()
+            self.assertTrue(np.array_equal(expected['plus'], actual['plus']))
+            self.assertTrue(np.array_equal(expected['cross'], actual['cross']))
 
     def test_no_source_model_given(self):
         self.waveform_generator.time_domain_source_model = None
@@ -194,7 +203,7 @@ class TestTimeDomainStrainMethod(unittest.TestCase):
     def setUp(self):
         self.waveform_generator = \
             tupak.gw.waveform_generator.WaveformGenerator(1, 4096,
-                                                          time_domain_source_model=gaussian_time_domain_strain_2)
+                                                          time_domain_source_model=dummy_func_dict_return_value)
         self.simulation_parameters = dict(amplitude=1e-21, mu=100, sigma=1,
                                           ra=1.375,
                                           dec=-1.2108,
@@ -211,32 +220,47 @@ class TestTimeDomainStrainMethod(unittest.TestCase):
             self.waveform_generator.time_domain_strain()
 
     def test_time_domain_source_model_call(self):
-        self.waveform_generator.time_domain_source_model = MagicMock(return_value=3)
-        self.assertEqual(3, self.waveform_generator.time_domain_strain())
+        self.waveform_generator.parameters = self.simulation_parameters
+        expected = self.waveform_generator.time_domain_source_model(self.waveform_generator.time_array,
+                                                                    self.simulation_parameters['amplitude'],
+                                                                    self.simulation_parameters['mu'],
+                                                                    self.simulation_parameters['sigma'],
+                                                                    self.simulation_parameters['ra'],
+                                                                    self.simulation_parameters['dec'],
+                                                                    self.simulation_parameters['geocent_time'],
+                                                                    self.simulation_parameters['psi'])
+        actual = self.waveform_generator.time_domain_strain()
+        self.assertTrue(np.array_equal(expected['plus'], actual['plus']))
+        self.assertTrue(np.array_equal(expected['cross'], actual['cross']))
 
     def test_frequency_domain_source_model_call_with_ndarray(self):
         self.waveform_generator.time_domain_source_model = None
-        self.waveform_generator.frequency_domain_source_model = MagicMock(return_value=np.array([1, 2, 3]))
+        self.waveform_generator.frequency_domain_source_model = dummy_func_array_return_value
+        self.waveform_generator.parameters = self.simulation_parameters
 
         def side_effect(value, value2):
             return value
 
         with mock.patch('tupak.core.utils.infft') as m:
             m.side_effect = side_effect
-            self.assertTrue(np.array_equal(np.array([1, 2, 3]), self.waveform_generator.time_domain_strain()))
+            expected = self.waveform_generator.frequency_domain_strain()
+            actual = self.waveform_generator.time_domain_strain()
+            self.assertTrue(np.array_equal(expected, actual))
 
     def test_frequency_domain_source_model_call_with_dict(self):
         self.waveform_generator.time_domain_source_model = None
-        self.waveform_generator.frequency_domain_source_model = MagicMock(return_value=dict(plus=1, cross=2))
+        self.waveform_generator.frequency_domain_source_model = dummy_func_dict_return_value
+        self.waveform_generator.parameters = self.simulation_parameters
 
         def side_effect(value, value2):
-            return value, value2
+            return value
 
         with mock.patch('tupak.core.utils.infft') as m:
             m.side_effect = side_effect
-            self.assertDictEqual(dict(plus=(1, self.waveform_generator.sampling_frequency),
-                                      cross=(2, self.waveform_generator.sampling_frequency)),
-                                 self.waveform_generator.time_domain_strain())
+            expected = self.waveform_generator.frequency_domain_strain()
+            actual = self.waveform_generator.time_domain_strain()
+            self.assertTrue(np.array_equal(expected['plus'], actual['plus']))
+            self.assertTrue(np.array_equal(expected['cross'], actual['cross']))
 
     def test_no_source_model_given(self):
         self.waveform_generator.time_domain_source_model = None
@@ -245,7 +269,9 @@ class TestTimeDomainStrainMethod(unittest.TestCase):
             self.waveform_generator.time_domain_strain()
 
     def test_key_popping(self):
-        self.waveform_generator.parameter_conversion = MagicMock(return_value=(dict(a=1e-21, m=100, s=1,
+        self.waveform_generator.parameter_conversion = MagicMock(return_value=(dict(amplitude=1e-2,
+                                                                                    mu=100,
+                                                                                    sigma=1,
                                                                                     ra=1.375, dec=-1.2108,
                                                                                     geocent_time=1126259642.413,
                                                                                     psi=2.659, c=None, d=None),
@@ -255,7 +281,7 @@ class TestTimeDomainStrainMethod(unittest.TestCase):
         except RuntimeError:
             pass
         self.assertListEqual(sorted(self.waveform_generator.parameters.keys()),
-                             sorted(['a', 'm', 's', 'ra', 'dec', 'geocent_time', 'psi']))
+                             sorted(['amplitude', 'mu', 'sigma', 'ra', 'dec', 'geocent_time', 'psi']))
 
 
 if __name__ == '__main__':