def gas_image(dir='./', snapnum=0, band_ids = [9,10,11], this_file=None, center=None, edge_on=False, cosmo_wrap = False, \
massmap=False, projaxis=0, sfr_image=False, unit_length_mpc=False, include_lighting=True, **kwargs):
snap_ext = str(snapnum).zfill(3)
# LZK
# SNAP_PREFIX = "/snapshot_"
SNAP_PREFIX = "/snap_"
if this_file == None:
print "setting snapshot"
this_file = dir + SNAP_PREFIX + snap_ext
if not os.path.exists(this_file + '.hdf5'):
this_file = dir + '/snapdir_' + snap_ext + SNAP_PREFIX + snap_ext
print this_file
gas_data = gas_image_data(this_file, **kwargs)
if cosmo_wrap:
boxsize = np.max(gas_data.gas_xyz) - np.min(gas_data.gas_xyz)
if not (center is None):
gas_data.gas_xyz[:, 0] -= np.mean(center[0])
gas_data.gas_xyz[:, 1] -= np.mean(center[1])
gas_data.gas_xyz[:, 2] -= np.mean(center[2])
else:
bh_xyz = np.array(ws.read_block(this_file, 'POS ', parttype=5))
gas_data.gas_xyz[:, 0] -= np.mean(bh_xyz[:, 0])
gas_data.gas_xyz[:, 1] -= np.mean(bh_xyz[:, 1])
gas_data.gas_xyz[:, 2] -= np.mean(bh_xyz[:, 2])
if cosmo_wrap:
gas_data.gas_xyz[gas_data.gas_xyz[:, 0] > boxsize / 2.0, 0] -= boxsize
gas_data.gas_xyz[gas_data.gas_xyz[:, 1] > boxsize / 2.0, 1] -= boxsize
gas_data.gas_xyz[gas_data.gas_xyz[:, 2] > boxsize / 2.0, 2] -= boxsize
gas_data.gas_xyz[gas_data.gas_xyz[:, 0] < -1.0 * boxsize / 2.0,
0] += boxsize
gas_data.gas_xyz[gas_data.gas_xyz[:, 1] < -1.0 * boxsize / 2.0,
1] += boxsize
gas_data.gas_xyz[gas_data.gas_xyz[:, 2] < -1.0 * boxsize / 2.0,
2] += boxsize
gas_map_temperature_cuts = np.array([300., 2.0e4, 3.0e5])
kernel_widths = np.array([0.8, 0.3, 0.6])
if unit_length_mpc:
KAPPA_UNITS = 2.0885 * np.array([1.1, 2.0, 1.5]) / 1000.0 / 1000.0
else:
KAPPA_UNITS = 2.0885 * np.array([1.1, 2.0, 1.5])
if massmap:
color_weights = np.ones_like(gas_data.gas_mass)
else:
# this just returns temperature so that we can determine colors
color_weights = units.springel_units.gas_code_to_temperature(
gas_data.gas_u, gas_data.gas_nume)
if sfr_image:
weights = gas_data.gas_sfr
else:
weights = gas_data.gas_mass
h_all = gas_data.gas_hsml
try:
xrange = kwargs['xrange']
yrange = kwargs['yrange']
if xrange != None and yrange != None:
fov = np.max([xrange, yrange])
print np.min(gas_data.gas_xyz[:, 0]), np.max(gas_data.gas_xyz[:,
0])
ok = (gas_data.gas_xyz[:,0] > -1.1*fov) & (gas_data.gas_xyz[:,0] < 1.1*fov) & \
(gas_data.gas_xyz[:,1] > -1.1*fov) & (gas_data.gas_xyz[:,1] < 1.1*fov) & \
(gas_data.gas_xyz[:,2] > -1.1*fov) & (gas_data.gas_xyz[:,2] < 1.1*fov)
print "Found xrange/yrange arguments in gas_images() call. Clipping down to {:d} particles from {:d}".format(
np.sum(ok), len(ok))
gas_data.gas_xyz = gas_data.gas_xyz[ok, :]
color_weights = color_weights[ok]
weights = weights[ok]
h_all = h_all[ok]
except:
print "failed to clip"
try:
ok = weights > 0
if np.sum(ok) < len(ok):
print "Found particles with negative weights in gas_images() call. Clipping down to {:d} particles from {:d}".format(
np.sum(ok), len(ok))
gas_data.gas_xyz = gas_data.gas_xyz[ok, :]
color_weights = color_weights[ok]
weights = weights[ok]
h_all = h_all[ok]
except:
print "failed to clip based on zero-weight values"
sys.exit()
if edge_on:
x_index = 0
y_index = 2
z_index = 1
elif projaxis == 0:
x_index = 0
y_index = 1
z_index = 2
elif projaxis == 1:
x_index = 2
y_index = 0
z_index = 1
elif projaxis == 2:
x_index = 1
y_index = 2
z_index = 0
else:
x_index = 0
y_index = 1
z_index = 2
if massmap or sfr_image:
image, out_u, out_g, out_r = projection.raytrace_projection_compute( gas_data.gas_xyz[:,x_index], gas_data.gas_xyz[:,y_index], gas_data.gas_xyz[:,z_index], \
h_all[:], weights[:], weights[:], weights[:], weights[:], \
0.0, 0.0, 0.0, TRIM_PARTICLES=0, \
**kwargs )
image24, massmap = makepic.make_threeband_image_process_bandmaps( out_r,out_g,out_u, \
**kwargs)
return np.log10(image), np.log10(image)
#def raytrace_projection_compute( x, y, z, hsml, mass, wt1, wt2, wt3, \
# kappa_1, kappa_2, kappa_3, xrange=0, yrange=0, zrange=0, pixels=720, \
# TRIM_PARTICLES=1 ):
else:
print "executing projection routine."
out_gas,out_u,out_g,out_r = projection.gas_raytrace_temperature( \
gas_map_temperature_cuts, \
gas_data.gas_xyz[:,x_index], gas_data.gas_xyz[:,y_index], gas_data.gas_xyz[:,z_index],
color_weights[:], weights[:], h_all[:], \
kernel_width=kernel_widths,\
isosurfaces = 1,\
KAPPA_UNITS = KAPPA_UNITS,\
add_temperature_weights = 0,\
**kwargs)
print "mapping u/g/r to three band image."
image24, massmap = makepic.make_threeband_image_process_bandmaps(
out_r, out_g, out_u, **kwargs)
print "layering three band image."
image24 = makepic.layer_band_images(image24,
massmap) # layer_band_images...
if (include_lighting == 1 and np.sum(image24[:, :, :2]) > 0):
print "adding lighting"
image24 = makepic.include_lighting(image24, massmap)
if False: #(include_lighting==1 and np.sum(image24[:,:,:2]) > 0 ):
print "Lighting is being included!!! "
light = viscolors.CustomLightSource(azdeg=0, altdeg=65)
if (len(massmap.shape) > 2):
## do some clipping to regulate the lighting:
elevation = massmap.sum(axis=2)
minden = maxden / dynrange
print " "
print elevation.min(), elevation.max(), elevation.mean()
print minden, maxden
print " "
elevation = (elevation - minden) / (maxden - minden)
elevation[elevation < 0.] = 0.
elevation[elevation > 1.] = 1.
elevation *= maxden
grad_max = maxden / 5.
grad_max = maxden / 6.
#image24_lit = light.shade_rgb(image24, massmap.sum(axis=2))
image24_lit = light.shade_rgb(image24,
elevation,
vmin=-grad_max,
vmax=grad_max)
else:
image24_lit = light.shade(massmap, matplotlib.cm.get_cmap(
'hot')) #reversed args # ptorrey -- important
image24 = image24_lit
else:
print "Lighting is not being included :( "
return image24, massmap
def stellar_image(dir='./', snapnum=0, band_ids = [9,10,11],
cosmo=False,
this_file=None, center=None,
snapbase='snapshot',
illustris_tng=False,
tng_subnr=None,
tng_groupnr=None,
cosmo_wrap=False,
massmap=False,
projaxis=0,
**kwargs):
if this_file==None:
snap_ext = "000"+str(snapnum)
snap_ext = snap_ext[-3:]
this_file = dir+'/'+snapbase+'_'+snap_ext
if not os.path.exists(this_file+'.hdf5'):
this_file = dir+'/snapdir_'+snap_ext+'/'+snapbase+'_'+snap_ext
print "setting snapshot to {:s}".format( this_file )
if illustris_tng==True:
print "Making a TNG stellar image."
star_data = tng_stellar_image_data(this_file, fof_num=tng_groupnr, sub_num=tng_subnr, **kwargs)
else:
star_data = stellar_image_data(this_file, **kwargs)
if cosmo_wrap:
boxsize = star_data.head.boxsize #np.max( star_data.gas_xyz ) - np.min( star_data.gas_xyz )
if illustris_tng and isinstance( tng_subnr, int ):
image_center = readsubf.subfind_catalog( dir, snapnum, subcat=True, grpcat=False,
keysel=['SubhaloPos'] )
center = image_center.SubhaloPos[tng_subnr, :]
elif illustris_tng and isinstance( tng_groupnr, int ):
image_center = readsubf.subfind_catalog( dir, snapnum, subcat=True, grpcat=True,
keysel=['SubhaloPos', 'GroupFirstSub'] )
center = image_center.SubhaloPos[ image_center.GroupFirstSub[tng_groupnr], :]
elif center is not None:
center = center
else:
try:
bh_id = np.array(ws.read_block(this_file,'ID ', parttype=5))
bh_xyz = np.array(ws.read_block(this_file,'POS ', parttype=5))
center = [ np.mean(bh_xyz[0,0]), np.mean(bh_xyz[0,1]), np.mean(bh_xyz[0,2]) ]
except:
center = [0.0, 0.0, 0.0 ]
star_data.gas_xyz[:,0] -= np.mean(center[0])
star_data.gas_xyz[:,1] -= np.mean(center[1])
star_data.gas_xyz[:,2] -= np.mean(center[2])
star_data.star_xyz[:,0] -= np.mean(center[0])
star_data.star_xyz[:,1] -= np.mean(center[1])
star_data.star_xyz[:,2] -= np.mean(center[2])
if cosmo_wrap:
star_data.star_xyz[ star_data.star_xyz[:,0] > boxsize/2.0 ,0] -= boxsize
star_data.star_xyz[ star_data.star_xyz[:,1] > boxsize/2.0 ,1] -= boxsize
star_data.star_xyz[ star_data.star_xyz[:,2] > boxsize/2.0 ,2] -= boxsize
star_data.star_xyz[ star_data.star_xyz[:,0] < -1.0* boxsize/2.0 ,0] += boxsize
star_data.star_xyz[ star_data.star_xyz[:,1] < -1.0* boxsize/2.0 ,1] += boxsize
star_data.star_xyz[ star_data.star_xyz[:,2] < -1.0* boxsize/2.0 ,2] += boxsize
star_data.gas_xyz[ star_data.gas_xyz[:,0] > boxsize/2.0 ,0] -= boxsize
star_data.gas_xyz[ star_data.gas_xyz[:,1] > boxsize/2.0 ,1] -= boxsize
star_data.gas_xyz[ star_data.gas_xyz[:,2] > boxsize/2.0 ,2] -= boxsize
star_data.gas_xyz[ star_data.gas_xyz[:,0] < -1.0* boxsize/2.0 ,0] += boxsize
star_data.gas_xyz[ star_data.gas_xyz[:,1] < -1.0* boxsize/2.0 ,1] += boxsize
star_data.gas_xyz[ star_data.gas_xyz[:,2] < -1.0* boxsize/2.0 ,2] += boxsize
try:
xrange = kwargs['xrange']
yrange = kwargs['yrange']
if xrange != None and yrange != None:
fov = np.max( [xrange, yrange] )
ok = (star_data.star_xyz[:,0] > -1.1*fov) & (star_data.star_xyz[:,0] < 1.1*fov) & \
(star_data.star_xyz[:,1] > -1.1*fov) & (star_data.star_xyz[:,1] < 1.1*fov) & \
(star_data.star_xyz[:,2] > -1.1*fov) & (star_data.star_xyz[:,2] < 1.1*fov) & \
(star_data.star_bt > 0)
if np.sum(ok) > 10:
star_data.star_xyz = star_data.star_xyz[ok,:]
star_data.star_bt = star_data.star_bt[ok]
star_data.star_z = star_data.star_z[ok]
star_data.star_mass= star_data.star_mass[ok]
else:
print "WARN: I found 10 or less particles in the FOV. Expect an error."
except:
print "Didnt find xrange/yrange values. Not clipping. Using some default range..."
if cosmo:
star_age = np.array([ units.age_from_a(star_data.star_bt[index], a0=star_data.head.time) for index in np.arange(star_data.star_bt.shape[0]) ])
else:
star_age = star_data.head.time - star_data.star_bt
star_hsml = calc_hsml.get_particle_hsml( star_data.star_xyz[:,0], star_data.star_xyz[:,1], star_data.star_xyz[:,2] , **kwargs )
print "min/max of star_hslm : {:f}|{:f}".format(np.min(star_hsml), np.max(star_hsml))
if projaxis==0:
x_axis_index = 0
y_axis_index = 1
z_axis_index = 2
if projaxis==1:
x_axis_index = 2
y_axis_index = 0
z_axis_index = 1
if projaxis==2:
x_axis_index = 1
y_axis_index = 2
z_axis_index = 0
print x_axis_index, y_axis_index, z_axis_index
# This calls "stellar raytrace" which itself is a wrapper for "raytrace projection compute".
# out_gas should have units of mass/area; out_u/g/r have funky units, but are ~luminosity/area
# NO CLIPPING TO THIS POINT!
out_gas,out_u,out_g,out_r = projection.stellar_raytrace( star_data.star_xyz[:,x_axis_index], star_data.star_xyz[:,y_axis_index], star_data.star_xyz[:,z_axis_index], \
star_data.star_mass[:], star_age[:], star_data.star_z[:], star_hsml[:], \
star_data.gas_xyz[:,x_axis_index], star_data.gas_xyz[:,y_axis_index], star_data.gas_xyz[:,z_axis_index], star_data.gas_mass, \
star_data.gas_z, star_data.gas_hsml, \
band_ids=band_ids, \
**kwargs)
print "within stellar_image the u,g,r images have sizes and min/max values of:"
print out_gas.shape
print out_u.shape
print out_g.shape
print out_r.shape
print np.min(out_u), np.max(out_u)
print np.min(out_g), np.max(out_g)
print np.min(out_r), np.max(out_r)
if(np.array(out_gas).size<=1):
out_gas=out_u=out_g=out_r=np.zeros((pixels,pixels))
image24=massmap=np.zeros((pixels,pixels,3))
print "SIZE OF OUTPUT IMAGE IS ZERO. RETURNING IMAGES WITH JUST ZEROS."
else:
## make the resulting maps into an image
# THIS IS WHERE CLIPPING SHOULD HAPPEN USING maxden and dynrange
image24, massmap = makepic.make_threeband_image_process_bandmaps( out_r,out_g,out_u, **kwargs)
print "Making a three-band image from the R-G-B maps."
print "within stellar_image image24 has a size:"
print image24.shape
print "and a min/max in the first channel of"
print np.min(image24[:,:,0]), np.max(image24[:,:,0])
include_lighting=0
if (include_lighting==1):
#light = matplotlib.colors.LightSource(azdeg=0,altdeg=65)
light = viscolors.CustomLightSource(azdeg=0,altdeg=65)
if (len(massmap.shape)>2):
## do some clipping to regulate the lighting:
elevation = massmap.sum(axis=2)
minden = maxden / dynrange
elevation = (elevation - minden) / (maxden - minden)
elevation[elevation < 0.] = 0.
elevation[elevation > 1.] = 1.
elevation *= maxden
grad_max = maxden / 5.
grad_max = maxden / 6.
#image24_lit = light.shade_rgb(image24, massmap.sum(axis=2))
image24_lit = light.shade_rgb(image24, elevation, vmin=-grad_max, vmax=grad_max)
else:
image24_lit = light.shade(image24, massmap)
image24 = image24_lit
return image24
def sz_gas_image(dir='./',
snapnum=0,
this_file=None,
center=None,
dynrange=None,
maxden=None,
**kwargs):
snap_ext = str(snapnum).zfill(3)
if this_file == None:
this_file = dir + '/snapshot_' + snap_ext
if not os.path.exists(this_file + '.hdf5'):
this_file = dir + '/snapdir_' + snap_ext + '/snapshot_' + snap_ext
this_file = this_file + '.hdf5'
print this_file
data = gas_image_data(this_file, **kwargs)
try:
bh_xyz = np.array(ws.read_block(this_file, 'POS ', parttype=5))
except:
bh_xyz = np.array([0, 0, 0])
if center != None:
data.xyz[:, 0] -= np.mean(bh_xyz[:, 0])
data.xyz[:, 1] -= np.mean(bh_xyz[:, 1])
data.xyz[:, 2] -= np.mean(bh_xyz[:, 2])
effective_mass = data.gas_nume * data.gas_mass * data.gas_temperature
weights = effective_mass
h_all = data.gas_hsml
try:
xrange = kwargs['xrange']
yrange = kwargs['yrange']
except:
xrange = [-1e10, 1e10]
yrange = [-1e10, 1e10]
fov = np.max([xrange, yrange])
ok = (data.gas_xyz[:,0] > -1.1*fov) & (data.gas_xyz[:,0] < 1.1*fov) & \
(data.gas_xyz[:,1] > -1.1*fov) & (data.gas_xyz[:,1] < 1.1*fov) & \
(data.gas_xyz[:,2] > -1.1*fov) & (data.gas_xyz[:,2] < 1.1*fov)
data.gas_xyz = data.gas_xyz[ok, :]
weights = weights[ok]
h_all = h_all[ok]
ok = np.arange(h_all.shape[0])
image, out_u, out_g, out_r = projection.raytrace_projection_compute(
data.gas_xyz[:, 0],
data.gas_xyz[:, 1],
data.gas_xyz[:, 2],
h_all,
weights,
weights,
weights,
weights,
0,
0,
0,
TRIM_PARTICLES=0)
print "image max/min is {:g}/{:g}".format(np.min(image), np.max(image))
if ((maxden != None) and (dynrange != None)):
image[image > maxden] = maxden
image[image < maxden / dynrange] = maxden / dynrange
print "image clipping is {:g}/{:g}".format(maxden, maxden / dynrange)
else:
print "no clipping (maxden/dynrange not set)..."
return np.log10(image)
def __init__(self,
this_file,
next_file=None,
interp_coef=0.0,
interp=True,
**kwargs):
if next_file == None or interp == False: # no interp
print "have a single file, no interp"
# self.head = ws.snapshot_header(this_file)
self.gas_xyz = np.array(
ws.read_block(this_file, 'POS ', parttype=0))
self.gas_mass = np.array(
ws.read_block(this_file, 'MASS', parttype=0))
self.gas_u = np.array(ws.read_block(this_file, 'U ', parttype=0))
try:
vol = np.array(ws.read_block(this_file, 'VOL ', parttype=0))
self.gas_hsml = (vol * 3.0 / (4.0 * 3.14159))**0.333
except:
print "failed to load gas hsml via volumes. Trying hsml values directly."
try:
self.gas_hsml = np.array(
ws.read_block(this_file, 'HSML', parttype=0))
except:
print "failed to load gas hsml via hsml (directly). Loading via mass and density."
try:
self.gas_hsml = 3.0 * (
3.0 / (4.0 * 3.14159) * self.gas_mass / np.array(
ws.read_block(this_file, 'RHO ',
parttype=0)))**0.33333
except:
print "failed to load gas hsml values for a third time via mass and density. stopping."
sys.exit()
try:
self.gas_sfr = np.array(
ws.read_block(this_file, 'SFR ', parttype=0))
except:
self.gas_sfr = np.zeros_like(gas_mass)
try:
self.gas_nume = np.array(
ws.read_block(this_file, 'NE ', parttype=0))
except:
print "failed to load gas nume."
self.gas_nume = np.ones_like(self.gas_mass)
try:
self.gas_z = np.array(
ws.read_block(this_file, 'Z ', parttype=0))
except:
try:
self.gas_z = np.array(
ws.read_block(this_file, 'GZ ', parttype=0))
except:
print "failed to load metallicities. Setting to solar."
self.gas_z = np.ones_like(self.gas_mass) * 0.02
else:
# head = ws.snapshot_header(this_file)
gas_id1 = np.array(ws.read_block(this_file, 'ID ', parttype=0))
order = np.argsort(gas_id1)
gas_id1 = gas_id1[order]
gas_xyz1 = np.array(ws.read_block(this_file, 'POS ',
parttype=0))[order, :]
gas_hsml1 = np.array(ws.read_block(this_file, 'HSML',
parttype=0))[order]
gas_mass1 = np.array(ws.read_block(this_file, 'MASS',
parttype=0))[order]
gas_u1 = np.array(ws.read_block(this_file, 'U ',
parttype=0))[order]
gas_nume1 = np.array(ws.read_block(this_file, 'NE ',
parttype=0))[order]
try:
gas_z1 = np.array(ws.read_block(this_file, 'Z ', parttype=0))
gas_z1 = np.sum(gas_z1[order, 2:], axis=1)
except:
gas_z1 = np.array(ws.read_block(this_file, 'GZ ',
parttype=0))[order]
gas_id2 = np.array(ws.read_block(next_file, 'ID ', parttype=0))
order = np.argsort(gas_id2)
gas_id2 = gas_id2[order]
gas_xyz2 = np.array(ws.read_block(next_file, 'POS ',
parttype=0))[order, :]
gas_hsml2 = np.array(ws.read_block(next_file, 'HSML',
parttype=0))[order]
gas_mass2 = np.array(ws.read_block(next_file, 'MASS',
parttype=0))[order]
gas_u2 = np.array(ws.read_block(next_file, 'U ',
parttype=0))[order]
gas_nume2 = np.array(ws.read_block(next_file, 'NE ',
parttype=0))[order]
try:
gas_z2 = np.array(ws.read_block(next_file, 'Z ', parttype=0))
gas_z2 = np.sum(gas_z2[order, 2:], axis=1)
except:
gas_z2 = np.array(ws.read_block(next_file, 'GZ ',
parttype=0))[order]
gas_id2 = np.array(ws.read_block(next_file, 'ID ', parttype=0))
i1 = np.in1d(gas_id1, gas_id2)
gas_xyz1 = gas_xyz1[i1, :]
gas_hsml1 = gas_hsml1[i1]
gas_mass1 = gas_mass1[i1]
gas_u1 = gas_u1[i1]
gas_nume1 = gas_nume1[i1]
gas_z1 = gas_z1[i1]
i1 = np.in1d(gas_id2, gas_id1)
gas_xyz2 = gas_xyz2[i1, :]
gas_hsml2 = gas_hsml2[i1]
gas_mass2 = gas_mass2[i1]
gas_u2 = gas_u2[i1]
gas_nume2 = gas_nume2[i1]
gas_z2 = gas_z2[i1]
self.gas_xyz = (1.0 -
interp_coef) * gas_xyz1 + interp_coef * gas_xyz2
self.gas_hsml = (1.0 -
interp_coef) * gas_hsml1 + interp_coef * gas_hsml2
self.gas_mass = (1.0 -
interp_coef) * gas_mass1 + interp_coef * gas_mass2
self.gas_u = (1.0 - interp_coef) * gas_u1 + interp_coef * gas_u2
self.gas_nume = (1.0 -
interp_coef) * gas_nume1 + interp_coef * gas_nume2
self.gas_z = (1.0 - interp_coef) * gas_z1 + interp_coef * gas_z2
rundir = '/n/home01/ptorrey/Runs/ISM_Boxes/'
snapnr = 50
little_h = 0.6774
print_phase_fractions = True
for snapnr in [11]:
this_scalefactor = scalefactors[snapnr]
this_redshift = 1.0 / this_scalefactor - 1.0
for runname in ['explicit_feedback_256_soft']: #[455, 910, 1820]:
base = rundir + runname + '/'
path = base + 'output/snapdir_' + str(snapnr).zfill(
3) + '/snapshot_' + str(snapnr).zfill(3)
print "density loading"
rho = np.array(ws.read_block(path, 'RHO ', 0)).astype(
np.float32) / 1000.0**3.0
print "Internal Energy loading"
u = np.array(ws.read_block(path, 'U ', 0)).astype(np.float32)
print "Electron Abundance loading"
ne = np.array(ws.read_block(path, 'NE ', 0)).astype(np.float32)
print "Temperature Loading"
t = np.array(units.gas_code_to_temperature(u, ne)).astype(np.float32)
del u
del ne
print "Mass Loading"
m = np.array(ws.read_block(path, 'MASS', 0)).astype(np.float32)
ms = np.array(ws.read_block(path, 'MASS', 4)).astype(np.float32)
rho = np.array(units.gas_code_to_cgs_density(rho)).astype(
np.float32) * little_h * little_h / this_scalefactor**3
def old_repo_gas_image(dir='./',
snapnum=0,
band_ids=[9, 10, 11],
this_file=None,
center=None,
edge_on=False,
cosmo_wrap=False,
massmap=False,
**kwargs):
if this_file == None:
snap_ext = str(snapnum).zfill(3)
this_file = dir + '/snapshot_' + snap_ext
if not os.path.exists(this_file + '.hdf5'):
this_file = dir + '/snapdir_' + snap_ext + '/snapshot_' + snap_ext
else:
this_file = this_file + '.hdf5'
print this_file
data = gas_image_data(this_file, **kwargs)
try:
bh_xyz = np.array(ws.read_block(this_file, 'POS ', parttype=5))
except:
bh_xyz = np.array([0, 0, 0])
if center != None:
data.xyz[:, 0] -= np.mean(bh_xyz[:, 0])
data.xyz[:, 1] -= np.mean(bh_xyz[:, 1])
data.xyz[:, 2] -= np.mean(bh_xyz[:, 2])
weights = data.gas_mass
h_all = data.gas_hsml
try:
xrange = kwargs['xrange']
yrange = kwargs['yrange']
except:
xrange = [-1e10, 1e10]
yrange = [-1e10, 1e10]
fov = np.max([xrange, yrange])
ok = (data.gas_xyz[:,0] > -1.1*fov) & (data.gas_xyz[:,0] < 1.1*fov) & \
(data.gas_xyz[:,1] > -1.1*fov) & (data.gas_xyz[:,1] < 1.1*fov) & \
(data.gas_xyz[:,2] > -1.1*fov) & (data.gas_xyz[:,2] < 1.1*fov)
data.gas_xyz = data.gas_xyz[ok, :]
weights = weights[ok]
h_all = h_all[ok]
ok = np.arange(h_all.shape[0])
image, out_u, out_g, out_r = projection.raytrace_projection_compute(
data.gas_xyz[:, 0],
data.gas_xyz[:, 1],
data.gas_xyz[:, 2],
h_all,
weights,
weights,
weights,
weights,
0,
0,
0,
TRIM_PARTICLES=0)
image[image > maxden] = maxden
image[image < maxden / dynrange] = maxden / dynrange
return np.log10(image)
def write_one_snapshot(self, snap_index, **kwargs):
this_file = self.snap_list[snap_index]
header = ws.snapshot_header(this_file)
self.values['time'][snap_index] = header.time
print this_file
gas_xyz = np.array(ws.read_block(this_file, 'POS ', parttype=0))
gas_sfr = np.array(ws.read_block(this_file, 'SFR ', parttype=0))
gas_z = np.array(ws.read_block(this_file, 'Z ', parttype=0))
gas_z = gas_z[:, 0]
gas_mass = np.array(ws.read_block(this_file, 'MASS', parttype=0))
gas_ids = np.array(ws.read_block(this_file, 'ID ', parttype=0))
gas_u = np.array(ws.read_block(this_file, 'U ', parttype=0))
gas_ne = np.array(ws.read_block(this_file, 'NE ', parttype=0))
gas_t = units.gas_code_to_temperature(gas_u, gas_ne)
for index, value in enumerate(gas_ids): # IDs, add init metallicity
#print index, value, gas_z.shape, self.init_metallicity.shape
try:
gas_z[index] += self.init_metallicity[
value] # all that matters here is that init metallicty can be indexed by ID
except:
try:
gas_z[index] += self.init_metallicity[value & 2**30 - 1]
except:
gas_z[index] += 0.0 # confused...
sys.exit()
try:
star_xyz = np.array(ws.read_block(this_file, 'POS ', parttype=4))
star_mass = np.array(ws.read_block(this_file, 'MASS', parttype=4))
star_ages = np.array(ws.read_block(this_file, 'AGE ', parttype=4))
star_ages = header.time - star_ages
except:
star_xyz = np.zeros((10, 3))
star_mass = np.zeros(10)
star_ages = np.zeros(10)
bh_id = np.array(ws.read_block(this_file, 'ID ', parttype=5))
bh_xyz = np.array(ws.read_block(this_file, 'POS ', parttype=5))
bh_order = np.argsort(bh_id)
bh_xyz = bh_xyz[bh_order, :]
if bh_xyz.shape[0] == 1:
self.values['sep'][snap_index] = 0
else:
self.values['sep'][snap_index] = np.sqrt(
(bh_xyz[0, 0] - bh_xyz[1, 0])**2 +
(bh_xyz[0, 1] - bh_xyz[1, 1])**2 +
(bh_xyz[0, 2] - bh_xyz[1, 2])**2)
dr = np.zeros((bh_xyz.shape[0], gas_xyz.shape[0]))
for iii in np.arange(bh_xyz.shape[0]):
dx2 = (gas_xyz[:, 0] - bh_xyz[iii, 0])**2
dy2 = (gas_xyz[:, 1] - bh_xyz[iii, 1])**2
dz2 = (gas_xyz[:, 2] - bh_xyz[iii, 2])**2
dr[iii, :] = np.sqrt(
dx2 + dy2 + dz2) # offset from this BH; and then you have it.
min_dr = np.min(
dr, axis=0) # now the min_dr, regardless of n_bh, (right?!)
all_molec_gas_mass = gas_mass[gas_t < 100]
all_molec_gas_min_dr = min_dr[gas_t < 100]
all_molec_gas_dr = dr[:, gas_t < 100]
all_atomic_gas_mass = gas_mass[(gas_t > 100) & (gas_t < 1000)]
all_atomic_gas_min_dr = min_dr[(gas_t > 100) & (gas_t < 1000)]
all_atomic_gas_dr = dr[:, (gas_t > 100) & (gas_t < 1000)]
all_warm_gas_mass = gas_mass[(gas_t > 1000) & (gas_t < 1e5)]
all_warm_gas_min_dr = min_dr[(gas_t > 1000) & (gas_t < 1e5)]
all_warm_gas_dr = dr[:, (gas_t > 1000) & (gas_t < 1e5)]
all_hot_gas_mass = gas_mass[(gas_t > 1e5)]
all_hot_gas_min_dr = min_dr[(gas_t > 1e5)]
all_hot_gas_dr = dr[:, (gas_t > 1e5)]
for iii in np.arange(bh_xyz.shape[0]):
for jjj in np.arange(self.rad_bins.shape[0]):
this_index = (dr[iii, :] < self.rad_bins[jjj]) & (np.equal(
dr[iii, :], min_dr))
self.values['sfr'][iii, jjj,
snap_index] = np.sum(gas_sfr[this_index])
self.values['z'][iii, jjj,
snap_index] = np.median(gas_z[this_index])
self.values['mgas'][iii, jjj,
snap_index] = np.sum(gas_mass[this_index])
this_index = (
all_molec_gas_dr[iii, :] < self.rad_bins[jjj]) & (np.equal(
all_molec_gas_dr[iii, :], all_molec_gas_min_dr))
self.values['m_molec_gas'][iii, jjj, snap_index] = np.sum(
all_molec_gas_mass[this_index])
this_index = (all_atomic_gas_dr[iii, :] < self.rad_bins[jjj]
) & (np.equal(all_atomic_gas_dr[iii, :],
all_atomic_gas_min_dr))
self.values['m_atomic_gas'][iii, jjj, snap_index] = np.sum(
all_atomic_gas_mass[this_index])
this_index = (all_warm_gas_dr[iii, :] < self.rad_bins[jjj]) & (
np.equal(all_warm_gas_dr[iii, :], all_warm_gas_min_dr))
self.values['m_warm_gas'][iii, jjj, snap_index] = np.sum(
all_warm_gas_mass[this_index])
this_index = (all_hot_gas_dr[iii, :] < self.rad_bins[jjj]) & (
np.equal(all_hot_gas_dr[iii, :], all_hot_gas_min_dr))
self.values['m_hot_gas'][iii, jjj, snap_index] = np.sum(
all_hot_gas_mass[this_index])
dr = np.zeros((bh_xyz.shape[0], star_xyz.shape[0]))
for iii in np.arange(bh_xyz.shape[0]):
dx2 = (star_xyz[:, 0] - bh_xyz[iii, 0])**2
dy2 = (star_xyz[:, 1] - bh_xyz[iii, 1])**2
dz2 = (star_xyz[:, 2] - bh_xyz[iii, 2])**2
dr[iii, :] = np.sqrt(
dx2 + dy2 + dz2) # offset from this BH; and then you have it.
min_dr = np.min(
dr, axis=0) # now the min_dr, regardless of n_bh, (right?!)
for iii in np.arange(bh_xyz.shape[0]):
for jjj in np.arange(self.rad_bins.shape[0]):
this_index = (dr[iii, :] < self.rad_bins[jjj]) & (np.equal(
dr[iii, :], min_dr))
if np.sum(this_index) > 0:
self.values['stellar_ages'][iii, jjj,
snap_index] = np.median(
star_ages[this_index])
self.values['mstar'][iii, jjj, snap_index] = np.sum(
star_mass[this_index])
this_index = (dr[iii, :] < self.rad_bins[jjj]) & (np.equal(
dr[iii, :], min_dr)) & (star_ages < 0.01)
if np.sum(this_index) > 0:
self.values['sfr_sm'][iii, jjj, snap_index] = np.sum(
star_mass[this_index]) / 0.01 * 10.0
this_index = (dr[iii, :] < self.rad_bins[jjj]) & (np.equal(
dr[iii, :], min_dr)) & (star_ages < 0.1)
if np.sum(this_index) > 0:
self.values['sfr_sm2'][iii, jjj, snap_index] = np.sum(
star_mass[this_index]) / 0.1 * 10.0
def dm_image(dir='./', snapnum=0, this_file=None, center=None, cosmo_wrap=False,
maxden=None, dynrange=None, **kwargs):
snap_ext = str(snapnum).zfill(3)
if this_file==None:
print "setting snapshot"
this_file = dir+'/snapshot_'+snap_ext
if not os.path.exists(this_file+'.hdf5'):
this_file = dir+'/snapdir_'+snap_ext+'/snapshot_'+snap_ext
# this_file = this_file+'.hdf5'
print this_file
data = dm_image_data(this_file, **kwargs)
try:
bh_xyz = np.array(ws.read_block(this_file,'POS ', parttype=5))
except:
bh_xyz = np.array( [0, 0, 0] )
if cosmo_wrap:
boxsize = np.max( data.xyz ) - np.min( data.xyz )
if center != None:
data.xyz[:,0] -= np.mean(center[0])
data.xyz[:,1] -= np.mean(center[1])
data.xyz[:,2] -= np.mean(center[2])
else:
data.xyz[:,0] -= np.mean(bh_xyz[:,0])
data.xyz[:,1] -= np.mean(bh_xyz[:,1])
data.xyz[:,2] -= np.mean(bh_xyz[:,2])
if cosmo_wrap:
data.xyz[ data.xyz[:,0] > boxsize/2.0 ,0] -= boxsize
data.xyz[ data.xyz[:,1] > boxsize/2.0 ,1] -= boxsize
data.xyz[ data.xyz[:,2] > boxsize/2.0 ,2] -= boxsize
data.xyz[ data.xyz[:,0] < -1.0* boxsize/2.0 ,0] += boxsize
data.xyz[ data.xyz[:,1] < -1.0* boxsize/2.0 ,1] += boxsize
data.xyz[ data.xyz[:,2] < -1.0* boxsize/2.0 ,2] += boxsize
# color_weights = units.springel_units.gas_code_to_temperature( gas_data.gas_u, gas_data.gas_nume)
color_weights = 1e4
weights = data.mass
h_all = data.hsml
try:
xrange = kwargs['xrange']
yrange = kwargs['yrange']
if xrange != None and yrange != None:
fov = np.max( [xrange, yrange] )
ok = (data.xyz[:,0] > -1.1*fov) & (data.xyz[:,0] < 1.1*fov) & \
(data.xyz[:,1] > -1.1*fov) & (data.xyz[:,1] < 1.1*fov) & \
(data.xyz[:,2] > -1.1*fov) & (data.xyz[:,2] < 1.1*fov)
data.xyz = data.xyz[ok,:]
# color_weights = color_weights[ok]
weights = weights[ok]
h_all = h_all[ok]
except:
print "failed to clip"
try:
zrange = kwargs['zrange']
if zrange != None:
ok = (data.xyz[:,2] > np.min(zrange) ) & (data.xyz[:,2] < np.max(zrange) )
data.xyz = data.xyz[ok,:]
weights = weights[ok]
h_all = h_all[ok]
except:
print "failed to clip"
sys.exit()
ok = np.arange(h_all.shape[0])
#weights = h_all
print np.min(weights), np.max(weights)
print data.xyz.shape
print h_all.shape
print weights.shape
image, out_u, out_g, out_r = projection.raytrace_projection_compute(
data.xyz[:,0], data.xyz[:,1], data.xyz[:,2], h_all, weights,
weights, weights, weights, 0, 0, 0, TRIM_PARTICLES=0,
**kwargs)
print image.shape
print "image max/min is {:g}/{:g}".format(np.min(image), np.max(image) )
if((maxden != None) and (dynrange != None)):
image[image > maxden] = maxden
image[image < maxden / dynrange ] = maxden / dynrange
print "image clipping is {:g}/{:g}".format(maxden, maxden / dynrange )
else:
print "no clipping (maxden/dynrange not set)..."
#gas_raytrace_temperature( \
# gas_map_temperature_cuts, \
# gas_data.gas_xyz[ok,0], gas_data.gas_xyz[ok,1], gas_data.gas_xyz[ok,2],
# color_weights[ok], weights[ok], h_all[ok], \
# kernel_width=kernel_widths,\
# isosurfaces = 1,\
# KAPPA_UNITS = 2.0885*np.array([1.1,2.0,1.5]),\
# add_temperature_weights = 0,\
# **kwargs)
# image24, massmap = makepic.make_threeband_image_process_bandmaps( out_r,out_g,out_u, \
# **kwargs)
print np.min(image), np.max(image)
return np.log10(image)
h2 = ws.snapshot_header(ic2) # load header 2
npart2 = h2.nall
nall = h1.nall + h2.nall
all_blocks1 = ws.return_tags(ic1)
f = ws.openfile('new_ics.hdf5')
my_header = h1
last_ID = 1
for parttype in np.arange(6):
this_part_blocks = all_blocks1[parttype]
for block in this_part_blocks:
print parttype, block
data = None
if npart1[parttype] > 0: data1 = ws.read_block(ic1, block, parttype=parttype)
else: data1 = np.zeros( 1 )
if npart2[parttype] > 0: data2 = ws.read_block(ic1, block, parttype=parttype)
else: data2 = np.zeros( 1 )
if npart1[parttype] > 0 and npart2[parttype] > 0:
if parttype==5:
data = data1
my_header.nall[5] = 1
else:
data = np.append(data1, data2, axis=0)
elif npart1[parttype] > 0: data = data1
elif npart2[parttype] > 0: data = data2
else: print "[error]: You shouldnt be here. expect an error."
right=1.0,
hspace=0,
wspace=0)
ax1_.axis('off')
ax2_.axis('off')
ax3_.axis('off')
import units.springel_units as units
header = ws.snapshot_header('../' + run_base +
'/output/snapdir_' +
str(snapnum).zfill(3) +
'/snapshot_' +
str(snapnum).zfill(3))
star_pos = ws.read_block(
'../' + run_base + '/output/snapdir_' +
str(snapnum).zfill(3) + '/snapshot_' +
str(snapnum).zfill(3), "POS ", 4)
star_bt = ws.read_block(
'../' + run_base + '/output/snapdir_' +
str(snapnum).zfill(3) + '/snapshot_' +
str(snapnum).zfill(3), "GAGE", 4)
star_hsml = ws.read_block(
'../' + run_base + '/output/snapdir_' +
str(snapnum).zfill(3) + '/snapshot_' +
str(snapnum).zfill(3), "STH ", 4)
star_cp1 = ws.read_block(
'../' + run_base + '/output/snapdir_' +
str(snapnum).zfill(3) + '/snapshot_' +
str(snapnum).zfill(3), "CP1 ", 4)
star_cp2 = ws.read_block(
'../' + run_base + '/output/snapdir_' +
def make_one_frame(self,
iii,
frame_type=0,
snapnr_label=False,
overwrite=True,
show_time_label=False,
**kwargs):
frame_name = "./plots/frames/" + self.savetag + "/frame_" + self.frametag + self.savetag + "_" + str(
self.all_snapnrs[iii]).zfill(4) + ".png"
if (overwrite or not os.path.exists(frame_name)):
this_time = iii * self.dt_frames
time_diff = self.snap_timing - this_time - self.snap_timing[0]
try:
index = np.where(
np.abs(time_diff) == np.min(np.abs(time_diff)))[0][0]
except:
print np.where(np.abs(time_diff) == np.min(np.abs(time_diff)))
sys.exit()
if time_diff[
index] < 0 and self.interp == True: # on the negative side, get the positive side
this_file1 = self.snap_list[index]
this_file2 = self.snap_list[index + 1]
interp_coef = (this_time - self.snap_timing[index]) / (
self.snap_timing[index + 1] - self.snap_timing[index])
elif time_diff[index] > 0 and self.interp == True:
this_file1 = self.snap_list[index - 1]
this_file2 = self.snap_list[index]
interp_coef = (this_time - self.snap_timing[index - 1]) / (
self.snap_timing[index + 1] - self.snap_timing[index])
else:
this_file1 = self.snap_list[index]
this_file2 = self.snap_list[index]
interp_coef = 0.0
fov_min = (1.0 - self.zoom_factor) * self.fov
fov_max = self.fov
log_fac = np.log10(fov_max / fov_min)
fac = (3.1416 / 2.0 + np.arctan(
(1.0 * iii /
(1.0 * self.n_frames)) * 2 * 3.14159)) / (3.1416) #0.0 -> 1.0
this_fov = fov_max * 10.0**(-1.0 * fac * log_fac)
if self.proj_depth > 0:
this_thick = self.proj_depth
else:
this_thick = this_fov
print " this_fov and thickness: "
print this_fov, this_thick
if frame_type == 0: # single gas projection
image = gasimages.gas_image('',
0,
this_file=this_file1,
next_file=this_file2,
interp_coef=interp_coef,
xrange=[-self.fov, self.fov],
yrange=[-self.fov, self.fov],
dynrange=self.dynrange,
maxden=self.maxden,
pixels=self.npixels,
interp=self.interp,
**kwargs)
fig = plt.figure(figsize=(5, 5))
ax = fig.add_subplot(1, 1, 1)
fig.subplots_adjust(left=0.0,
right=1.0,
top=1.0,
bottom=0.0,
wspace=0.0,
hspace=0.0)
imgplot = ax.imshow(
image,
origin='lower',
extent=[-self.fov, self.fov, -self.fov, self.fov])
ax.set_xlim([-self.fov, self.fov])
ax.set_ylim([-self.fov, self.fov])
fig.savefig(frame_name, dpi=self.npixels)
elif frame_type == 1:
image1 = gasimages.gas_image('',
0,
this_file=this_file1,
next_file=this_file2,
interp_coef=interp_coef,
xrange=[-self.fov, self.fov],
yrange=[-self.fov, self.fov],
dynrange=self.dynrange,
maxden=self.maxden,
pixels=self.npixels,
interp=self.interp,
**kwargs)
image2 = gasimages.gas_image('',
0,
this_file=this_file1,
next_file=this_file2,
interp_coef=interp_coef,
xrange=[-self.fov, self.fov],
yrange=[-self.fov, self.fov],
dynrange=self.dynrange,
maxden=self.maxden,
pixels=self.npixels,
interp=self.interp,
**kwargs)
fig = plt.figure(figsize=(10, 5))
ax = fig.add_subplot(1, 2, 1)
fig.subplots_adjust(left=0.0,
right=1.0,
top=1.0,
bottom=0.0,
wspace=0.0,
hspace=0.0)
imgplot = ax.imshow(
image1,
origin='lower',
extent=[-self.fov, self.fov, -self.fov, self.fov])
ax.set_xlim([-self.fov, self.fov])
ax.set_ylim([-self.fov, self.fov])
ax = fig.add_subplot(1, 2, 2)
fig.subplots_adjust(left=0.0,
right=1.0,
top=1.0,
bottom=0.0,
wspace=0.0,
hspace=0.0)
imgplot = ax.imshow(
image2,
origin='lower',
extent=[-self.fov, self.fov, -self.fov, self.fov])
ax.set_xlim([-self.fov, self.fov])
ax.set_ylim([-self.fov, self.fov])
fig.savefig(frame_name, dpi=self.npixels)
elif frame_type == 2:
image = stellarimages.stellar_image(
'',
0,
this_file=this_file1,
next_file=this_file2,
interp_coef=interp_coef,
xrange=[-self.fov, self.fov],
yrange=[-self.fov, self.fov],
dynrange=self.dynrange,
maxden=self.maxden,
pixels=self.npixels,
interp=self.interp,
**kwargs)
# massmap = np.transpose( massmap )
pic = np.zeros_like(image)
for jjj in np.arange(image.shape[2]):
pic[:, :, jjj] = np.transpose(image[:, :, jjj])
fig = plt.figure(figsize=(1, 1))
ax = fig.add_subplot(1, 1, 1)
fig.subplots_adjust(left=0.0,
right=1.0,
top=1.0,
bottom=0.0,
wspace=0.0,
hspace=0.0)
extent = [-this_fov, this_fov, -this_fov, this_fov]
imgplot = ax.imshow(pic, origin='lower', extent=extent)
fig.savefig(frame_name, dpi=self.npixels)
elif frame_type == 3:
print kwargs.get('center_pos', [0, 0, 0])[0]
print self.center_pos
print self.arepo
print this_fov
print this_thick
show_gasstarxray = 'gas'
if self.plottype == 'gas': show_gasstarxray = 'gas'
if self.plottype == 'sfr': show_gasstarxray = 'sfr'
if self.IFU or self.plottype == 'sfr':
self.threecolor = False
if this_thick > 0:
zrange = [-this_thick, this_thick]
else:
zrange = 0
image, massmap = pfh_image_maker.image_maker(
self.sdir,
self.all_snapnrs[iii],
snapdir_master='',
outdir_master='',
filename_set_manually='tmp_remove',
center_pos=self.center_pos,
#center_on_bh=0,
xrange=[-this_fov, this_fov],
zrange=zrange,
# dynrange=self.dynrange, maxden=self.maxden,
threecolor=self.threecolor,
pixels=self.npixels,
show_gasstarxray=show_gasstarxray,
theta=self.theta,
arepo=self.arepo,
**kwargs)
print "image values"
print image
print image.shape
print massmap
print " "
massmap = np.transpose(massmap)
pic = np.zeros_like(image)
for jjj in np.arange(image.shape[2]):
pic[:, :, jjj] = np.transpose(image[:, :, jjj])
if self.IFU:
manga_mask = self.create_ifu_mask(pic)
#star_xyz = ws.read_block(this_file1, 'POS ', 4)
fig = plt.figure(figsize=(1, 1))
ax = fig.add_subplot(1, 1, 1)
fig.subplots_adjust(left=0.0,
right=1.0,
top=1.0,
bottom=0.0,
wspace=0.0,
hspace=0.0)
extent = [-this_fov, this_fov, -this_fov, this_fov]
imgplot = ax.imshow(pic, origin='lower', extent=extent)
if self.IFU:
imgplot = ax.imshow(np.log10(manga_mask),
origin='lower',
extent=extent,
cmap='Greys',
vmin=-1,
vmax=1,
alpha=0.5)
if self.plot_stars:
try:
star_xyz = ws.read_block(this_file1, 'POS ', 4)
ax.scatter(star_xyz[:, 0] -
kwargs.get('center_pos', [0, 0, 0])[0],
star_xyz[:, 1] -
kwargs.get('center_pos', [0, 0, 0])[1],
s=0.1,
lw=0,
facecolor='r')
except:
print "no stars to scatter"
ax.set_xlim([-1.0 * this_fov, this_fov])
ax.set_ylim([-1.0 * this_fov, this_fov])
plt.axis('off')
if snapnr_label:
ax.text(0.75 * this_fov,
0.85 * this_fov,
str(iii).zfill(3),
fontsize=3,
path_effects=[
PathEffects.withStroke(linewidth=1,
foreground="w")
]) #color='w' )
if show_time_label:
ax.text(0.05,
0.95,
"t={:.2f}Gyr".format(iii * self.dt_frames),
fontsize=2,
ha='left',
transform=ax.transAxes,
path_effects=[
PathEffects.withStroke(linewidth=0.2,
foreground="w")
])
if self.IFU:
png_file = "./plots/frames/frame_ifu_" + self.savetag + "_" + str(
self.all_snapnrs[iii]).zfill(4) + ".png"
hdf5_file = "./plots/frames/frame_ifu_" + self.savetag + "_" + str(
self.all_snapnrs[iii]).zfill(4) + ".hdf5"
fig.savefig(png_file, dpi=self.npixels)
f = hdf5lib.OpenFile(hdf5_file, mode='w')
group = hdf5lib.CreateGroup(f, "GasSurfaceDensity")
hdf5lib.CreateArray(f, group, 'data', massmap.flatten())
group = hdf5lib.CreateGroup(f, "MangaMap")
hdf5lib.CreateArray(f, group, 'manga_mask',
-1.0 * manga_mask.flatten())
f.close()
else:
fig.savefig(frame_name, dpi=self.npixels)
elif frame_type == 4:
image, massmap = pfh_image_maker.image_maker(
self.sdir,
self.all_snapnrs[iii],
snapdir_master='',
outdir_master='',
show_time_label=0,
filename_set_manually='tmp_remove',
center_on_bh=0,
xrange=[-this_fov, this_fov],
dynrange=self.dynrange,
maxden=self.maxden,
pixels=self.npixels,
**kwargs)
pic1 = np.zeros_like(image)
for jjj in np.arange(3):
pic1[:, :, jjj] = np.transpose(image[:, :, jjj])
image, massmap = pfh_image_maker.image_maker(
self.sdir,
self.all_snapnrs[iii],
snapdir_master='',
outdir_master='',
show_time_label=0,
filename_set_manually='tmp_remove',
center_on_bh=0,
xrange=[-this_fov, this_fov],
dynrange=self.dynrange,
maxden=self.maxden,
pixels=self.npixels,
theta=90,
**kwargs)
pic2 = np.zeros_like(image)
for jjj in np.arange(3):
pic2[:, :, jjj] = np.transpose(image[:, :, jjj])
fig = plt.figure(figsize=(2, 1))
ax = fig.add_subplot(1, 2, 1)
fig.subplots_adjust(left=0.0,
right=1.0,
top=1.0,
bottom=0.0,
wspace=0.0,
hspace=0.0)
imgplot = ax.imshow(
pic1,
origin='lower',
extent=[-this_fov, this_fov, -this_fov, this_fov])
ax.set_xlim([-1.0 * this_fov, this_fov])
ax.set_ylim([-1.0 * this_fov, this_fov])
plt.axis('off')
ax = fig.add_subplot(1, 2, 2)
fig.subplots_adjust(left=0.0,
right=1.0,
top=1.0,
bottom=0.0,
wspace=0.0,
hspace=0.0)
imgplot = ax.imshow(
pic2,
origin='lower',
extent=[-this_fov, this_fov, -this_fov, this_fov])
ax.set_xlim([-1.0 * this_fov, this_fov])
ax.set_ylim([-1.0 * this_fov, this_fov])
plt.axis('off')
fig.savefig(frame_name, dpi=self.npixels)
else:
fig = plt.figure(figsize=(10, 5))
fig.subplots_adjust(left=0.0,
right=1.0,
top=1.0,
bottom=0.0,
wspace=0.0,
hspace=0.0)
ax = fig.add_subplot(1, 2, 1)
imgplot = ax.imshow(
image1,
origin='lower',
extent=[-self.fov, self.fov, -self.fov, self.fov])
ax.set_xlim([-self.fov, self.fov])
ax.set_ylim([-self.fov, self.fov])
ax.axes.xaxis.set_ticklabels([])
ax.axes.yaxis.set_ticklabels([])
ax = fig.add_subplot(1, 2, 2)
imgplot = ax.imshow(image2,
origin='lower',
extent=[
-self.fov / 10.0, self.fov / 10.0,
-self.fov / 10.0, self.fov / 10.0
])
ax.set_xlim([-self.fov / 10.0, self.fov / 10.0])
ax.set_ylim([-self.fov / 10.0, self.fov / 10.0])
ax.axes.xaxis.set_ticklabels([])
ax.axes.yaxis.set_ticklabels([])
fig.savefig(frame_name, dpi=self.npixels)
def __init__(self,
path,
min_time=None,
max_time=None,
n_frames=None,
fov=15.0,
npixels=256,
tag='',
frametag='',
interp=False,
zoom_factor=0.0,
IFU=False,
plottype='gas',
threecolor=True,
dynrange=0.6e5,
maxden=0.01,
plot_stars=False,
theta=0,
proj_depth=0,
arepo=0,
set_bh_center_and_leave=0,
center_pos=[0, 0, 0],
**kwargs):
print path
self.sdir = path
snap_list = np.array(glob.glob(path + '/snap*')) #_*hdf5'))
print snap_list
try:
all_snapnrs = np.array([
int(file[file.index('snapshot_') + 9:file.index('.hdf5')])
for file in snap_list
],
dtype=int)
except:
try:
all_snapnrs = np.array([
int(file[file.index('snap_') + 5:file.index('.hdf5')])
for file in snap_list
],
dtype=int)
except:
all_snapnrs = np.array([
int(file[file.index('snapdir_') + 8:])
for file in snap_list
],
dtype=int)
all_snapnrs = np.sort(all_snapnrs)
# all_snapnrs = all_snapnrs[ all_snapnrs > 0 ]
print all_snapnrs
first_snapnr = all_snapnrs[0]
new_snap_list = [None] * all_snapnrs.shape[0] #(snap_list.shape[0]
self.all_snapnrs = all_snapnrs
snap_timing = np.zeros(all_snapnrs.shape[0])
for index in np.arange(all_snapnrs.shape[0]):
try:
this_file = path + "snapshot_" + str(
index + int(first_snapnr)).zfill(3) + ".hdf5"
new_snap_list[index] = this_file
# print this_file
head = ws.snapshot_header(this_file[:this_file.index('.hdf5')])
snap_timing[index] = head.time
except:
this_file = path + "snapdir_" + str(
index + int(first_snapnr)).zfill(3) + "/snapshot_" + str(
index + int(first_snapnr)).zfill(3) #+".0.hdf5"
new_snap_list[index] = this_file
# print this_file
head = ws.snapshot_header(
this_file) #[:this_file.index('.hdf5')])
snap_timing[index] = head.time
# snap_timing[index] = -1.0
# print "snapshot failed..."
self.interp = interp
self.savetag = tag
self.frametag = frametag
self.snap_list = np.array(new_snap_list)
self.snap_timing = snap_timing
self.fov = fov
self.npixels = npixels
self.maxden = maxden
self.dynrange = dynrange
self.zoom_factor = zoom_factor
self.IFU = IFU # True / False
self.plottype = plottype # gas, sfr, stellar XYZ band
self.threecolor = threecolor # True / False
self.plot_stars = plot_stars
self.theta = theta
self.proj_depth = proj_depth
self.arepo = arepo
print "snap timing:"
print snap_timing
if (not os.path.exists("./plots/frames/" + self.savetag)):
os.makedirs("./plots/frames/" + self.savetag)
if min_time == None:
self.min_movie_time = np.min(snap_timing)
else:
self.min_movie_time = min_time
if max_time == None:
self.max_movie_time = np.max(snap_timing)
else:
self.max_movie_time = max_time
if n_frames != None:
self.n_frames = n_frames
else:
self.n_frames = snap_list.shape[0]
self.dt_frames = (self.max_movie_time -
self.min_movie_time) / (1.0 * self.n_frames - 1.0)
if center_pos == [0, 0, 0]:
print "here!"
print set_bh_center_and_leave
if set_bh_center_and_leave == 1:
bhpos = ws.read_block(self.snap_list[0], 'POS ', 5)[0]
self.center_pos = bhpos
else:
self.center_pos = np.array([0, 0, 0])
else:
self.center_pos = center_pos
def __init__(self, this_file, next_file=None, interp_coef=0.0, interp=True, crude=False, ra_range=None, dec_range=None,
ra_center=None, dec_center=None, zrange=None, convert_to_sg=False, \
**kwargs):
if next_file==None or interp==False: # no interp
print "have a single file, no interp"
self.head = ws.snapshot_header(this_file)
self.xyz = np.array(ws.read_block(this_file,'POS ', parttype=1))
self.mass = np.array(ws.read_block(this_file,'MASS', parttype=1))
npart = np.shape(self.xyz)[0]*1.0
if crude:
self.hsml = np.zeros_like( self.mass ) + self.head.boxsize / ( npart**0.333 ) * 1.0
else:
sys.exit()
self.hsml = calc_hsml.get_particle_hsml( self.xyz[:,0], self.xyz[:,1], self.xyz[:,2] , **kwargs )
if True: #ra_center is not None and dec_center is not None: # need to do some rotation
if True: # this likely makes the most sense. MW is at the center.
tmp_x = self.xyz[:,0] - self.head.boxsize/2.0
tmp_y = self.xyz[:,1] - self.head.boxsize/2.0
tmp_z = self.xyz[:,2] - self.head.boxsize/2.0
else: # this could be possible, but puts the MW at the origin, at the edge of the box.
tmp_x = self.xyz[:,0]
tmp_y = self.xyz[:,1]
tmp_z = self.xyz[:,2]
tmp_x[ tmp_x > self.head.boxsize/2.0 ] -= self.head.boxsize
tmp_y[ tmp_y > self.head.boxsize/2.0 ] -= self.head.boxsize
tmp_z[ tmp_z > self.head.boxsize/2.0 ] -= self.head.boxsize
tmp_r = np.sqrt( tmp_x**2 + tmp_y**2 + tmp_z**2 )
dec = np.arcsin( tmp_z / tmp_r ) * 180.0 / 3.14159
ra = np.arctan2( tmp_y, tmp_x ) * 180.0 / 3.14159
c = SkyCoord(ra=ra*u.degree, dec=dec*u.degree)
if convert_to_sg:
new_coords = c.transform_to('supergalactic')
self.xyz[:,0] = new_coords.cartesian.x * tmp_r + self.head.boxsize/2.0
self.xyz[:,1] = new_coords.cartesian.y * tmp_r + self.head.boxsize/2.0
self.xyz[:,2] = new_coords.cartesian.z * tmp_r + self.head.boxsize/2.0
else:
self.xyz[:,0] = c.cartesian.x * tmp_r + self.head.boxsize/2.0
self.xyz[:,1] = c.cartesian.y * tmp_r + self.head.boxsize/2.0
self.xyz[:,2] = c.cartesian.z * tmp_r + self.head.boxsize/2.0
new_r = np.sqrt( self.xyz[:,0] **2 + self.xyz[:,1] **2 + self.xyz[:,2] **2 )
if ra_range is not None and dec_range is not None:
tmp_x = self.xyz[:,0] - self.head.boxsize/2.0
tmp_y = self.xyz[:,1] - self.head.boxsize/2.0
tmp_z = self.xyz[:,2] - self.head.boxsize/2.0
tmp_r = np.sqrt( tmp_x**2 + tmp_y**2 + tmp_z**2 )
dec = np.arcsin( tmp_z / tmp_r )
ra = np.arctan2( tmp_y, tmp_x )
in_box_index = ( ra > np.min(ra_range) ) & ( ra < np.max(ra_range) ) & \
( dec> np.min(dec_range)) & (dec < np.max(dec_range))
print "Found {:d} out of {:d} particles in our projection wedge".format( np.sum(in_box_index), len(in_box_index) )
# sys.exit()
self.xyz = self.xyz[in_box_index,:]
self.mass = self.mass[in_box_index]
self.hsml = self.hsml[in_box_index]
else:
# TODO: FINISH THIS BLOCK
head = ws.snapshot_header(this_file)
gas_id1 = np.array(ws.read_block(this_file,'ID ', parttype=0))
order = np.argsort(gas_id1)
gas_id1 = gas_id1[order]
gas_xyz1 = np.array(ws.read_block(this_file,'POS ', parttype=0))[order,:]
gas_hsml1 = np.array(ws.read_block(this_file,'HSML', parttype=0))[order]
gas_mass1 = np.array(ws.read_block(this_file,'MASS', parttype=0))[order]
gas_u1 = np.array(ws.read_block(this_file,'U ', parttype=0))[order]
gas_nume1 = np.array(ws.read_block(this_file,'NE ', parttype=0))[order]
try:
gas_z1 = np.array(ws.read_block(this_file,'Z ', parttype=0))
gas_z1 = np.sum( gas_z1[order,2:], axis=1)
except:
gas_z1 = np.array(ws.read_block(this_file,'GZ ', parttype=0))[order]
gas_id2 = np.array(ws.read_block(next_file,'ID ', parttype=0))
order = np.argsort(gas_id2)
gas_id2 = gas_id2[order]
gas_xyz2 = np.array(ws.read_block(next_file,'POS ', parttype=0))[order,:]
gas_hsml2 = np.array(ws.read_block(next_file,'HSML', parttype=0))[order]
gas_mass2 = np.array(ws.read_block(next_file,'MASS', parttype=0))[order]
gas_u2 = np.array(ws.read_block(next_file,'U ', parttype=0))[order]
gas_nume2 = np.array(ws.read_block(next_file,'NE ', parttype=0))[order]
try:
gas_z2 = np.array(ws.read_block(next_file,'Z ', parttype=0))
gas_z2 = np.sum( gas_z2[order,2:], axis=1)
except:
gas_z2 = np.array(ws.read_block(next_file,'GZ ', parttype=0))[order]
gas_id2 = np.array(ws.read_block(next_file,'ID ', parttype=0))
i1 = np.in1d( gas_id1, gas_id2 )
gas_xyz1 = gas_xyz1[i1,:]
gas_hsml1= gas_hsml1[i1]
gas_mass1= gas_mass1[i1]
gas_u1 = gas_u1[i1]
gas_nume1= gas_nume1[i1]
gas_z1 = gas_z1[i1]
i1 = np.in1d( gas_id2, gas_id1 )
gas_xyz2 = gas_xyz2[i1,:]
gas_hsml2= gas_hsml2[i1]
gas_mass2= gas_mass2[i1]
gas_u2 = gas_u2[i1]
gas_nume2= gas_nume2[i1]
gas_z2 = gas_z2[i1]
self.gas_xyz = (1.0 - interp_coef) * gas_xyz1 + interp_coef * gas_xyz2
self.gas_hsml= (1.0 - interp_coef) * gas_hsml1 + interp_coef * gas_hsml2
self.gas_mass= (1.0 - interp_coef) * gas_mass1 + interp_coef * gas_mass2
self.gas_u = (1.0 - interp_coef) * gas_u1 + interp_coef * gas_u2
self.gas_nume= (1.0 - interp_coef) * gas_nume1 + interp_coef * gas_nume2
self.gas_z = (1.0 - interp_coef) * gas_z1 + interp_coef * gas_z2
def __init__(self, this_file, next_file=None, interp_coef=0.0, interp=True, dust=True, **kwargs):
if next_file==None or interp==False: # no interp
print "have a single file, no interp"
print this_file
self.head = ws.snapshot_header(this_file)
if dust:
print "loading dust data blocks"
self.gas_xyz = np.array(ws.read_block(this_file,'POS ', parttype=0))
try:
self.gas_hsml = np.array(ws.read_block(this_file,'HSML', parttype=0))
except:
print "failed to load HSML values directly. Try using cell volumes instead"
rho = np.array(ws.read_block(this_file,'RHO ', parttype=0))
mass = np.array(ws.read_block(this_file,'MASS', parttype=0))
vol = mass / rho
self.gas_hsml = (0.75 * vol / 3.14159 ) ** 0.33333 * 3
self.gas_mass = np.array(ws.read_block(this_file,'MASS', parttype=0))
self.gas_u = np.array(ws.read_block(this_file,'U ', parttype=0))
self.gas_rho = np.array(ws.read_block(this_file,'RHO ', parttype=0))
self.gas_numh = np.array(ws.read_block(this_file,'NH ', parttype=0))
self.gas_nume = np.array(ws.read_block(this_file,'NE ', parttype=0))
try:
self.gas_z = np.array(ws.read_block(this_file,'Z ', parttype=0))
except:
self.gas_z = np.array(ws.read_block(this_file,'GZ ', parttype=0))
else:
self.gas_xyz = np.zeros( (10,3) )
self.gas_hsml = np.zeros( 10 )
self.gas_mass = np.zeros( 10 )
self.gas_u = np.zeros( 10 )
self.gas_rho = np.zeros( 10 )
self.gas_numh = np.zeros( 10 )
self.gas_nume = np.zeros( 10 )
self.gas_z = np.zeros( 10 )
self.star_xyz = np.array(ws.read_block(this_file,'POS ', parttype=4))
self.star_mass = np.array(ws.read_block(this_file,'MASS', parttype=4))
try:
self.star_bt = np.array(ws.read_block(this_file,'AGE ',parttype=4))
self.star_z = np.array(ws.read_block(this_file,'Z ',parttype=4))
except:
try:
self.star_bt = np.array(ws.read_block(this_file,'GAGE',parttype=4))
self.star_z = np.array(ws.read_block(this_file,'GZ ',parttype=4))
except:
print "failed to load stars properly. Setting to zero to avoid failure, but this will not produce an image."
self.star_xyz = np.zeros( (10,3) )
self.star_bt = np.zeros( (10) )
self.star_z = np.zeros( (10) )
self.star_mass = np.zeros( (10) )
def __init__(self, path, tag, factor=1, **kwargs):
snap_list = np.array(glob.glob(path + '/snap*_*hdf5'))
print path, snap_list
all_snapnrs = np.sort(
np.array([
int(file[file.index('snapshot_') + 9:file.index('.hdf5')])
for file in snap_list
],
dtype=int))
first_snapnr = all_snapnrs[0]
new_snap_list = [None] * snap_list.shape[0]
snap_timing = np.zeros(snap_list.shape[0])
for index in np.arange(snap_list.shape[0]):
this_file = path + "snapshot_" + str(
index * factor + int(first_snapnr)).zfill(3) + ".hdf5"
print this_file
new_snap_list[index] = this_file
try:
head = ws.snapshot_header(this_file[:this_file.index('.hdf5')])
snap_timing[index] = head.time
except:
snap_timing[index] = -1.0
self.first_snapnr = first_snapnr
self.savetag = tag
self.snap_list = np.array(new_snap_list)
self.snap_timing = snap_timing
self.n_snaps = self.snap_list.shape[0]
print " "
print tag
print " "
self.alloc_loc_arrays()
this_file = path + "snapshot_" + str(
int(first_snapnr)).zfill(3) + ".hdf5"
print this_file
gas_xyz = np.array(ws.read_block(this_file, 'POS ', parttype=0))
star_xyz = np.append(
np.array(ws.read_block(this_file, 'POS ', parttype=2)),
np.array(ws.read_block(this_file, 'POS ', parttype=3)),
axis=0)
star_mass = np.append(
np.array(ws.read_block(this_file, 'MASS', parttype=2)),
np.array(ws.read_block(this_file, 'MASS', parttype=3)),
axis=0) * 1e10
bh_id = np.array(ws.read_block(this_file, 'ID ', parttype=5))
bh_xyz = np.array(ws.read_block(this_file, 'POS ', parttype=5))
bh_order = np.argsort(bh_id)
bh_xyz = bh_xyz[bh_order, :]
# bh_xyz = np.array(ws.read_block(this_file,'POS ', parttype=5))
# calculate stellar offsets, to determine each galaxy's initial stellar mass (for MZ rleaiton)
dr = np.zeros((bh_xyz.shape[0], star_xyz.shape[0]))
for iii in np.arange(bh_xyz.shape[0]):
dx2 = (star_xyz[:, 0] - bh_xyz[iii, 0])**2
dy2 = (star_xyz[:, 1] - bh_xyz[iii, 1])**2
dz2 = (star_xyz[:, 2] - bh_xyz[iii, 2])**2
dr[iii, :] = np.sqrt(
dx2 + dy2 + dz2) # offset from this BH; and then you have it.
min_dr = np.min(dr, axis=0) # now the min_dr, regardless of n_bh
print bh_xyz
print dr
print dr.shape
print min_dr.shape
m_gal_1 = 0
m_gal_2 = 0
for iii in np.arange(min_dr.shape[0]):
if min_dr[iii] == dr[0, iii]:
m_gal_1 += star_mass[iii]
else:
m_gal_2 += star_mass[iii]
print m_gal_1, m_gal_2
print np.log10(m_gal_1), np.log10(m_gal_2)
x1 = np.log10(m_gal_1)
x2 = np.log10(m_gal_2)
z10_tmp = 27.7911 - 6.94493 * x1 + 0.808097 * x1 * x1 - 0.0301508 * x1 * x1 * x1
z20_tmp = 27.7911 - 6.94493 * x2 + 0.808097 * x2 * x2 - 0.0301508 * x2 * x2 * x2
print " Initial Metallicity Summary: "
print " We think galaxy 1 has a mass of {:.3f} and a metallicity of {:.3f}".format(
x1, z10_tmp)
print " We think galaxy 2 has a mass of {:.3f} and a metallicity of {:.3f}".format(
x2, z20_tmp)
# calculate gas offsets, to determine each galaxy's metallicity gradient
dr = np.zeros((bh_xyz.shape[0], gas_xyz.shape[0]))
for iii in np.arange(bh_xyz.shape[0]):
dx2 = (gas_xyz[:, 0] - bh_xyz[iii, 0])**2
dy2 = (gas_xyz[:, 1] - bh_xyz[iii, 1])**2
dz2 = (gas_xyz[:, 2] - bh_xyz[iii, 2])**2
dr[iii, :] = np.sqrt(
dx2 + dy2 + dz2) # offset from this BH; and then you have it.
min_dr = np.min(
dr, axis=0) # now the min_dr, regardless of n_bh, (right?!)
tmp_ids = np.array(ws.read_block(this_file, 'ID ', parttype=0))
self.init_ids = np.zeros(tmp_ids.max() + 1)
self.init_metallicity = np.zeros(tmp_ids.max() + 1)
for iii in np.arange(min_dr.shape[0]):
if min_dr[iii] == dr[0, iii]:
x = np.log10(m_gal_1)
else:
x = np.log10(m_gal_2)
oh_12 = 27.7911 - 6.94493 * x + 0.808097 * x * x - 0.0301508 * x * x * x
z_zsolar = 10.0**(
oh_12 - 12) / 0.0004 # central metallicity in solar units.
this_id = tmp_ids[iii]
this_z = 0.0127 * z_zsolar * 10.0**(-0.1 * min_dr[iii])
self.init_ids[this_id] = this_id
self.init_metallicity[this_id] = this_z
def __init__(self, this_file, next_file=None, interp_coef=0.0, interp=True, dust=True, fof_num=-1, sub_num=-1, **kwargs):
self.head = ws.snapshot_header(this_file)
if dust:
print "loading dust data blocks"
self.gas_xyz = np.array(ws.read_block(this_file,'POS ', parttype=0))
try:
self.gas_hsml = np.array(ws.read_block(this_file,'HSML', parttype=0))
except:
print "failed to load HSML values directly. Try using cell volumes instead"
rho = np.array(ws.read_block(this_file,'RHO ', parttype=0))
mass = np.array(ws.read_block(this_file,'MASS', parttype=0))
vol = mass / rho
self.gas_hsml = (0.75 * vol / 3.14159 ) ** 0.33333 * 3
self.gas_mass = np.array(ws.read_block(this_file,'MASS', parttype=0))
self.gas_u = np.array(ws.read_block(this_file,'U ', parttype=0))
self.gas_rho = np.array(ws.read_block(this_file,'RHO ', parttype=0))
self.gas_numh = np.array(ws.read_block(this_file,'NH ', parttype=0))
self.gas_nume = np.array(ws.read_block(this_file,'NE ', parttype=0))
try:
self.gas_z = np.array(ws.read_block(this_file,'Z ', parttype=0))
except:
self.gas_z = np.array(ws.read_block(this_file,'GZ ', parttype=0))
else:
self.gas_xyz = np.zeros( (10,3) )
self.gas_hsml = np.zeros( 10 )
self.gas_mass = np.zeros( 10 )
self.gas_u = np.zeros( 10 )
self.gas_rho = np.zeros( 10 )
self.gas_numh = np.zeros( 10 )
self.gas_nume = np.zeros( 10 )
self.gas_z = np.zeros( 10 )
import simread.readhaloHDF5 as hr
print this_file
snapdir = this_file[:this_file.index('output')+7]
run=this_file[this_file.index('Runs/')+5:this_file.index('output') ].replace('/','')
print run
snapnum = int(this_file[-3:])
halo_reader = hr.HaloReader( snapdir, snapnum, run=run )
self.star_xyz = halo_reader.read( 'POS ', 4, fof_num, sub_num )
self.star_bt = halo_reader.read( 'GAGE', 4, fof_num, sub_num )
self.star_z = halo_reader.read( 'GZ ', 4, fof_num, sub_num )
self.star_mass = halo_reader.read( 'MASS', 4, fof_num, sub_num )
print "loaded stellar data from "+this_file
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
