def h5distread(step):
basename = fname_cc(step, 'input')
ccall = {k: [] for k in vars_cc_min}
for f in glob.glob(basename + '#*'):
ccf = h5_read_dict(f)
for k in vars_cc_min:
ccall[k].append(ccf[k])
return {k: np.concatenate(ccall[k]) for k in vars_cc_min}
def CMF(outfile,
M1,
M2,
s1=False,
disrupt=None,
returnUnevolved=False,
cc=None,
returnZeta0=False,
A=None,
idx_m21=None):
"""Generate cores array m/M for mass function plot.
Arguments:
outfile {string} -- HDF5 core catalog file with 'coredata' dir.
M1, M2 {float} -- Host halos with M in mass range [M1, M2], where M is central's `infall_mass` from core catalog.
cc {dict} -- If given, `outfile` is not read and `cc` is used instead.
A {float} -- Fitting parameter used for fast mass evolution if `returnZeta0`.
disrupt {None, int, or list} -- If given, applies core disruption criteria defined in SHMLM when filtering substructure. (default: {None})
Keyword Arguments:
s1 {bool} -- If true, consider only 1st order substructure (i.e. subhalos). (default: {False})
returnUnevolved {bool} -- If true, use unevolved m (i.e. core `infall_mass`) (default: {False})
returnZeta0 {bool} -- If true, does fast mass evolution for the case of zeta=0, A=`A`. (default: {False})
Returns:
np 1d array -- m/M array with appropriate mask
integer -- number of host halos (M)
"""
if not cc:
# Load extended core catalog
cc = h5_read_dict(outfile, 'coredata')
if 'x' in cc.keys():
idx_filteredsatcores, M, X, Y, Z, nHalo = SHMLM.core_mask(
cc, M1, M2, s1=s1, disrupt=disrupt, idx_m21=idx_m21)
else:
idx_filteredsatcores, M, nHalo = SHMLM.core_mask(cc,
M1,
M2,
s1=s1,
disrupt=disrupt,
idx_m21=idx_m21)
# m/M array for CMF
if returnUnevolved:
plot_arr = cc['infall_mass'][idx_filteredsatcores] / M
elif returnZeta0:
plot_arr = SHMLM.fast_m_evolved(
cc['infall_mass'][idx_filteredsatcores] / M,
cc['infall_step'][idx_filteredsatcores], A)
else:
plot_arr = cc['m_evolved'][idx_filteredsatcores] / M
return plot_arr, nHalo
#from cc_generate_fitting_phasemerge import m_evolved_col, A_arr, zeta_arr
from cc_generate_fitting_dt_LJDS import m_evolved_col, A_arr, zeta_arr
from scipy import spatial #KD Tree for subhalo-core matching
from scipy.stats import norm # Gaussian fitting
import itertools as it
massbinsize = 0.5
step = 499
# step = 247
z = SHMLM.step2z[step]
# cc = h5_read_dict('/home/isultan/projects/halomassloss/core_catalog_mevolved/output_merg_fof_fitting_localhost_mergedcoretag/09_03_2019.AQ.{}.corepropertiesextend.hdf5'.format(step), 'coredata')
# cc = h5_read_dict('/home/isultan/projects/halomassloss/core_catalog_mevolved/output_merg_fof_fitting_localhost_mergedcoretag_parallel/09_03_2019.AQ.{}.corepropertiesextend.hdf5'.format(step), 'coredata')
#cc = h5_read_dict('/home/isultan/projects/halomassloss/core_catalog_mevolved/output_merg_fof_fitting_localhost_phasemerge/09_03_2019.AQ.{}.corepropertiesextend.hdf5'.format(step), 'coredata')
cc = h5_read_dict(
'/home/isultan/projects/halomassloss/core_catalog_mevolved/output_LJDS_localhost_dtfactor_0.5_fitting2/m000p-{}.corepropertiesextend.hdf5'
.format(step), 'coredata')
# cc = h5_read_dict('/home/isultan/projects/halomassloss/core_catalog_mevolved/output_LJDS_localhost_dtfactor_0.5/m000p-{}.corepropertiesextend.hdf5'.format(step), 'coredata')
#cc = h5_read_dict('/home/isultan/projects/halomassloss/core_catalog_mevolved/output_merg_fof_fitting_localhost_dtfactor_0.5_fiducial_A_0.9_zeta_0.005_FIDUCIALPARTICLECUTMASS_PARTICLES100MASS_spline/09_03_2019.AQ.{}.corepropertiesextend.hdf5'.format(step), 'coredata')
#cc = h5_read_dict('/home/isultan/projects/halomassloss/core_catalog_mevolved/output_merg_fof_fitting_localhost_dtfactor_0.5_fiducial_A_0.9_zeta_0.005_FIDUCIALPARTICLECUTMASS_PARTICLES100MASS_spline_zbin_1.5_2/09_03_2019.AQ.{}.corepropertiesextend.hdf5'.format(step), 'coredata')
#cc = h5_read_dict('/home/isultan/projects/halomassloss/core_catalog_mevolved/output_merg_fof_fitting_localhost_dtfactor_0.5_fiducial_A_0.9_zeta_0.005_FIDUCIALPARTICLECUTMASS_PARTICLES100MASS_spline_match_zbin_1.5_2/09_03_2019.AQ.{}.corepropertiesextend.hdf5'.format(step), 'coredata')
#cc_pm = loadpickle('phasemerge_{}.pkl'.format(step))
#_, idx1, idx2 = np.intersect1d(cc['core_tag'], cc_pm['core_tag'], assume_unique=False, return_indices=True)
#cc['phaseSpaceMerged'] = np.zeros_like(cc['central'])
#cc['phaseSpaceMerged'][idx1] = cc_pm['phaseSpaceMerged'][idx2]
# sh = gio_read_dict('/home/isultan/data/AlphaQ/subhalos/m000-{}.subhaloproperties'.format(step), ['fof_halo_tag','subhalo_mean_x','subhalo_mean_y','subhalo_mean_z','subhalo_mean_vx', 'subhalo_mean_vy', 'subhalo_mean_vz', 'subhalo_count', 'subhalo_tag', 'subhalo_mass'])
mt_cores = gio_read_dict(
'/home/isultan/data/LJDS/MergerTrees_updated/m000p-{}.treenodes'.format(
step), [
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
from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import subhalo_mass_loss_model as SHMLM
import genericio as gio
from tqdm import tqdm
from itk import hist, h5_read_dict, gio_read_dict, loadpickle, plt_latex, periodic_bcs, many_to_one
from cc_generate_fitting import m_evolved_col, A_arr, zeta_arr
from scipy import spatial #KD Tree for subhalo-core matching
from scipy.stats import norm # Gaussian fitting
import itertools as it
step = 499
z = SHMLM.step2z[step]
cc = h5_read_dict('/home/isultan/projects/halomassloss/core_catalog_mevolved/output_merg_fof_fitting_localhost_mergedcoretag/09_03_2019.AQ.{}.corepropertiesextend.hdf5'.format(step), 'coredata')
sh = gio_read_dict('/home/isultan/data/AlphaQ/subhalos/m000-{}.subhaloproperties'.format(step), ['fof_halo_tag','subhalo_mean_x','subhalo_mean_y','subhalo_mean_z','subhalo_mean_vx', 'subhalo_mean_vy', 'subhalo_mean_vz', 'subhalo_count', 'subhalo_tag', 'subhalo_mass'])
mt = gio_read_dict('/home/isultan/data/AlphaQ/updated_tree_nodes/09_03_2019.AQ.{}.treenodes'.format(step), ['fof_halo_tag', 'fof_halo_mass', 'fof_halo_center_x', 'fof_halo_center_y', 'fof_halo_center_z'])
M1, M2 = 10**9., 10**15.5
# Subhalo-core matching: to look at the subhalo mass:core evolved mass ratio
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 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'])
'''
#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:
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]
# 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)
# m = (cc[m_evolved_col(A, zeta)][satellites_mask] == 0)*cc['infall_tree_node_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_tree_node_mass'][c1] > cc['infall_tree_node_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 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.
# 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_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=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
