def generate_cores_kdtree(cc,
M1,
M2,
s1=False,
disrupt=None,
onlyFiltered=False,
giveFragmentsFofMass=False,
z=0):
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 generate_cores_kdtree(M1, M2, s1=False, disrupt=None):
idx_filteredsatcores, M, X, Y, Z, nHalo = SHMLM.core_mask(cc, M1, M2, s1=s1, disrupt=disrupt)
cc_filtered = { k:cc[k][idx_filteredsatcores].copy() for k in cc.keys() }
cc_filtered['M'] = M
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)
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 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
sh['Z'] = mt['fof_halo_center_z'][idx_m21_sh]
### 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()}
# Use ALL cores in matching
cores_tree, cc_filtered = spatial.cKDTree( np.vstack((cc['x'], cc['y'], cc['z'])).T ), cc
### End: Which cores to match subhalos to ###
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)
distance_mask = SHMLM.dist(subhalo_mean_x, subhalo_mean_y, subhalo_mean_z, sh['X'], sh['Y'], sh['Z']) >= SHMLM.SUBHALOHOSTVIRIALRADIUSFACTOR*SHMLM.getRvir(sh['M'], z)
sh_mask = (sh['subhalo_tag']!=0)&(M1<=sh['M'])&(sh['M']<=M2)#&(distance_mask) #&(sh['subhalo_mass']>=SHMLM.SUBHALOMASSCUT)
print '{}% satellite subhalos in given mass range remain after mask.'.format(np.sum(sh_mask)/np.sum((sh['subhalo_tag']!=0)&(M1<=sh['M'])&(sh['M']<=M2))*100.)
sh_arr = np.vstack((subhalo_mean_x[sh_mask], subhalo_mean_y[sh_mask], subhalo_mean_z[sh_mask])).T
distance_upper_bound = SHMLM.SUBHALOVIRIALRADIUSFACTOR * sh['rvir'][sh_mask]
# for SUBHALOVIRIALRADIUSFACTOR in [1,2,3,4]:
# distance_upper_bound = SUBHALOVIRIALRADIUSFACTOR * sh['rvir'][sh_mask]
# mostMassiveCore(A=1.1, zeta=0.01, plotFlag=False)
'''
# Search radius of 2*rvir around each subhalo, only look for 1:1 matches
dist, idx = [], []
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
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()
if rank == 0:
numCnts = len(cnts)
else: