Added option to plot spectra in velocity space

ladsantos · Apr 6, 2018 · f819e7c · f819e7c
1 parent 164d8f7
commit f819e7c
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 16 deletions.
diff --git a/sunburn/analysis.py b/sunburn/analysis.py
@@ -327,7 +327,6 @@ def _integrate(wl_range):
                 wavelength_range = ds_range / light_speed * central_wl + \
                     central_wl
                 # Compute the integrated flux for the line
-                print(wavelength_range)
                 _integrate(wavelength_range)
 
             # If more than one line is requested, then the integrated fluxes
@@ -347,7 +346,6 @@ def _integrate(wl_range):
                     wavelength_range = ds_range / light_speed * central_wl + \
                         central_wl
                     # Compute the integrated flux for the line
-                    print(wavelength_range)
                     _integrate(wavelength_range)
                     count += 1
 

diff --git a/sunburn/hst_observation.py b/sunburn/hst_observation.py
@@ -68,7 +68,7 @@ def __init__(self, dataset_name, instrument, good_pixel_limits=None,
                                           'yet.')
 
     # Plot all the spectra in a wavelength range
-    def plot_spectra(self, wavelength_range=None, chip_index=None,
+    def plot_spectra(self, wavelength_range=None, chip_index=None, ref_wl=None,
                      uncertainties=False, figure_sizes=(9.0, 6.5),
                      axes_font_size=18, legend_font_size=13, rotate_x_ticks=30):
         """
@@ -81,6 +81,9 @@ def plot_spectra(self, wavelength_range=None, chip_index=None,
 
             chip_index():
 
+            ref_wl (``float``, optional): Reference wavelength used to plot the
+                spectra in Doppler velocity space.
+
             uncertainties (``bool``, optional): If ``True``, then plot the
                 spectra with their respective uncertainties. Default is
                 ``False``.
@@ -122,16 +125,24 @@ def plot_spectra(self, wavelength_range=None, chip_index=None,
                                          wavelength_range[0])
             max_wl = tools.nearest_index(self.orbit[i].wavelength[ind],
                                          wavelength_range[1])
+
+            if isinstance(ref_wl, float):
+                x_axis = \
+                    (self.orbit[i].wavelength[ind][min_wl:max_wl] - ref_wl) / \
+                    ref_wl * c.c.to(u.km / u.s).value
+                x_label = r'Velocity (km s$^{-1}$)'
+            else:
+                x_axis = self.orbit[i].wavelength[ind][min_wl:max_wl]
+                x_label = r'Wavelength ($\mathrm{\AA}$)'
+
             if uncertainties is False:
-                plt.plot(self.orbit[i].wavelength[ind][min_wl:max_wl],
-                         self.orbit[i].flux[ind][min_wl:max_wl],
+                plt.plot(x_axis, self.orbit[i].flux[ind][min_wl:max_wl],
                          label=label)
             else:
-                plt.errorbar(self.orbit[i].wavelength[ind][min_wl:max_wl],
-                             self.orbit[i].flux[ind][min_wl:max_wl],
+                plt.errorbar(x_axis, self.orbit[i].flux[ind][min_wl:max_wl],
                              yerr=self.orbit[i].error[ind][min_wl:max_wl],
                              fmt='.', label=label)
-        plt.xlabel(r'Wavelength ($\mathrm{\AA}$)')
+        plt.xlabel(x_label)
         plt.ylabel(r'Flux (erg s$^{-1}$ cm$^{-2}$ $\mathrm{\AA}^{-1}$)')
         plt.legend(fontsize=legend_font_size)
         if rotate_x_ticks is not None:
@@ -351,7 +362,7 @@ def integrated_flux(self, wavelength_range,
 
     # Plot the spectrum
     def plot_spectrum(self, wavelength_range=None, chip_index=None,
-                      plot_uncertainties=False, rotate_x_ticks=30):
+                      plot_uncertainties=False, rotate_x_ticks=30, ref_wl=None):
         """
         Plot the spectrum, with the option of selecting a specific wavelength
         range or the red or blue chips of the detector. In order to visualize
@@ -374,6 +385,9 @@ def plot_spectrum(self, wavelength_range=None, chip_index=None,
                 ``False``.
 
             rotate_x_ticks ():
+
+            ref_wl (``float``, optional): Reference wavelength used to plot the
+                spectra in Doppler velocity space.
         """
 
         if wavelength_range is not None:
@@ -383,18 +397,23 @@ def plot_spectrum(self, wavelength_range=None, chip_index=None,
             max_wl = tools.nearest_index(self.wavelength[ind],
                                          wavelength_range[1])
 
+            if isinstance(ref_wl, float):
+                x_axis = c.c.to(u.km / u.s).value * \
+                    (self.wavelength[ind][min_wl:max_wl] - ref_wl) / ref_wl
+                x_label = r'Velocity (km s$^{-1}$)'
+            else:
+                x_axis = self.wavelength[ind][min_wl:max_wl]
+                x_label = r'Wavelength ($\mathrm{\AA}$)'
+
             # Finally plot it
             if plot_uncertainties is False:
-                plt.plot(self.wavelength[ind][min_wl:max_wl],
-                         self.flux[ind][min_wl:max_wl],
+                plt.plot(x_axis, self.flux[ind][min_wl:max_wl],
                          label=self.start_JD.value)
             else:
-                plt.errorbar(self.wavelength[ind][min_wl:max_wl],
-                             self.flux[ind][min_wl:max_wl],
+                plt.errorbar(x_axis, self.flux[ind][min_wl:max_wl],
                              yerr=self.error[ind][min_wl:max_wl],
-                             fmt='.',
-                             label=self.start_JD.value)
-            plt.xlabel(r'Wavelength ($\mathrm{\AA}$)')
+                             fmt='.', label=self.start_JD.value)
+            plt.xlabel(x_label)
             plt.ylabel(r'Flux (erg s$^{-1}$ cm$^{-2}$ $\mathrm{\AA}^{-1}$)')
 
         elif chip_index is not None: