示例#1
0
def read_block_vars(fname, ptype):
    if ptype == 0: print "Processing GAS PARTICLES"
    if ptype == 1: print "Processing DM PARTICLES"
    if ptype == 4: print "Processing STARS+WIND PARTICLES"
    if ptype == 5: print "Processing BH PARTICLES"
    # position
    rrr = rs.read_block(fname, "POS ", parttype=ptype)
    # potential
    #pot = rs.read_block(fname, "POT ", parttype = ptype)
    mmm = rs.read_block(fname, "MASS", parttype=ptype)
    return rrr, mmm
示例#2
0
def get_disk_grid_3D(disk, NR, Nphi, Box, **kwargs):

    if (disk.double == 0):
        dt = "float32"
    if (disk.double == 1):
        dt = "float64"

    mesh_option = kwargs.get("from_file")

    #Ntheta is obtained from the values of NR and Nphi

    R = sample_radial_positions(disk, NR)
    phi = sample_azimuthal_angles(disk, Nphi)

    Ntheta1 = disk.aspect_ratio / 2.0 / math.sqrt(2.0 * math.pi) * Nphi
    Ntheta0 = 32
    theta = sample_polar_angles(disk, Ntheta0, Ntheta1)
    Ntheta = theta.shape[0]

    pos = np.zeros([NR * Nphi * Ntheta, 3], dtype=dt)

    if (mesh_option == None):
        offsetX = Box.X / 2.0
        offsetY = Box.Y / 2.0
        offsetZ = Box.Z / 2.0
        ip = 0
        for i in range(0, NR):
            for j in range(0, Nphi):
                for k in range(0, Ntheta):
                    pos[ip,
                        0] = R[i] * math.cos(phi) * math.sin(theta) + offsetX
                    pos[ip,
                        1] = R[i] * math.sin(phi) * math.sin(theta) + offsetY
                    pos[ip, 2] = R[i] * math.cos(theta) + offsetZ
                    ip += 1

    else:
        import readsnapHDF5 as rs
        pos = rs.read_block(mesh_option, "POS ", parttype=0)

    pos[:, 0] = np.where(
        (np.sqrt((pos[:, 0] - BoxSizeX / 2.0)**2 +
                 (pos[:, 1] - BoxSizeY / 2.0)**2) < disk.outer_radius) &
        (np.sqrt((pos[:, 0] - BoxSizeX / 2.0)**2 +
                 (pos[:, 1] - BoxSizeY / 2.0)**2) > disk.inner_radius),
        pos[:, 0], -1)

    ind = np.where(pos[:, 0] >= 0)[0]
    newpos = np.zeros([len(ind), 3], dtype=dt)
    newpos[:, 0] = pos[ind, 0]
    newpos[:, 1] = pos[ind, 1]
    newpos[:, 2] = pos[ind, 2]

    ids = np.arange(1, len(ind) + 1, dtype="int32")

    return newpos, ids
示例#3
0
def find_vt_tracer_ids(snap_num,sub_id,base,scale_factor,gal_radfac):
	print "Finding velocity tracer IDs..."
	cat = readsubfHDF5.subfind_catalog(base, snap_num)
	
	# I assume that the subhalo is the primary subhalo in its group
	sub_list = cat.GroupFirstSub
	grp_id = np.argmin( np.abs(sub_id-sub_list) )

	sub_pos = cat.SubhaloPos[sub_id]
	sub_mass = cat.SubhaloMass[sub_id]
	sub_Rvir = cat.Group_R_Crit200[grp_id]

	gal_rad = sub_Rvir * gal_radfac #* 0.1

	#snapname = base + "snapdir_"+str(snap_num).zfill(3)+"/snap_"+str(snap_num).zfill(3)
	snapname = base + "snap_"+str(snap_num).zfill(3)

	#Find all velocity tracers within 0.1 Rvir of the subhalo
	vt_pos = readsnapHDF5.read_block(snapname,"POS ",parttype=2)
	vt_vel = readsnapHDF5.read_block(snapname,"VEL ",parttype=2)	
	vt_ids = readsnapHDF5.read_block(snapname,"ID  ",parttype=2)
	#readsnapHDF5.list_blocks(snapname+".hdf5")

	vt_x = np.logical_and(vt_pos[:,0] > sub_pos[0]-gal_rad, vt_pos[:,0] < sub_pos[0]+gal_rad)
	vt_y = np.logical_and(vt_pos[:,1] > sub_pos[1]-gal_rad, vt_pos[:,1] < sub_pos[1]+gal_rad)
	vt_z = np.logical_and(vt_pos[:,2] > sub_pos[2]-gal_rad, vt_pos[:,2] < sub_pos[2]+gal_rad)


	vt_ind = np.logical_and(np.logical_and(vt_x,vt_y),vt_z)
	gal_vt_pos = vt_pos[vt_ind] - sub_pos
	#gal_vt_vel = vt_vel[vt_ind]
	

	gal_vt_rad = np.sqrt(gal_vt_pos[:,0]**2 + gal_vt_pos[:,1]**2 + gal_vt_pos[:,2]**2)
	vt_ind2 = gal_vt_rad < gal_rad

	gal_vt_pos = gal_vt_pos[vt_ind2]
	#gal_vt_vel = gal_vt_vel[vt_ind2]
	gal_vt_ids = vt_ids[vt_ind][vt_ind2]
	
	print "Done finding velocity tracer IDs!"
	return gal_vt_ids
示例#4
0
def get_disk_grid_3D(disk,NR, Nphi, Box ,**kwargs):

    if (disk.double == 0):
        dt="float32"
    if (disk.double == 1):
        dt="float64"

    mesh_option = kwargs.get("from_file")

    #Ntheta is obtained from the values of NR and Nphi
    
    R = sample_radial_positions(disk, NR)
    phi = sample_azimuthal_angles(disk,Nphi)

    Ntheta1 = disk.aspect_ratio/2.0/math.sqrt(2.0*math.pi) * Nphi
    Ntheta0 = 32
    theta = sample_polar_angles(disk,Ntheta0,Ntheta1)
    Ntheta = theta.shape[0]
    
    pos=np.zeros([NR*Nphi*Ntheta, 3], dtype=dt)

    if (mesh_option == None):
        offsetX=Box.X/2.0
        offsetY=Box.Y/2.0
        offsetZ=Box.Z/2.0
        ip = 0
        for i in range(0,NR):
            for j in range(0,Nphi):
                for k in range(0,Ntheta):
                    pos[ip,0]= R[i] * math.cos(phi) *  math.sin(theta) +offsetX
                    pos[ip,1]= R[i] * math.sin(phi)*  math.sin(theta) +offsetY
                    pos[ip,2]= R[i] * math.cos(theta)+ offsetZ
                    ip+=1

    else:
        import readsnapHDF5 as rs
        pos = rs.read_block(mesh_option,"POS ",parttype=0)

        
    pos[:,0]=np.where( (np.sqrt((pos[:,0]-BoxSizeX/2.0)**2+(pos[:,1]-BoxSizeY/2.0)**2) < disk.outer_radius) &
                       (np.sqrt((pos[:,0]-BoxSizeX/2.0)**2+(pos[:,1]-BoxSizeY/2.0)**2) > disk.inner_radius), pos[:,0],-1)
   

    ind=np.where(pos[:,0] >= 0)[0]
    newpos=np.zeros([len(ind), 3], dtype=dt)
    newpos[:,0]=pos[ind,0]
    newpos[:,1]=pos[ind,1]
    newpos[:,2]=pos[ind,2]
 
    ids=np.arange(1,len(ind)+1, dtype="int32")

    return newpos,ids
示例#5
0
# (Remove any file that is not a data file for a specific snapshot)


post_setup = time.time()
#===================================================================
#===================================================================
# Get the tracer IDs!
snapname = base + "/snap_"+str(latest_snap).zfill(3)
redshift = readsnapHDF5.snapshot_header(snapname).redshift
scale_factor = 1./(1.+redshift)
gal_radfac =  0.2#scale_factor * 2.5

# Read in from snapshot once:
pre_read = time.time()
# GAS #
gas_pos = readsnapHDF5.read_block(snapname,"POS ",parttype=0)
gas_ids = readsnapHDF5.read_block(snapname,"ID  ",parttype=0)
# STARS #
star_pos = readsnapHDF5.read_block(snapname,"POS ",parttype=4)
star_ids = readsnapHDF5.read_block(snapname,"ID  ",parttype=4)
# TRACERS #
parent_ids = readsnapHDF5.read_block(snapname,"TRPI",parttype=3)
tracer_ids = readsnapHDF5.read_block(snapname,"TFID",parttype=3)
post_read = time.time()

print "reading time: ",post_read-pre_read

i = 0
for sub_id in sub_list:
	i = i+1
示例#6
0
def Illustris_halo(snapshot_root,
                   snapnum,
                   halo_number,
                   TREECOOL_file,
                   fout,
                   ptype=0):

    # find snapshot name and read header
    snapshot = snapshot_root + 'snapdir_%03d/snap_%03d' % (snapnum, snapnum)
    header = rs.snapshot_header(snapshot)
    redshift = header.redshift
    BoxSize = header.boxsize / 1e3  #Mpc/h
    filenum = header.filenum
    Omega_m = header.omega0
    Omega_L = header.omegaL
    h = header.hubble
    massarr = header.massarr * 1e10  #Msun/h

    print '\nBoxSize         = %.1f Mpc/h' % BoxSize
    print 'Number of files = %d' % filenum
    print 'Omega_m         = %.3f' % Omega_m
    print 'Omega_l         = %.3f' % Omega_L
    print 'redshift        = %.3f' % redshift

    # read number of particles in halos and subhalos and number of subhalos
    halos = groupcat.loadHalos(
        snapshot_root,
        snapnum,
        fields=['GroupLenType', 'GroupPos', 'GroupMass'])
    halo_len = halos['GroupLenType'][:, ptype]
    halo_pos = halos['GroupPos'] / 1e3
    halo_mass = halos['GroupMass'] * 1e10
    del halos

    # find where the halo starts and ends in the file
    begin = np.sum(halo_len[:halo_number], dtype=np.int64)
    end = begin + halo_len[halo_number]
    print begin, end

    # do a loop over all snapshot subfiles
    f = h5py.File(fout, 'w')
    pos_f = f.create_dataset('pos', (0, 3), maxshape=(None, 3))
    vel_f = f.create_dataset('vel', (0, 3), maxshape=(None, 3))
    if ptype == 0:
        mass_f = f.create_dataset('mass', (0, ), maxshape=(None, ))
        MHI_f = f.create_dataset('M_HI', (0, ), maxshape=(None, ))
        radii_f = f.create_dataset('radii', (0, ), maxshape=(None, ))
    if ptype == 1:
        radii_f = f.create_dataset('radii', (0, ), maxshape=(None, ))
        mass_f = f.create_dataset('mass_c', (0, ), maxshape=(None, ))

    begin_subfile, particles = 0, 0
    for i in xrange(filenum):

        # find subfile name and read the number of particles in it
        snapshot = snapshot_root + 'snapdir_%03d/snap_%03d.%d' % (snapnum,
                                                                  snapnum, i)
        header = rs.snapshot_header(snapshot)
        npart = header.npart

        end_subfile = begin_subfile + npart[ptype]

        # if all particles in the halo has been read exit loop
        if end < begin_subfile: break

        # if the subfile does not contain any particle move to next subfile
        if begin > end_subfile:
            begin_subfile = end_subfile
            continue

        print 'Working with subfile %03d' % i
        pos = rs.read_block(snapshot, 'POS ', parttype=ptype,
                            verbose=False) / 1e3
        pos = pos.astype(np.float32)
        vel = rs.read_block(snapshot, 'VEL ', parttype=ptype,
                            verbose=False) / np.sqrt(1.0 + redshift)  #km/s

        if ptype == 0:
            MHI = rs.read_block(snapshot, 'NH  ', parttype=0,
                                verbose=False)  #HI/H
            mass = rs.read_block(snapshot, 'MASS', parttype=0,
                                 verbose=False) * 1e10
            SFR = rs.read_block(snapshot, 'SFR ', parttype=0, verbose=False)
            indexes = np.where(SFR > 0.0)[0]
            del SFR

            # find the metallicity of star-forming particles
            metals = rs.read_block(snapshot, 'GZ  ', parttype=0, verbose=False)
            metals = metals[indexes] / 0.0127

            # find densities of star-forming particles: units of h^2 Msun/Mpc^3
            rho = rs.read_block(snapshot, 'RHO ', parttype=0,
                                verbose=False) * 1e19
            Volume = mass / rho  #(Mpc/h)^3
            radii = (Volume / (4.0 * np.pi / 3.0))**(1.0 / 3.0)  #Mpc/h

            # find density and radius of star-forming particles
            radii_SFR = radii[indexes]
            rho = rho[indexes]

            # find HI/H fraction for star-forming particles
            MHI[indexes] = HIL.Rahmati_HI_Illustris(rho,
                                                    radii_SFR,
                                                    metals,
                                                    redshift,
                                                    h,
                                                    TREECOOL_file,
                                                    Gamma=None,
                                                    fac=1,
                                                    correct_H2=True)  #HI/H
            MHI *= (0.76 * mass)

        if ptype == 1:
            radii = rs.read_block(snapshot, 'SFHS', parttype=1,
                                  verbose=False) / 1e3  #Mpc/h
            mass = np.ones(len(radii)) * massarr[1]

        # find the indexes of current subfile that contribute to halo
        begin_array = begin - begin_subfile
        end_array = begin_array + (end - begin)

        if end > end_subfile:
            end_array = end_subfile - begin_subfile
            begin = end_subfile

        new_size = particles + (end_array - begin_array)

        pos_f.resize((new_size, 3))
        pos_f[particles:] = pos[begin_array:end_array]
        vel_f.resize((new_size, 3))
        vel_f[particles:] = vel[begin_array:end_array]

        if ptype == 0:
            mass_f.resize((new_size, ))
            mass_f[particles:] = mass[begin_array:end_array]

            MHI_f.resize((new_size, ))
            MHI_f[particles:] = MHI[begin_array:end_array]

            radii_f.resize((new_size, ))
            radii_f[particles:] = radii[begin_array:end_array]

        if ptype == 1:
            radii_f.resize((new_size, ))
            radii_f[particles:] = radii[begin_array:end_array]

            mass_f.resize((new_size, ))
            mass_f[particles:] = mass[begin_array:end_array]

        particles = new_size
        begin_subfile = end_subfile

    f.close()
    print 'Halo mass = %.3e' % halo_mass[halo_number]
    print 'Halo pos  =', halo_pos[halo_number]
    print 'Number of particles in the halo = %ld' % particles
