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("")