IRATE-format / irate /

Full commit

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 = []
        for line in f:
            if linen == 0:
                columns = line.lstrip('#')
                columns = columns.split()
                linen = 1
            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
                        #What to do if it's a line about units
                        thiscom = thiscom.lstrip('Units: ').split(' in ')
                        mycom = [thiscom[0]+' Units',thiscom[1]]
            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)):

def read_rockstar_binary(fname,verbose=False):
    print "Sorry, haven't gotten around to writing this yet.  Please use an ASCII file instead."