Commits

Shea Garrison-Kimmel  committed b61ef2f

Initial import of rockster2irate. Currently only has an ascii option

  • Participants
  • Parent commits f208c5c

Comments (0)

Files changed (2)

File irate/rockstar.py

+#!/bin/python
+
+from numpy import fromstring
+import h5py
+
+def read_rockstar_ascii(fname,verbose=False):
+	f = open('fname')
+	linen = 0
+	halos = []
+	try:
+		for line in f:
+			if linen == 0:
+				columns = line.lstrip('#')
+				columns.split(' ')
+				linen = 1
+				continue
+			if line[0] == "#":
+				print "Haven't yet written the code to go through units and statistics.  Hope to get to it sometime."
+				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])
+	J = 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]]
+		J[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,j,coldata[17],coldata[18],coldata[19],coldata[20]]
+	
+	return [mycoldata,myheaders]
+		
+
+def write_halodata(halocols,headers,outname):
+    #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)

File scripts/rockstar2irate

+#!/bin/python
+
+import sys,os
+from optparse import OptionParser
+from irate.rockstar import read_rockstar_ascii,read_rockstar_binary,write_halodata
+import h5py
+
+usage = "usage: %prog <input file> <format (-a or -b)> [output file (defaults to input file + '-irate.hdf5')]"
+
+parser = OptionParser(usage=usage)
+
+parser.add_option("-b",action="store_true",dest="binary",help="Specifies that input file is a Rockstar binary file.  Defaut = %default because not currently supported",default=False)
+parser.add_option("-a",action="store_false",dest="binary",help="Specifies that the input file is a Rockstar ascii file.  Default = True, because currently the only supported format.")
+parser.add_option('-o',action='store_true',dest='overwrite',help='Overwrite existing IRATE file, if one exists.  Default = %default',default=False)
+parser.add_option('-q',action='store_true',dest='quiet',help='Run with minimal output.  Default = %default',default=False)
+
+(ops,args) = parser.parse_args()
+
+if len(sys.argv) == 1:
+	parser.print_help()
+	sys.exit(1337)
+
+if len(args) == 1:
+	inname = args[0]
+	outname = inname+'-irate.hdf5'
+elif len(args) == 2:
+	inname = args[0]
+	outname = args[1]
+else:
+	parser.print_help()
+	sys.exit(1337)
+
+if not outname.endswith('.hdf5'):  outname = outname + ".hdf5"
+
+if os.path.isfile(outname):
+    if ops.overwrite:
+        print "Overwriting "+outname+" with new file."
+        os.remove(outname)
+    else:
+        print "Output file already exists; please choose a different output file or run with -o."
+        sys.exit(1337)
+verbose = not ops.quiet
+
+if binary:
+	print "Reading data from Rockstar binary file "+inname
+	[halodata,haloheaders] = read_rockstar_binary(inname,verbose)
+else:
+	print "Reading data from Rockstar ascii file "+inname
+	[halodata,haloheaders] = read_rockstar_ascii(inname,verbose)
+	
+print "Writing data to "+outname
+write_halodata(halodata,haloheaders,outname)
+print "Successfully created IRATE file "+outname