from readgadget import *
import numpy as np
import os,sys
sp = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0,'%s/../../classes' % sp)
sys.path.insert(0,'%s/../../particlesearch' % sp)
from halo import Halo
def read_rockstar(obj):
[docs] obj.logOutput('READING ROCKSTAR')
RSBIN = obj.locs.RS_BINFILE
RSPARENT = obj.locs.RS_PARENTFILE
f = open(RSPARENT,'r')
nlines = len(f.readlines())
f.close()
comments = 0
ncolumns = 0
f = open(RSPARENT,'r')
for i in range(0,nlines):
line = f.readline()
if line[0] == '#':
comments += 1
else:
ncolumns = len(line.split())
break
f.close()
num_halos = nlines - comments
obj.num_halos = num_halos
if num_halos == 0:
return 1
obj.logOutput('reading halos and creating objects...',indent=1)
## PYTHON PyGadgetReader
hdata = readrockstargalaxies(RSBIN,'halos',**obj.rto.pygr_OPTS)
HID = hdata['id']
hcount = hdata['num_p']
halo_mass = hdata['m'].astype(np.float64)
halo_r = hdata['r'].astype(np.float64)
hcm_x = hdata['pos'][:,0].astype(np.float64)
hcm_y = hdata['pos'][:,1].astype(np.float64)
hcm_z = hdata['pos'][:,2].astype(np.float64)
h_Jx = hdata['J'][:,0].astype(np.float64)
h_Jy = hdata['J'][:,1].astype(np.float64)
h_Jz = hdata['J'][:,2].astype(np.float64)
spin = hdata['spin'].astype(np.float64)
energy = hdata['energy'].astype(np.float64)
vmax = hdata['vmax'].astype(np.float64)
rvmax = hdata['rvmax'].astype(np.float64)
vrms = hdata['vrms'].astype(np.float64)
parentHID,parenthparent = np.genfromtxt(RSPARENT,usecols=(0,ncolumns-1),unpack=True)
bin_sortargs = np.argsort(HID)
parent_sortargs = np.argsort(parentHID)
hparent = np.zeros(len(HID),dtype=np.int32)
hparent.fill(-2)
# check for single objects
# convert to array
if isinstance(HID,float):
HID = np.asarray([HID])
hcount = np.asarray([hcount])
halo_mass = np.asarray([halo_mass])
halo_r = np.asarray([halo_r])
hcm_x = np.asarray([hcm_x])
hcm_y = np.asarray([hcm_y])
hcm_z = np.asarray([hcm_z])
h_Jx = np.asarray([h_Jx])
h_Jy = np.asarray([h_Jy])
h_Jz = np.asarray([h_Jz])
hparent = np.asarray([hparent])
bin_sortargs = np.asarray([bin_sortargs])
spin = np.asarray([spin])
energy = np.asarray([energy])
vmax = np.asarray([vmax])
rvmax = np.asarray([rvmax])
vrms = np.asarray([vrms])
if isinstance(parentHID,float):
parentHID = np.asarray([parentHID])
parenthparent = np.asarray([parenthparent])
parent_sortargs = np.asarray([parent_sortargs])
for i in range(0,len(HID)):
if HID[bin_sortargs[i]] == parentHID[parent_sortargs[i]]:
hparent[bin_sortargs[i]] = parenthparent[parent_sortargs[i]]
parentHID = None
parenthparent = None
#bin_sortargs = None
si = bin_sortargs
## order halos via HID + unit conversion
HID = HID[si]
hcount = hcount[si]
halo_mass = halo_mass[si] / obj.h
halo_r = halo_r[si] / obj.h
hcm_x = hcm_x[si] / obj.h * 1000.
hcm_y = hcm_y[si] / obj.h * 1000.
hcm_z = hcm_z[si] / obj.h * 1000.
h_Jx = h_Jx[si] * 1000./obj.h**2
h_Jy = h_Jy[si] * 1000./obj.h**2
h_Jz = h_Jz[si] * 1000./obj.h**2
hparent = hparent[si]
spin = spin[si]
energy = energy[si]
vmax = vmax[si]
rvmax = rvmax[si] / obj.h
vrms = vrms[si]
## build dictionary of halo IDs to map to indexes
obj.halo_index_translator = {}
for i in range(0,len(HID)):
obj.halo_index_translator[int(HID[i])] = i
## tranlate from halo IDs to halo_indexes
for i in range(0,len(hparent)):
if hparent[i] > -1:
hparent[i] = obj.halo_index_translator[hparent[i]]
obj.num_halos = num_halos
for i in range(0,len(hcm_x)):
obj.halos.append(Halo(hcm_x[i],hcm_y[i],hcm_z[i],
halo_r[i],halo_mass[i],hparent[i],
h_Jx[i],h_Jy[i],h_Jz[i]))
obj.halos[i].particle_count = hcount[i]
obj.halos[i].sub_halo_list = []
obj.halos[i].sub_halos = []
obj.halos[i].spin = spin[i]
obj.halos[i].energy = energy[i]
obj.halos[i].vmax = vmax[i]
obj.halos[i].rvmax = rvmax[i]
obj.halos[i].vrms = vrms[i]
##############
## PARTICLE ##
##############
obj.logOutput('sorting RS particles...',indent=1)
hdata = readrockstargalaxies(RSBIN,'particles',**obj.rto.pygr_OPTS)
RSDMPIDS = hdata[:,0]
RSHALOIDS = hdata[:,1]
hdata = None
RSHALOIDS = RSHALOIDS.astype(np.int32)
glist = np.zeros(obj.nparts[0],dtype=np.int32)
dmlist = np.zeros(obj.nparts[1],dtype=np.int32)
slist = np.zeros(obj.nparts[4],dtype=np.int32)
nlowres = 0
if obj.rto.zoomSim:
for i in range(0,6):
if (1<<i) & (obj.rto.ZOOM_LOWRES):
nlowres += obj.nparts[i]
lrlist = np.zeros(nlowres,dtype=np.int32)
lrlist.fill(-1)
else:
lrlist = np.zeros(0,dtype=np.int32)
glist.fill(-1)
dmlist.fill(-1)
slist.fill(-1)
## CYTHON ME?
for i in range(0,len(RSDMPIDS)):
index = RSDMPIDS[i]
if index < obj.nparts[0]:
glist[index] = RSHALOIDS[i]
elif index >= obj.nparts[0] and index < obj.nparts[0]+obj.nparts[1]:
index -= obj.nparts[0]
dmlist[index] = RSHALOIDS[i]
elif index >= obj.nparts[0]+obj.nparts[1] and index < obj.nparts[0]+obj.nparts[1]+obj.nparts[4]:
index -= obj.nparts[0] + obj.nparts[1]
slist[index] = RSHALOIDS[i]
elif obj.rto.zoomSim:
index -= obj.nparts[0] + obj.nparts[1] + obj.nparts[4]
lrlist[index] = RSHALOIDS[i]
obj.particleLists.halo_glist = glist
obj.particleLists.halo_dmlist = dmlist
obj.particleLists.halo_slist = slist
obj.particleLists.halo_lrlist = lrlist
glist = None
dmlist = None
slist = None
lrlist = None
HID = None
hcount = None
halo_mass = None
halo_r = None
hcm_x = None
hcm_y = None
hcm_z = None
h_Jx = None
h_Jy = None
h_Jz = None
hparent = None
spin = None
energy = None
vmax = None
rvmax = None
vrms = None
RSDMPIDS = None
RSHALOIDS = None
return 0