ramses tform derived quantity using part2birth

ramses_tform.py

import pynbody

part2birth = '/home/itp/roskar/ramses/galaxy_formation/part2birth'

@pynbody.ramses.RamsesSnap.derived_quantity
def tform(self) : 
    """Generate a tform array that is compatible with pynbody unit system 
    and will allow one to get the ages of stars. 
    
    make sure the 'part2birth' above corresponds to the location of the executable on your machine
    """
    
    from scipy.io.numpyio import fread

    top = self
    while hasattr(top,'base') : top = self.base

    ncpu = top._info['ncpu']
    nstar = len(top.s)

    top.s['tform'] = -1.0
    done = 0
    for i in range(ncpu) : 
        try : 
            f = open('%s/birth/birth_%s.out%05d'%(top.filename[:-12],top._timestep_id,i+1))
        except IOError : 
            import os
            
            os.system("cd %s; mkdir birth; %s -inp output_%s; cd .."%(top.filename[:-12],part2birth,top._timestep_id))
            f = open('%s/birth/birth_%s.out%05d'%(top.filename[:-12],top._timestep_id,i+1))

        n = fread(f,1,'i')
        n /= 8
        ages = fread(f,n,'d')
        new = np.where(ages > 0)[0]
        top.s['tform'][done:done+len(new)] = ages[new]
        done += len(new)
        f.close()
    top.s['tform'].units = 'Gyr'

    return self.s['tform']