def auxilliary_fn(
dfid,
df,
id_finder,
halo_data_dir,
main_halo_id,
):
'''Function for retrieving additional quantities to store in ptracks.
'''
# Get the tidal tensor data
tidal_df = data_products.tidal_tensor_data_grudic(
id_finder.snum,
ids = dfid.index,
)
dfid = pd.concat( [ dfid, tidal_df ], axis=1 )
# Get the enclosed mass data
h_data = halo_data.HaloData(
data_dir = halo_data_dir,
mt_kwargs = { 'index': 'snum', 'tag': 'smooth' },
)
positions = np.array([
dfid['P0'],
dfid['P1'],
dfid['P2'],
]).transpose()
dfid['M_enc'] = h_data.get_enclosed_mass(
positions = positions,
snum = id_finder.snum,
hubble_param = id_finder.attrs['hubble'],
mt_halo_id = main_halo_id,
)
return dfid
def auxilliary_fn(
dfid,
df,
id_finder,
halo_data_dir,
main_halo_id,
):
'''Function for retrieving additional quantities to store in ptracks.
'''
# Get the radius from halo center
h_data = halo_data.HaloData(
data_dir = halo_data_dir,
mt_kwargs = { 'index': 'snum', 'tag': 'smooth' },
)
r = np.zeros( df['P0'].shape, )
pos_keys = [ 'P0', 'P1', 'P2' ]
for i, h_pos_key in enumerate( [ 'Xc', 'Yc', 'Zc' ] ):
origin = (
h_data.get_mt_data(
h_pos_key,
snums = [ id_finder.snum ],
a_power = 1.
)[0] /
id_finder.attrs['hubble']
)
xi = df[pos_keys[i]] - origin
r += xi**2.
r = np.sqrt( r )
# Identify stars in the galaxy
galaxy_radius = 4. * h_data.get_mt_data(
'Rstar0.5',
snums = [ id_finder.snum ],
mt_halo_id = main_halo_id,
a_power = 1.,
)[0] / id_finder.attrs['hubble']
in_galaxy = r < galaxy_radius
is_star = df['PType'] == 4
# Get the galaxy stellar mass
m_star = df['M'][in_galaxy & is_star].sum()
# Format (have to repeat for each particle....)
dfid['M_star'] = np.full( dfid['P0'].shape, m_star )
return dfid
def read_data(self):
'''Read the input data.
Modifies:
self.ptrack (h5py file) : Loaded tracked particle data.
self.halo_data (HaloData instance): For the halo data.
'''
# Get the tag for particle tracking.
if self.ptracks_tag is default:
self.ptracks_tag = self.tag
# Load the particle track data
ptrack_filename = 'ptracks_{}.hdf5'.format(self.ptracks_tag)
self.ptrack_filepath = os.path.join(self.out_dir, ptrack_filename)
self.ptrack = h5py.File(self.ptrack_filepath, 'r')
if self.halo_data_dir is default:
self.halo_data_dir = self.ptrack['parameters'].attrs['sdir']
# Load the halo data
self.halo_data = halo_data.HaloData(self.halo_data_dir)
def __init__(
self,
particle_positions,
redshift,
snum,
hubble,
galaxy_cut,
length_scale,
mt_length_scale = None,
halo_length_scale = 'Rvir',
particle_masses = None,
minimum_criteria = 'n_star',
minimum_value = 10,
ids_to_return = None,
ids_with_supplementary_data = [],
supplementary_data_keys = [],
halo_data = None,
halo_finder = 'AHF',
halo_data_dir = None,
halo_file_tag = None,
mtree_halos_index = None,
main_mt_halo_id = None,
low_memory_mode = True,
memory_mode_divisions = 10,
):
'''Initialize.
Args:
particle_positions (array-like) :
Positions with dimensions (n_particles, 3).
If given a np.ndarray, assumes pkpc units.
Also accepts a unyt array.
redshift (float) :
Redshift the particles are at.
snum (int) :
Snapshot the particles correspond to.
hubble (float) :
Cosmological hubble parameter (little h)
galaxy_cut (float) :
The fraction of the length scale a particle must be inside to be
counted as part of a galaxy.
length_scale (str) :
Anything within galaxy_cut*length_scale is counted as being
inside the galaxy.
mt_length_scale (str) :
Same as length scale, but for merger tree halos. Defaults to
length_scale.
halo_length_scale (str) :
Anything within halo_length_scale is counted as being
inside the halo.
particle_masses (np.ndarray, optional) :
Masses of particles.
minimum_criteria (str) :
Options...
'n_star' -- halos must contain a minimum number of stars to
count as containing a galaxy.
'M_star' -- halos must contain a minimum stellar mass to count
as containing a galaxy.
minimum_value (int or float) :
The minimum amount of something (specified in minimum criteria)
in order for a galaxy to count as hosting a halo.
ids_to_return (list of strs, optional) :
The types of id you want to get out.
ids_with_supplementary_data (list of strs, optional) :
What types of IDs should include supplementary data pulled
from the halo files.
supplementary_data_keys (list of strs, optional) :
What data keys in the halo files should be accessed and
included as part of supplementary data.
halo_data (HaloData object, optional) :
An instance of an object that retrieves halo data
If not given initiate one using the halo_data_dir in kwargs
halo_data_dir (str, optional) :
Directory the AHF data is in. Necessary if halo_data is not
provided.
halo_finder (str, optional) :
What code were the halos found with? Examples are 'AHF',
'Rockstar'.
halo_file_tag (int, optional) :
What identifying tag to use for the merger tree files? Necessary
if using merger tree information.
mtree_halos_index (str or int, optional) :
The index argument to pass to AHFReader.get_mtree_halos().
For most cases this should be the final snapshot number, but see
AHFReader.get_mtree_halos's documentation.
main_mt_halo_id (int, optional) :
Index of the main merger tree halo.
low_memory_mode (bool) :
If True, use less memory at the cost of reduced performance.
memory_mode_divisions (int) :
When low_memory_mode is True, we reduce the memory cost by doing
less work at once. This divides certain
arrays into a number of divisions equal to
memory_mode_divisions, and then the results are calculated for
each division before bringing everything together.
'''
if self.mt_length_scale is None:
self.mt_length_scale = length_scale
# Setup the default halo_data
if halo_data is None:
self.halo_data = h_data.HaloData(
self.halo_data_dir,
halo_finder = self.halo_finder,
mt_kwargs = {
'index' : self.mtree_halos_index,
'tag' : self.halo_file_tag,
},
)
if not isinstance( particle_positions, unyt.array.unyt_array ):
particle_positions = copy.copy( particle_positions ) * unyt.kpc
# In the case of a minimum stellar mass, we need to divide the minimum
# value by 1/h when getting its values out.
if self.minimum_criteria in [ 'M_star', 'Mvir']:
self.min_conversion_factor = self.hubble
else:
self.min_conversion_factor = 1
# Derived properties
self.n_particles = self.particle_positions.shape[0]
def setUp(self):
self.halo_data = halo_data.HaloData(
sdir,
mt_kwargs={'tag': 'smooth'},
)
def setUp(self):
self.halo_data = halo_data.HaloData(sdir)