Merge pull request #12 from chummels/flux_error

Adding flux_error field to output spectra

trident/absorption_spectrum/absorption_spectrum.py

@@ -67,6 +67,7 @@ def __init__(self, lambda_min, lambda_max, n_lambda):
                                     float(n_lambda - 1), "angstrom")
         self.line_list = []
         self.continuum_list = []
+        self.snr = 100  # default signal to noise ratio for error estimation
 
     def add_line(self, label, field_name, wavelength,
                  f_value, gamma, atomic_mass,
@@ -195,6 +196,7 @@ def make_spectrum(self, input_file, output_file=None,
            spectrum generation.
            Default: "auto"
         """
+        self.snr = 100
         if line_list_file is not None:
             mylog.info("'line_list_file' keyword is deprecated. Please use " \
                        "'output_absorbers_file'.")
@@ -250,7 +252,7 @@ def make_spectrum(self, input_file, output_file=None,
 
         if output_file is None:
             pass
-        elif output_file.endswith('.h5'):
+        elif output_file.endswith('.h5') or output_file.endswith('.hdf5'):
             self._write_spectrum_hdf5(output_file)
         elif output_file.endswith('.fits'):
             self._write_spectrum_fits(output_file)
@@ -660,10 +662,12 @@ def _write_spectrum_ascii(self, filename):
         """
         mylog.info("Writing spectrum to ascii file: %s.", filename)
         f = open(filename, 'w')
-        f.write("# wavelength[A] tau flux\n")
+        f.write("# wavelength[A] tau flux flux_error\n")
         for i in range(self.lambda_field.size):
-            f.write("%e %e %e\n" % (self.lambda_field[i],
-                                    self.tau_field[i], self.flux_field[i]))
+            f.write("%e %e %e %e\n" % (self.lambda_field[i],
+                                    self.tau_field[i], 
+                                    self.flux_field[i],
+                                    self._error_func(self.flux_field[i])))
         f.close()
 
     @parallel_root_only
@@ -675,7 +679,8 @@ def _write_spectrum_fits(self, filename):
         col1 = pyfits.Column(name='wavelength', format='E', array=self.lambda_field)
         col2 = pyfits.Column(name='tau', format='E', array=self.tau_field)
         col3 = pyfits.Column(name='flux', format='E', array=self.flux_field)
-        cols = pyfits.ColDefs([col1, col2, col3])
+        col4 = pyfits.Column(name='flux_error', format='E', array=self._error_func(self.flux_field))
+        cols = pyfits.ColDefs([col1, col2, col3, col4])
         tbhdu = pyfits.BinTableHDU.from_columns(cols)
         tbhdu.writeto(filename, overwrite=True)
 
@@ -690,4 +695,21 @@ def _write_spectrum_hdf5(self, filename):
         output.create_dataset('wavelength', data=self.lambda_field)
         output.create_dataset('tau', data=self.tau_field)
         output.create_dataset('flux', data=self.flux_field)
+        output.create_dataset('flux_error', data=self._error_func(self.flux_field))
         output.close()
+
+    def _error_func(self, flux):
+        """
+        Many observational analysis programs require a flux error channel
+        in addition to a flux channel.  So we create a zeroth order 
+        approximation of the flux error, simply by taking the square root
+        of the flux.  Unfortunately, with flux normalized to be < 1, this
+        would result in errors larger than the flux values themselves,
+        so we normalize by an arbitrary signal-to-noise ratio, which by default
+        is set to 100.  This yields a typical error for a normalized spectrum of
+        sqrt(1.0*100)/100 = 0.1.  This assures our flux errors are smaller 
+        than our fluxes for most flux reasonable flux values.  Note that 
+        when a signal to noise ratio is specified for adding gaussian noise,
+        it uses this updated value for estimating the errors.
+        """
+        return np.sqrt(flux*self.snr)/self.snr

trident/spectrum_generator.py

@@ -496,6 +496,9 @@ def add_milky_way_foreground(self, flux_field=None,
         MW_spectrum = self._get_milky_way_foreground(filename=filename)
         flux_field *= MW_spectrum
 
+        # Negative fluxes don't make sense, so clip
+        np.clip(flux_field, 0, np.inf, out=flux_field)
+
     def add_qso_spectrum(self, flux_field=None,
                          emitting_redshift=None,
                          observing_redshift=None,
@@ -558,6 +561,9 @@ def add_qso_spectrum(self, flux_field=None,
                                               filename=filename)
         flux_field *= qso_spectrum
 
+        # Negative fluxes don't make sense, so clip
+        np.clip(flux_field, 0, np.inf, out=flux_field)
+
     def add_gaussian_noise(self, snr, seed=None):
         """
         Postprocess a spectrum to add gaussian random noise of a given SNR.
@@ -598,11 +604,15 @@ def add_gaussian_noise(self, snr, seed=None):
         >>> sg.add_gaussian_noise(10)
         >>> sg.plot_spectrum('spec_noise.png')
         """
+        self.snr = snr
         np.random.seed(seed)
         noise = np.random.normal(loc=0.0, scale=1/float(snr),
                                  size=self.flux_field.size)
         self.add_noise_vector(noise)
 
+        # Negative fluxes don't make sense, so clip
+        np.clip(self.flux_field, 0, np.inf, out=self.flux_field)
+
     def add_noise_vector(self, noise):
         """
         Add an array of noise to the spectrum.
@@ -634,6 +644,7 @@ def add_noise_vector(self, noise):
                 "Flux (%s) and noise (%s) vectors must have same shape." %
                 (self.flux_field.shape, noise.shape))
         self.flux_field += noise
+        self.snr = 1 / np.std(noise)
 
     def apply_lsf(self, function=None, width=None, filename=None):
         """
@@ -702,6 +713,9 @@ def apply_lsf(self, function=None, width=None, filename=None):
         from astropy.convolution import convolve
         self.flux_field = convolve(self.flux_field, lsf.kernel)
 
+        # Negative fluxes don't make sense, so clip
+        np.clip(self.flux_field, 0, np.inf, out=self.flux_field)
+
     def load_spectrum(self, lambda_field=None, tau_field=None, flux_field=None):
         """
         Load data arrays into an existing spectrum object.
@@ -876,9 +890,9 @@ def save_spectrum(self, filename='spectrum.h5', format=None):
         >>> sg.plot_spectrum('temp.png')
         """
         if format is None:
-            if filename.endswith('.h5'):
+            if filename.endswith('.h5') or filename.endswith('hdf5'):
                 self._write_spectrum_hdf5(filename)
-            elif filename.endswith('.fits'):
+            elif filename.endswith('.fits') or filename.endswith('FITS'):
                 self._write_spectrum_fits(filename)
             else:
                 self._write_spectrum_ascii(filename)