Source

IRATE-format / irate / rockstar.py

#!/bin/python

from numpy import fromstring,empty,loadtxt
import h5py

def read_rockstar_ascii(fname,verbose):
    """Reads a Rockstar ASCII file and collimates it.  It should also read the
        parameters that Rockstar gives you at the start of the ASCII file, but
        I haven't had time to do that yet.  Returned are the columns in the
        file and the comments that head the columns.
    :param fname:  The input ASCII file to read and collimate.
    """
    f = open(fname)
    linen = 0
    comments = []
    try:
        for line in f:
            if linen == 0:
                columns = line.lstrip('#')
                columns = columns.split()
                linen = 1
                continue
            elif line[0] == "#":
                thisl = line.lstrip('#').rstrip('\n')
                thisl = thisl.split(';')                    
                for j in range(len(thisl)):
                    thiscom = thisl[j]
                    if not thiscom.startswith('Units'):
                        #So, not a units line
                        thiscom = thiscom.split(':')
                        if len(thiscom) == 1:   #Then it didn't split on a :
                            thiscom = thiscom[0]
                            thiscom = thiscom.split('=')
                            for k in range(len(thiscom)):
                                thiscom[k] = thiscom[k].strip(' ')
                        #Ok, now thiscom should be 2 things for sure
                        comments.append(thiscom)
                    else:
                        #What to do if it's a line about units
                        thiscom = thiscom.lstrip('Units: ').split(' in ')
                        mycom = [thiscom[0]+' Units',thiscom[1]]
                        comments.append(mycom)
                        
            else:  break
        
    finally:  f.close()
    
    odens = columns[2].lstrip('m')       #The overdensity critera used
    
    coldata = loadtxt(fname,unpack=True,comments='#')
    nhalos = len(coldata[0])
    
    if verbose:  print "Read {0} columns of data for {1} halos.  Now collimating.".format(len(columns),nhalos)
    
    #Now to do the collimation; each column should be a property, except for the few that are vectors
    #coldata = empty([len(halos[0]),nhalos])
    #for j in range(nhalos):
    #    for i in range(len(halos[0])):
    #        coldata[i][j] = halos[j][i]
    
    cen = empty([nhalos,3])
    vel = empty([nhalos,3])
    angmom = empty([nhalos,3])
    
    #id num_p mvir mbound_vir rvir vmax rvmax vrms x y z vx vy vz Jx Jy Jz E Spin PosUncertainty VelUncertainty
    myheaders = ['id','npart','M'+odens,'mbound_'+odens,'R'+odens,'Vmax','Rmax','vrms','Center','Velocity','J','E','Spin','PositionUncertainty','VelocityUncertainty']
    #Making some capitalization changes to match with AHF for ease of use
    
    for i in range(len(coldata[0])):
        cen[i] = [coldata[8][i],coldata[9][i],coldata[10][i]]
        vel[i] = [coldata[11][i],coldata[12][i],coldata[13][i]]
        angmom[i] = [coldata[14][i],coldata[15][i],coldata[16][i]]
    
    mycoldata = [coldata[0],coldata[1],coldata[2],coldata[3],coldata[4],coldata[5],coldata[6],coldata[7],cen,vel,angmom,coldata[17],coldata[18],coldata[19],coldata[20]]
    
    return [mycoldata,myheaders,comments]
        

def write_halodata(halocols,headers,comments,outname):
    """Writes the data from a Rockstar file into an IRATE file.  The
    strings given in headers determines the names given to the datasets in
    the IRATE file.
    
    :param halocols:  The data to be written into the IRATE file.  Each entry
        in the list should be a column from the ASCII file, which turns into
        a dataset in the IRATE file.
    :param headers:  The names to be given to the datasets.  Ought to be in the
        same order as halocols; otherwise, output file won't make any sense.
    :param outname:  The IRATE file to create.
    """
    #Write colimated data with headers as an IRATE file
    f = h5py.File(outname,'a')
    fileheader = f.create_group('CatalogHeader')
    header = fileheader.create_group('RockstarHeader')
    for i in range(len(comments)):
        header.attrs[comments[i][0]] = comments[i][1]

    cat = f.create_group('Catalog')
    #Now to create the datasets:
    for i in range(len(headers)):
        cat.create_dataset(headers[i],data=halocols[i])
    f.close()
    


def read_rockstar_binary(fname,verbose=False):
    print "Sorry, haven't gotten around to writing this yet.  Please use an ASCII file instead."
    sys.exit(1337)
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.