Star formation for Gadget

from pylab import *
from scipy import integrate
from yt.mods import *
import yt.visualization.eps_writer as eps
from yt.analysis_modules.star_analysis.api import StarFormationRate

#
# Cosmology Stuff
#

# Standard cosmology params
# Comes from Planck 2015. Table 4  last column of:
# http://arxiv.org/pdf/1502.01589v2.pdf
Omega_m = 0.3089
Omega_L = 0.6911
H_0 = 67.74
c = 3.0e5  # km/s
h = H_0/100.
Omega_r = 4.165e-5/(h*h)

# Calculate E(z)
def E(z):
    return sqrt(Omega_r*(1.0+z)**4+Omega_m*(1.0+z)**3+Omega_L)

# Get t from z(t)
# Eq. 3 & 30 of http://arxiv.org/pdf/astro-ph/9905116v4.pdf
# Return value in years
def tl(z):
    f = lambda x: 1.0/((1.0+x)*E(x))
    return integrate.quad(f,z,inf,full_output=1)[0]*9.78e9/h/1.0e6
vtl = vectorize(tl)                             # Vectorized version

#
# YT Stuff
#

# Load in Gadget snapshot
ds = load('snapshot_024.hdf5')
sp = ds.all_data()

# Get star masses in Msun
mass = sp[('PartType4', 'Masses')].in_units('Msun')

#  Get creation time of each star particle in Myr
act = sp[('PartType4','StellarFormationTime')]  # Gadget creation time in terms of a(t)
zct = array(1.0/act - 1.0)                      # Convert a(t) to z(t)
ct = vtl(zct)                                   # Convert z(t) to t in Myr
ct = YTArray(ct,'Myr')                          # Assign to YTArray with proper units

# Get star formation rate with known masses and creation time
sfr = StarFormationRate(ds, data_source=sp,star_mass=mass,star_creation_time=ct)
sfr.write_out(name="StarFormationRate.out")