Ability to add noise to spectrum with a given snr value.

spectrum_overload/Spectrum.py

@@ -137,6 +137,12 @@ def wav_select(self, wav_min, wav_max):
             self.xaxis = x_org
             raise
 
+    def add_noise(self, snr):
+        """Add noise level of snr to the flux of the spectrum."""
+        sigma = self.flux / snr
+        # Add normal distributed noise at the SNR level.
+        self.flux += np.random.normal(0, sigma)
+
     def doppler_shift(self, RV):
         """Function to compute a wavelength shift due to radial velocity.
 

spectrum_overload/test/test_Spectrum.py

@@ -379,3 +379,12 @@ def test_spline_interpolation_when_given_a_ndarray():
 
 
 # test_doppler_shift_with_hypothesis()
+@pytest.mark.parametrize('snr', [50, 100, 200, 1000])
+def test_add_noise(snr):
+    """Test addition of noise."""
+    x = np.linspace(2000, 2200, 10000000)
+    y = np.ones_like(x)
+    spec = Spectrum(xaxis=x, flux=y)
+    spec.add_noise(snr)
+
+    assert np.isclose(np.std(spec.flux), 1. / snr, atol=1e-5)