def generate_cores_kdtree(M1,
M2,
s1=False,
disrupt=None,
onlyFiltered=False,
giveFragmentsFofMass=False):
idx_filteredsatcores, M, X, Y, Z, nHalo = SHMLM.core_mask(cc,
M1=M1,
M2=M2,
s1=s1,
disrupt=disrupt,
z=z)
cc_filtered = {k: cc[k][idx_filteredsatcores].copy() for k in cc.keys()}
cc_filtered['M'] = M
if giveFragmentsFofMass:
fragmask = cc_filtered['fof_halo_tag'] < 0
frag_realfht = np.bitwise_and(
cc_filtered['fof_halo_tag'][fragmask] * -1, 0xffffffffffff)
idx_m21 = many_to_one(frag_realfht, mt_cores['fof_halo_tag'])
cc_filtered['M'][fragmask] = mt_cores['fof_halo_mass'][idx_m21]
cc_filtered['x'] = SHMLM.periodic_bcs(cc_filtered['x'], X, SHMLM.BOXSIZE)
cc_filtered['y'] = SHMLM.periodic_bcs(cc_filtered['y'], Y, SHMLM.BOXSIZE)
cc_filtered['z'] = SHMLM.periodic_bcs(cc_filtered['z'], Z, SHMLM.BOXSIZE)
Mmask = (M1 <= cc_filtered['M']) & (cc_filtered['M'] <= M2)
cc_filtered = {k: cc_filtered[k][Mmask].copy() for k in cc_filtered.keys()}
if onlyFiltered:
return cc_filtered
else:
return spatial.cKDTree(
np.vstack((cc_filtered['x'], cc_filtered['y'],
cc_filtered['z'])).T), cc_filtered, nHalo
def sh_cores_match(cc, sh, mt, z):
cores_tree, cc_filtered, _ = generate_cores_kdtree(cc,
M1=10**9,
M2=10**16,
s1=False,
disrupt=None,
z=z)
sh['rvir'] = SHMLM.getRvir(sh['subhalo_mass'], z)
idx_m21_sh = many_to_one(sh['fof_halo_tag'], mt['fof_halo_tag'])
sh['M'] = mt['fof_halo_mass'][idx_m21_sh]
sh['X'] = mt['fof_halo_center_x'][idx_m21_sh]
sh['Y'] = mt['fof_halo_center_y'][idx_m21_sh]
sh['Z'] = mt['fof_halo_center_z'][idx_m21_sh]
subhalo_mean_x = SHMLM.periodic_bcs(sh['subhalo_mean_x'], sh['X'],
SHMLM.BOXSIZE)
subhalo_mean_y = SHMLM.periodic_bcs(sh['subhalo_mean_y'], sh['Y'],
SHMLM.BOXSIZE)
subhalo_mean_z = SHMLM.periodic_bcs(sh['subhalo_mean_z'], sh['Z'],
SHMLM.BOXSIZE)
sh_mask = (sh['subhalo_tag'] != 0)
sh_arr = np.vstack((subhalo_mean_x[sh_mask], subhalo_mean_y[sh_mask],
subhalo_mean_z[sh_mask])).T
distance_upper_bound = 2 * sh['rvir'][sh_mask]
# Search radius of 2*rvir around each subhalo, only look for 1:1 matches
dist, idx = [], []
for i in range(len(distance_upper_bound)):
dv, iv = cores_tree.query(sh_arr[i],
k=2,
distance_upper_bound=distance_upper_bound[i])
dist.append(dv)
idx.append(iv)
dist = np.array(dist)
idx = np.array(idx)
f1, f2 = (dist != np.inf).T
fmask = f1 ^ f2
percentmatch = np.sum(fmask) / np.sum(sh_mask) * 100
f1i = f1[np.invert(fmask)]
percentmany = np.sum(f1i) / len(f1i) * 100
percentnone = np.sum(np.invert(f1i)) / len(f1i) * 100
print '{}% of masked subhalos have 1:1 core match. Of the unmatched subhalos, {}% have mutliple cores and {}% have no core.'.format(
percentmatch, percentmany, percentnone)
return np.flatnonzero(sh_mask)[fmask], idx[:, 0][fmask], cc_filtered
def core_mask(cc, M1, M2, s1=False, disrupt=None, z=0, idx_m21=None):
"""Given a core catalog and filtering criteria, returns indices in cc of filtered satellite cores.
Arguments:
cc {dict} -- core catalog
M1, M2 {float} -- Host halos with M in mass range [M1, M2], where M is central's `infall_mass` from core catalog.
disrupt {None, int, or list} -- If given, applies core filtering criteria defined in SHMLM when filtering substructure. (default: {None})
z {int} -- Redshift (default: {0})
Keyword Arguments:
s1 {bool} -- If true, consider only 1st order substructure (i.e. subhalos). (default: {False})
Returns:
np 1d array -- indices of satellite filtered cores in `cc`
np 1d array -- M array (corresponds with satellite filtered cores)
integer -- number of host halos (M)
"""
satellites_mask = cc['central'] == 0
centrals_mask = cc['central'] == 1
# Create M array (corresponds with cc[satellites_mask]) to be host tree node mass of each satellite
if idx_m21 is None:
idx_m21 = many_to_one(cc['tree_node_index'][satellites_mask],
cc['tree_node_index'][centrals_mask])
M = cc['infall_mass'][centrals_mask][idx_m21]
if 'x' in cc.keys():
X, Y, Z = cc['x'][centrals_mask][idx_m21], cc['y'][centrals_mask][
idx_m21], cc['z'][centrals_mask][idx_m21]
mask = np.flatnonzero(
(M1 <= M) &
(M <=
M2)) #& (cc['infall_mass'][satellites_mask] >= PARTICLES100MASS) )
if s1:
Coretag = cc['core_tag'][centrals_mask][idx_m21]
mask = np.intersect1d(
mask, np.flatnonzero(cc['host_core'][satellites_mask] == Coretag))
if disrupt is not None:
cc_satellites = {k: cc[k][satellites_mask] for k in cc.keys()}
mask = np.intersect1d(
mask,
np.flatnonzero(
disruption_mask(cc_satellites, disrupt, M, X, Y, Z, z)))
nHalo = np.sum((M1 <= cc['infall_mass'][centrals_mask])
& (cc['infall_mass'][centrals_mask] <= M2))
if 'x' in cc.keys():
return np.flatnonzero(
satellites_mask)[mask], M[mask], X[mask], Y[mask], Z[mask], nHalo
else:
return np.flatnonzero(satellites_mask)[mask], M[mask], nHalo
def SHMF(M1, M2, partcut):
"""Returns 'subhalo mass function' m/M for M in [`M1`,`M2`]."""
subhalos_read()
assert np.all(
np.isin(sh['fof_halo_tag'], mt['fof_halo_tag'])
), 'Every subhalo does not have corresponding fof_halo_tag in merger tree.'
assert len(np.unique(mt['fof_halo_tag'])) == len(
mt['fof_halo_tag']), 'Duplicate fof_halo_tag(s) found in merger tree.'
idx_m21 = many_to_one(sh['fof_halo_tag'], mt['fof_halo_tag'])
# Initialize M array (corresponds with sh) to be host halo fof mass of each subhalo
M = mt['fof_halo_mass'][idx_m21]
# m/M with mask: M in [M1,M2] with central subhalos (subhalo tag == 0) removed
mask = (M1 <= M) & (M <= M2) & (sh['subhalo_tag'] != 0)
if partcut:
mask = mask & (sh['subhalo_mass'] >= SHMLM.PARTICLES100MASS)
plot_arr = (sh['subhalo_mass'] / M)[mask]
# nH = np.sum( np.isin(mt['fof_halo_tag'], sh['fof_halo_tag'][mask]) ) #all halos with at least 1 subhalo
nH = np.sum((M1 <= mt['fof_halo_mass']) & (mt['fof_halo_mass'] <= M2))
return plot_arr, nH
def create_core_catalog_mevolved(virialFlag, writeOutputFlag=True, useLocalHost=False):
"""
Appends mevolved to core catalog and saves output in HDF5.
Works by computing mevolved for step+1 at each step and saving that in memory.
"""
if writeOutputFlag:
print 'Reading data from {} and writing output to {}'.format(cc_data_dir, cc_output_dir)
cc = {}
cc_prev = {}
for step in tqdm(steps):
# Read in cc for step
cc = { v:gio.gio_read(fname_cc(step, 'input'), v)[0] for v in vars_cc }
if virialFlag:
# Convert all mass to virial
cc['infall_mass'] = SHMLM.m_vir(cc['infall_mass'], step)
satellites_mask = cc['central'] == 0
centrals_mask = cc['central'] == 1
# assert np.sum(satellites_mask) + np.sum(centrals_mask) == len(cc['infall_mass']), 'central flag not 0 or 1.'
numSatellites = np.sum(satellites_mask)
# Verify there are no satellites at first step
if step == steps[0]:
assert numSatellites == 0, 'Satellites found at first step.'
# Add column for m_evolved and initialize to 0
for A in A_arr:
for zeta in zeta_arr:
cc[m_evolved_col(A, zeta)] = np.zeros_like(cc['infall_mass'])
# If there are satellites (not applicable for first step)
if numSatellites != 0:
# Match satellites to prev step cores using core tag.
vals, idx1, idx2 = np.intersect1d( cc['core_tag'][satellites_mask], cc_prev['core_tag'], return_indices=True )
# assert len(vals) == len(cc['core_tag'][satellites_mask]), 'All cores from prev step did not carry over.'
# assert len(cc['core_tag'][satellites_mask]) == len(np.unique(cc['core_tag'][satellites_mask])), 'Satellite core tags not unique for this time step.'
for A in A_arr:
for zeta in zeta_arr:
# Set m_evolved of all satellites that have core tag match on prev step to next_m_evolved of prev step.
# DOUBLE MASK ASSIGNMENT: cc['m_evolved'][satellites_mask][idx1] = cc_prev['next_m_evolved'][idx2]
cc[m_evolved_col(A, zeta)][ np.flatnonzero(satellites_mask)[idx1] ] = cc_prev[m_evolved_col(A, zeta, next=True)][idx2]
# Initialize m array (corresponds with cc[satellites_mask]) to be either infall_mass (step after infall) or m_evolved (subsequent steps).
m = (cc[m_evolved_col(A, zeta)][satellites_mask] == 0)*cc['infall_mass'][satellites_mask] + (cc[m_evolved_col(A, zeta)][satellites_mask] != 0)*cc[m_evolved_col(A, zeta)][satellites_mask]
# Initialize m array (corresponds with cc[satellites_mask]) to be either virial_infall_mass (step after infall) or m_evolved (subsequent steps).
#m = (cc[m_evolved_col(A, zeta)][satellites_mask] == 0)*SHMLM.m_vir(cc['infall_mass'][satellites_mask], step) + (cc[m_evolved_col(A, zeta)][satellites_mask] != 0)*cc[m_evolved_col(A, zeta)][satellites_mask]
# Set m_evolved of satellites with m_evolved=0 to infall mass.
cc[m_evolved_col(A, zeta)][satellites_mask] = m
# Initialize M array (corresponds with cc[satellites_mask]) to be host tree node mass of each satellite
M = cc['infall_mass'][centrals_mask][ many_to_one( cc['tree_node_index'][satellites_mask], cc['tree_node_index'][centrals_mask] ) ]
# assert len(M) == len(m) == len(cc['m_evolved'][satellites_mask]), 'M, m, cc[satellites_mask] lengths not equal.'
# wasInFragment: persistent flag that is False by default and becomes True when core is inside a fragment halo
cc['wasInFragment'] = cc['fof_halo_tag']<0
if step != steps[0]:
_, i1, i2 = np.intersect1d( cc['core_tag'], cc_prev['core_tag'], return_indices=True )
cc['wasInFragment'][i1] = np.logical_or( cc['wasInFragment'][i1], cc_prev['wasInFragment'][i2] )
if writeOutputFlag:
# Write output to disk
h5_write_dict( fname_cc(step, 'output'), cc, 'coredata' )
# Compute m_evolved of satellites according to SHMLModel for NEXT time step and save as cc_prev['next_m_evolved'] in memory.
# Mass loss assumed to begin at step AFTER infall detected.
if step != steps[-1]:
cc_prev = { 'core_tag':cc['core_tag'].copy(), 'wasInFragment':cc['wasInFragment'].copy() }
if numSatellites != 0:
localhost_m21 = many_to_one( cc['host_core'][satellites_mask], cc['core_tag'] )
for A in A_arr:
for zeta in zeta_arr:
cc_prev[m_evolved_col(A, zeta, next=True)] = np.zeros_like(cc['infall_mass'])
if numSatellites != 0: # If there are satellites (not applicable for first step)
m = cc[m_evolved_col(A, zeta)][satellites_mask]
if useLocalHost:
Mlocal = cc[m_evolved_col(A, zeta)][localhost_m21]
M_A_zeta = (Mlocal==0)*M + (Mlocal!=0)*Mlocal
cc_prev[m_evolved_col(A, zeta, next=True)][satellites_mask] = SHMLM.m_evolved(m0=m, M0=M_A_zeta, step=steps[steps.index(step)+1], step_prev=step, A=A, zeta=zeta)
else:
cc_prev[m_evolved_col(A, zeta, next=True)]][satellites_mask] = SHMLM.m_evolved(m0=m, M0=M, step=steps[steps.index(step)+1], step_prev=step, A=23.7, zeta=0.36)
return cc
def create_core_catalog_mevolved(virialFlag,
writeOutputFlag=True,
useLocalHost=False):
"""
Appends mevolved to core catalog and saves output in HDF5.
Works by computing mevolved for step+1 at each step and saving that in memory.
"""
if writeOutputFlag:
print 'Reading data from {} and writing output to {}'.format(
cc_data_dir, cc_output_dir)
global cc, coresMaskedArray, distance_upper_bound, pool_mergedCoreTag, coresKDTree, KDTree_idx
cc = {}
cc_prev = {}
for step in tqdm(steps):
# Read in cc for step
cc = {v: gio.gio_read(fname_cc(step, 'input'), v)[0] for v in vars_cc}
if virialFlag:
# Convert all mass to virial
cc['infall_mass'] = SHMLM.m_vir(cc['infall_mass'], step)
satellites_mask = cc['central'] == 0
centrals_mask = cc['central'] == 1
# assert np.sum(satellites_mask) + np.sum(centrals_mask) == len(cc['infall_mass']), 'central flag not 0 or 1.'
numSatellites = np.sum(satellites_mask)
# Verify there are no satellites at first step
if step == steps[0]:
assert numSatellites == 0, 'Satellites found at first step.'
# Add column for m_evolved and initialize to 0
for A in A_arr:
for zeta in zeta_arr:
cc[m_evolved_col(A, zeta)] = np.zeros_like(cc['infall_mass'])
cc['mergedCoreTag'] = np.zeros_like(cc['core_tag'])
cc['mergedCoreStep'] = np.zeros_like(cc['infall_step'])
# wasInFragment: persistent flag that is False by default and becomes True when core is inside a fragment halo
cc['wasInFragment'] = cc['fof_halo_tag'] < 0
if step != steps[0]:
_, i1, i2 = np.intersect1d(cc['core_tag'],
cc_prev['core_tag'],
return_indices=True)
cc['wasInFragment'][i1] = np.logical_or(
cc['wasInFragment'][i1], cc_prev['wasInFragment'][i2])
cc['mergedCoreTag'][i1] = cc_prev['mergedCoreTag'][i2]
cc['mergedCoreStep'][i1] = cc_prev['mergedCoreStep'][i2]
# If there are satellites (not applicable for first step)
if numSatellites != 0:
# Match satellites to prev step cores using core tag.
vals, idx1, idx2 = np.intersect1d(cc['core_tag'][satellites_mask],
cc_prev['core_tag'],
return_indices=True)
# assert len(vals) == len(cc['core_tag'][satellites_mask]), 'All cores from prev step did not carry over.'
# assert len(cc['core_tag'][satellites_mask]) == len(np.unique(cc['core_tag'][satellites_mask])), 'Satellite core tags not unique for this time step.'
for A in A_arr:
for zeta in zeta_arr:
# Set m_evolved of all satellites that have core tag match on prev step to next_m_evolved of prev step.
# DOUBLE MASK ASSIGNMENT: cc['m_evolved'][satellites_mask][idx1] = cc_prev['next_m_evolved'][idx2]
cc[m_evolved_col(A, zeta)][np.flatnonzero(satellites_mask)
[idx1]] = cc_prev[m_evolved_col(
A, zeta, next=True)][idx2]
# Initialize m array (corresponds with cc[satellites_mask]) to be either infall_mass (step after infall) or m_evolved (subsequent steps).
m = (cc[m_evolved_col(A, zeta)][satellites_mask]
== 0) * cc['infall_mass'][satellites_mask] + (
cc[m_evolved_col(A, zeta)][satellites_mask] !=
0) * cc[m_evolved_col(A, zeta)][satellites_mask]
# Set m_evolved of satellites with m_evolved=0 to infall mass.
cc[m_evolved_col(A, zeta)][satellites_mask] = m
# Initialize M array (corresponds with cc[satellites_mask]) to be host tree node mass of each satellite
idx_m21 = many_to_one(cc['tree_node_index'][satellites_mask],
cc['tree_node_index'][centrals_mask])
M, X, Y, Z = (cc[k][centrals_mask][idx_m21]
for k in ['infall_mass', 'x', 'y', 'z'])
KDTreemask = cc['mergedCoreTag'][satellites_mask] == 0
x = SHMLM.periodic_bcs(cc['x'][satellites_mask][KDTreemask],
X[KDTreemask], SHMLM.BOXSIZE)
y = SHMLM.periodic_bcs(cc['y'][satellites_mask][KDTreemask],
Y[KDTreemask], SHMLM.BOXSIZE)
z = SHMLM.periodic_bcs(cc['z'][satellites_mask][KDTreemask],
Z[KDTreemask], SHMLM.BOXSIZE)
# coresMaskedArray = np.vstack((x, y, z)).T
coresMaskedArray_Tree = np.vstack(
(np.r_[x,
cc['x'][centrals_mask]], np.r_[y,
cc['y'][centrals_mask]],
np.r_[z, cc['z'][centrals_mask]])).T
coresMaskedArray = coresMaskedArray_Tree
KDTree_idx = np.r_[np.flatnonzero(satellites_mask)[KDTreemask],
np.flatnonzero(centrals_mask)]
coresKDTree = spatial.cKDTree(coresMaskedArray_Tree)
# Set mergedCoreTag as central core of KDTreemask cores with no mergedCoreTag that are `merged`
'''
disruptMask = cc['merged'][satellites_mask][KDTreemask]|(cc['radius'][satellites_mask][KDTreemask] >= SHMLM.CORERADIUSCUT)
cc['mergedCoreTag'][ np.flatnonzero(satellites_mask)[KDTreemask] ] += disruptMask*cc['core_tag'][centrals_mask][idx_m21][KDTreemask]
cc['mergedCoreStep'][ np.flatnonzero(satellites_mask)[KDTreemask] ] += disruptMask*step
'''
### core merging detection
# rvir_KDTreemask = SHMLM.getRvir(cc[m_evolved_col(1.1, 0.068)][satellites_mask][KDTreemask], SHMLM.step2z[step]) #infall should be exact m_evolved in production
# rvir_KDTreemask = SHMLM.getRvir(cc['infall_mass'][satellites_mask][KDTreemask], SHMLM.step2z[step]) #infall should be exact m_evolved in production
# distance_upper_bound = SHMLM.COREVIRIALRADIUSFACTOR * rvir_KDTreemask
distance_upper_bound = SHMLM.COREVIRIALRADIUSFACTOR * SHMLM.getRvir(
cc['infall_mass'][KDTree_idx], SHMLM.step2z[step])
import ctypes
# pool_mergedCoreTag = { i:np.full(len(KDTree_idx), 0, dtype=np.int64) for i in range(1,pool_size+1) }
pool_mergedCoreTag = {
i: multiprocessing.RawArray(
ctypes.c_int64, np.zeros(len(KDTree_idx), dtype=np.int64))
for i in range(1, pool_size + 1)
}
# Search radius of 2*rvir around each core
from datetime import datetime
tstart = datetime.now()
work = np.arange(len(distance_upper_bound))
t1 = datetime.now() - tstart
print "work=", t1
tstart = datetime.now()
p = Pool(pool_size)
t2 = datetime.now() - tstart
print "Pool=", t2
tstart = datetime.now()
p.map(query, np.array_split(work, pool_size))
t3 = datetime.now() - tstart
print "mapp=", t3
tstart = datetime.now()
p.close()
t4 = datetime.now() - tstart
print "close=", t4
tstart = datetime.now()
p.join()
t5 = datetime.now() - tstart
print "join=", t5
tstart = datetime.now()
# mergedCoreTag = np.full_like(pool_mergedCoreTag[1], 0)
mergedCoreTag = np.zeros(len(KDTree_idx), dtype=np.int64)
# pool_mergedCoreTag = {k:np.array(pool_mergedCoreTag[k]) for k in pool_mergedCoreTag.keys()}
for i in range(1, pool_size + 1):
mCT = np.array(pool_mergedCoreTag[i])
newMergedMask = mCT != 0
# print np.sum(newMergedMask)
mergedCoreTag[newMergedMask] = mCT[newMergedMask]
t6 = datetime.now() - tstart
newMergedMask = mergedCoreTag != 0
cc['mergedCoreTag'][
KDTree_idx[newMergedMask]] = mergedCoreTag[newMergedMask]
cc['mergedCoreStep'][KDTree_idx[newMergedMask]] = step
print np.sum(cc['mergedCoreTag'] != 0)
t7 = datetime.now() - tstart
print t1, t2, t3, t4, t5, t6, t7
###
print 'Merged centrals: ', np.sum(
cc['mergedCoreTag'][centrals_mask] != 0)
cc['mergedCoreTag'][centrals_mask] = 0
cc['mergedCoreStep'][centrals_mask] = 0
if writeOutputFlag:
# Write output to disk
h5_write_dict(fname_cc(step, 'output'), cc, 'coredata')
# Compute m_evolved of satellites according to SHMLModel for NEXT time step and save as cc_prev['next_m_evolved'] in memory.
# Mass loss assumed to begin at step AFTER infall detected.
if step != steps[-1]:
# cc_prev = { 'core_tag':cc['core_tag'].copy(), 'wasInFragment':cc['wasInFragment'].copy() }
cc_prev = {
k: cc[k].copy()
for k in [
'core_tag', 'wasInFragment', 'mergedCoreTag',
'mergedCoreStep'
]
}
if numSatellites != 0:
localhost_m21 = many_to_one(cc['host_core'][satellites_mask],
cc['core_tag'])
for A in A_arr:
for zeta in zeta_arr:
cc_prev[m_evolved_col(A, zeta, next=True)] = np.zeros_like(
cc['infall_mass'])
if numSatellites != 0: # If there are satellites (not applicable for first step)
m = cc[m_evolved_col(A, zeta)][satellites_mask]
if useLocalHost:
Mlocal = cc[m_evolved_col(A, zeta)][localhost_m21]
M_A_zeta = (Mlocal
== 0) * M + (Mlocal != 0) * Mlocal
cc_prev[m_evolved_col(
A, zeta,
next=True)][satellites_mask] = SHMLM.m_evolved(
m0=m,
M0=M_A_zeta,
step=steps[steps.index(step) + 1],
step_prev=step,
A=A,
zeta=zeta)
else:
cc_prev[m_evolved_col(
A, zeta,
next=True)][satellites_mask] = SHMLM.m_evolved(
m0=m,
M0=M,
step=steps[steps.index(step) + 1],
step_prev=step,
A=A,
zeta=zeta)
return cc
def create_core_catalog_mevolved(virialFlag=False,
writeOutputFlag=True,
useLocalHost=False):
"""
Appends mevolved to core catalog and saves output in HDF5.
Works by computing mevolved for step+1 at each step and saving that in memory.
"""
if writeOutputFlag:
print 'Reading data from {} and writing output to {}'.format(
cc_data_dir, cc_output_dir)
cc = {}
cc_prev = {}
for step in tqdm(steps):
# Read in cc for step
cc = {v: gio.gio_read(fname_cc(step, 'input'), v)[0] for v in vars_cc}
if virialFlag:
# Convert all mass to virial
cc['infall_mass'] = SHMLM.m_vir(cc['infall_mass'], step)
satellites_mask = cc['central'] == 0
centrals_mask = cc['central'] == 1
# assert np.sum(satellites_mask) + np.sum(centrals_mask) == len(cc['infall_mass']), 'central flag not 0 or 1.'
numSatellites = np.sum(satellites_mask)
# Verify there are no satellites at first step
if step == steps[0]:
assert numSatellites == 0, 'Satellites found at first step.'
# Add column for m_evolved and initialize to 0
for A in A_arr:
for zeta in zeta_arr:
cc[m_evolved_col(A, zeta)] = np.zeros_like(cc['infall_mass'])
'''
#cc['mergedCoreTag'] = np.zeros_like(cc['core_tag'])
#cc['mergedCoreStep'] = np.zeros_like(cc['infall_step'])
cc['phaseSpaceMerged'] = np.zeros_like(cc['central'])
# wasInFragment: persistent flag that is False by default and becomes True when core is inside a fragment halo
cc['wasInFragment'] = cc['fof_halo_tag']<0
if step != steps[0]:
_, i1, i2 = np.intersect1d( cc['core_tag'], cc_prev['core_tag'], return_indices=True )
cc['wasInFragment'][i1] = np.logical_or( cc['wasInFragment'][i1], cc_prev['wasInFragment'][i2] )
#cc['mergedCoreTag'][i1] = cc_prev['mergedCoreTag'][i2]
#cc['mergedCoreStep'][i1] = cc_prev['mergedCoreStep'][i2]
cc['phaseSpaceMerged'][i1] = cc_prev['phaseSpaceMerged'][i2]
'''
# If there are satellites (not applicable for first step)
if numSatellites != 0:
# Match satellites to prev step cores using core tag.
vals, idx1, idx2 = np.intersect1d(cc['core_tag'][satellites_mask],
cc_prev['core_tag'],
return_indices=True)
# assert len(vals) == len(cc['core_tag'][satellites_mask]), 'All cores from prev step did not carry over.'
# assert len(cc['core_tag'][satellites_mask]) == len(np.unique(cc['core_tag'][satellites_mask])), 'Satellite core tags not unique for this time step.'
# Initialize M array (corresponds with cc[satellites_mask]) to be host tree node mass of each satellite
idx_m21 = many_to_one(cc['tree_node_index'][satellites_mask],
cc['tree_node_index'][centrals_mask])
M, X, Y, Z = (cc[k][centrals_mask][idx_m21]
for k in ['infall_mass', 'x', 'y', 'z'])
for A in A_arr:
for zeta in zeta_arr:
# Set m_evolved of all satellites that have core tag match on prev step to next_m_evolved of prev step.
# DOUBLE MASK ASSIGNMENT: cc['m_evolved'][satellites_mask][idx1] = cc_prev['next_m_evolved'][idx2]
cc[m_evolved_col(A, zeta)][np.flatnonzero(satellites_mask)
[idx1]] = cc_prev[m_evolved_col(
A, zeta, next=True)][idx2]
# Initialize m array (corresponds with cc[satellites_mask]) to be either infall_mass (step after infall) or m_evolved (subsequent steps).
initMask = cc[m_evolved_col(A, zeta)][satellites_mask] == 0
minfall = cc['infall_mass'][satellites_mask][initMask]
cc[m_evolved_col(A, zeta)][np.flatnonzero(
satellites_mask)[initMask]] = SHMLM.m_evolved(
m0=minfall,
M0=M[initMask],
step=step,
step_prev=steps[steps.index(step) - 1],
A=A,
zeta=zeta,
dtFactorFlag=True)
# m = (cc[m_evolved_col(A, zeta)][satellites_mask] == 0)*cc['infall_mass'][satellites_mask] + (cc[m_evolved_col(A, zeta)][satellites_mask] != 0)*cc[m_evolved_col(A, zeta)][satellites_mask]
# Set m_evolved of satellites with m_evolved=0 to infall mass.
# cc[m_evolved_col(A, zeta)][satellites_mask] = m
'''
x, y, z = cc['x'].copy(), cc['y'].copy(), cc['z'].copy()
x[satellites_mask] = SHMLM.periodic_bcs(cc['x'][satellites_mask], X, SHMLM.BOXSIZE)
y[satellites_mask] = SHMLM.periodic_bcs(cc['y'][satellites_mask], Y, SHMLM.BOXSIZE)
z[satellites_mask] = SHMLM.periodic_bcs(cc['z'][satellites_mask], Z, SHMLM.BOXSIZE)
nonMergeMask = np.flatnonzero(cc['phaseSpaceMerged']==0)
c1, c2 = np.array( list(combinations(nonMergeMask, 2)) ).T #c1, c2 are indices relative of cc
Delta_x = ((x[c1]-x[c2])**2 + (y[c1]-y[c2])**2 + (z[c1]-z[c2])**2)**0.5
Delta_v = ((cc['vx'][c1]-cc['vx'][c2])**2 + (cc['vy'][c1]-cc['vy'][c2])**2 + (cc['vz'][c1]-cc['vz'][c2])**2)**0.5
mass1bigger = cc['infall_mass'][c1] > cc['infall_mass'][c2]
massbiggeridx = np.where(mass1bigger, c1, c2)
masssmalleridx = np.where(mass1bigger, c2, c1)
sigma_x = 3*cc['radius'][massbiggeridx]
sigma_v = 3*cc['vel_disp'][massbiggeridx]
mergeFound = ((Delta_x/sigma_x + Delta_v/sigma_v) < 2)
cc['phaseSpaceMerged'][np.unique(masssmalleridx[mergeFound])] = 1
print "Step", step
print "Merged pairs found: {} out of {} pairs ({}%)".format( np.sum(mergeFound), len(mergeFound), np.sum(mergeFound)/len(mergeFound)*100. )
print "Total number of cores marked as merged:", np.sum(cc['phaseSpaceMerged']==1)
print "Number of central cores marked as merged:", np.sum(cc['phaseSpaceMerged'][centrals_mask]==1)
cc['phaseSpaceMerged'][centrals_mask] = 0
'''
if writeOutputFlag:
# Write output to disk
h5_write_dict(fname_cc(step, 'output'), cc, 'coredata')
# Compute m_evolved of satellites according to SHMLModel for NEXT time step and save as cc_prev['next_m_evolved'] in memory.
# Mass loss assumed to begin at step AFTER infall detected.
if step != steps[-1]:
# cc_prev = { k:cc[k].copy() for k in ['core_tag', 'wasInFragment', 'phaseSpaceMerged'] }
cc_prev = {'core_tag': cc['core_tag'].copy()}
if numSatellites != 0:
localhost_m21 = many_to_one(cc['host_core'][satellites_mask],
cc['core_tag'])
for A in A_arr:
for zeta in zeta_arr:
cc_prev[m_evolved_col(A, zeta, next=True)] = np.zeros_like(
cc['infall_mass'])
if numSatellites != 0: # If there are satellites (not applicable for first step)
m = cc[m_evolved_col(A, zeta)][satellites_mask]
if useLocalHost:
Mlocal = cc[m_evolved_col(A, zeta)][localhost_m21]
M_A_zeta = (Mlocal
== 0) * M + (Mlocal != 0) * Mlocal
cc_prev[m_evolved_col(
A, zeta,
next=True)][satellites_mask] = SHMLM.m_evolved(
m0=m,
M0=M_A_zeta,
step=steps[steps.index(step) + 1],
step_prev=step,
A=A,
zeta=zeta)
else:
cc_prev[m_evolved_col(
A, zeta,
next=True)][satellites_mask] = SHMLM.m_evolved(
m0=m,
M0=M,
step=steps[steps.index(step) + 1],
step_prev=step,
A=A,
zeta=zeta)
return cc
shmassmask = sh['subhalo_mass'][sh_mask][iarr] >= (SHMLM.PARTICLES100MASS *
10)
plt.hist((sh['subhalo_mass'][sh_mask][iarr][shmassmask] -
cc_filtered[m_evolved_col(A, zeta)][carr][shmassmask]) /
sh['subhalo_mass'][sh_mask][iarr][shmassmask],
range=(-5, 5),
bins=100) #, alpha=.5)
plt.axvline(x=0, c='k')
if __name__ == "__main__":
plt_latex()
sh['rvir'] = SHMLM.getRvir(sh['subhalo_mass'], z)
idx_m21_sh = many_to_one(sh['fof_halo_tag'], mt['fof_halo_tag'])
sh['M'] = mt['fof_halo_mass'][idx_m21_sh]
sh['X'] = mt['fof_halo_center_x'][idx_m21_sh]
sh['Y'] = mt['fof_halo_center_y'][idx_m21_sh]
sh['Z'] = mt['fof_halo_center_z'][idx_m21_sh]
# M1, M2 = 10**12., 10**15.5
### Which cores to match subhalos to ###
# Only include satellite cores in matching
# cores_tree, cc_filtered, nHalo = generate_cores_kdtree(M1=M1, M2=M2, s1=False, disrupt=None)
# _, _, nHalo = generate_cores_kdtree(M1=M1, M2=M2, s1=False, disrupt=None)
# Attempt to remove 'true' FOF central cores from list of all cores
# cc_mask = np.invert( np.isin(cc['x'], sh['X'])&(np.isin(cc['y'], sh['Y']))&(np.isin(cc['z'], sh['Z'])) )
# cores_tree, cc_filtered = spatial.cKDTree( np.vstack((cc['x'][cc_mask], cc['y'][cc_mask], cc['z'][cc_mask])).T ), {k:cc[k][cc_mask].copy() for k in cc.keys()}
def create_core_catalog_mevolved(writeOutputFlag=True,
useLocalHost=False,
save_cc_prev=False,
resumeStep=None,
useGPU=False):
"""
Appends mevolved to core catalog and saves output in HDF5.
Works by computing mevolved for step+1 at each step and saving that in memory.
"""
if writeOutputFlag:
print('Reading data from {} and writing output to {}'.format(
cc_data_dir, cc_output_dir))
cc = {}
cc_prev = {}
for step in tqdm(steps):
# Start at step `resumeStep`. Assumes code finished running and saved ccprev for step `resumeStep-1`.
if resumeStep is not None:
if step < resumeStep:
continue
elif step == resumeStep:
prevstep = steps[steps.index(step) - 1]
fname_ccprev_prevstep = fname_cc(prevstep, 'output_ccprev')
print(
f'Resuming at step {step} using {fname_ccprev_prevstep}.')
start = time.time()
cc_prev = h5_read_dict(fname_ccprev_prevstep, 'coredata')
cc_prev['core_tag'] = h5distread(prevstep)['core_tag']
print(
f'Finished reading {fname_ccprev_prevstep} and input hdf5 in {time.time()-start} seconds.'
)
print(f'Beginning step {step}. Reading hdf5 files..')
start = time.time()
# Read in cc for step
cc = h5distread(step)
print(f'Finished reading hdf5 files in {time.time()-start} seconds.')
satellites_mask = cc['central'] == 0
centrals_mask = cc['central'] == 1
# assert np.sum(satellites_mask) + np.sum(centrals_mask) == len(cc['infall_tree_node_mass']), 'central flag not 0 or 1.'
numSatellites = np.sum(satellites_mask)
# Verify there are no satellites at first step
if step == steps[0]:
assert numSatellites == 0, 'Satellites found at first step.'
# Add column for m_evolved and initialize to 0
cc[m_evolved_col(A, zeta)] = np.zeros_like(cc['infall_tree_node_mass'])
# If there are satellites (not applicable for first step)
if numSatellites != 0:
# Match satellites to prev step cores using core tag.
print('intersect1d cc to cc_prev...')
start = time.time()
if useGPU:
vals, idx1, idx2 = intersect1d_GPU(
cc['core_tag'][satellites_mask],
cc_prev['core_tag'],
return_indices=True)
else:
vals, idx1, idx2 = np.intersect1d(
cc['core_tag'][satellites_mask],
cc_prev['core_tag'],
return_indices=True)
print(f'Matches found between cc and cc_prev: {len(vals)}.')
print(
f'Finished intersect1d cc to cc_prev in {time.time()-start} seconds.'
)
print('satellites:centrals m21...')
start = time.time()
if useGPU:
idx_m21 = many_to_one_GPU(
cc['tree_node_index'][satellites_mask],
cc['tree_node_index'][centrals_mask])
else:
idx_m21 = many_to_one(cc['tree_node_index'][satellites_mask],
cc['tree_node_index'][centrals_mask],
verbose=True,
assert_x0_unique=False,
assert_x1_in_x0=False)
print(
f'Finished satellites:centrals m21 in {time.time()-start} seconds.'
)
M = cc['infall_tree_node_mass'][centrals_mask][idx_m21]
# Set m_evolved of all satellites that have core tag match on prev step to next_m_evolved of prev step.
cc[m_evolved_col(A, zeta)][np.flatnonzero(satellites_mask)
[idx1]] = cc_prev[m_evolved_col(
A, zeta, next=True)][idx2]
# Initialize m array (corresponds with cc[satellites_mask]) to be either infall_mass (step after infall) or m_evolved (subsequent steps).
print('SHMLM (new satellites)...')
start = time.time()
initMask = cc[m_evolved_col(A, zeta)][satellites_mask] == 0
minfall = cc['infall_tree_node_mass'][satellites_mask][initMask]
cc[m_evolved_col(A, zeta)][np.flatnonzero(
satellites_mask)[initMask]] = SHMLM.m_evolved(
m0=minfall,
M0=M[initMask],
step=step,
step_prev=steps[steps.index(step) - 1],
A=A,
zeta=zeta,
dtFactorFlag=True)
print(
f'Finished SHMLM (new satellites) in {time.time()-start} seconds.'
)
if writeOutputFlag and (step == 499 or step == 247):
# Write output to disk
cc_save = {
# 'fof_halo_tag':cc['fof_halo_tag'],
# 'central':cc['central'],
m_evolved_col(A, zeta):
cc[m_evolved_col(A, zeta)]
# 'M':M
}
h5_write_dict(fname_cc(step, 'output'), cc_save, 'coredata')
# Compute m_evolved of satellites according to SHMLModel for NEXT time step and save as cc_prev['next_m_evolved'] in memory.
# Mass loss assumed to begin at step AFTER infall detected.
if step != steps[-1]:
cc_prev = {'core_tag': cc['core_tag'].copy()}
if numSatellites != 0:
print('host_core:core_tag m21...')
start = time.time()
if useGPU:
localhost_m21 = many_to_one_GPU(
cc['host_core'][satellites_mask], cc['core_tag'])
else:
localhost_m21 = many_to_one(
cc['host_core'][satellites_mask],
cc['core_tag'],
verbose=True,
assert_x0_unique=False,
assert_x1_in_x0=False)
print(
f'Finished host_core:core_tag m21 in {time.time()-start} seconds.'
)
cc_prev[m_evolved_col(A, zeta, next=True)] = np.zeros_like(
cc['infall_tree_node_mass'])
if numSatellites != 0: # If there are satellites (not applicable for first step)
print('SHMLM (next step)...')
start = time.time()
m = cc[m_evolved_col(A, zeta)][satellites_mask]
if useLocalHost:
Mlocal = cc[m_evolved_col(A, zeta)][localhost_m21]
M_A_zeta = (Mlocal == 0) * M + (Mlocal != 0) * Mlocal
cc_prev[m_evolved_col(
A, zeta,
next=True)][satellites_mask] = SHMLM.m_evolved(
m0=m,
M0=M_A_zeta,
step=steps[steps.index(step) + 1],
step_prev=step,
A=A,
zeta=zeta)
else:
cc_prev[m_evolved_col(
A, zeta,
next=True)][satellites_mask] = SHMLM.m_evolved(
m0=m,
M0=M,
step=steps[steps.index(step) + 1],
step_prev=step,
A=A,
zeta=zeta)
print(
f'Finished SHMLM (next step) in {time.time()-start} seconds.'
)
if save_cc_prev:
print('writing ccprev hdf5...')
start = time.time()
h5_write_dict(
fname_cc(step, 'output_ccprev'), {
m_evolved_col(A, zeta, next=True):
cc_prev[m_evolved_col(A, zeta, next=True)]
}, 'coredata')
# if os.path.exists( fname_cc(steps[steps.index(step)-1], 'output_ccprev') ):
# os.remove( fname_cc(steps[steps.index(step)-1], 'output_ccprev') )
print(
f'Finished writing ccprev hdf5 in {time.time()-start} seconds.'
)
return cc
def create_core_catalog_mevolved(virialFlag=False,
writeOutputFlag=True,
useLocalHost=False,
checkpointRestart=None):
"""
Appends mevolved to core catalog and saves output in HDF5.
Works by computing mevolved for step+1 at each step and saving that in memory.
"""
if writeOutputFlag:
print 'Reading data from {} and writing output to {}'.format(
cc_data_dir, cc_output_dir)
cc = {}
cc_prev = {}
if checkpointRestart is None:
itersteps = steps
else:
itersteps = steps[steps.index(checkpointRestart):]
for step in tqdm(itersteps):
# Read in cc for step
if checkpointRestart is None:
cc = {
v: gio.gio_read(fname_cc(step, 'input'), v)[0]
for v in vars_cc
}
if virialFlag:
# Convert all mass to virial
cc['infall_tree_node_mass'] = SHMLM.m_vir(
cc['infall_tree_node_mass'], step)
else:
cc = h5_read_dict(fname_cc(checkpointRestart, 'output'),
'coredata')
satellites_mask = cc['central'] == 0
centrals_mask = cc['central'] == 1
# assert np.sum(satellites_mask) + np.sum(centrals_mask) == len(cc['infall_tree_node_mass']), 'central flag not 0 or 1.'
numSatellites = np.sum(satellites_mask)
# Verify there are no satellites at first step
if step == steps[0]:
assert numSatellites == 0, 'Satellites found at first step.'
# Add column for m_evolved and initialize to 0
if checkpointRestart is None:
for A in A_arr:
for zeta in zeta_arr:
cc[m_evolved_col(A, zeta)] = np.zeros_like(
cc['infall_tree_node_mass'])
# If there are satellites (not applicable for first step)
if numSatellites != 0:
if checkpointRestart is None:
# Match satellites to prev step cores using core tag.
vals, idx1, idx2 = np.intersect1d(
cc['core_tag'][satellites_mask],
cc_prev['core_tag'],
return_indices=True)
# assert len(vals) == len(cc['core_tag'][satellites_mask]), 'All cores from prev step did not carry over.'
# assert len(cc['core_tag'][satellites_mask]) == len(np.unique(cc['core_tag'][satellites_mask])), 'Satellite core tags not unique for this time step.'
# Initialize M array (corresponds with cc[satellites_mask]) to be host tree node mass of each satellite
idx_m21 = many_to_one(cc['tree_node_index'][satellites_mask],
cc['tree_node_index'][centrals_mask])
M, X, Y, Z = (cc[k][centrals_mask][idx_m21]
for k in ['infall_tree_node_mass', 'x', 'y', 'z'])
if checkpointRestart is None:
for A in A_arr:
for zeta in zeta_arr:
# Set m_evolved of all satellites that have core tag match on prev step to next_m_evolved of prev step.
# DOUBLE MASK ASSIGNMENT: cc['m_evolved'][satellites_mask][idx1] = cc_prev['next_m_evolved'][idx2]
cc[m_evolved_col(A, zeta)][np.flatnonzero(
satellites_mask)[idx1]] = cc_prev[m_evolved_col(
A, zeta, next=True)][idx2]
initMask = cc[m_evolved_col(
A, zeta)][satellites_mask] == 0
minfall = cc['infall_tree_node_mass'][satellites_mask][
initMask]
cc[m_evolved_col(A, zeta)][np.flatnonzero(
satellites_mask)[initMask]] = SHMLM.m_evolved(
m0=minfall,
M0=M[initMask],
step=step,
step_prev=steps[steps.index(step) - 1],
A=A,
zeta=zeta,
dtFactorFlag=True)
if writeOutputFlag and (checkpointRestart is None):
# Write output to disk
h5_write_dict(fname_cc(step, 'output'), cc, 'coredata')
#cc should be checkpointRestart step file if not None
if checkpointRestart is not None:
checkpointRestart = None
# Compute m_evolved of satellites according to SHMLModel for NEXT time step and save as cc_prev['next_m_evolved'] in memory.
if step != steps[-1]:
cc_prev = {'core_tag': cc['core_tag'].copy()}
if numSatellites != 0:
localhost_m21 = many_to_one(cc['host_core'][satellites_mask],
cc['core_tag'])
for A in A_arr:
for zeta in zeta_arr:
cc_prev[m_evolved_col(A, zeta, next=True)] = np.zeros_like(
cc['infall_tree_node_mass'])
if numSatellites != 0: # If there are satellites (not applicable for first step)
m = cc[m_evolved_col(A, zeta)][satellites_mask]
if useLocalHost:
Mlocal = cc[m_evolved_col(A, zeta)][localhost_m21]
M_A_zeta = (Mlocal
== 0) * M + (Mlocal != 0) * Mlocal
cc_prev[m_evolved_col(
A, zeta,
next=True)][satellites_mask] = SHMLM.m_evolved(
m0=m,
M0=M_A_zeta,
step=steps[steps.index(step) + 1],
step_prev=step,
A=A,
zeta=zeta)
else:
cc_prev[m_evolved_col(
A, zeta,
next=True)][satellites_mask] = SHMLM.m_evolved(
m0=m,
M0=M,
step=steps[steps.index(step) + 1],
step_prev=step,
A=A,
zeta=zeta)
return cc
def create_core_catalog_mevolved(writeOutputFlag=True, useLocalHost=False):
"""
Appends mevolved to core catalog and saves output in HDF5.
Works by computing mevolved for step+1 at each step and saving that in memory.
"""
if writeOutputFlag:
print 'Reading data from {} and writing output to {}'.format(
cc_data_dir, cc_output_dir)
cc = {}
cc_prev = {}
for step in tqdm(steps):
# Read in cc for step
cc = {
v: gio.gio_read(fname_cc(step, 'input'), v)[0]
for v in vars_cc(step)
}
satellites_mask = cc['central'] == 0
centrals_mask = cc['central'] == 1
# assert np.sum(satellites_mask) + np.sum(centrals_mask) == len(cc['infall_tree_node_mass']), 'central flag not 0 or 1.'
numSatellites = np.sum(satellites_mask)
# Verify there are no satellites at first step
if step == steps[0]:
assert numSatellites == 0, 'Satellites found at first step.'
# Add column for m_evolved and initialize to 0
cc[m_evolved_col(A, zeta)] = np.zeros_like(cc['infall_tree_node_mass'])
# If there are satellites (not applicable for first step)
if numSatellites != 0:
# Match satellites to prev step cores using core tag.
vals, idx1, idx2 = np.intersect1d(cc['core_tag'][satellites_mask],
cc_prev['core_tag'],
return_indices=True)
idx_m21 = many_to_one(cc['tree_node_index'][satellites_mask],
cc['tree_node_index'][centrals_mask])
M = cc['infall_tree_node_mass'][centrals_mask][idx_m21]
# Set m_evolved of all satellites that have core tag match on prev step to next_m_evolved of prev step.
cc[m_evolved_col(A, zeta)][np.flatnonzero(satellites_mask)
[idx1]] = cc_prev[m_evolved_col(
A, zeta, next=True)][idx2]
# Initialize m array (corresponds with cc[satellites_mask]) to be either infall_mass (step after infall) or m_evolved (subsequent steps).
initMask = cc[m_evolved_col(A, zeta)][satellites_mask] == 0
minfall = cc['infall_tree_node_mass'][satellites_mask][initMask]
cc[m_evolved_col(A, zeta)][np.flatnonzero(
satellites_mask)[initMask]] = SHMLM.m_evolved(
m0=minfall,
M0=M[initMask],
step=step,
step_prev=steps[steps.index(step) - 1],
A=A,
zeta=zeta,
dtFactorFlag=True)
if writeOutputFlag and (step == 499 or step == 247):
# Write output to disk
h5_write_dict(fname_cc(step, 'output'), cc, 'coredata')
# Compute m_evolved of satellites according to SHMLModel for NEXT time step and save as cc_prev['next_m_evolved'] in memory.
# Mass loss assumed to begin at step AFTER infall detected.
if step != steps[-1]:
cc_prev = {'core_tag': cc['core_tag'].copy()}
if numSatellites != 0:
localhost_m21 = many_to_one(cc['host_core'][satellites_mask],
cc['core_tag'])
cc_prev[m_evolved_col(A, zeta, next=True)] = np.zeros_like(
cc['infall_tree_node_mass'])
if numSatellites != 0: # If there are satellites (not applicable for first step)
m = cc[m_evolved_col(A, zeta)][satellites_mask]
if useLocalHost:
Mlocal = cc[m_evolved_col(A, zeta)][localhost_m21]
M_A_zeta = (Mlocal == 0) * M + (Mlocal != 0) * Mlocal
cc_prev[m_evolved_col(
A, zeta,
next=True)][satellites_mask] = SHMLM.m_evolved(
m0=m,
M0=M_A_zeta,
step=steps[steps.index(step) + 1],
step_prev=step,
A=A,
zeta=zeta)
else:
cc_prev[m_evolved_col(
A, zeta,
next=True)][satellites_mask] = SHMLM.m_evolved(
m0=m,
M0=M,
step=steps[steps.index(step) + 1],
step_prev=step,
A=A,
zeta=zeta)
return cc
def create_core_catalog_mevolved(virialFlag,
writeOutputFlag=True,
useLocalHost=False):
"""
Appends mevolved to core catalog and saves output in HDF5.
Works by computing mevolved for step+1 at each step and saving that in memory.
"""
if writeOutputFlag:
print 'Reading data from {} and writing output to {}'.format(
cc_data_dir, cc_output_dir)
cc = {}
cc_prev = {}
for step in tqdm(steps):
# Read in cc for step
cc = {v: gio.gio_read(fname_cc(step, 'input'), v)[0] for v in vars_cc}
if virialFlag:
# Convert all mass to virial
cc['infall_mass'] = SHMLM.m_vir(cc['infall_mass'], step)
satellites_mask = cc['central'] == 0
centrals_mask = cc['central'] == 1
# assert np.sum(satellites_mask) + np.sum(centrals_mask) == len(cc['infall_mass']), 'central flag not 0 or 1.'
numSatellites = np.sum(satellites_mask)
# Verify there are no satellites at first step
if step == steps[0]:
assert numSatellites == 0, 'Satellites found at first step.'
# Add column for m_evolved and initialize to 0
for A in A_arr:
for zeta in zeta_arr:
cc[m_evolved_col(A, zeta)] = np.zeros_like(cc['infall_mass'])
cc['mergedCoreTag'] = np.zeros_like(cc['core_tag'])
cc['mergedCoreStep'] = np.zeros_like(cc['infall_step'])
# wasInFragment: persistent flag that is False by default and becomes True when core is inside a fragment halo
cc['wasInFragment'] = cc['fof_halo_tag'] < 0
if step != steps[0]:
_, i1, i2 = np.intersect1d(cc['core_tag'],
cc_prev['core_tag'],
return_indices=True)
cc['wasInFragment'][i1] = np.logical_or(
cc['wasInFragment'][i1], cc_prev['wasInFragment'][i2])
cc['mergedCoreTag'][i1] = cc_prev['mergedCoreTag'][i2]
cc['mergedCoreStep'][i1] = cc_prev['mergedCoreStep'][i2]
# If there are satellites (not applicable for first step)
if numSatellites != 0:
# Match satellites to prev step cores using core tag.
vals, idx1, idx2 = np.intersect1d(cc['core_tag'][satellites_mask],
cc_prev['core_tag'],
return_indices=True)
# assert len(vals) == len(cc['core_tag'][satellites_mask]), 'All cores from prev step did not carry over.'
# assert len(cc['core_tag'][satellites_mask]) == len(np.unique(cc['core_tag'][satellites_mask])), 'Satellite core tags not unique for this time step.'
for A in A_arr:
for zeta in zeta_arr:
# Set m_evolved of all satellites that have core tag match on prev step to next_m_evolved of prev step.
# DOUBLE MASK ASSIGNMENT: cc['m_evolved'][satellites_mask][idx1] = cc_prev['next_m_evolved'][idx2]
cc[m_evolved_col(A, zeta)][np.flatnonzero(satellites_mask)
[idx1]] = cc_prev[m_evolved_col(
A, zeta, next=True)][idx2]
# Initialize m array (corresponds with cc[satellites_mask]) to be either infall_mass (step after infall) or m_evolved (subsequent steps).
m = (cc[m_evolved_col(A, zeta)][satellites_mask]
== 0) * cc['infall_mass'][satellites_mask] + (
cc[m_evolved_col(A, zeta)][satellites_mask] !=
0) * cc[m_evolved_col(A, zeta)][satellites_mask]
# Set m_evolved of satellites with m_evolved=0 to infall mass.
cc[m_evolved_col(A, zeta)][satellites_mask] = m
# Initialize M array (corresponds with cc[satellites_mask]) to be host tree node mass of each satellite
idx_m21 = many_to_one(cc['tree_node_index'][satellites_mask],
cc['tree_node_index'][centrals_mask])
M, X, Y, Z = (cc[k][centrals_mask][idx_m21]
for k in ['infall_mass', 'x', 'y', 'z'])
KDTreemask = cc['mergedCoreTag'][satellites_mask] == 0
x = SHMLM.periodic_bcs(cc['x'][satellites_mask][KDTreemask],
X[KDTreemask], SHMLM.BOXSIZE)
y = SHMLM.periodic_bcs(cc['y'][satellites_mask][KDTreemask],
Y[KDTreemask], SHMLM.BOXSIZE)
z = SHMLM.periodic_bcs(cc['z'][satellites_mask][KDTreemask],
Z[KDTreemask], SHMLM.BOXSIZE)
coresMaskedArray = np.vstack((x, y, z)).T
coresMaskedArray_Tree = np.vstack(
(np.r_[x,
cc['x'][centrals_mask]], np.r_[y,
cc['y'][centrals_mask]],
np.r_[z, cc['z'][centrals_mask]])).T
KDTree_idx = np.r_[np.flatnonzero(satellites_mask)[KDTreemask],
np.flatnonzero(centrals_mask)]
coresKDTree = spatial.cKDTree(coresMaskedArray_Tree)
# Set mergedCoreTag as central core of KDTreemask cores with no mergedCoreTag that are `merged`
disruptMask = cc['merged'][satellites_mask][KDTreemask] | (
cc['radius'][satellites_mask][KDTreemask] >=
SHMLM.CORERADIUSCUT)
cc['mergedCoreTag'][np.flatnonzero(satellites_mask)
[KDTreemask]] += disruptMask * cc['core_tag'][
centrals_mask][idx_m21][KDTreemask]
cc['mergedCoreStep'][np.flatnonzero(satellites_mask)
[KDTreemask]] += disruptMask * step
### core merging detection
# rvir_KDTreemask = SHMLM.getRvir(cc[m_evolved_col(1.1, 0.068)][satellites_mask][KDTreemask], SHMLM.step2z[step]) #infall should be exact m_evolved in production
rvir_KDTreemask = SHMLM.getRvir(
cc['infall_mass'][satellites_mask][KDTreemask], SHMLM.
step2z[step]) #infall should be exact m_evolved in production
distance_upper_bound = SHMLM.COREVIRIALRADIUSFACTOR * rvir_KDTreemask
# Search radius of 2*rvir around each core
for i in tqdm(range(len(distance_upper_bound))):
qres = coresKDTree.query_ball_point(coresMaskedArray[i],
r=distance_upper_bound[i])
assert len(qres) > 0, 'Merge with self not detected!'
if len(qres) > 1:
# idxmax = qres[ np.argmax(cc['infall_mass'][satellites_mask][KDTreemask][qres]) ]
idxmax = qres[np.argmax(
cc['infall_mass'][KDTree_idx][qres])]
qres = np.array(qres)
qres2 = qres[qres != idxmax]
# cc['mergedCoreTag'][ np.flatnonzero(satellites_mask)[KDTreemask][qres2] ] = cc['core_tag'][satellites_mask][KDTreemask][idxmax]
# cc['mergedCoreStep'][ np.flatnonzero(satellites_mask)[KDTreemask][qres2] ] = step
cc['mergedCoreTag'][
KDTree_idx[qres2]] = cc['core_tag'][KDTree_idx][idxmax]
cc['mergedCoreStep'][KDTree_idx[qres2]] = step
###
if writeOutputFlag:
# Write output to disk
h5_write_dict(fname_cc(step, 'output'), cc, 'coredata')
# Compute m_evolved of satellites according to SHMLModel for NEXT time step and save as cc_prev['next_m_evolved'] in memory.
# Mass loss assumed to begin at step AFTER infall detected.
if step != steps[-1]:
# cc_prev = { 'core_tag':cc['core_tag'].copy(), 'wasInFragment':cc['wasInFragment'].copy() }
cc_prev = {
k: cc[k].copy()
for k in [
'core_tag', 'wasInFragment', 'mergedCoreTag',
'mergedCoreStep'
]
}
if numSatellites != 0:
localhost_m21 = many_to_one(cc['host_core'][satellites_mask],
cc['core_tag'])
for A in A_arr:
for zeta in zeta_arr:
cc_prev[m_evolved_col(A, zeta, next=True)] = np.zeros_like(
cc['infall_mass'])
if numSatellites != 0: # If there are satellites (not applicable for first step)
m = cc[m_evolved_col(A, zeta)][satellites_mask]
if useLocalHost:
Mlocal = cc[m_evolved_col(A, zeta)][localhost_m21]
M_A_zeta = (Mlocal
== 0) * M + (Mlocal != 0) * Mlocal
cc_prev[m_evolved_col(
A, zeta,
next=True)][satellites_mask] = SHMLM.m_evolved(
m0=m,
M0=M_A_zeta,
step=steps[steps.index(step) + 1],
step_prev=step,
A=A,
zeta=zeta)
else:
cc_prev[m_evolved_col(
A, zeta,
next=True)][satellites_mask] = SHMLM.m_evolved(
m0=m,
M0=M,
step=steps[steps.index(step) + 1],
step_prev=step,
A=A,
zeta=zeta)
return cc
centrals_mask = cc['central'] == 1
numSatellites = np.sum(satellites_mask)
# Verify there are no satellites at first step
if step == steps[0]:
assert numSatellites == 0, 'Satellites found at first step.'
#cc['mergedCoreTag'] = np.zeros_like(cc['core_tag'])
#cc['mergedCoreStep'] = np.zeros_like(cc['infall_step'])
# cc['phaseSpaceMerged'] = np.zeros_like(cc['central'])
# If there are satellites (not applicable for first step)
assert numSatellites != 0
# Initialize M array (corresponds with cc[satellites_mask]) to be host tree node mass of each satellite
idx_m21 = many_to_one( cc['tree_node_index'][satellites_mask], cc['tree_node_index'][centrals_mask] )
M, X, Y, Z = (cc[k][centrals_mask][idx_m21] for k in ['infall_mass', 'x', 'y', 'z'])
x, y, z = cc['x'].copy(), cc['y'].copy(), cc['z'].copy()
x[satellites_mask] = SHMLM.periodic_bcs(cc['x'][satellites_mask], X, SHMLM.BOXSIZE)
y[satellites_mask] = SHMLM.periodic_bcs(cc['y'][satellites_mask], Y, SHMLM.BOXSIZE)
z[satellites_mask] = SHMLM.periodic_bcs(cc['z'][satellites_mask], Z, SHMLM.BOXSIZE)
srt_fht = np.argsort(cc['fof_halo_tag'])
cc_srt = { k:cc[k][srt_fht].copy() for k in cc.keys() }
cc_srt['x'] = x[srt_fht]
cc_srt['y'] = y[srt_fht]
cc_srt['z'] = z[srt_fht]
_, idx, cnts = np.unique(cc_srt['fof_halo_tag'], return_index=True, return_counts=True)
comm.barrier()
