create_spectrum.py

from __future__ import print_function
import numpy as np
from yt.analysis_modules.absorption_spectrum.api import AbsorptionSpectrum
import yt
import h5py

sim_dir = "./"
sim_dim = np.array([4.0, 2.0, 0.0])

def main():
    # Load data
    ds = yt.load(sim_dir + "DD0030/DD0030")

    # Define ray. Start and end positions are normalized units (?).
    start_pos, end_pos = np.array([0.0, 1.0, 0.0]), np.array([4.0, 1.0, 0.0])
    ray_vec = end_pos - start_pos
    ray_len = np.linalg.norm(ray_vec)
    ray = ds.ray(start_pos, end_pos)

    # Grab relevant information from Ray
    density = ray["HI_Density"]
    temperature = ray["Temperature"]
    # Assuming the length of the simulation box is 400.0 EnzoUnits.
    dl = ray["dts"] * ray_len
    redshift = np.zeros(density.shape)
    # Calculate line-of-sight velocity
    los_velocity = np.dot(ray_vec / ray_len,
                          np.array([ray["x-velocity"],
                                    ray["y-velocity"],
                                    np.zeros(ray["x-velocity"].shape)]))

    # Test for reasonable limits
    print("HI_Density:  {}, {}".format(density.min(), density.max()))
    print("Temperature: {}, {}".format(temperature.min(), temperature.max()))
    print("dl:          {}, {}".format(dl.min(), dl.max()))

    # Write values to an hdf5 file
    h5f = h5py.File('lightray.h5', 'w')
    h5f['density'] = density
    h5f['temperature'] = temperature
    h5f['dl'] = dl
    h5f['redshift'] = redshift
    h5f['los_velocity'] = los_velocity
    h5f.close()

    # Create our absorption spectrum object
    sp = AbsorptionSpectrum(900.0, 1800.0, 10000)

    # Create absorption line
    my_label = 'HI Lya'
    field = 'density'
    wavelength = 1215.6700 # Angstroms
    f_value = 4.164E-01
    gamma = 6.265e+08
    mass = 1.00794
    # Add the line
    sp.add_line(my_label, field, wavelength, f_value, gamma, mass, label_threshold=1.e10)

    # Make the spectrum, returning the wave length and the flux
    wavelength, flux = sp.make_spectrum('lightray.h5', output_file='spectrum.h5',
                                        line_list_file='lines.txt',
                                        use_peculiar_velocity=False)

    # Check to see if there really are any features in the spectrum
    print("min: {}, max: {}".format(min(flux), max(flux)))


if __name__ == '__main__':
    main()