Ejemplo n.º 1
0
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
Ejemplo n.º 2
0
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("")