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_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_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