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, 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=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
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