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_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_create_spectrum_H_O_lines():
"""
Test that we can create a basic spectrum with hydrogen and oxygen lines
"""
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)
sg.make_spectrum(ray, lines=['H', 'O'])
sg.plot_spectrum(os.path.join(dirpath, 'spec.png'))
shutil.rmtree(dirpath)
def test_create_spectrum_all_lines():
"""
Test that we can create a basic spectrum with all available lines
"""
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)
sg.make_spectrum(ray, lines='all')
sg.plot_spectrum(os.path.join(dirpath, 'spec.png'))
shutil.rmtree(dirpath)
assert True
def test_create_spectrum_H_lines_no_continuum():
"""
Test that we can create a basic spectrum with H I lines but no Lyman
continuum
"""
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)
sg.make_spectrum(ray, lines=['H I'], ly_continuum=False)
sg.plot_spectrum(os.path.join(dirpath, 'spec.png'))
shutil.rmtree(dirpath)
def test_make_onezone_ray():
"""
Tests the make_onezone_ray infrastructure by creating a ray and making a
spectrum from it.
"""
dirpath = tempfile.mkdtemp()
ray_filename = os.path.join(dirpath, 'ray.h5')
image_filename = os.path.join(dirpath, 'spec.png')
ray = make_onezone_ray(column_densities={'H_p0_number_density': 1e21},
filename=ray_filename)
sg_final = SpectrumGenerator(lambda_min=1200, lambda_max=1300, dlambda=0.5)
sg_final.make_spectrum(ray, lines=['Ly a'])
sg_final.plot_spectrum(image_filename)
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 test_make_onezone_dataset():
"""
Tests the make_onezone_dataset infrastructure by generating a one_zone
dataset and then creating a ray and spectrum from it.
"""
dirpath = tempfile.mkdtemp()
ray_filename = os.path.join(dirpath, 'ray.h5')
image_filename = os.path.join(dirpath, 'spec.png')
ds = make_onezone_dataset()
ray = make_simple_ray(ds,
start_position=ds.domain_left_edge,
end_position=ds.domain_right_edge,
fields=['density', 'temperature', 'metallicity'],
data_filename=ray_filename)
sg = SpectrumGenerator('COS')
sg.make_spectrum(ray)
sg.plot_spectrum(image_filename)
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("")