def test_apply_wcs_to_pixels(self):
expected_result = [9719.456, 16292.345, 24981.703]
sp = spectro.Spectrum(TestSpectrum.adhdu)
result = [sp.wlen[0], sp.wlen[1000], sp.wlen[-1]]
assert_almost_equal(result[0], expected_result[0], 3)
assert_almost_equal(result[1], expected_result[1], 3)
assert_almost_equal(result[2], expected_result[2], 3)
def test_get_wunit2(self):
expected_result = u.Angstrom
sp = spectro.Spectrum(TestSpectrum.apfhdu)
result = sp.wunit
assert_equal(result, expected_result)
def test_get_wcs_from_hdu2(self):
expected_result = TestSpectrum.wcs_string
sp = spectro.Spectrum(TestSpectrum.apfhdu, wunit=u.Angstrom)
result = sp.wcs.wcs.to_header()
assert_equal(result, expected_result)
def test_get_pixel_array_from_hdu2(self):
expected_result = np.arange(TestSpectrum.apfhdu.header['NAXIS1'])
sp = spectro.Spectrum(TestSpectrum.apfhdu, wunit=u.Angstrom)
result = sp.pix
assert_array_equal(result, expected_result)
def test_get_counts_array_from_hdu2(self):
# test passing an astropy hdu
expected_result = TestSpectrum.adhdu.data
sp = spectro.Spectrum(TestSpectrum.apfhdu, wunit=u.Angstrom)
result = sp.counts
assert_array_equal(result, expected_result)
def specplot(hdulist,
spec_ext,
var_ext,
annotations=None,
ylimits=None,
output_plot_name=None):
"""
Plot a 1-D spectrum and annotates.
This is utility function to do most of the work for plotting a spectrum.
The axes are labeled automatically using information from the extension
in the input HDU list.
Lines can be annotated provide that the list is defined
in the spectro.py module. A redshift can be applied to the line list.
The user can reset the y-axis limits. Work in progress is the ability
to draw band limits to identify where the signal is not good due to
atmospheric absorption. The plot can be saved as PNG.
Parameters
----------
hdulist : HDUList
hdulist can be either from astropy.io.fits or from pyfits.
Unfortunately the two versions are giving incompatible hdulists.
The code here uses a workaround that allows both, until the
old pyfits is deprecated.
spec_ext : int or str
The extension that contains the spectrum. The extension identifier
can be an int or a string representation with extname and extver,
eg. 'sci,1'.
var_ext : int or str
The extension that contains the variance. The extension identifier
can be an int or a string representation with extname and extver,
eg. 'var,1'.
annotations : SpecPlotAnnotation, optional
An instance of SpecPlotAnnotation containing annotation information
for plot, eg. title, line list names, redshift, whether to draw the
band limits.
ylimits : list, optional
The lower and upper limits for the y-axis. Individual values
are int or float.
output_plot_name : str, optional
If set, the plot will be saved to a file of that name. The .png
filename extension is optional, it will be added appropriately.
The plot will be display on screen as well as being saved to disk.
Returns
_______
No return values. A plot is produced on screen, and it can be saved
to disk.
Raises
______
NotImplementedError
Some features are not yet implemented.
See Also
--------
splot : An app that uses this function and is callable from the shell.
Examples
--------
TODO : add examples for specplot()
"""
# ----- Get the information needed to produce the plot
# Get the spectrum.
# print 'debug - specplot - Extension parsed as:',
# get_valid_extension(spec_ext)
# print 'debug - specplot - The hdulist is:', hdulist.info()
spectrum = spectro.Spectrum(hdulist[get_valid_extension(spec_ext)])
if var_ext is not None:
error = spectro.Spectrum(hdulist[get_valid_extension(var_ext)])
error.counts = np.sqrt(error.counts)
else:
error = None
# To simplify the rest of the scripts, create an instance of
# SpecPlotAnnotations that has everything set to False if no
# annotations were defined.
if annotations is None:
# create an instance that has everything set to False
annotations = SpecPlotAnnotations()
# Get the line list and redshift it.
if annotations.annotate_lines:
linelist = spectro.LineList(annotations.line_list_name,
annotations.redshift)
else:
linelist = None
# Get the band limits.
if annotations.draw_bands_limits:
errmsg = 'Draw bands limits feature not implemented.'
raise NotImplementedError(errmsg)
# bandlist = spectro.BandList(spectrum.wlen[0], spectrum.wlen[-1])
# ----- START PLOTTING
#
# Plot the spectrum and set y-axis limits
plot = plottools.SpPlot(title=annotations.title)
plot.plot_spectrum(spectrum)
if error is not None:
plot.plot_spectrum(error, color='g')
if ylimits is not None:
plot.adjust_ylimits(ylimits[0], ylimits[1])
# Annotate the lines.
if linelist is not None:
lines_to_plot = []
for line in linelist.lines:
lines_to_plot.append(
(line.obswlen.to(spectrum.wunit).value, line.name))
plot.annotate_lines(lines_to_plot)
# Draw the band limits
# if annotations.draw_band_limits:
# plot.draw_band_limits()
# Save the plot to disk.
if output_plot_name is not None:
plot.write_png(output_plot_name)
return
