Source

IRATE-format / irate / rockstar.py

#!/bin/python

from numpy import fromstring,empty
import h5py

def read_rockstar_ascii(fname):
	"""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
	halos = []
	showflag = True
	try:
		for line in f:
			if linen == 0:
				columns = line.lstrip('#')
				columns.split(' ')
				linen = 1
				continue
			if line[0] == "#":
				if showflag:
                                    print "Haven't yet written the code to go through units and statistics.  Hope to get to it sometime."
				    showflag = False
				continue
			halos.append(fromstring(line,dtype='float',sep=' '))
		
	finally:  f.close()
	
	nhalos = len(halos)
	
	print "Read {0} columns of data for {1} halos.  Now collimating.".format(len(halos[0]),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','Mvir','mbound_vir','Rvir','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]
		

def write_halodata(halocols,headers,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')
    head = fileheader.create_group('Rockstar')

    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."
	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.