示例#7
0
all_ids = np.array([],dtype="uint32")
types = np.array([],dtype="uint32")
mass = np.array([],dtype="float64")
pos = np.array([],dtype="float64")
vel = np.array([],dtype="float64")
u = np.array([],dtype="float64")
T = np.array([],dtype="float64")
rho = np.array([],dtype="float64")
sfr = np.array([],dtype="float64")
age = np.array([],dtype="float64")
pot = np.array([],dtype="float64")


for parttype in parttype_list:
	ids = rs.read_block(snapname,"ID  ",parttype=parttype)
	all_ids = np.append(all_ids, ids)
	types = np.append(types, np.zeros(len(ids))+parttype)
	
ngas = (types == 0).sum()
ndm = (types == 1).sum()
nstar = (types == 4).sum()

maxid = np.max(all_ids)

all_bound_ids = readsubf.subf_ids(base,snapnum,0, 0, long_ids=True,read_all=True).SubIDs

rev = np.zeros(maxid+1) - 1
rev[all_bound_ids] = np.arange(len(all_bound_ids),dtype="uint32")

r = rev[all_ids]
示例#8
0
def sw_err(N_part,N_gb):

	import matplotlib.pyplot as plt
	import readsnapHDF5 as rs
	from numpy import *
	
	print "Calculating sound-wave error for: "
	print "N = "+str(N_part)
	print "Ngb = "+str(N_gb)
	print ""
	
	#N_part = 1000
	#N_gb = 7
	snap_total = 6-1
	time_out = 1.
	#snapnum=0
	ga=1.4
	probn = 5
	
	xmin=0.0#10.0
	xmax=40.0#30.0
	
	wave_cent = array([])
	
	Lx = xmax-xmin
	wave_width = Lx/10.
	
	
	######################################################
	# "Exact" wave speed:
	P = 300.0
	rho = 1.0
	del_rho = 1.*10**-3
	#del_P = del_rho * 1.0
	
	cs = sqrt(ga * P/rho)
	#print cs
	vel_exact = 0.1 * cs
	
	######################################################
	L1_r_wave_arr = array([])
	L1_v_arr = array([])
	L1_v_wave_arr = array([])
	
	
	for snapnum in arange(snap_total+1):
		print "\n"
		plt.close('all')
		
		#base = "/n/hernquistfs1/jsuresh/SPH_tests/02-2011-re-py/"
		base = "/n/hernquistfs1/jsuresh/SPH_tests/04-2011-1D/"
		prob_folder = "prob"+str(probn)+"/"
		probdir = base+prob_folder
		Ndir = probdir+"N_"+str(N_part)
		outputdir = Ndir+"/Ngb_"+str(N_gb)
		outdir = outputdir + "/outfiles"
		snapsdir = outputdir + "/snaps"
		# filename="./shock_tube/snapshot_"+str(num).zfill(3)		
		filename= snapsdir +"/snapshot_"+ str(snapnum).zfill(3) # +".hdf5"
		# print filename
		
		r=rs.read_block(filename,"RHO ",parttype=0)
		x=rs.read_block(filename,"POS ",parttype=0)[:,0]
		v=rs.read_block(filename,"VEL ",parttype=0)[:,0]
		u=rs.read_block(filename,"U   ",parttype=0)
		#P=co.GetPressure(u, r, gamma)
		#A=co.GetEntropy(u, r, gamma)
		m=rs.read_block(filename,"MASS",parttype=0)
		
		P = (ga-1)*u*r
		S = P/(r**ga)
		
		#Find wave center:
		current_wave_cent = x[argmax(r)]
		wave_cent = append(wave_cent, current_wave_cent)
		
		# Demarcate wave:
		N = float(len(v))
		
		left_edge = current_wave_cent - wave_width/2.
		right_edge = current_wave_cent + wave_width/2.
		
		argleft = argmin(abs(x-left_edge))
		argright = argmin(abs(x-right_edge))
		
		wave_r = r[argright:argleft] #indices switched because x is listed in decreasing order
		wave_v = v[argright:argleft]
		wave_x = x[argright:argleft]
		
		N_wave = float(len(wave_v))
		
		L1_v_wave = 0.
		for v_part in wave_v:
			L1_v_wave = L1_v_wave + 1/N_wave * abs(v_part-vel_exact)/abs(vel_exact)
		
		L1_v_wave_arr = append(L1_v_wave_arr,L1_v_wave)
		#"Wave L1_v: "+str(L1_v_wave)
		
		####
		L1_v = 0.
		for v_part in v:
			L1_v = L1_v + 1/N * abs(v_part-vel_exact)/abs(vel_exact)
			
		L1_v_arr = append(L1_v_arr,L1_v)
			
		#print "L1_v: " +str(L1_v)
		
		if L1_v_wave > L1_v:
			zz = 1
			#print "Wave error > Total error"
			# [it's ok: difference comes from N_wave as opposed to N]
			
			
		##############################
		# Now calculate the density errors:
		wave_r_exact = array([])
		
		L1_r_wave = 0.
		
		for i in arange(N_wave):
			r_part = wave_r[i]
			x_part = wave_x[i]
			
			f = (((x_part-current_wave_cent)/wave_width)**2 - 1.0)**4
			rho_temp = rho + del_rho * f
		
			#print "f = "+str(f)
			#print "pos: " + str(x_part)
		
			L1_r_wave = L1_r_wave + 1/N * abs(r_part-rho_temp)/abs(rho_temp)
		
		L1_r_wave_arr = append(L1_r_wave_arr,L1_r_wave)
		
		#print "Wave L1_r: "+str(L1_r_wave)
		
		print "v: "+str(L1_v_wave)
		print "r: "+str(L1_r_wave)
		print ""
		
	#Find wave speed:
	wave_vel = (wave_cent[snap_total]-wave_cent[0])/(snap_total*time_out)
	
	err = (wave_vel-vel_exact)/vel_exact
	print "Fractional error in wave speed: "+str(err)
	
	plt.close('all') 
	
	'''
	fig = plt.figure(1, figsize=(7.0,7.0))
	ax = fig.add_subplot(1,1,1)
	ax.plot(arange(snap_total+1),wave_cent,'-')
	#ax.axis([xmin,xmax,0.9*r.min(), 1.1*r.max()])
	#ax.axis([xmin,xmax,r.min(), r.max()])
	#ax.set_xlabel('x', fontsize=15)
	#ax.set_ylabel(r'$\mathrm{\rho}$', fontsize=15)
	fig.show()
	
	fig = plt.figure(2, figsize=(7.0,7.0))
	ydat = log10(L1_v_wave_arr)
	ax = fig.add_subplot(1,1,1)
	ax.plot(arange(snap_total+1),log10(L1_v_wave_arr),'-')
	#ax.axis([xmin,xmax,0.9*r.min(), 1.1*r.max()])
	ax.axis([1,6,ydat.min(), ydat.max()])
	#ax.set_xlabel('x', fontsize=15)
	#ax.set_ylabel(r'$\mathrm{\rho}$', fontsize=15)
	fig.show()
	
	fig = plt.figure(3, figsize=(7.0,7.0))
	ydat = log10(L1_r_wave_arr)
	ax = fig.add_subplot(1,1,1)
	ax.plot(arange(snap_total+1),log10(L1_r_wave_arr),'-')
	#ax.axis([xmin,xmax,0.9*r.min(), 1.1*r.max()])
	ax.axis([1,6,ydat.min(), ydat.max()])
	#ax.set_xlabel('x', fontsize=15)
	#ax.set_ylabel(r'$\mathrm{\rho}$', fontsize=15)
	fig.show()'''
	
	return [L1_v_wave, L1_v, L1_r_wave]
