from yt.mods import *
from yt.analysis_modules.star_analysis.api import *
from yt.analysis_modules.halo_finding.api import *

pf = load("output_0108")

#loading halos
haloes = LoadHaloes(pf, "DD0108")

i=0
for halo in haloes:
  # Get the pertinent arrays.
  rv = halo.virial_radius()
  center = halo.center_of_mass()
  #sp = halo.get_sphere()
  sp = pf.h.sphere(center,rv)
  if i >= 20: break
  # This puts the particle data for *all* the particles in the region re
  # into the arrays sm and ct.
  sm = sp["ParticleMassMsun"]
  ct = sp["creation_time"]
  # First pick out only stars.
  stars = (ct > 0)
  ct = ct[stars]
  sm = sm[stars]
  # Pick out only old stars using Numpy array fancy indexing.
  # 100 is a time in code units.
  #sm_old = sm[ct < 100]
  #ct_old = ct[ct < 100]
  sfr = StarFormationRate(pf, star_mass=sm, star_creation_time=ct, \
    volume=sp.volume('mpc'))
  sfr.write_out(name="StarFormationRate_virial_%02i.out"%i)
  i=i+1