def test_save_load_spectrum_fits():
"""
Test that we can save and load spectra in the FITS format
"""
dirpath = tempfile.mkdtemp()
filename = os.path.join(dirpath, 'ray.h5')
ray = make_onezone_ray(column_densities={'H_p0_number_density': 1e21},
filename=filename)
sg = SpectrumGenerator(lambda_min=1200, lambda_max=1300, dlambda=0.5)
sg.make_spectrum(ray, lines=['Ly a'])
sg.save_spectrum(os.path.join(dirpath, 'spec.fits'))
del (sg)
sg = load_spectrum(os.path.join(dirpath, 'spec.fits'))
sg.plot_spectrum(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 verify(save=False):
"""
Verify that the bulk of Trident's functionality is working. First, it
ensures that the user has a configuration file and ion table datafile,
and creates/downloads these files if they do not exist. Next, it
creates a single-cell grid-based dataset in memory, generates a ray
by sending a sightline through that dataset, then makes a spectrum from
the ray object. It saves all data to a tempdir before deleting it.
**Parameters**
:save: boolean, optional
By default, verify saves all of its outputs to a temporary directory
and then removes it upon completion. If you would like to see the
resulting data from verify(), set this to be True and it will save
a light ray, and raw and processed spectra in the current working
directory.
Default: False
**Example**
Verify Trident works.
>>> import trident
>>> trident.verify()
"""
parse_config()
from trident.spectrum_generator import SpectrumGenerator
from trident.ray_generator import make_simple_ray
print("")
print("Creating single-cell dataset")
print("----------------------------")
print("")
try:
ds = make_onezone_dataset()
except BaseException:
print("Failed to create single-cell dataset")
raise
print("")
print("Creating ray object through single-cell dataset")
print("-----------------------------------------------")
print("")
if save:
tempdir = '.'
else:
tempdir = tempfile.mkdtemp()
try:
ray = make_simple_ray(ds,
start_position=ds.domain_left_edge,
end_position=ds.domain_right_edge,
data_filename=os.path.join(tempdir, 'ray.h5'),
fields=['density', 'temperature', 'metallicity'])
except BaseException:
print("Failed to create ray object")
raise
print("")
print("Create spectrum with Lyman alpha, Mg II, and O VI lines")
print("-------------------------------------------------------")
print("")
sg = SpectrumGenerator('COS')
sg.make_spectrum(ray, lines=['Ly a', 'Mg II', 'O VI'])
sg.save_spectrum(os.path.join(tempdir, 'spec_raw.h5'))
sg.plot_spectrum(os.path.join(tempdir, 'spec_raw.png'))
# Test other post processing of the spectrum
sg.add_qso_spectrum()
sg.add_milky_way_foreground()
sg.apply_lsf()
sg.add_gaussian_noise(30)
sg.save_spectrum(os.path.join(tempdir, 'spec_final.h5'))
sg.plot_spectrum(os.path.join(tempdir, 'spec_final.png'))
if not save:
print("Removing all temporary data files...")
shutil.rmtree(tempdir)
print("")
print("Congratulations, you have verified that Trident is installed correctly.")
print("Now let's science!")
print("")
