def compute_veldisp(path, snapnum):
# try loading snapshot
sn = arepo.Snapshot(path + 'output/', snapnum, combineFiles=True)
time = sn.Time.as_unit(arepo.u.d).value * u.d
time = time.to_value(u.Myr)
firstpart = True
for i in [2, 3, 4]:
numpart = sn.NumPart_Total[i]
if numpart == 0:
continue
part = getattr(sn, 'part' + str(i))
pos_in_kpc = part.pos.as_unit(arepo.u.kpc).value
vel_in_kms = part.vel.as_unit(arepo.u.kms).value
x, y, z = pos_in_kpc[:, 0], pos_in_kpc[:, 1], pos_in_kpc[:, 2]
vx, vy, vz = vel_in_kms[:, 0], vel_in_kms[:, 1], vel_in_kms[:, 2]
R = np.sqrt(np.add(np.square(x), np.square(y)))
phi = np.arctan2(y, x)
sphi = np.sin(phi)
cphi = np.cos(phi)
vR = np.add(np.multiply(vx, cphi), np.multiply(vy, sphi))
vphi = np.subtract(np.multiply(vy, cphi), np.multiply(vx, sphi))
this_pos = np.transpose([R, phi, z])
this_vel = np.transpose([vR, vphi, vz])
if sn.MassTable[i].value == 0:
this_mass = part.mass.as_unit(arepo.u.msol).value
else:
this_mass = np.full(numpart,
sn.MassTable[i].as_unit(arepo.u.msol).value)
keys = np.logical_and(R > Rmin, R < Rmax)
if firstpart:
pos = np.copy(this_pos[keys])
vel = np.copy(this_vel[keys])
mass = np.copy(this_mass[keys])
firstpart = False
else:
pos = np.concatenate((pos, this_pos[keys]))
vel = np.concatenate((vel, this_vel[keys]))
mass = np.concatenate((mass, this_mass[keys]))
avevR, stdvR = weighted_avg_and_std(vel[:, 0], mass)
avevphi, stdvphi = weighted_avg_and_std(vel[:, 1], mass)
avevz, stdvz = weighted_avg_and_std(vel[:, 2], mass)
return time, stdvR, stdvphi, stdvz
def read_snap(path,
idx,
parttype=[0],
fields=['Coordinates', 'Masses', 'Velocities', 'ParticleIDs']):
fname = path + '/output'
return arepo.Snapshot(fname,
idx,
parttype=parttype,
fields=fields,
combineFiles=True)
def compute_SFR(path, name, output_dir='data/'):
# get the last snapshot
files = glob.glob(path + '/output/snapdir_*/snapshot_*.0.hdf5')
indices = np.array([int(re.findall('\d?\d\d\d', f)[-1]) for f in files])
last_key = np.argmax(indices)
file_last = files[last_key]
last_snap = arepo.Snapshot(file_last, combineFiles=True)
star_birthtime = last_snap.GFM_StellarFormationTime * time_conv
star_mini = last_snap.GFM_InitialMass.as_unit(arepo.u.msol).value
pickle.dump((star_birthtime, star_mini),
open(output_dir + 'SFR_' + name + '.p', 'wb'))
def read_snapshot(base, number, return_time=False):
snap = arepo.Snapshot(base, number, combineFiles=True)
for i, npart in enumerate(snap.NumPart_Total):
if npart == 0:
continue
part = getattr(snap, 'part'+str(i))
if return_time:
time = snap.Time.as_unit(arepo.u.d).value * u.d
time = time.to_value(u.Myr)
return snap, time
else:
return snap
def compute_gas_flow(path,
snapnum,
Rmin=0.0,
Rmax=30.0,
nbins=60,
logspace=False,
center=None):
# try loading snapshot
try:
sn = arepo.Snapshot(path + '/output/',
snapnum,
combineFiles=True,
parttype=[0],
fields=['Coordinates', 'Masses', 'Velocities'])
except:
print("unable to load path:" + path, " snapnum: ", snapnum)
return None
firstpart = True
pos = sn.part0.pos
if center is not None:
pos = np.subtract(pos, center)
vel = sn.part0.vel
mass = sn.part0.mass
Rlist, flow_rate = gas_flow(pos,
mass,
vel,
Rmin,
Rmax,
nbins=nbins,
logspace=logspace)
time = sn.Time.as_unit(arepo.u.d).value * u.d
time = time.to_value(u.Myr)
out = {}
out['Rlist'] = Rlist
out['flow_rate'] = flow_rate
out['time'] = time
return out
def compute_SFR_profile(path, snapnum, name, output_dir='data/', center=[200, 200, 200],
Rmin=0.1, Rmax=20, nbins=20):
# try loading snapshot
sn = arepo.Snapshot(path+'/output/', snapnum, combineFiles=True)
#get pos of gas
pos = np.subtract(sn.part0.pos, center)
R = np.linalg.norm(pos[:,:2], axis=1)
sfr = np.copy(sn.part0.sfr.value)
#do binning
# bins = np.logspace(np.log10(Rmin), np.log10(Rmax), nbins)
bins = np.linspace(Rmin, Rmax, nbins)
R_binned, tot_sfr_binned = my_bin(R, bins, sfr, np.sum)
# now go through and divide by area of bin
for i in range(len(R_binned)):
surf_area = np.pi * (bins[i+1]**2 - bins[i]**2)
tot_sfr_binned /= surf_area
pickle.dump((R_binned, tot_sfr_binned), open(output_dir+'SFR_profile_'+name+'_snap'+str(snapnum)+'.p', 'wb'))
def compute_potential(path, snapnum, name, center=np.array([0, 0, 0]), output_dir='data/'):
# try loading snapshot
sn = arepo.Snapshot(path+'output/', snapnum, combineFiles=True)
return sn
def make_projection_snap(path,
snapnum,
parttype=[0, 2, 3, 4],
center=np.array([200, 200, 200]),
width=30.,
nres=256):
sn = arepo.Snapshot(path + '/output',
snapnum,
parttype=parttype,
combineFiles=True,
fields=['Coordinates', 'Masses'])
time = sn.Time
range_xy = [[center[0] - width / 2.0, center[0] + width / 2.0],
[center[1] - width / 2.0, center[1] + width / 2.0]]
range_xz = [[center[0] - width / 2.0, center[0] + width / 2.0],
[center[2] - width / 2.0, center[2] + width / 2.0]]
range_yz = [[center[1] - width / 2.0, center[1] + width / 2.0],
[center[2] - width / 2.0, center[2] + width / 2.0]]
surf = (width / nres)**(2.0)
heatmap_xy_out = []
heatmap_xz_out = []
heatmap_yz_out = []
for pt in parttype:
if sn.NumPart_Total[pt] == 0:
heatmap_xy_out.append(np.zeros((nres, nres)))
heatmap_xz_out.append(np.zeros((nres, nres)))
heatmap_yz_out.append(np.zeros((nres, nres)))
continue
part = getattr(sn, 'part' + str(pt))
x = part.pos[:, 0]
y = part.pos[:, 1]
z = part.pos[:, 2]
xbool = np.logical_and(x > center[0] - width / 2.0,
x < center[0] + width / 2.0)
ybool = np.logical_and(y > center[1] - width / 2.0,
y < center[1] + width / 2.0)
zbool = np.logical_and(z > center[2] - width / 2.0,
z < center[2] + width / 2.0)
keys = np.logical_and(np.logical_and(xbool, ybool), zbool)
if sn.MassTable[pt] > 0:
weights = None
postfac = sn.MassTable[pt] / surf
else:
weights = part.Masses[keys] / surf
postfac = 1.0
heatmap_xy, _, _ = np.histogram2d(x[keys],
y[keys],
bins=(nres, nres),
range=range_xy,
weights=weights)
heatmap_xz, _, _ = np.histogram2d(x[keys],
z[keys],
bins=(nres, nres),
range=range_xz,
weights=weights)
heatmap_yz, _, _ = np.histogram2d(y[keys],
z[keys],
bins=(nres, nres),
range=range_yz,
weights=weights)
heatmap_xy *= postfac
heatmap_xz *= postfac
heatmap_yz *= postfac
heatmap_xy_out.append(heatmap_xy)
heatmap_xz_out.append(heatmap_xz)
heatmap_yz_out.append(heatmap_yz)
return heatmap_xy_out, heatmap_xz_out, heatmap_xy_out, time
def make_projection_snap(path, snapnum, parttype=[0, 2, 3, 4],
center=np.array([200, 200, 200]), width=30., nres=256,
corot=False, bar_angle=None, fourier_n=5):
sn = arepo.Snapshot(path+'/output', snapnum, parttype=parttype,
combineFiles=True, fields=['Coordinates', 'Masses'])
time = sn.Time
ba, _ = bar_angle['poly_eval'][fourier_n]
ba = ba[snapnum]
range_xy = [[center[0] - width/2.0, center[0] + width/2.0], [center[1] - width/2.0, center[1] + width/2.0]]
range_xz = [[center[0] - width/2.0, center[0] + width/2.0], [center[2] - width/2.0, center[2] + width/2.0]]
range_yz = [[center[1] - width/2.0, center[1] + width/2.0], [center[2] - width/2.0, center[2] + width/2.0]]
surf = (width/nres)**(2.0)
heatmap_xy_out = []
heatmap_xz_out = []
heatmap_yz_out = []
for pt in parttype:
if sn.NumPart_Total[pt] == 0:
heatmap_xy_out.append(np.zeros((nres, nres)))
heatmap_xz_out.append(np.zeros((nres, nres)))
heatmap_yz_out.append(np.zeros((nres, nres)))
continue
part = getattr(sn, 'part'+str(pt))
x = part.pos[:,0].value
y = part.pos[:,1].value
z = part.pos[:,2].value
if corot:
x -= center[0]
y -= center[1]
phi = np.arctan2(y, x)
R = np.sqrt(x**2 + y**2)
phi -= ba
x = R * np.cos(phi)
y = R * np.sin(phi)
x += center[0]
y += center[1]
xbool = np.logical_and(x > center[0] - width/2.0, x < center[0] + width/2.0)
ybool = np.logical_and(y > center[1] - width/2.0, y < center[1] + width/2.0)
zbool = np.logical_and(z > center[2] - width/2.0, z < center[2] + width/2.0)
keys = np.logical_and(np.logical_and(xbool, ybool), zbool)
if sn.MassTable[pt] > 0:
weights = None
postfac = sn.MassTable[pt] / surf
else:
weights = part.Masses[keys] / surf
postfac = 1.0
heatmap_xy, _, _ = np.histogram2d(x[keys], y[keys], bins=(nres, nres), range=range_xy, weights=weights)
heatmap_xz, _, _ = np.histogram2d(x[keys], z[keys], bins=(nres, nres), range=range_xz, weights=weights)
heatmap_yz, _, _ = np.histogram2d(y[keys], z[keys], bins=(nres, nres), range=range_yz, weights=weights)
heatmap_xy *= postfac
heatmap_xz *= postfac
heatmap_yz *= postfac
heatmap_xy_out.append(heatmap_xy)
heatmap_xz_out.append(heatmap_xz)
heatmap_yz_out.append(heatmap_yz)
return heatmap_xy_out, heatmap_xz_out, heatmap_xy_out, time
def compute_fourier_component(path,
snapnum,
Rmin=0.0,
Rmax=30.0,
nbins=60,
logspace=False,
center=None):
# try loading snapshot
try:
sn = arepo.Snapshot(path + '/output/',
snapnum,
combineFiles=True,
parttype=[2, 3, 4],
fields=['Coordinates', 'Masses'])
except:
print("unable to load path:" + path, " snapnum: ", snapnum)
return None
firstpart = True
for i, npart in enumerate(sn.NumPart_Total):
if i not in [2, 3]:
continue
if npart == 0:
continue
part = getattr(sn, 'part' + str(i))
# compute the center of mass
this_mass = sn.MassTable[i].as_unit(arepo.u.msol).value
this_pos = part.pos.as_unit(arepo.u.kpc).value
if center is not None:
this_pos = np.subtract(this_pos, center)
# if mass is zero, then we need to load each individual mass
if this_mass == 0:
this_mass = part.mass.as_unit(arepo.u.msol).value
else:
this_mass = np.full(npart, this_mass)
# now concatenate if needed
if firstpart:
mass = np.copy(this_mass)
pos = np.copy(this_pos)
firstpart = False
else:
mass = np.concatenate((mass, this_mass))
pos = np.concatenate((pos, this_pos))
Rlist, A0, _ = fourier_component(pos,
mass,
0,
Rmin,
Rmax,
nbins=nbins,
logspace=logspace)
Rlist, A1r, A1i = fourier_component(pos,
mass,
1,
Rmin,
Rmax,
nbins=nbins,
logspace=logspace)
Rlist, A2r, A2i = fourier_component(pos,
mass,
2,
Rmin,
Rmax,
nbins=nbins,
logspace=logspace)
Rlist, A3r, A3i = fourier_component(pos,
mass,
3,
Rmin,
Rmax,
nbins=nbins,
logspace=logspace)
Rlist, A4r, A4i = fourier_component(pos,
mass,
4,
Rmin,
Rmax,
nbins=nbins,
logspace=logspace)
Rlist, A5r, A5i = fourier_component(pos,
mass,
5,
Rmin,
Rmax,
nbins=nbins,
logspace=logspace)
Rlist, A6r, A6i = fourier_component(pos,
mass,
6,
Rmin,
Rmax,
nbins=nbins,
logspace=logspace)
Rlist, A7r, A7i = fourier_component(pos,
mass,
7,
Rmin,
Rmax,
nbins=nbins,
logspace=logspace)
Rlist, A8r, A8i = fourier_component(pos,
mass,
8,
Rmin,
Rmax,
nbins=nbins,
logspace=logspace)
Rlist, A9r, A9i = fourier_component(pos,
mass,
9,
Rmin,
Rmax,
nbins=nbins,
logspace=logspace)
Rlist, A10r, A10i = fourier_component(pos,
mass,
10,
Rmin,
Rmax,
nbins=nbins,
logspace=logspace)
time = sn.Time.as_unit(arepo.u.d).value * u.d
time = time.to_value(u.Myr)
out = {}
out['Rlist'] = Rlist
out['A0'] = A0
out['A1r'], out['A1i'] = A1r, A1i
out['A2r'], out['A2i'] = A2r, A2i
out['A3r'], out['A3i'] = A3r, A3i
out['A4r'], out['A4i'] = A4r, A4i
out['A5r'], out['A5i'] = A5r, A5i
out['A6r'], out['A6i'] = A6r, A6i
out['A7r'], out['A7i'] = A7r, A7i
out['A8r'], out['A8i'] = A8r, A8i
out['A9r'], out['A9i'] = A9r, A9i
out['A10r'], out['A10i'] = A10r, A10i
out['time'] = time
return out
def compute_acc_alltypes(path,
snapnum,
center=np.array([0., 0., 0.]),
bar_Lz=400,
theta=0.35,
num_threads=1):
# try loading snapshot
try:
sn = arepo.Snapshot(path + '/output/',
snapnum,
combineFiles=True,
fields=['Coordinates', 'Masses', 'Velocities'])
except:
print("unable to load path:" + path, " snapnum: ", snapnum)
return None
acc_out = {}
# now separate into bar/not bar for disk
cyl = get_cyl(sn, 2, center=center)
R = cyl[:, 0]
vphi = cyl[:, 4]
bar_key = R * vphi < bar_Lz
disk_key = np.logical_not(bar_key)
bar_pos = sn.part2.pos.value[bar_key]
for i, npart in enumerate(sn.NumPart_Total):
if npart == 0:
continue
part = getattr(sn, 'part' + str(i))
soft = getattr(sn.parameters, 'SofteningComovingType' + str(i))
# compute the center of mass
mass_val = sn.MassTable[i].value
tree_pos = part.pos.value
if center is not None:
tree_pos = np.subtract(tree_pos, center)
# if mass is zero, then we need to load each individual mass
if mass_val == 0:
tree_mass = part.mass.value
else:
tree_mass = np.full(npart, mass_val)
if i == 2:
tree_pos_b = tree_pos[bar_key]
tree_mass_b = tree_mass[bar_key]
tree_pos_d = tree_pos[disk_key]
tree_mass_d = tree_mass[disk_key]
acc_b = compute_acc(bar_pos,
tree_pos_b,
tree_mass_b,
theta,
soft,
num_threads=num_threads)
acc_d = compute_acc(bar_pos,
tree_pos_d,
tree_mass_d,
theta,
soft,
num_threads=num_threads)
acc_out['acc_bar'] = acc_b
acc_out['acc_disk'] = acc_d
else:
if i == 0:
name = 'acc_gas'
maxnode_fac = 20.0
elif i == 1:
name = 'acc_halo'
maxnode_fac = 1.5
elif i == 3:
name = 'acc_bulge'
maxnode_fac = 1.5
elif i == 4:
name == 'acc_star'
maxnode_fac = 1.5
print(i, name, maxnode_fac)
acc = compute_acc(bar_pos,
tree_pos,
tree_mass,
theta,
soft,
num_threads=num_threads,
maxnode_fac=maxnode_fac)
print('done')
acc_out[name] = acc
acc_out['time'] = sn.Time.value
return acc_out
def compute_REPLACE(path, snapnum, name, output_dir='data/'):
# try loading snapshot
sn = arepo.Snapshot(path + 'output/', snapnum, combineFiles=True)
return sn