def test_input_types():
"""
Test that spectra can be generated from ray file, dataset, or
data container.
"""
dirpath = tempfile.mkdtemp()
filename = os.path.join(dirpath, 'ray.h5')
ray = make_onezone_ray(filename=filename)
sg = SpectrumGenerator(lambda_min=1200, lambda_max=1300, dlambda=0.5)
spectra = []
# from file
sg.make_spectrum(filename, lines=['H I'], ly_continuum=False)
spectra.append(sg.flux_field[:])
# from dataset
sg.make_spectrum(ray, lines=['H I'], ly_continuum=False)
spectra.append(sg.flux_field[:])
assert (spectra[0] == spectra[1]).all()
# from data container
sg.make_spectrum(ray.all_data(), lines=['H I'], ly_continuum=False)
spectra.append(sg.flux_field[:])
assert (spectra[0] == spectra[2]).all()
plot_spectrum(sg.lambda_field,
spectra,
filename=os.path.join(dirpath, 'spec.png'))
shutil.rmtree(dirpath)
def test_plot_multiple_spectra():
"""
Test that multiple spectra can be plotted on top of each other. Example
from plot_spectrum docstrings.
"""
dirpath = tempfile.mkdtemp()
filename = os.path.join(dirpath, 'ray.h5')
ray = make_onezone_ray(column_densities={'H_p0_number_density': 1e21},
filename=filename)
sg_final = SpectrumGenerator(lambda_min=1200, lambda_max=1300, dlambda=0.5)
sg_final.make_spectrum(ray, lines=['Ly a'])
sg_final.save_spectrum(os.path.join(dirpath, 'spec_raw.h5'))
sg_final.add_gaussian_noise(10)
sg_raw = load_spectrum(os.path.join(dirpath, 'spec_raw.h5'))
plot_spectrum([sg_raw.lambda_field, sg_final.lambda_field],
[sg_raw.flux_field, sg_final.flux_field],
stagger=0,
step=[False, True],
label=['Raw', 'Noisy'],
filename=os.path.join(dirpath, 'raw_and_noise.png'))
shutil.rmtree(dirpath)
assert True
def plot_spectrum(self,
filename="spectrum.png",
lambda_limits=None,
flux_limits=None,
step=False,
title=None,
label=None,
figsize=None,
features=None,
axis_labels=None):
"""
Plot the current spectrum and save to disk.
This is a convenience method that wraps the
:class:`~trident.plot_spectrum` standalone function for use with the
data from the :class:`~trident.SpectrumGenerator` itself.
**Parameters**
:filename: string, optional
Output filename of the plotted spectrum. Will be a png file.
Default: 'spectrum.png'
:lambda_limits: tuple or list of floats, optional
The minimum and maximum of the lambda range (x-axis) for the plot
in angstroms. If specified as None, will use whole lambda range
of spectrum. Example: (1200, 1400) for 1200-1400 Angstroms.
Default: None
:flux_limits: tuple or list of floats, optional
The minimum and maximum of the flux range (y-axis) for the plot.
If specified as None, limits are automatically from
[0, 1.1*max(flux)]. Example: (0, 1) for normal flux range before
postprocessing.
Default: None
:step: boolean, optional
Plot the spectrum as a series of step functions. Appropriate for
plotting processed and noisy data.
:title: string, optional
Optional title for plot
Default: None
:label: string, optional
Label for spectrum to be plotted. Will automatically trigger a
legend to be generated.
Default: None
:features: dict, optional
Include vertical lines with labels to represent certain spectral
features. Each entry in the dictionary consists of a key string to
be overplot and the value float as to where in wavelength space it
will be plot as a vertical line with the corresponding label.
Example: features={'Ly a' : 1216, 'Ly b' : 1026}
Default: None
:axis_labels: tuple of strings, optional
Optionally set the axis labels directly. If set to None, defaults to
('Wavelength [$\\rm\\AA$]', 'Relative Flux').
Default: None
**Example**
Create a one-zone ray, and generate a COS spectrum from that ray. Plot
the resulting spectrum highlighting the Lyman alpha feature.
>>> import trident
>>> ray = trident.make_onezone_ray()
>>> sg = trident.SpectrumGenerator('COS')
>>> sg.make_spectrum(ray)
>>> sg.plot_spectrum('spec_raw.png', features={'Ly a' : 1216})
"""
plot_spectrum(self.lambda_field,
self.flux_field,
filename=filename,
lambda_limits=lambda_limits,
flux_limits=flux_limits,
title=title,
label=label,
figsize=figsize,
step=step,
features=features,
axis_labels=axis_labels)