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