示例#9
0
def tsw_err(probn,N_part,N_gb):

	import matplotlib.pyplot as plt
	import readsnapHDF5 as rs
	from numpy import *
	from scipy.interpolate import interp1d
	
	verbose = 0
	
	N_part = 100
	N_gb = 33
	snap_total = 100
	time_out = 1.
	ga=1.0
	probn = 6
	
	xmin=0.0
	xmax=1.0
	
	
	######################################################
	
	plt.close('all')
	
	base = "/n/hernquistfs1/jsuresh/SPH_tests/06-2011-tsw/"
	prob_folder = "prob"+str(probn)+"/"
	probdir = base+prob_folder
	Ndir = probdir+"N_"+str(N_part)
	outputdir = Ndir+"/Ngb_"+str(N_gb)
	outdir = outputdir + "/outfiles"
	snapsdir = outputdir + "/snaps"	
	#filename= snapsdir +"/snapshot_"+ str(snapnum).zfill(3)
	
	first_name = snapsdir +"/snapshot_000"
	last_name = snapsdir +"/snapshot_"+ str(snap_total).zfill(3)
	
	
	r_first=rs.read_block(first_name,"RHO ",parttype=0)
	x_first=rs.read_block(first_name,"POS ",parttype=0)[:,0]
	v_first=rs.read_block(first_name,"VEL ",parttype=0)[:,0]
	u_first=rs.read_block(first_name,"U   ",parttype=0)
	#m_first=rs.read_block(first_name,"MASS",parttype=0)
	
	r_last=rs.read_block(last_name,"RHO ",parttype=0)
	x_last=rs.read_block(last_name,"POS ",parttype=0)[:,0]
	v_last=rs.read_block(last_name,"VEL ",parttype=0)[:,0]
	u_last=rs.read_block(last_name,"U   ",parttype=0)
	#m_last=rs.read_block(last_name,"MASS",parttype=0)
	
	#P = (ga-1)*u*r
	#S = P/(r**ga)
		
		
	
	##############################
	# ERROR CALCULATION
	
	if verbose == 1:
		print min(x_last)
		print min(x_first)
		print max(x_last)
		print max(x_first)
	
	# First create splines of the first snapshot:
	#r_interp = interp1d(x_first,r_first)
	r_interp = interp1d(x_first[::-1],r_first[::-1])
	v_interp = interp1d(x_first[::-1],v_first[::-1])
	
	#L1_r = sum(abs(r_last-r_interp(x_last)))
	#L1_v = sum(abs(v_last-v_interp(x_last)))
	
	#L1_r = 0.
	#L1_v = 0.
	r_diff = zeros(len(x_last))
	v_diff = zeros(len(v_last))
	
	for i in arange(len(x_last)):
		if x_last[i] < min(x_first):
			r_temp = r_last[argmin(x_last)]
			v_temp = v_last[argmin(x_last)]
		elif x_last[i] > max(x_first):
			r_temp = r_last[argmax(x_last)]
			v_temp = v_last[argmax(x_last)]
		else :
			if verbose == 1:
				print "x_last[i]" + str(x_last[i])
			r_temp = r_interp(x_last[i])
			v_temp = v_interp(x_last[i])
			
		r_diff[i] = r_last[i]-r_temp
		v_diff[i] = v_last[i]-v_temp
		#L1_r = L1_r + abs(r_last[i]-r_temp)
		#L1_v = L1_v + abs(v_last[i]-v_temp)
	
	L1_r = sum(abs(r_diff))/N_part
	L1_v = sum(abs(v_diff))/N_part
	print "r: "+str(L1_r)
	print "v: "+str(L1_v)
	
	if verbose == 1:
		fig = plt.figure(1, figsize=(7.0,7.0))
		ax = fig.add_subplot(1,1,1)
		ax.plot(x_last,r_diff,'-')
		#ax.axis([xmin,xmax,0.9*r.min(), 1.1*r.max()])
		#ax.axis([xmin,xmax,r.min(), r.max()])
		#ax.set_xlabel('x', fontsize=15)
		#ax.set_ylabel(r'$\mathrm{\rho}$', fontsize=15)
		fig.show()
		
		fig = plt.figure(2, figsize=(7.0,7.0))
		ax = fig.add_subplot(1,1,1)
		ax.plot(x_last,v_diff,'-')
		#ax.axis([xmin,xmax,0.9*r.min(), 1.1*r.max()])
		#ax.axis([xmin,xmax,r.min(), r.max()])
		#ax.set_xlabel('x', fontsize=15)
		#ax.set_ylabel(r'$\mathrm{\rho}$', fontsize=15)
		fig.show()
	
	return [L1_r,L1_v]
	'''
示例#10
0
def get_disk_grid_2D(disk, N1, N2, BoxSizeX, BoxSizeY, **kwargs):

    if (disk.double == 0):
        dt = "float32"
    if (disk.double == 1):
        dt = "float64"

    mesh_option = kwargs.get("from_file")

    Ntot = N1 * N2
    if (disk.mesh_type == "ring"):
        Ntot = 0
        delta = (disk.outer_radius - disk.inner_radius) / N1
        R = np.linspace(disk.inner_radius + 0.5 * delta, disk.outer_radius, N1)
        for i in range(0, N1):
            Ntot += int(2 * np.pi * R[i] / delta)

    pos = np.zeros([Ntot, 3], dtype=dt)

    if (mesh_option == None):
        if (disk.mesh_type == "cartesian"):
            offsetX = BoxSizeX / 2.0 - disk.outer_radius
            offsetY = BoxSizeY / 2.0 - disk.outer_radius
            deltax = float(2.0 * disk.outer_radius / N1)
            deltay = float(2.0 * disk.outer_radius / N2)
            for i in range(0, N1):
                for j in range(0, N2):
                    pos[i + j * N1,
                        0] = float(i * deltax + deltax / 2.0) + offsetX
                    pos[i + j * N1,
                        1] = float(j * deltay + deltay / 2.0) + offsetY
                    pos[i + j * N1, 2] = 0.0

        if (disk.mesh_type == "polar"):
            offsetX = BoxSizeX / 2.0
            offsetY = BoxSizeY / 2.0
            deltalog = (math.log(disk.outer_radius) -
                        math.log(disk.inner_radius)) / N1
            logR = (np.arange(0,N1, dtype=dt)+0.5)/N1 * (math.log(disk.outer_radius)-math.log(disk.inner_radius))+ \
                   math.log(disk.inner_radius)

            ip = 0
            for i in range(0, N1):
                R = math.exp(logR[i])
                for j in range(0, N2):
                    phi = float((j + 0.5) * 2.0 * math.pi / N2)
                    if (disk.mesh_alignment == "interleaved"):
                        phi += (-1)**i * 0.25 * 2.0 * math.pi / N2
                    pos[ip, 0] = R * math.cos(phi) + offsetX
                    pos[ip, 1] = R * math.sin(phi) + offsetY
                    pos[ip, 2] = 0.0
                    ip += 1

        if (disk.mesh_type == "ring"):
            offsetX = BoxSizeX / 2.0
            offsetY = BoxSizeY / 2.0
            delta = (disk.outer_radius - disk.inner_radius) / N1
            R = np.linspace(disk.inner_radius + 0.5 * delta, disk.outer_radius,
                            N1)
            ip = 0
            for i in range(0, N1):
                for j in range(0, int(2 * np.pi * R[i] / delta)):
                    phi = float((j + 0.5) * 2.0 * math.pi /
                                int(2 * np.pi * R[i] / delta))
                    if (disk.mesh_alignment == "interleaved"):
                        phi += (-1)**i * 0.25 * 2.0 * math.pi / int(
                            2 * np.pi * R[i] / delta)
                    pos[ip, 0] = R[i] * math.cos(phi) + offsetX
                    pos[ip, 2] = 0.0
                    pos[ip, 1] = R[i] * math.sin(phi) + offsetY
                    ip += 1

        if (disk.mesh_type == "unstructured"):
            offsetX = BoxSizeX / 2.0
            offsetY = BoxSizeY / 2.0
            ip = 0
            for i in range(0, N1):
                R = float((i + 0.5) * (disk.outer_radius - disk.inner_radius) /
                          N1) + disk.inner_radius
                delta_R = float(
                    (i + 0.5) * (disk.outer_radius - disk.inner_radius) /
                    N1) + disk.inner_radius - R
                for j in range(0, N2):
                    phi = float((j + 0.5) * 2.0 * math.pi / N2)
                    pos[ip, 0] = R * math.cos(phi) + offsetX
                    pos[ip, 1] = R * math.sin(phi) + offsetY
                    pos[ip, 2] = 0.0
                    ip += 1

    else:
        import readsnapHDF5 as rs
        pos = rs.read_block(mesh_option, "POS ", parttype=0)

    pos[:, 0] = np.where(
        (np.sqrt((pos[:, 0] - BoxSizeX / 2.0)**2 +
                 (pos[:, 1] - BoxSizeY / 2.0)**2) < disk.outer_radius) &
        (np.sqrt((pos[:, 0] - BoxSizeX / 2.0)**2 +
                 (pos[:, 1] - BoxSizeY / 2.0)**2) > disk.inner_radius),
        pos[:, 0], -1)

    ind = np.where(pos[:, 0] >= 0)[0]
    newpos = np.zeros([len(ind), 3], dtype=dt)
    newpos[:, 0] = pos[ind, 0]
    newpos[:, 1] = pos[ind, 1]
    newpos[:, 2] = pos[ind, 2]

    ids = np.arange(1, len(ind) + 1, dtype="int32")

    return newpos, ids
示例#11
0
def get_disk_grid_2D(disk,N1, N2,BoxSizeX,BoxSizeY,**kwargs):

    if (disk.double == 0):
        dt="float32"
    if (disk.double == 1):
        dt="float64"

    mesh_option = kwargs.get("from_file")

    Ntot = N1 * N2
    if(disk.mesh_type == "ring"):
      Ntot = 0
      delta = (disk.outer_radius - disk.inner_radius)/N1
      R = np.linspace(disk.inner_radius + 0.5 * delta,disk.outer_radius,N1) 
      for i in range(0,N1):
	Ntot+= int(2*np.pi*R[i]/delta)
    
    pos=np.zeros([Ntot, 3], dtype=dt)

    if (mesh_option == None):
        if (disk.mesh_type == "cartesian"):
            offsetX=BoxSizeX/2.0-disk.outer_radius
            offsetY=BoxSizeY/2.0-disk.outer_radius
            deltax=float(2.0 * disk.outer_radius / N1)
            deltay=float(2.0 * disk.outer_radius / N2)
            for i in range(0,N1):
                for j in range(0,N2):
                    pos[i+j*N1,0]=float(i*deltax+deltax/2.0)+offsetX
                    pos[i+j*N1,1]=float(j*deltay+deltay/2.0)+offsetY
                    pos[i+j*N1,2]=0.0
        
        if (disk.mesh_type == "polar"):
            offsetX=BoxSizeX/2.0
            offsetY=BoxSizeY/2.0
            deltalog = (math.log(disk.outer_radius)-math.log(disk.inner_radius))/ N1
            logR = (np.arange(0,N1, dtype=dt)+0.5)/N1 * (math.log(disk.outer_radius)-math.log(disk.inner_radius))+ \
                   math.log(disk.inner_radius)
            
            ip = 0
            for i in range(0,N1):
                R = math.exp(logR[i])
                for j in range(0,N2):
                    phi=float((j+0.5)*2.0*math.pi/N2)
                    if (disk.mesh_alignment == "interleaved"):
                        phi+=(-1)**i*0.25*2.0*math.pi/N2
                    pos[ip,0]= R * math.cos(phi)+offsetX
                    pos[ip,1]= R * math.sin(phi)+offsetY
                    pos[ip,2]=0.0
                    ip+=1
            
        if (disk.mesh_type == "ring"):
            offsetX=BoxSizeX/2.0
            offsetY=BoxSizeY/2.0
            delta = (disk.outer_radius - disk.inner_radius)/N1
            R = np.linspace(disk.inner_radius + 0.5 * delta,disk.outer_radius,N1) 
            ip = 0
            for i in range(0,N1): 
                for j in range(0,int(2*np.pi*R[i]/delta)):
                    phi=float((j+0.5)*2.0*math.pi/int(2*np.pi*R[i]/delta))
                    if (disk.mesh_alignment == "interleaved"):
                        phi+=(-1)**i*0.25*2.0*math.pi/int(2*np.pi*R[i]/delta)
                    pos[ip,0]= R[i] * math.cos(phi)+offsetX
                    pos[ip,2]=0.0
                    pos[ip,1]= R[i] * math.sin(phi)+offsetY
                    ip+=1

        if (disk.mesh_type == "unstructured"):
            offsetX=BoxSizeX/2.0
            offsetY=BoxSizeY/2.0
            ip = 0
            for i in range(0,N1):
                R = float((i+0.5)*(disk.outer_radius-disk.inner_radius)/N1) + disk.inner_radius
                delta_R = float((i+0.5)*(disk.outer_radius-disk.inner_radius)/N1) + disk.inner_radius - R
                for j in range(0,N2):
                    phi=float((j+0.5)*2.0*math.pi/N2)
                    pos[ip,0]= R * math.cos(phi)+offsetX
                    pos[ip,1]= R * math.sin(phi)+offsetY
                    pos[ip,2]=0.0
                    ip+=1     

    else:
        import readsnapHDF5 as rs
        pos = rs.read_block(mesh_option,"POS ",parttype=0)

        
    pos[:,0]=np.where( (np.sqrt((pos[:,0]-BoxSizeX/2.0)**2+(pos[:,1]-BoxSizeY/2.0)**2) < disk.outer_radius) &
                       (np.sqrt((pos[:,0]-BoxSizeX/2.0)**2+(pos[:,1]-BoxSizeY/2.0)**2) > disk.inner_radius), pos[:,0],-1)
   

    ind=np.where(pos[:,0] >= 0)[0]
    newpos=np.zeros([len(ind), 3], dtype=dt)
    newpos[:,0]=pos[ind,0]
    newpos[:,1]=pos[ind,1]
    newpos[:,2]=pos[ind,2]
 
    ids=np.arange(1,len(ind)+1, dtype="int32")

    return newpos,ids
        if (('128' in directory) | ('128' in current_dir)):
            radial_zones = 128
        elif (('256' in directory) | ('256' in current_dir)):
            radial_zones = 256
        elif (('512' in directory) | ('512' in current_dir)):
            radial_zones = 512
        elif (('1024' in directory) | ('1024' in current_dir)):
            radial_zones = 1024

        print "SNAPSHOT #", num
        #open the snapshot header
        filename = directory + base + "snap_" + str(num).zfill(3)
        header = rs.snapshot_header(filename)
        time = header.time / 2.0 / math.pi
        BoxX, BoxY = header.boxsize, header.boxsize
        pos = rs.read_block(filename, "POS ", parttype=0)
        dens = rs.read_block(filename, "RHO ", parttype=0)
        vel = rs.read_block(filename, "VEL ", parttype=0)

        x, y = pos[:, 0], pos[:, 1]
        vx, vy = vel[:, 0], vel[:, 1]

        # define regular grid spatially covering input data
        n = 1024
        xg = np.linspace(rangeX[0], rangeX[1], n)
        yg = np.linspace(rangeY[0], rangeY[1], n)
        delta_x = np.diff(xg).mean()
        delta_y = np.diff(yg).mean()
        X, Y = np.meshgrid(xg, yg)

        #data to interpolate
示例#13
0
        # define the M_HI and sigma_HI array
        M_ptype      = np.zeros(len(halo_len), dtype=np.float64)
        sigma2_ptype = np.zeros(len(halo_len), dtype=np.float64)


        # do a loop over all files
        Omega_ptype, done = 0.0, False
        #start, end, end_gal = 0, end_halos[0], end_all_galaxies[0]
        #pars = [Number, start_h, end_h, start_g, end_g, halo_num, gal_num,
        #        gal_in_local_halo]
        pars = np.array([0, 0, halo_len[0], 0], dtype=np.int64)
        for i in xrange(filenum):

            snapshot = snapshot_root + 'snapdir_%03d/snap_%03d.%d'%(num,num,i)
        
            vel = rs.read_block(snapshot, 'VEL ', parttype=ptype, verbose=False)/np.sqrt(1.0+redshift) #km/s                         

            mass  = rs.read_block(snapshot, 'MASS', parttype=ptype, verbose=False)*1e10
            mass = mass.astype(np.float32)

            # compute the HI mass in halos and galaxies
            if not(done):
                done = HIL.sigma_HI_halos(pars, done, halo_len, halo_vel,
                                          mass, vel, sigma2_ptype, M_ptype)
                             
            Omega_ptype += np.sum(mass,  dtype=np.float64)

            print '\ni = %03d: z=%.3f'%(i,redshift)
            print 'Omega(%d,z=%d) = %.6e'%(ptype,round(redshift),Omega_ptype/(BoxSize**3*rho_crit))
            
            if done: break                                
    vy = np.sqrt(1 - eb * eb) * np.cos(ecc_anom) / (1 - eb * np.cos(ecc_anom))
    x2, x1 = -qb / (1 + qb) * x, 1.0 / (1 + qb) * x
    y2, y1 = -qb / (1 + qb) * y, 1.0 / (1 + qb) * y

    print "Looping over snapshots..."
    force1x_disk, force1y_disk, force2x_disk, force2y_disk = np.zeros(
        len(time)), np.zeros(len(time)), np.zeros(len(time)), np.zeros(
            len(time))
    force1x_cav, force1y_cav, force2x_cav, force2y_cav = np.zeros(
        len(time)), np.zeros(len(time)), np.zeros(len(time)), np.zeros(
            len(time))
    for i, snap in enumerate(range(init_snap, final_snap + 1)):
        print "SNAPSHOT #", snap
        filename = directory + base + snap_base + str(snap).zfill(3)
        header = rs.snapshot_header(filename)
        pos = rs.read_block(filename, "POS ", parttype=0)
        mass = rs.read_block(filename, "MASS", parttype=0)
        acc = rs.read_block(filename, "ACCE", parttype=0)
        ids = rs.read_block(filename, "ID  ", parttype=0)

        r1 = np.sqrt((pos[:,0] - (x1[i] + 0.5  * header.boxsize))**2 +\
                     (pos[:,1] - (y1[i] + 0.5  * header.boxsize))**2)
        r2 = np.sqrt((pos[:,0] - (x2[i] + 0.5  * header.boxsize))**2 +\
                     (pos[:,1] - (y2[i] + 0.5  * header.boxsize))**2)
        r =  np.sqrt((pos[:,0] - 0.5  * header.boxsize)**2 +\
                     (pos[:,1] - 0.5  * header.boxsize)**2)

        ind = (ids >= -2) & (r > (1 + eb))
        force1x_disk[i] = (-mass[ind] / r1[ind] / r1[ind] / r1[ind] * (
            (x1[i] + 0.5 * header.boxsize) - pos[ind, 0])).sum()
        force1y_disk[i] = (-mass[ind] / r1[ind] / r1[ind] / r1[ind] * (
示例#15
0
def find_MC_tracer_ids(snap_num,sub_id,base,scale_factor,gal_radfac):
	print "Finding MC tracer IDs..."
	cat = readsubfHDF5.subfind_catalog(base, snap_num)
	
	# I assume that the subhalo is the primary subhalo in its group
	sub_list = cat.GroupFirstSub
	grp_id = np.argmin( np.abs(sub_id-sub_list) )

	sub_pos = cat.SubhaloPos[sub_id]
	sub_mass = cat.SubhaloMass[sub_id]
	sub_Rvir = cat.Group_R_Crit200[grp_id]

	gal_rad = sub_Rvir * gal_radfac # * 0.1

	#snapname = base + "snapdir_"+str(snap_num).zfill(3)+"/snap_"+str(snap_num).zfill(3)
	snapname = base + "snap_"+str(snap_num).zfill(3)

	#Find all gas and star particles within 0.1 Rvir of the subhalo
	# GAS #
	gas_pos = readsnapHDF5.read_block(snapname,"POS ",parttype=0)
	gas_ids = readsnapHDF5.read_block(snapname,"ID  ",parttype=0)
	#readsnapHDF5.list_blocks(snapname+".hdf5")

	gas_x = np.logical_and(gas_pos[:,0] > sub_pos[0]-gal_rad, gas_pos[:,0] < sub_pos[0]+gal_rad)
	gas_y = np.logical_and(gas_pos[:,1] > sub_pos[1]-gal_rad, gas_pos[:,1] < sub_pos[1]+gal_rad)
	gas_z = np.logical_and(gas_pos[:,2] > sub_pos[2]-gal_rad, gas_pos[:,2] < sub_pos[2]+gal_rad)

	gas_xyz_ind = np.logical_and(np.logical_and(gas_x,gas_y),gas_z)
	gas_pos = gas_pos[gas_xyz_ind]
	gas_pos = gas_pos - sub_pos

	gas_r_ind = np.sqrt(gas_pos[:,0]**2 + gas_pos[:,1]**2 + gas_pos[:,2]**2) < gal_rad
	gal_gas_ids = gas_ids[gas_xyz_ind][gas_r_ind]
	print "len(gal_gas_ids )",len(gal_gas_ids)

	# Clear gas data:
	gas_pos = 0.
	gas_ids = 0.
	gas_x = 0.
	gas_y = 0.
	gas_z = 0.
	gas_xyz_ind = 0.
	gas_r_ind = 0.


	# STARS #
	star_pos = readsnapHDF5.read_block(snapname,"POS ",parttype=4)
	star_ids = readsnapHDF5.read_block(snapname,"ID  ",parttype=4)
	#readsnapHDF5.list_blocks(snapname+".hdf5")

	star_x = np.logical_and(star_pos[:,0] > sub_pos[0]-gal_rad, star_pos[:,0] < sub_pos[0]+gal_rad)
	star_y = np.logical_and(star_pos[:,1] > sub_pos[1]-gal_rad, star_pos[:,1] < sub_pos[1]+gal_rad)
	star_z = np.logical_and(star_pos[:,2] > sub_pos[2]-gal_rad, star_pos[:,2] < sub_pos[2]+gal_rad)

	star_xyz_ind = np.logical_and(np.logical_and(star_x,star_y),star_z)
	star_pos = star_pos[star_xyz_ind]
	star_pos = star_pos - sub_pos

	star_r_ind = np.sqrt(star_pos[:,0]**2 + star_pos[:,1]**2 + star_pos[:,2]**2) < gal_rad
	gal_star_ids = star_ids[star_xyz_ind][star_r_ind]
	print "len(gal_star_ids )",len(gal_star_ids)

	# Clear star data:
	star_pos = 0.
	star_ids = 0.
	star_x = 0.
	star_y = 0.
	star_z = 0.
	star_xyz_ind = 0.
	star_r_ind = 0.


	# Now find all MC tracers associated with the gas and star particles we have selected:
	parent_ids = readsnapHDF5.read_block(snapname,"TRPI",parttype=3)
	tracer_ids = readsnapHDF5.read_block(snapname,"TFID",parttype=3)
	temp1 = tracer_ids[np.in1d(parent_ids,gal_gas_ids)]
	temp2 = tracer_ids[np.in1d(parent_ids,gal_star_ids)]

	print "len(temp1)",len(temp1)

	gal_MC_ids = np.append(temp1,temp2)
	print "len(gal_MC_ids )",len(gal_MC_ids)
	
	print "Done finding MC tracer IDs..."
	return gal_MC_ids
示例#16
0
		old_subid = new_subid
		old_partIDs = new_partIDs
		old_subpos = new_subpos
		old_submass = new_submass
		old_subvel = new_subvel

		#new_snapname = base + "snapdir_"+str(current_snap).zfill(3)+"/snap_"+str(current_snap).zfill(3)
		new_snapname = base + "/snap_"+str(current_snap).zfill(3)
		new_cat = readsubfHDF5.subfind_catalog(base, current_snap)

		redshift = readsnapHDF5.snapshot_header(new_snapname).redshift
		scale_factor = 1./(1.+redshift)

		readsnapHDF5.list_blocks(new_snapname+".hdf5")
		# Velocity tracer information:
		vt_ids = readsnapHDF5.read_block(new_snapname,"ID  ",parttype=2)
		vt_ind = np.in1d(vt_ids,gal_vt_ids)
		vt_ids = 0.
		vt_pos = readsnapHDF5.read_block(new_snapname,"POS ",parttype=2)
		gal_vt_pos = vt_pos[vt_ind]
		vt_pos = 0.
		vt_vel = readsnapHDF5.read_block(new_snapname,"VEL ",parttype=2)
		gal_vt_vel = vt_vel[vt_ind]
		vt_vel = 0.

		time15 = time.time()
		# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
		# MC tracer information:
		MC_ids = readsnapHDF5.read_block(new_snapname,"TFID",parttype=3)
		MC_ind = np.in1d(MC_ids,gal_MC_ids,assume_unique=True)
		MC_ids = 0.
示例#17
0
HI_mass_shell = np.zeros((halos, bins), dtype=np.float64)
part_in_halo  = np.zeros(halos,         dtype=np.int64)

# find the numbers each cpu will work on
array   = np.arange(filenum)
numbers = np.where(array%nprocs==myrank)[0]

# do a loop over each subsnapshot
for i in numbers:

    # find subfile name and read the number of particles in it
    snapshot = snapshot_root + 'snapdir_%03d/snap_%03d.%d'%(snapnum, snapnum, i)
    header = rs.snapshot_header(snapshot)
    npart  = header.npart 

    pos  = rs.read_block(snapshot, 'POS ', parttype=0, verbose=False)/1e3
    pos  = pos.astype(np.float32)
    MHI  = rs.read_block(snapshot, 'NH  ', parttype=0, verbose=False)#HI/H
    mass = rs.read_block(snapshot, 'MASS', parttype=0, verbose=False)*1e10
    SFR  = rs.read_block(snapshot, 'SFR ', parttype=0, verbose=False)
    indexes = np.where(SFR>0.0)[0];  del SFR

    # find the metallicity of star-forming particles
    metals = rs.read_block(snapshot, 'GZ  ', parttype=0, verbose=False)
    metals = metals[indexes]/0.0127

    # find densities of star-forming particles: units of h^2 Msun/Mpc^3
    rho = rs.read_block(snapshot, 'RHO ', parttype=0, verbose=False)*1e19
    Volume = mass/rho                            #(Mpc/h)^3
    radii  = (Volume/(4.0*np.pi/3.0))**(1.0/3.0) #Mpc/h 
def Illustris_region(snapshot_root, snapnum, TREECOOL_file, x_min, x_max, 
                     y_min, y_max, z_min, z_max, padding, fout,
                     redshift_space=False, axis=0):


    # read snapshot and find number of subfiles
    snapshot = snapshot_root + 'snapdir_%03d/snap_%03d'%(snapnum,snapnum)
    header   = rs.snapshot_header(snapshot)
    nall     = header.nall
    redshift = header.redshift
    BoxSize  = header.boxsize/1e3 #Mpc/h
    filenum  = header.filenum
    Omega_m  = header.omega0
    Omega_L  = header.omegaL
    h        = header.hubble
    Hubble = 100.0*np.sqrt(Omega_m*(1.0+redshift)**3+Omega_L) #km/s/(Mpc/h)

    if myrank==0:
        print '\n'
        print 'BoxSize         = %.3f Mpc/h'%BoxSize
        print 'Number of files = %d'%filenum
        print 'Omega_m         = %.3f'%Omega_m
        print 'Omega_l         = %.3f'%Omega_L
        print 'redshift        = %.3f'%redshift

    # find the numbers each cpu will work on
    array   = np.arange(0, filenum)
    numbers = np.where(array%nprocs==myrank)[0]

    # do a loop over the different realizations
    particles = 0
    for i in numbers:

        snapshot = snapshot_root + 'snapdir_%03d/snap_%03d.%d'%(snapnum,snapnum,i)
        pos = rs.read_block(snapshot, 'POS ', parttype=0, verbose=False)/1e3
        pos = pos.astype(np.float32)

        # read velocities and displace particle positions
        if redshift_space:
            vel = rs.read_block(snapshot, 'VEL ', parttype=0, verbose=False)/np.sqrt(1.0+redshift) #km/s
            RSL.pos_redshift_space(pos, vel, BoxSize, Hubble, redshift, axis)

        # check if particles are in the region
        indexes_region = np.where((pos[:,0]>=x_min-padding) & (pos[:,0]<=x_max+padding) &\
                                  (pos[:,1]>=y_min-padding) & (pos[:,1]<=y_max+padding) &\
                                  (pos[:,2]>=z_min-padding) & (pos[:,2]<=z_max+padding))[0]

        # if particles are not in the region continue
        local_particles = indexes_region.shape[0]
        print 'Myrank = %d ---> num = %d ---> part = %ld'%(myrank,i,local_particles)
        if local_particles==0:  continue

        # find radii, HI and gas masses
        MHI  = rs.read_block(snapshot, 'NH  ', parttype=0, verbose=False)#HI/H
        mass = rs.read_block(snapshot, 'MASS', parttype=0, verbose=False)*1e10
        SFR  = rs.read_block(snapshot, 'SFR ', parttype=0, verbose=False)
        indexes = np.where(SFR>0.0)[0];  del SFR

        # find the metallicity of star-forming particles
        metals = rs.read_block(snapshot, 'GZ  ', parttype=0, verbose=False)
        metals = metals[indexes]/0.0127

        # find densities of star-forming particles: units of h^2 Msun/Mpc^3
        rho    = rs.read_block(snapshot, 'RHO ', parttype=0, verbose=False)*1e19
        Volume = mass/rho                            #(Mpc/h)^3
        radii  = (Volume/(4.0*np.pi/3.0))**(1.0/3.0) #Mpc/h 
        rho    = rho[indexes]                        #h^2 Msun/Mpc^3
        Volume = Volume[indexes]                     #(Mpc/h)^3

        # find volume and radius of star-forming particles
        radii_SFR  = (Volume/(4.0*np.pi/3.0))**(1.0/3.0) #Mpc/h 
            
        # find HI/H fraction for star-forming particles
        MHI[indexes] = HIL.Rahmati_HI_Illustris(rho, radii_SFR, metals, redshift, 
                                                h, TREECOOL_file, Gamma=None,
                                                fac=1, correct_H2=True) #HI/H
        MHI *= (0.76*mass)
            

        # select the particles belonging to the region
        pos   = pos[indexes_region]
        MHI   = MHI[indexes_region]
        radii = radii[indexes_region]
        mass  = mass[indexes_region]

        # write partial files        
        new_size = particles + local_particles    

        if particles==0:
            f = h5py.File(fout[:-5]+'_%d.hdf5'%myrank, 'w')
            f.create_dataset('pos',   data=pos,   maxshape=(None,3))
            f.create_dataset('M_HI',  data=MHI,   maxshape=(None,))
            f.create_dataset('radii', data=radii, maxshape=(None,))
            f.create_dataset('mass',  data=mass,  maxshape=(None,))
        else:
            f = h5py.File(fout[:-5]+'_%d.hdf5'%myrank, 'a')
            pos_f   = f['pos'];    pos_f.resize((new_size,3))
            M_HI_f  = f['M_HI'];   M_HI_f.resize((new_size,))
            radii_f = f['radii'];  radii_f.resize((new_size,))
            mass_f  = f['mass'];   mass_f.resize((new_size,))
            pos_f[particles:]   = pos
            M_HI_f[particles:]  = MHI
            radii_f[particles:] = radii
            mass_f[particles:]  = mass
        f.close()
        particles += local_particles
                
    comm.Barrier()

    # sum the particles found in each cpu
    All_particles = 0 
    All_particles = comm.reduce(particles, op=MPI.SUM, root=0)

    # Master will merge partial files into a file one
    if myrank==0:

        print 'Found %d particles'%All_particles
        f = h5py.File(fout,'w')
        
        f1 = h5py.File(fout[:-5]+'_0.hdf5','r')
        pos   = f1['pos'][:]
        M_HI  = f1['M_HI'][:]
        radii = f1['radii'][:]
        mass  = f1['mass'][:]
        f1.close()

        particles = mass.shape[0]
        pos_f   = f.create_dataset('pos',   data=pos,   maxshape=(None,3))
        M_HI_f  = f.create_dataset('M_HI',  data=M_HI,  maxshape=(None,))
        radii_f = f.create_dataset('radii', data=radii, maxshape=(None,))
        mass_f  = f.create_dataset('mass',  data=mass,  maxshape=(None,))

        for i in xrange(1,nprocs):
            f1 = h5py.File(fout[:-5]+'_%d.hdf5'%i,'r')
            pos   = f1['pos'][:]
            M_HI  = f1['M_HI'][:]
            radii = f1['radii'][:]
            mass  = f1['mass'][:]
            f1.close()
            
            size = mass.shape[0]
            
            pos_f.resize((particles+size,3));  pos_f[particles:] = pos
            M_HI_f.resize((particles+size,));  M_HI_f[particles:] = M_HI
            radii_f.resize((particles+size,)); radii_f[particles:] = radii
            mass_f.resize((particles+size,));  mass_f[particles:] = mass

            particles += size

        f.close()

        for i in xrange(nprocs):
            os.system('rm '+fout[:-5]+'_%d.hdf5'%i)
示例#19
0
def ray_tracing(filename,nframe,target,cam_pos,cam_orient,frame,near,far,xbins,ybins,Box,basename):
	cam_pos = np.array(cam_pos)
	cam_orient = np.array(cam_orient)	
	target = np.array(target)

	print cam_pos,cam_orient,target
	
    	head=rs.snapshot_header(filename)
	pos=rs.read_block(filename, "POS ", parttype=0)
	mass=rs.read_block(filename, "MASS", parttype=0)
        hsmlfac=2.75*1.1
        hsml=hsmlfac*(3.0/4.0/np.pi*rs.read_block(filename, "VOL ", parttype=0))**(1.0/3.0)
	
	b,l,t,r = frame[0], frame[1], frame[2], frame[3]
	
        #translate box so that center is at Box/2.0
	centerx,centery,centerz = Box/2.0,Box/2.0,Box/2.0
        pos[:,0]=Box/2.0 + (pos[:,0]-centerx)
        pos[:,1]=Box/2.0 + (pos[:,1]-centery)
        pos[:,2]=Box/2.0 + (pos[:,2]-centerz)

        for k in range(0,3):
            ind=pos[:,k]>Box
            pos[ind,k]=pos[ind,k]-Box
            ind=pos[:,k]<0
            pos[ind,k]=pos[ind,k]+Box
            print "min/max coord: ",k,pos[:,k].min(), pos[:,k].max()

        #save original values
        x_orig=pos[:,0]
        y_orig=pos[:,1]
        z_orig=pos[:,2]
        hsml_orig=hsml
        mass_orig=mass
        
	#construct homog. transformation matrix
	PS=np.matrix([[(2*near)/(r-l),0,(r+l)/(r-l),0],[0,(2*near)/(t-b),(t+b)/(t-b),0],[0,0,-(far+near)/(far-near),-(2*far*near)/(far-near)],[0,0,-1,0]])
	nvec=-(cam_pos-target)*(-1.0)/np.sqrt((cam_pos[0]-target[0])**2.0 + (cam_pos[1]-target[1])**2.0 + (cam_pos[2]-target[2])**2.0)  #-1 
	temp=np.cross(cam_orient,nvec)
	rvec=temp/np.sqrt(temp[0]**2.0 + temp[1]**2.0 + temp[2]**2.0)
	uvec=np.cross(nvec, rvec) 
	R=np.matrix([[rvec[0],rvec[1],rvec[2],0],[uvec[0],uvec[1],uvec[2],0],[nvec[0],nvec[1],nvec[2],0],[0,0,0,1]])
	T=np.matrix([[1,0,0,-cam_pos[0]],[0,1,0,-cam_pos[1]],[0,0,1,-cam_pos[2]],[0,0,0,1]])

	PSRT=PS*R*T

	#PSRT tranformation: world coordinates -> camera coordinates
	x=PSRT[0,0]*x_orig + PSRT[0,1]*y_orig + PSRT[0,2]*z_orig + PSRT[0,3]*1
	y=PSRT[1,0]*x_orig + PSRT[1,1]*y_orig + PSRT[1,2]*z_orig + PSRT[1,3]*1
	z=PSRT[2,0]*x_orig + PSRT[2,1]*y_orig + PSRT[2,2]*z_orig + PSRT[2,3]*1
	w=PSRT[3,0]*x_orig + PSRT[3,1]*y_orig + PSRT[3,2]*z_orig + PSRT[3,3]*1

        hsml_x=PS[0,0]*hsml_orig + PS[0,1]*hsml_orig + PS[0,2]*hsml_orig + PS[0,3]*1
        hsml_y=PS[1,0]*hsml_orig + PS[1,1]*hsml_orig + PS[1,2]*hsml_orig + PS[1,3]*1
	
	#homog. scaling
	x/=w
	y/=w
	z/=w
	mass=mass_orig
	s=np.abs(w)
	hsml_x/=s
	hsml_y/=s
	hsml_o=hsml_orig

	#clipping in frustum (clip a bit larger for particle contributions outside of frustum)
	index=(np.abs(x) < 1.01)  & (np.abs(y) < 1.01) & (np.abs(z) < 1.01)
	x=x[index]
	y=y[index]
	z=z[index]
	hsml_x=hsml_x[index]
	hsml_y=hsml_y[index]
	hsml_o=hsml_o[index]
	mass=mass[index]

	print "Number of particles in frustum: ", mass.shape[0]
	
	#sort descending according to pseudo-depth
	index=np.argsort(z)[::-1]
	x=x[index]
	y=y[index]
	z=z[index]
	hsml_x=hsml_x[index]
	hsml_y=hsml_y[index]
        hsml_o=hsml_o[index]
	mass=mass[index]

	#avoid single pixel flickering
        pixfac=0.5
        hsml_x[hsml_x<pixfac*2.0/xbins]=0
        hsml_y[hsml_y<pixfac*2.0/ybins]=0

	#now ray-trace
	print "start render..."

        image=rc.Render(x, y, np.repeat(Ap, x.shape[0]).astype("float32"), hsml_x, hsml_y, hsml_o/(hsmlfac*(3.0/4.0/np.pi)**(1.0/3.0)), mass, xbins, ybins, hsmlfac)
	print "done."

	#save file
	fd=open(basename+str(nframe).zfill(4)+".dat", "wb")
	image.astype("float64").tofile(fd)
	fd.close()
	print image.shape	
	#clean image
	image=image*0.0	
示例#20
0
import numpy as np
import readsnapHDF5
import time

base = "/n/home07/dnelson/sims.tracers/256_20Mpc/output/"
current_snap = 234
new_snapname = base + "/snap_"+str(current_snap).zfill(3)

pre_read = time.time()
MC_ids = readsnapHDF5.read_block(new_snapname,"TFID",parttype=3) # May want to fix this for very large # of tracers (read in 1 snapshot at a time)
post_read = time.time()

ids_argsort = np.argsort(MC_ids)
post_sort = time.time()

ids_dblargsort = np.argsort(ids_argsort)
post_dblsrot = time.time()

#MC_ind = np.in1d(MC_ids,gal_MC_ids,assume_unique=True) #This step takes a long time!
#MC_ids = 0.
# now match the tracers with their parent particles:
#parent_ids = readsnapHDF5.read_block(new_snapname,"TRPI",parttype=3)[MC_ind]




print "done w part 1"



示例#21
0
def analyze_st(probn,N_part,N_gb,err_norm_mode):
# This code carries out the full analysis of a given shock-tube test run.
# 02-19-2011: DOES NOT CONTAIN BINNING ROUTINES

	import readsnapHDF5 as rs
	import matplotlib.pyplot as plt
	from numpy import *
	from scipy.interpolate import interp1d
	import pylab
	
	from err_norm import *
	
	err_var = "r"
	plot_mode = 2
	
	'''L1_mode = 2
	plot_mode = 1'''

	print "Analyzing shocktube from problem "+str(probn)+"..."
	print "N_particles = " + str(N_part)
	print "N_gb = " + str(N_gb)

	base = "/n/hernquistfs1/jsuresh/SPH_tests/04-2011-box/"
	probdir = base+"prob"+str(probn)+"/"
	exact_name = "p"+str(probn)+"_exact.npz"
	exact_file = probdir + exact_name

	##########################################################
	# IMPORTS THE .NPZ FILE WITH THE EXACT SOLUTION
	npzfile = load(exact_file)
	x_ex = npzfile['x_vec']-10.0
	r_ex = npzfile['d_vec']
	v_ex = npzfile['v_vec']
	p_ex = npzfile['p_vec']
	S_ex = npzfile['S_vec']
	
	gamma=1.4

	##########################################################
	# IMPORTS THE SPH-CALCULATED QUANTITIES OF THE GIVEN SHOCKTUBE
	'''if probn == 1:
		num=6
	if probn == 3:
		num=1

	xmin=10.0
	xmax=30.0'''

	num = 6
	
	Ndir = probdir+"N_"+str(N_part)
	outputdir = Ndir+"/Ngb_"+str(N_gb)
	snapsdir = outputdir+"/snaps/"
	snapname="snapshot_"+str(num).zfill(3)
	filename = snapsdir+snapname#+".HDF5"


	# Read HDF5 block to get fluid quantities:
	r_file=rs.read_block(filename,"RHO ",parttype=0)
	x_file=rs.read_block(filename,"POS ",parttype=0)[:,0]
	v_file=rs.read_block(filename,"VEL ",parttype=0)[:,0]
	u_file=rs.read_block(filename,"U   ",parttype=0)
	m_file=rs.read_block(filename,"MASS",parttype=0)

	# Initialize arrays
	j = 0
	r_ca = array([])
	x_ca = array([])
	v_ca = array([])
	u_ca = array([])
	m_ca = array([])

	# Convert the blocks to arrays
	for i in arange(x_file.size):
		if x_file[i] <= 30. and x_file[i] >= 10.:
			r_ca = append(r_ca,r_file[i])
			x_ca = append(x_ca,x_file[i] - 10.)
			v_ca = append(v_ca,v_file[i])
			u_ca = append(u_ca,u_file[i])
			m_ca = append(m_ca,m_file[i])
			j = j+1


	# Calculate pressure and entropy
	P_ca = (gamma-1)*u_ca*r_ca
	S_ca = P_ca/(r_ca**gamma)

	eff_num_particles = r_ca.size*1.
	
	##########################################################
	# CALCULATE ERROR OF SPH CALCULATION
	
	if err_var == 'r':
		dat_ex = r_ex
		dat_ca = r_ca
		
	if err_var == 'v':
		dat_ex = v_ex
		dat_ca = v_ca
		
	if err_var == 'p':
		dat_ex = p_ex
		dat_ca = p_ca
		
	if err_var == 'S':
		dat_ex = S_ex
		dat_ca = S_ca

	ex_tuple = [x_ex,r_ex]
	ca_tuple = [x_ca,r_ca]
	
	L_err = err_norm(err_norm_mode,ex_tuple,ca_tuple)
	
	
	print "Using error norm L_"+str(err_norm_mode)
	print "Reading from: " +filename
	print "Effective number of particles: "+str(eff_num_particles)
	print "Tabulated error: "+str(L_err)
	print "\n\n"

	


	##########################################################
	# PLOTS


	if plot_mode > 0:
		fig = plt.figure(1, figsize=(10.0,15.0))
		ax = fig.add_subplot(3,1,1)
		ax.scatter(x_ex,dat_ex,facecolor='w',edgecolor='b',s=10)
		ax.axis([0.,20.,0.9*dat_ex.min(), 1.1*dat_ex.max()])
		ax.set_xlabel('x', fontsize=15)
		ax.set_ylabel(err_var, fontsize=15)
	
	
		ax = fig.add_subplot(3,1,2)
		ax.scatter(x_ca,dat_ca,facecolor='w',edgecolor='b',s=10)
		ax.axis([x_ca.min(),x_ca.max(),0.9*dat_ca.min(), 1.1*dat_ca.max()])
		ax.set_xlabel('x', fontsize=15)
		ax.set_ylabel(err_var, fontsize=15)
	
	
		ax = fig.add_subplot(3,1,3)
		ax.plot(x_ca,dat_ca,'+b',x_ex,dat_ex,'+r')
		ax.axis([x_ca.min(),x_ca.max(),0.9*dat_ca.min(), 1.1*dat_ca.max()])
		ax.set_xlabel('x', fontsize=15)
		ax.set_ylabel(err_var, fontsize=15)


	if plot_mode == 1: 
		fig.show()
	
	if plot_mode == 2:
		#fig.show()
		#print "flag1"
		plt.savefig(outputdir+"/"+err_var+"_dat.eps")
		plt.close(fig)
	##########################################################
	# RETURN
	
	return [err_norm_mode,L_err]
    ecc_anom = np.vectorize(KeplerEquation)(np.array(time)  + np.pi,eb)
    x = np.cos(ecc_anom) - eb
    y = np.sqrt(1 - eb * eb) * np.sin(ecc_anom)
    vx = -np.sin(ecc_anom) / (1 - eb * np.cos(ecc_anom))
    vy = np.sqrt(1 - eb * eb) * np.cos(ecc_anom)/ (1 - eb * np.cos(ecc_anom))
    x2, x1 = -qb/(1 + qb) * x, 1.0/(1 + qb) * x
    y2, y1 = -qb/(1 + qb) * y, 1.0/(1 + qb) * y

    print "Looping over snapshots..."
    force1x_disk,force1y_disk, force2x_disk, force2y_disk = np.zeros(len(time)),np.zeros(len(time)),np.zeros(len(time)),np.zeros(len(time))
    force1x_cav,force1y_cav, force2x_cav, force2y_cav = np.zeros(len(time)),np.zeros(len(time)),np.zeros(len(time)),np.zeros(len(time))
    for i,snap in enumerate(range(init_snap,final_snap+1)):
        print "SNAPSHOT #",snap
        filename=directory+base+snap_base+str(snap).zfill(3)
        header = rs.snapshot_header(filename)
        pos = rs.read_block(filename,"POS ", parttype=0)
        mass = rs.read_block(filename,"MASS", parttype=0)
        acc = rs.read_block(filename,"ACCE", parttype=0)
        ids = rs.read_block(filename,"ID  ", parttype=0)

        r1 = np.sqrt((pos[:,0] - (x1[i] + 0.5  * header.boxsize))**2 +\
                     (pos[:,1] - (y1[i] + 0.5  * header.boxsize))**2)
        r2 = np.sqrt((pos[:,0] - (x2[i] + 0.5  * header.boxsize))**2 +\
                     (pos[:,1] - (y2[i] + 0.5  * header.boxsize))**2)
        r =  np.sqrt((pos[:,0] - 0.5  * header.boxsize)**2 +\
                     (pos[:,1] - 0.5  * header.boxsize)**2)

        ind = (ids >= -2) & (r > (1 + eb))
        force1x_disk[i] = (-mass[ind] / r1[ind] / r1[ind] / r1[ind] * ((x1[i] + 0.5  * header.boxsize) - pos[ind,0])).sum()
        force1y_disk[i] = (-mass[ind] / r1[ind] / r1[ind] / r1[ind] * ((y1[i] + 0.5  * header.boxsize) - pos[ind,1])).sum()
        force2x_disk[i] = (-mass[ind] / r2[ind] / r2[ind] / r2[ind] * ((x2[i] + 0.5  * header.boxsize) - pos[ind,0])).sum()
def calculate_halo(sub_id): #submit job to queue

	# I assume that the subhalo is the primary subhalo in its group
	sub_list = cat.GroupFirstSub
	grp_id = np.argmin( np.abs(sub_id-sub_list) )
	#grp_id = cat.SubhaloParent[sub_id]

	sub_partIDs = get_subhalo_ids(base,snap_num,sub_id)
	sub_pos = cat.SubhaloPos[sub_id]
	sub_mass = cat.SubhaloMass[sub_id]
	sub_vel = cat.SubhaloVel[sub_id]
	sub_Rvir = cat.Group_R_Crit200[grp_id]




	frame_num = 0
	for current_snap in snap_nums:
		# Find the halo in this snapshot that corresponds to the halo in our original snapshot:
		print "on snapshot number ",current_snap


		old_match_flag = new_match_flag
		old_subid = new_subid
		old_partIDs = new_partIDs
		old_subpos = new_subpos
		old_submass = new_submass
		old_subvel = new_subvel

		#new_snapname = base + "snapdir_"+str(current_snap).zfill(3)+"/snap_"+str(current_snap).zfill(3)
		new_snapname = base + "/snap_"+str(current_snap).zfill(3)
		new_cat = readsubfHDF5.subfind_catalog(base, current_snap)

		redshift = readsnapHDF5.snapshot_header(new_snapname).redshift
		scale_factor = 1./(1.+redshift)

		#readsnapHDF5.list_blocks(new_snapname+".hdf5")

		# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
		# MC tracer information:
		MC_ids = readsnapHDF5.read_block(new_snapname,"TFID",parttype=3) # May want to fix this for very large # of tracers (read in 1 snapshot at a time)
		MC_ind = np.in1d(MC_ids,gal_MC_ids,assume_unique=True) #This step takes a long time!
		MC_ids = 0.
		print "MC_ind.sum() ",MC_ind.sum()
		# now match the tracers with their parent particles:
		parent_ids = readsnapHDF5.read_block(new_snapname,"TRPI",parttype=3)[MC_ind]
		#parent_ids = np.unique(parent_ids)

		parent_ids.sort()
		(unique_parent_ids,unique_arg) = np.unique(parent_ids,return_index=True)
		foo = parent_ids.size-unique_arg[-1] # the number of reps for the very last unique entry in parent_ids

		rep_count = np.append(np.diff(unique_arg),foo)

		gas_ids = readsnapHDF5.read_block(new_snapname,"ID  ",parttype=0)
		gas_ind = np.in1d(gas_ids,unique_parent_ids,assume_unique=True)
		gas_ids = 0.
		gas_pos = readsnapHDF5.read_block(new_snapname,"POS ",parttype=0)
		pos1 = gas_pos[gas_ind]
		gas_pos = 0.
		gas_vel = readsnapHDF5.read_block(new_snapname,"VEL ",parttype=0)
		vel1 = gas_vel[gas_ind]
		gas_vel = 0.

		# now find which ones are stars (and get their properties):
		star_ids = readsnapHDF5.read_block(new_snapname,"ID  ",parttype=4)
		star_ind = np.in1d(star_ids,unique_parent_ids,assume_unique=True)
		star_ids = 0.
		star_pos = readsnapHDF5.read_block(new_snapname,"POS ",parttype=4)
		pos2 = star_pos[star_ind]
		star_pos = 0.
		star_vel = readsnapHDF5.read_block(new_snapname,"VEL ",parttype=4)
		vel2 = star_vel[star_ind]
		star_vel = 0.

		pos = np.concatenate((pos1,pos2))
		pos1 = 0.
		pos2 = 0.
		vel = np.concatenate((vel1,vel2))
		vel1 = 0.
		vel2 = 0.

		gal_MC_pos = np.repeat(pos,rep_count,axis=0)
		gal_MC_vel = np.repeat(vel,rep_count,axis=0)

		print "len(gal_MC_pos) ",len(gal_MC_pos)
		print "len(gal_MC_vel) ",len(gal_MC_vel)
		# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~


		# Generate image:
		#generate_tracer_image()
		print "not generating image "
	                        
		# Generate savefile: # save: snap_num, # of vts, gal_vt_ids at z=0, sub_id at that snap, 
		# subhalo_pos at that snap, Rvir at that time, 
		filename = save_dir + str(current_snap).zfill(3)+".dat"
		f = open(filename,'wb')

		primary_halo_flag = np.uint32(primary_halo_flag)
		gal_MC_posx = gal_MC_pos[:,0]
		gal_MC_posy = gal_MC_pos[:,1]
		gal_MC_posz = gal_MC_pos[:,2]
		gal_MC_velx = gal_MC_vel[:,0]
		gal_MC_vely = gal_MC_vel[:,1]
		gal_MC_velz = gal_MC_vel[:,2]

		print "len(gal_MC_posx) ",len(gal_MC_posx)
		print "len(gal_MC_velz) ",len(gal_MC_velz)


		# Saving stuff to file:
		# general stuff
		current_snap.astype("uint32").tofile(f)
		redshift.astype("float64").tofile(f)
		# MC tracer data
		n_MC.astype("uint32").tofile(f)
		gal_MC_ids.astype("uint64").tofile(f)
		gal_MC_posx.astype("float64").tofile(f)
		gal_MC_posy.astype("float64").tofile(f)
		gal_MC_posz.astype("float64").tofile(f)
		gal_MC_velx.astype("float64").tofile(f)
		gal_MC_vely.astype("float64").tofile(f)
		gal_MC_velz.astype("float64").tofile(f)
		# halo data
		primary_halo_flag.astype("uint32").tofile(f)
		new_subid.astype("uint32").tofile(f)
		new_subpos.astype("float64").tofile(f)
		new_subvel.astype("float64").tofile(f)
		new_subRvir.astype("float64").tofile(f)
		if False:
			if primary_halo_flag == 1: pass
			elif primary_halo_flag == 0: 
				primary_halo_pos.astype("float64").tofile(f)
				primary_halo_vel.astype("float64").tofile(f)		
		
		f.close()

		frame_num = frame_num + 1
		time4 = time.time()

	print "done!"
	return
示例#24
0
for option in sys.argv:
  if option == "128" or option == "256": nPoints = int(option)
  if option.find("device=") != -1: useDevice = int(option[-1]) 
 

import readsnapHDF5
snapname = dataDirectory + 'L25n128/output/snapdir_135/snap_135'
box_size = 25000.0  # kpc
particle_resolution = 128
linking_length = 0.2 * box_size / particle_resolution

# Read data
start = time.time()
print 'Reading data...'
#0 = gas particle, 1 = dark matter (DM), 4 = star particles
posParticles_dm = readsnapHDF5.read_block(snapname, "POS ", parttype=1)
posParticles_gas = readsnapHDF5.read_block(snapname, "POS ", parttype=0)
posParticles_star = readsnapHDF5.read_block(snapname, "POS ", parttype=4)
nParticles = len(posParticles_dm)
massParticles_dm = readsnapHDF5.read_block(snapname, "MASS", parttype=1)
massParticles_gas = readsnapHDF5.read_block(snapname, "MASS", parttype=0)
massParticles_star = readsnapHDF5.read_block(snapname, "MASS", parttype=4)
dm_particle_mass = massParticles_dm[0]
print 'Time:', time.time() - start, "\n"

posAll = { "dm":posParticles_dm, "gas":posParticles_gas, "star":posParticles_star }
massAll = { "dm":massParticles_dm, "gas":massParticles_gas, "star":massParticles_star }


nWidth = nPoints
nHeight = nPoints
x=linspace(1,30,200)
vphiprofilesq = vecvelprofilesq(x)
DvphiDRprofile = vecDvphiDRprofile(x)
vphiprofilesqgrad = np.gradient(vphiprofilesq,np.diff(x).mean())
#plt.plot(x,vphiprofilesqgrad,'ro',mew=0.0,ms=2.0)
#plt.plot(x,DvphiDRprofile,'b')
#plt.show()

#open the snapshot header
filename=base+"snap_"+str(num).zfill(3)
header = rs.snapshot_header(filename)
time = header.time
BoxX, BoxY = header.boxsize, header.boxsize


pos = rs.read_block(filename,"POS ",parttype=0)
vel  = rs.read_block(filename,"VEL ",parttype=0)
mass  = rs.read_block(filename,"MASS",parttype=0)
dens  = rs.read_block(filename,"RHO ",parttype=0)
vol   = rs.read_block(filename,"VOL ",parttype=0)
u   = rs.read_block(filename,"U   ",parttype=0)
ids  = rs.read_block(filename,"ID  ",parttype=0)

radius = np.sqrt((pos[:,0] - 0.5 * BoxX)**2 + (pos[:,1] - 0.5 * BoxY)**2)        
phi = np.arctan2((pos[:,1] - 0.5 * BoxY),(pos[:,0] - 0.5 * BoxX))

vphi = -np.sin(phi) * vel[:,0] + np.cos(phi) * vel[:,1]
vr   =  np.cos(phi) * vel[:,0] + np.sin(phi) * vel[:,1]

densprofile = vecdiskdens(radius[np.argsort(radius)])
vphiprofilesq = vecvelprofilesq(radius[np.argsort(radius)])
示例#26
0
	print "\n"
	plt.close('all')
	
	#base = "/n/hernquistfs1/jsuresh/SPH_tests/02-2011-re-py/"
	base = "/n/hernquistfs1/jsuresh/SPH_tests/04-2011-1D/"
	prob_folder = "prob"+str(probn)+"/"
	probdir = base+prob_folder
	Ndir = probdir+"N_"+str(N_part)
	outputdir = Ndir+"/Ngb_"+str(N_gb)
	outdir = outputdir + "/outfiles"
	snapsdir = outputdir + "/snaps"
	# filename="./shock_tube/snapshot_"+str(num).zfill(3)		
	filename= snapsdir +"/snapshot_"+ str(snapnum).zfill(3) # +".hdf5"
	# print filename
	
	r=rs.read_block(filename,"RHO ",parttype=0)
	x=rs.read_block(filename,"POS ",parttype=0)[:,0]
	v=rs.read_block(filename,"VEL ",parttype=0)[:,0]
	u=rs.read_block(filename,"U   ",parttype=0)
	#P=co.GetPressure(u, r, gamma)
	#A=co.GetEntropy(u, r, gamma)
	m=rs.read_block(filename,"MASS",parttype=0)
	
	P = (ga-1)*u*r
	S = P/(r**ga)
	
	#Find wave center:
	current_wave_cent = x[argmax(r)]
	wave_cent = append(wave_cent, current_wave_cent)
	
	# Demarcate wave:
x = linspace(1, 30, 200)
vphiprofilesq = vecvelprofilesq(x)
DvphiDRprofile = vecDvphiDRprofile(x)
vphiprofilesqgrad = np.gradient(vphiprofilesq, np.diff(x).mean())
#plt.plot(x,vphiprofilesqgrad,'ro',mew=0.0,ms=2.0)
#plt.plot(x,DvphiDRprofile,'b')
#plt.show()

#open the snapshot header
filename = base + "snap_" + str(num).zfill(3)
header = rs.snapshot_header(filename)
time = header.time
BoxX, BoxY = header.boxsize, header.boxsize

pos = rs.read_block(filename, "POS ", parttype=0)
vel = rs.read_block(filename, "VEL ", parttype=0)
mass = rs.read_block(filename, "MASS", parttype=0)
dens = rs.read_block(filename, "RHO ", parttype=0)
vol = rs.read_block(filename, "VOL ", parttype=0)
u = rs.read_block(filename, "U   ", parttype=0)
ids = rs.read_block(filename, "ID  ", parttype=0)

radius = np.sqrt((pos[:, 0] - 0.5 * BoxX)**2 + (pos[:, 1] - 0.5 * BoxY)**2)
phi = np.arctan2((pos[:, 1] - 0.5 * BoxY), (pos[:, 0] - 0.5 * BoxX))

vphi = -np.sin(phi) * vel[:, 0] + np.cos(phi) * vel[:, 1]
vr = np.cos(phi) * vel[:, 0] + np.sin(phi) * vel[:, 1]

densprofile = vecdiskdens(radius[np.argsort(radius)])
vphiprofilesq = vecvelprofilesq(radius[np.argsort(radius)])
示例#28
0
        # define the M_HI and sigma_HI array
        M_ptype = np.zeros(len(halo_len), dtype=np.float64)
        sigma2_ptype = np.zeros(len(halo_len), dtype=np.float64)

        # do a loop over all files
        Omega_ptype, done = 0.0, False
        #start, end, end_gal = 0, end_halos[0], end_all_galaxies[0]
        #pars = [Number, start_h, end_h, start_g, end_g, halo_num, gal_num,
        #        gal_in_local_halo]
        pars = np.array([0, 0, halo_len[0], 0], dtype=np.int64)
        for i in xrange(filenum):

            snapshot = snapshot_root + 'snapdir_%03d/snap_%03d.%d' % (num, num,
                                                                      i)

            vel = rs.read_block(
                snapshot, 'VEL ', parttype=ptype, verbose=False) / np.sqrt(
                    1.0 + redshift)  #km/s

            mass = rs.read_block(
                snapshot, 'MASS', parttype=ptype, verbose=False) * 1e10
            mass = mass.astype(np.float32)

            # compute the HI mass in halos and galaxies
            if not (done):
                done = HIL.sigma_HI_halos(pars, done, halo_len, halo_vel, mass,
                                          vel, sigma2_ptype, M_ptype)

            Omega_ptype += np.sum(mass, dtype=np.float64)

            print '\ni = %03d: z=%.3f' % (i, redshift)
            print 'Omega(%d,z=%d) = %.6e' % (ptype, round(redshift),
示例#29
0
# define the arrays containing the mass and number of particles in each spherical shell
mass_shell_CDM = np.zeros((halos, bins), dtype=np.float64)
part_in_halo_CDM = np.zeros(halos, dtype=np.int64)

# do a loop over each subsnapshot
numbers = np.where(np.arange(filenum) % nprocs == myrank)[0]
for i in numbers:

    # find subfile name and read the number of particles in it
    snapshot = snapshot_root + 'snapdir_%03d/snap_%03d.%d' % (snapnum, snapnum,
                                                              i)
    header = rs.snapshot_header(snapshot)
    npart = header.npart

    ###################### CDM #######################
    pos = rs.read_block(snapshot, 'POS ', parttype=1, verbose=False) / 1e3
    pos = pos.astype(np.float32)
    mass = rs.read_block(snapshot, 'MASS', parttype=1, verbose=False) * 1e10
    mass = mass.astype(np.float32)

    # sort the positions of the particles
    data = SL.sort_3D_pos(pos,
                          BoxSize,
                          cell_size,
                          return_indexes=True,
                          return_offset=True)
    pos = data.pos_sorted
    mass = mass[data.indexes]
    offset_CDM = data.offset

    HIL.HI_profile(halo_pos, halo_R, halo_id, pos, mass, offset_CDM,
示例#30
0
    # define the M_HI and sigma_HI array
    M_HI      = np.zeros(len(halo_len), dtype=np.float64)
    sigma2_HI = np.zeros(len(halo_len), dtype=np.float64)


    # do a loop over all files
    Omega_HI, done = 0.0, False
    #start, end, end_gal = 0, end_halos[0], end_all_galaxies[0]
    #pars = [Number, start_h, end_h, start_g, end_g, halo_num, gal_num,
    #        gal_in_local_halo]
    pars = np.array([0, 0, halo_len[0], 0], dtype=np.int64)
    for i in xrange(filenum):

        snapshot = snapshot_root + 'snapdir_%03d/snap_%03d.%d'%(num,num,i)
        
        vel = rs.read_block(snapshot, 'VEL ', parttype=0, verbose=False)/np.sqrt(1.0+redshift) #km/s                         

        MHI   = rs.read_block(snapshot, 'NH  ', parttype=0, verbose=False)#HI/H
        mass  = rs.read_block(snapshot, 'MASS', parttype=0, verbose=False)*1e10
        SFR   = rs.read_block(snapshot, 'SFR ', parttype=0, verbose=False)
        indexes = np.where(SFR>0.0)[0];  del SFR

        # find the metallicity of star-forming particles
        metals = rs.read_block(snapshot, 'GZ  ', parttype=0, verbose=False)
        metals = metals[indexes]/0.0127

        # find densities of star-forming particles: units of h^2 Msun/Mpc^3
        rho = rs.read_block(snapshot, 'RHO ', parttype=0, verbose=False)*1e19
        rho = rho[indexes]

        # find volume and radius of star-forming particles
示例#31
0
    def read(self, block_name, parttype, fof_num, sub_num):
        if sub_num < 0 and fof_num < 0:
            # Load all of the non-FoF particles.
            off = (self.group_offset[-1, parttype] +
                   self.cat.GroupLenType[-1, parttype])
            left = 1e9  # reads the rest.
            print(off, left)

        if sub_num >= 0 and fof_num < 0:
            off = self.halo_offset[sub_num, parttype]
            left = self.cat.SubhaloLenType[sub_num, parttype]

        if fof_num >= 0 and sub_num < 0:
            off = self.group_offset[fof_num, parttype]
            left = self.cat.GroupLenType[fof_num, parttype]

        if sub_num >= 0 and fof_num >= 0:
            real_sub_num = sub_num + self.cat.GroupFirstSub[fof_num]
            off = self.halo_offset[real_sub_num, parttype]
            left = self.cat.SubhaloLenType[real_sub_num, parttype]

        if left == 0:
            if self.verbose:
                print("READHALO: no particles of type...returning")
            return

        # Get first file that contains particles of required halo/fof/etc.
        findex = np.argmax(self.file_type_numbers[:, parttype] > off) - 1
        # np.argmax returns 0 when the offset corresponds to a particle
        # in the last file.
        if findex == -1:
            findex = self.file_num - 1

        # Convert the overall offset to an offset for just the file given by
        # findex by subtracting off the number of particles in previous files.
        for fnr in range(0, findex):
            off -= (self.file_type_numbers[fnr + 1, parttype] -
                    self.file_type_numbers[fnr, parttype])

        # Read data from file(s).
        first = True
        for fnr in range(findex, self.file_num):
            path = self.filenames[fnr]

            head = readsnapHDF5.snapshot_header(path)
            nloc = head.npart[parttype]

            if nloc > off:
                if self.verbose:
                    print("READHALO: data found in %s" % path)
                start = off
                if nloc - off > left:
                    # All remaining particles are in this file.
                    count = left
                else:
                    # Read to end of file.
                    count = nloc - off

                block = readsnapHDF5.read_block(path,
                                                block_name,
                                                parttype,
                                                slab_start=start,
                                                slab_len=count)
                if first:
                    data = block
                    first = False
                else:
                    data = np.append(data, block, axis=0)

                left -= count
                off += count
            if left == 0:
                break
            off -= nloc

        gc.collect()

        return data
示例#32
0
def Illustris_region(snapshot_root,
                     snapnum,
                     TREECOOL_file,
                     x_min,
                     x_max,
                     y_min,
                     y_max,
                     z_min,
                     z_max,
                     padding,
                     fout,
                     redshift_space=False,
                     axis=0):

    # read snapshot and find number of subfiles
    snapshot = snapshot_root + 'snapdir_%03d/snap_%03d' % (snapnum, snapnum)
    header = rs.snapshot_header(snapshot)
    nall = header.nall
    redshift = header.redshift
    BoxSize = header.boxsize / 1e3  #Mpc/h
    filenum = header.filenum
    Omega_m = header.omega0
    Omega_L = header.omegaL
    h = header.hubble
    Hubble = 100.0 * np.sqrt(Omega_m *
                             (1.0 + redshift)**3 + Omega_L)  #km/s/(Mpc/h)

    if myrank == 0:
        print '\n'
        print 'BoxSize         = %.3f Mpc/h' % BoxSize
        print 'Number of files = %d' % filenum
        print 'Omega_m         = %.3f' % Omega_m
        print 'Omega_l         = %.3f' % Omega_L
        print 'redshift        = %.3f' % redshift

    # find the numbers each cpu will work on
    array = np.arange(0, filenum)
    numbers = np.where(array % nprocs == myrank)[0]

    # do a loop over the different realizations
    particles = 0
    for i in numbers:

        snapshot = snapshot_root + 'snapdir_%03d/snap_%03d.%d' % (snapnum,
                                                                  snapnum, i)
        pos = rs.read_block(snapshot, 'POS ', parttype=0, verbose=False) / 1e3
        pos = pos.astype(np.float32)

        # read velocities and displace particle positions
        if redshift_space:
            vel = rs.read_block(snapshot, 'VEL ', parttype=0,
                                verbose=False) / np.sqrt(1.0 + redshift)  #km/s
            RSL.pos_redshift_space(pos, vel, BoxSize, Hubble, redshift, axis)

        # check if particles are in the region
        indexes_region = np.where((pos[:,0]>=x_min-padding) & (pos[:,0]<=x_max+padding) &\
                                  (pos[:,1]>=y_min-padding) & (pos[:,1]<=y_max+padding) &\
                                  (pos[:,2]>=z_min-padding) & (pos[:,2]<=z_max+padding))[0]

        # if particles are not in the region continue
        local_particles = indexes_region.shape[0]
        print 'Myrank = %d ---> num = %d ---> part = %ld' % (myrank, i,
                                                             local_particles)
        if local_particles == 0: continue

        # find radii, HI and gas masses
        MHI = rs.read_block(snapshot, 'NH  ', parttype=0, verbose=False)  #HI/H
        mass = rs.read_block(snapshot, 'MASS', parttype=0,
                             verbose=False) * 1e10
        SFR = rs.read_block(snapshot, 'SFR ', parttype=0, verbose=False)
        indexes = np.where(SFR > 0.0)[0]
        del SFR

        # find the metallicity of star-forming particles
        metals = rs.read_block(snapshot, 'GZ  ', parttype=0, verbose=False)
        metals = metals[indexes] / 0.0127

        # find densities of star-forming particles: units of h^2 Msun/Mpc^3
        rho = rs.read_block(snapshot, 'RHO ', parttype=0, verbose=False) * 1e19
        Volume = mass / rho  #(Mpc/h)^3
        radii = (Volume / (4.0 * np.pi / 3.0))**(1.0 / 3.0)  #Mpc/h
        rho = rho[indexes]  #h^2 Msun/Mpc^3
        Volume = Volume[indexes]  #(Mpc/h)^3

        # find volume and radius of star-forming particles
        radii_SFR = (Volume / (4.0 * np.pi / 3.0))**(1.0 / 3.0)  #Mpc/h

        # find HI/H fraction for star-forming particles
        MHI[indexes] = HIL.Rahmati_HI_Illustris(rho,
                                                radii_SFR,
                                                metals,
                                                redshift,
                                                h,
                                                TREECOOL_file,
                                                Gamma=None,
                                                fac=1,
                                                correct_H2=True)  #HI/H
        MHI *= (0.76 * mass)

        # select the particles belonging to the region
        pos = pos[indexes_region]
        MHI = MHI[indexes_region]
        radii = radii[indexes_region]
        mass = mass[indexes_region]

        # write partial files
        new_size = particles + local_particles

        if particles == 0:
            f = h5py.File(fout[:-5] + '_%d.hdf5' % myrank, 'w')
            f.create_dataset('pos', data=pos, maxshape=(None, 3))
            f.create_dataset('M_HI', data=MHI, maxshape=(None, ))
            f.create_dataset('radii', data=radii, maxshape=(None, ))
            f.create_dataset('mass', data=mass, maxshape=(None, ))
        else:
            f = h5py.File(fout[:-5] + '_%d.hdf5' % myrank, 'a')
            pos_f = f['pos']
            pos_f.resize((new_size, 3))
            M_HI_f = f['M_HI']
            M_HI_f.resize((new_size, ))
            radii_f = f['radii']
            radii_f.resize((new_size, ))
            mass_f = f['mass']
            mass_f.resize((new_size, ))
            pos_f[particles:] = pos
            M_HI_f[particles:] = MHI
            radii_f[particles:] = radii
            mass_f[particles:] = mass
        f.close()
        particles += local_particles

    comm.Barrier()

    # sum the particles found in each cpu
    All_particles = 0
    All_particles = comm.reduce(particles, op=MPI.SUM, root=0)

    # Master will merge partial files into a file one
    if myrank == 0:

        print 'Found %d particles' % All_particles
        f = h5py.File(fout, 'w')

        f1 = h5py.File(fout[:-5] + '_0.hdf5', 'r')
        pos = f1['pos'][:]
        M_HI = f1['M_HI'][:]
        radii = f1['radii'][:]
        mass = f1['mass'][:]
        f1.close()

        particles = mass.shape[0]
        pos_f = f.create_dataset('pos', data=pos, maxshape=(None, 3))
        M_HI_f = f.create_dataset('M_HI', data=M_HI, maxshape=(None, ))
        radii_f = f.create_dataset('radii', data=radii, maxshape=(None, ))
        mass_f = f.create_dataset('mass', data=mass, maxshape=(None, ))

        for i in xrange(1, nprocs):
            f1 = h5py.File(fout[:-5] + '_%d.hdf5' % i, 'r')
            pos = f1['pos'][:]
            M_HI = f1['M_HI'][:]
            radii = f1['radii'][:]
            mass = f1['mass'][:]
            f1.close()

            size = mass.shape[0]

            pos_f.resize((particles + size, 3))
            pos_f[particles:] = pos
            M_HI_f.resize((particles + size, ))
            M_HI_f[particles:] = M_HI
            radii_f.resize((particles + size, ))
            radii_f[particles:] = radii
            mass_f.resize((particles + size, ))
            mass_f[particles:] = mass

            particles += size

        f.close()

        for i in xrange(nprocs):
            os.system('rm ' + fout[:-5] + '_%d.hdf5' % i)
import numpy as np
import readsnapHDF5

# Run specifics:
base = "/n/home07/dnelson/sims.tracers/256_20Mpc/output/"
#base = "/n/home07/dnelson/sims.tracers/512_20Mpc/output/"
save_dir = "/n/hernquistfs1/jsuresh/AREPO_morph/Tracers_v2/argsort_files/256_20Mpc/"

# Which snapshots to run on?
latest_snap = 234
earliest_snap = 140

snap_nums = np.arange(earliest_snap, latest_snap+1)

for current_snap in snap_nums:
	print "argsorting snapshot "+str(current_snap)
	snapname = base + "snap_"+str(current_snap).zfill(3)

	MC_ids = readsnapHDF5.read_block(snapname,"TFID",parttype=3)
	args = np.argsort(MC_ids)

	savefile = save_dir + str(current_snap).zfill(3) + ".args"
	f = open(savefile,'wb')
	args.astype("uint64").tofile(f)
	f.close()
print len(catPos)
print len(catPos) - len(badHaloIndex)
print len(groupFirstSub)

#Creating output array (to be saved into numpy file)
outArray = np.zeros([len(groupFirstSub), 3+len(search_r_array)])
outArray[:,0] = groupFirstSub
outArray[:,1] = groupNum[groupFirstSub]
outArray[:,2] = subhalo_masses[groupFirstSub]

badHaloMass = np.zeros([len(badHaloPos),len(search_r_array)])

i = 0
while (i < len(partFileArray)):
	#loading in particle positional data
	partPos = rs.read_block(partFileArray[i], 'POS ', parttype=pNum, verbose=True).astype('float')
	partPos = partPos*(cc.h_little**-1.)
	#loading in particle mass data
	partMass = rs.read_block(partFileArray[i], 'MASS', parttype=pNum, verbose=True).astype('float')
	partMass = np.reshape(partMass, partMass.shape[0])*1e10*(cc.h_little**-1.)
	
	print "Number of Particles: " + str(len(partMass))
	
	
	#Setting up the the KD Tree
	kdTree = sp.cKDTree(partPos)

	j = 0
	while (j < len(groupFirstSub)):
		k = 0
		while (k < len(search_r_array)):
示例#35
0
M_HI_halo = np.zeros(halos, dtype=np.float64)

# find the numbers each cpu will work on
array   = np.arange(0, filenum)
numbers = np.where(array%nprocs==myrank)[0]

# do a loop over each subsnapshot
for i in numbers:

    # find subfile name and read the number of particles in it
    snapshot = snapshot_root + 'snapdir_%03d/snap_%03d.%d'%(snapnum, snapnum, i)
    header = rs.snapshot_header(snapshot)
    npart  = header.npart 

    print '\nWorking with subfile %03d : %d'%(i,myrank)
    pos  = rs.read_block(snapshot, 'POS ', parttype=0, verbose=False)/1e3
    pos  = pos.astype(np.float32)
    MHI  = rs.read_block(snapshot, 'NH  ', parttype=0, verbose=False)#HI/H
    mass = rs.read_block(snapshot, 'MASS', parttype=0, verbose=False)*1e10
    SFR  = rs.read_block(snapshot, 'SFR ', parttype=0, verbose=False)
    indexes = np.where(SFR>0.0)[0];  del SFR

    # find the metallicity of star-forming particles
    metals = rs.read_block(snapshot, 'GZ  ', parttype=0, verbose=False)
    metals = metals[indexes]/0.0127

    # find densities of star-forming particles: units of h^2 Msun/Mpc^3
    rho = rs.read_block(snapshot, 'RHO ', parttype=0, verbose=False)*1e19
    Volume = mass/rho                            #(Mpc/h)^3
    radii  = (Volume/(4.0*np.pi/3.0))**(1.0/3.0) #Mpc/h 
radial_bins=np.append(radial_bins1,radial_bins2)
#radial_bins=MasstoRadius(np.linspace(0.0,1.0,INTERPOL_BINS)*mass_bins_gas.max())
mass_bins_gas=vecGasMass(radial_bins)
MasstoRadius=InterpolatedUnivariateSpline(mass_bins_gas, radial_bins,k=3)


if not (load_mesh):
  print "    Inversion sampling..."
  #generate random positions gas
  #fully random
  radius_gas= MasstoRadius(rd.random_sample(N_gas)*mass_bins_gas.max())
  print "done."
else:
  #read-in, positions of gas cells
  print "        loading mesh..."
  pos = rs.read_block(mesh_file,"POS ",parttype=0)
  pos[:,0],pos[:,1], pos[:,2] = (pos[:,0]- 0.5 * BOX_WIDTH), (pos[:,1]- 0.5 * BOX_WIDTH), (pos[:,2]- 0.5 * BOX_HEIGHT)
  radius_gas = np.sqrt(pos[:,0]**2 + pos[:,1]**2)
  vertical_gas = pos[:,2]

  ind = (np.abs(vertical_gas) < z_max) & (radius_gas < R_max)
  radius_gas = radius_gas[ind]
  vertical_gas = vertical_gas[ind]
  
  print "done"
  

#first we bin our newly sampled list of radii
bin_inds=np.digitize(radius_gas,radial_bins)

#initialize some variables and functions
def Illustris_halo(snapshot_root, snapnum, halo_number, TREECOOL_file, fout,
                   ptype=0):

    # find snapshot name and read header
    snapshot = snapshot_root + 'snapdir_%03d/snap_%03d'%(snapnum, snapnum)
    header   = rs.snapshot_header(snapshot)
    redshift = header.redshift
    BoxSize  = header.boxsize/1e3 #Mpc/h
    filenum  = header.filenum
    Omega_m  = header.omega0
    Omega_L  = header.omegaL
    h        = header.hubble
    massarr  = header.massarr*1e10 #Msun/h

    print '\nBoxSize         = %.1f Mpc/h'%BoxSize
    print 'Number of files = %d'%filenum
    print 'Omega_m         = %.3f'%Omega_m
    print 'Omega_l         = %.3f'%Omega_L
    print 'redshift        = %.3f'%redshift

    # read number of particles in halos and subhalos and number of subhalos
    halos = groupcat.loadHalos(snapshot_root, snapnum, 
            fields=['GroupLenType','GroupPos','GroupMass'])
    halo_len  = halos['GroupLenType'][:,ptype]  
    halo_pos  = halos['GroupPos']/1e3
    halo_mass = halos['GroupMass']*1e10
    del halos


    # find where the halo starts and ends in the file
    begin = np.sum(halo_len[:halo_number], dtype=np.int64)
    end   = begin + halo_len[halo_number]
    print begin,end

    # do a loop over all snapshot subfiles
    f = h5py.File(fout,'w')
    pos_f   = f.create_dataset('pos',   (0,3),  maxshape=(None,3))
    vel_f   = f.create_dataset('vel',   (0,3),  maxshape=(None,3))
    if ptype==0:
        mass_f  = f.create_dataset('mass',  (0,),   maxshape=(None,))
        MHI_f   = f.create_dataset('M_HI',  (0,),   maxshape=(None,))
        radii_f = f.create_dataset('radii', (0,),   maxshape=(None,))
    if ptype==1:
        radii_f = f.create_dataset('radii',  (0,),   maxshape=(None,))
        mass_f  = f.create_dataset('mass_c', (0,),   maxshape=(None,))

    begin_subfile, particles = 0, 0
    for i in xrange(filenum):

        # find subfile name and read the number of particles in it
        snapshot = snapshot_root + 'snapdir_%03d/snap_%03d.%d'%(snapnum, snapnum, i)
        header = rs.snapshot_header(snapshot)
        npart  = header.npart 

        end_subfile = begin_subfile + npart[ptype]

        # if all particles in the halo has been read exit loop
        if end<begin_subfile:  break

        # if the subfile does not contain any particle move to next subfile
        if begin>end_subfile:
            begin_subfile = end_subfile;  continue


        print 'Working with subfile %03d'%i
        pos  = rs.read_block(snapshot, 'POS ', parttype=ptype, 
                             verbose=False)/1e3
        pos  = pos.astype(np.float32)
        vel  = rs.read_block(snapshot, 'VEL ', parttype=ptype, 
                             verbose=False)/np.sqrt(1.0+redshift) #km/s

        if ptype==0:
            MHI  = rs.read_block(snapshot, 'NH  ', parttype=0,
                                 verbose=False)#HI/H
            mass = rs.read_block(snapshot, 'MASS', parttype=0,
                                 verbose=False)*1e10
            SFR  = rs.read_block(snapshot, 'SFR ', parttype=0,
                                 verbose=False)
            indexes = np.where(SFR>0.0)[0];  del SFR

            # find the metallicity of star-forming particles
            metals = rs.read_block(snapshot, 'GZ  ', parttype=0, verbose=False)
            metals = metals[indexes]/0.0127

            # find densities of star-forming particles: units of h^2 Msun/Mpc^3
            rho = rs.read_block(snapshot, 'RHO ', parttype=0, 
                                verbose=False)*1e19
            Volume = mass/rho                            #(Mpc/h)^3
            radii  = (Volume/(4.0*np.pi/3.0))**(1.0/3.0) #Mpc/h 

            # find density and radius of star-forming particles
            radii_SFR = radii[indexes]    
            rho       = rho[indexes]

            # find HI/H fraction for star-forming particles
            MHI[indexes] = HIL.Rahmati_HI_Illustris(rho, radii_SFR, metals,
                                                    redshift, h, TREECOOL_file,
                                                    Gamma=None, fac=1,
                                                    correct_H2=True) #HI/H
            MHI *= (0.76*mass)
            
        if ptype==1:
            radii = rs.read_block(snapshot, 'SFHS', parttype=1,
                                  verbose=False)/1e3 #Mpc/h
            mass = np.ones(len(radii))*massarr[1]
            

        # find the indexes of current subfile that contribute to halo
        begin_array = begin - begin_subfile
        end_array   = begin_array + (end-begin)

        if end>end_subfile:
            end_array = end_subfile - begin_subfile
            begin     = end_subfile

        new_size = particles + (end_array - begin_array)

        pos_f.resize((new_size,3))
        pos_f[particles:] = pos[begin_array:end_array]
        vel_f.resize((new_size,3))
        vel_f[particles:] = vel[begin_array:end_array]

        if ptype==0:
            mass_f.resize((new_size,))
            mass_f[particles:] = mass[begin_array:end_array]

            MHI_f.resize((new_size,))
            MHI_f[particles:] = MHI[begin_array:end_array]

            radii_f.resize((new_size,))
            radii_f[particles:] = radii[begin_array:end_array]

        if ptype==1:
            radii_f.resize((new_size,))
            radii_f[particles:] = radii[begin_array:end_array]

            mass_f.resize((new_size,))
            mass_f[particles:] = mass[begin_array:end_array]

        particles = new_size
        begin_subfile = end_subfile


    f.close()
    print 'Halo mass = %.3e'%halo_mass[halo_number]
    print 'Halo pos  =',halo_pos[halo_number]
    print 'Number of particles in the halo = %ld'%particles