def test_find_galaxies_for_particle_tracks(self):
particle_track_gal_linker = galaxy_link.ParticleTrackGalaxyLinker(
**ptrack_gal_linker_kwargs)
particle_track_gal_linker.find_galaxies_for_particle_tracks()
f = h5py.File(self.savefile, 'r')
g = h5py.File(self.originalfile, 'r')
# What we expect (calculated using the positions of the particles in
# snum 500 )
particle_positions = g['P'][...][:, -1]
gal_linker = galaxy_linker.GalaxyLinker(particle_positions,
**gal_linker_kwargs)
expected = gal_linker.find_ids()
# Make the comparison
for key in expected.keys():
npt.assert_allclose(expected[key],
f[key][...][:, -1],
err_msg='Key {} failed'.format(key),
rtol=1e-4)
# What we expect (calculated using the positions of the particles in
# snum 550 )
gal_linker_kwargs_copy = dict(gal_linker_kwargs)
gal_linker_kwargs_copy['snum'] = g['snum'][1]
gal_linker_kwargs_copy['redshift'] = g['redshift'][1]
particle_positions = g['P'][...][:, 1]
gal_linker = galaxy_linker.GalaxyLinker(particle_positions,
**gal_linker_kwargs_copy)
# Make the comparison
expected = gal_linker.find_ids()
for key in expected.keys():
npt.assert_allclose(expected[key],
f[key][...][:, 1],
err_msg="Key '{}' failed".format(key),
rtol=1e-4)
# Make sure the main MT halo ID is the one we expect.
assert f.attrs['main_mt_halo_id'] == 0
# Make sure we have stored the data parameters too.
for key in ptrack_gal_linker_kwargs.keys():
if (key != 'ids_to_return') and (key != 'main_mt_halo_id'):
assert \
ptrack_gal_linker_kwargs[key] == f['parameters'].attrs[key]
def get_galaxy_and_halo_ids(args):
'''Get the galaxy and halo ids for a single snapshot.'''
particle_positions, kwargs = args
time_start = time.time()
# Find the galaxy for a given snapshot
gal_linker = galaxy_linker.GalaxyLinker(particle_positions,
**kwargs)
galaxy_and_halo_ids = gal_linker.find_ids()
time_end = time.time()
print( 'Snapshot {:>3} | redshift {:>7.3g} | done in {:.3g} seconds'\
.format(
kwargs['snum'],
kwargs['redshift'],
time_end - time_start
)
)
sys.stdout.flush()
# Try to avoid memory leaks
del kwargs
del gal_linker
gc.collect()
return galaxy_and_halo_ids
def get_enclosed_mass( self,
simulation_data_dir,
ptype,
galaxy_cut,
length_scale,
):
'''Get the mass inside galaxy_cut*length_scale for each Halo halo.
Args:
simulation_data_dir (str) :
Directory containing the raw particle data.
ptype (str) :
What particle type to get the mass for.
galaxy_cut (float) :
galaxy_cut*length_scale defines the radius around the center of the halo within which to get the mass.
length_scale (str) :
galaxy_cut*length_scale defines the radius around the center of the halo within which to get the mass.
Returns:
mass_inside_all_halos (np.ndarray) :
mass_inside_all_halos[i] is the mass of particle type ptype inside galaxy_cut*length scale around a galaxy.
'''
# Load the simulation data
s_data = particle_data.ParticleData(
simulation_data_dir,
self.halos_snum,
data_constants.PTYPES[ptype],
)
try:
particle_positions = s_data.data['P'].transpose()
# Case where there are no star particles at this redshift.
except KeyError:
return np.array( [ 0., ]*self.halos.index.size )
# Find the mass radii
galaxy_linker_kwargs = {
'particle_positions' : particle_positions,
'particle_masses' : s_data.data['M']*constants.UNITMASS_IN_MSUN,
'snum' : self.halos_snum,
'redshift' : s_data.redshift,
'hubble' : s_data.data_attrs['hubble'],
'galaxy_cut' : galaxy_cut,
'length_scale' : length_scale,
'halo_data' : self,
}
gal_linker = galaxy_linker.GalaxyLinker( **galaxy_linker_kwargs )
mass_inside_all_halos = gal_linker.mass_inside_all_halos
# Make sure to put hubble constant back in so we have consistent units.
mass_inside_all_halos *= s_data.data_attrs['hubble']
return mass_inside_all_halos
def get_mass_radii(
self,
mass_fractions,
simulation_data_dir,
galaxy_cut,
length_scale,
):
'''Get radii that enclose a fraction (mass_fractions[i]) of a halo's stellar mass.
Args:
mass_fractions (list of floats) :
Relevant mass fractions.
simulation_data_dir (str) :
Directory containing the raw particle data.
galaxy_cut (float) :
galaxy_cut*length_scale defines the radius around the center of the halo to look for stars.
length_scale (str) :
galaxy_cut*length_scale defines the radius around the center of the halo to look for stars.
Returns:
mass_radii (list of np.ndarrays) :
If M_sum_j = all mass inside galaxy_cut*length_scale for halo j, then mass_radii[i][j] is the radius that
contains a fraction mass_fractions[i] of M_sum_j.
'''
# Load the simulation data
s_data = particle_data.ParticleData(
simulation_data_dir,
self.halos_snum,
ptype = data_constants.PTYPES['star'],
)
try:
particle_positions = s_data.data['P'].transpose()
# Case where there are no star particles at this redshift.
except KeyError:
return [ np.array( [ np.nan, ]*self.halos.index.size ), ]*len( mass_fractions )
# Find the mass radii
galaxy_linker_kwargs = {
'particle_positions' : particle_positions,
'particle_masses' : s_data.data['M'],
'snum' : self.halos_snum,
'redshift' : s_data.redshift,
'hubble' : s_data.data_attrs['hubble'],
'galaxy_cut' : galaxy_cut,
'length_scale' : length_scale,
'halo_data' : self,
}
gal_linker = galaxy_linker.GalaxyLinker( **galaxy_linker_kwargs )
mass_radii = [ gal_linker.get_mass_radius( mass_fraction ) for mass_fraction in mass_fractions ]
return mass_radii
def get_galaxy_and_halo_ids(i):
'''Get the galaxy and halo ids for a single snapshot.'''
# Get the particle positions
particle_positions = self.ptrack['P'][...][:, i]
# Get the data parameters to pass to GalaxyLinker
kwargs = {
'halo_data': None,
'galaxy_cut': self.galaxy_cut,
'length_scale': self.length_scale,
'mt_length_scale': self.mt_length_scale,
'ids_to_return': self.ids_to_return,
'minimum_criteria': self.minimum_criteria,
'minimum_value': self.minimum_value,
'redshift': self.ptrack['redshift'][...][i],
'snum': self.ptrack['snum'][...][i],
'hubble': self.ptrack.attrs['hubble'],
'halo_data_dir': self.halo_data_dir,
'mtree_halos_index': self.mtree_halos_index,
'main_mt_halo_id': self.main_mt_halo_id,
'halo_file_tag': self.halo_file_tag,
}
time_start = time.time()
# Find the galaxy for a given snapshot
gal_linker = galaxy_linker.GalaxyLinker(particle_positions,
**kwargs)
galaxy_and_halo_ids = gal_linker.find_ids()
time_end = time.time()
print( 'Snapshot {:>3} | redshift {:>7.3g} | done in {:.3g} seconds'\
.format(
kwargs['snum'],
kwargs['redshift'],
time_end - time_start
)
)
sys.stdout.flush()
# Try to avoid memory leaks
del kwargs
del gal_linker
gc.collect()
return galaxy_and_halo_ids
def get_galaxy_identification_loop(self):
'''Loop over all snapshots and identify the galaxy in each.
Modifies:
self.ptrack_gal_ids (dict) : Where the galaxy IDs are stored.
'''
# Loop over each included snapshot.
n_snaps = self.ptrack['snum'][...].size
for i in range(n_snaps):
# Get the particle positions
particle_positions = self.ptrack['P'][...][:, i]
# Get the data parameters to pass to GalaxyLinker
kwargs = {
'halo_data': self.halo_data,
'galaxy_cut': self.galaxy_cut,
'length_scale': self.length_scale,
'mt_length_scale': self.mt_length_scale,
'ids_to_return': self.ids_to_return,
'ids_with_supplementary_data':
self.ids_with_supplementary_data,
'supplementary_data_keys': self.supplementary_data_keys,
'minimum_criteria': self.minimum_criteria,
'minimum_value': self.minimum_value,
'redshift': self.ptrack['redshift'][...][i],
'snum': self.ptrack['snum'][...][i],
'hubble': self.ptrack.attrs['hubble'],
'halo_data_dir': self.halo_data_dir,
'mtree_halos_index': self.mtree_halos_index,
'main_mt_halo_id': self.main_mt_halo_id,
'halo_file_tag': self.halo_file_tag,
}
time_start = time.time()
# Find the galaxy for a given snapshot
gal_linker = galaxy_linker.GalaxyLinker(particle_positions,
**kwargs)
galaxy_and_halo_ids = gal_linker.find_ids()
time_end = time.time()
print( 'Snapshot {:>3} | redshift {:>7.3g} | done in {:.3g} seconds'\
.format(
kwargs['snum'],
kwargs['redshift'],
time_end - time_start
)
)
sys.stdout.flush()
# Make the arrays to store the data in
if not hasattr(self, 'ptrack_gal_ids'):
self.ptrack_gal_ids = {}
for key in galaxy_and_halo_ids.keys():
dtype = type(galaxy_and_halo_ids[key][0])
self.ptrack_gal_ids[key] = np.empty(
(gal_linker.n_particles, n_snaps), dtype=dtype)
# Store the data in the primary array
for key in galaxy_and_halo_ids.keys():
self.ptrack_gal_ids[key][:, i] = galaxy_and_halo_ids[key]
# Try clearing up memory again, in case gal_linker is hanging around
del kwargs
del gal_linker
del galaxy_and_halo_ids
gc.collect()
def get_average_quantity_inside_galaxy( self,
data_key,
simulation_data_dir,
ptype,
galaxy_cut,
length_scale,
weight_data_key = 'M',
fill_value = np.nan,
):
'''Get the mass inside galaxy_cut*length_scale for each Halo halo.
Args:
data_key (str) :
Data key for the quantity to get the average of.
simulation_data_dir (str) :
Directory containing the raw particle data.
ptype (str) :
What particle type to get the mass for.
galaxy_cut (float) :
galaxy_cut*length_scale defines the radius around the center of the halo within which to get the mass.
length_scale (str) :
galaxy_cut*length_scale defines the radius around the center of the halo within which to get the mass.
weight_data_key (str) :
Data key for the weight to use when averaging.
fill_value (float) :
What value to use when the average quantity inside the galaxy is not resolved.
Returns:
average_quantity_inside_galaxy (np.ndarray) :
average_quantity_inside_galaxy[i] is the average value of the requested quantity for particle type ptype
inside galaxy_cut*length scale around a galaxy.
'''
# Load the simulation data
s_data = particle_data.ParticleData(
simulation_data_dir,
self.halos_snum,
data_constants.PTYPES[ptype],
# The following values need to be set, because they come into play when a galaxy is centered on halo finder
# data. That's obviously not the case here...
centered = True,
vel_centered = True,
hubble_corrected = True,
)
try:
particle_positions = s_data.data['P'].transpose()
# Case where there are no particles of the given ptype at this redshift.
except KeyError:
return np.array( [ fill_value, ]*self.halos.index.size )
# Find the mass radii
galaxy_linker_kwargs = {
'particle_positions' : particle_positions,
'snum' : self.halos_snum,
'redshift' : s_data.redshift,
'hubble' : s_data.data_attrs['hubble'],
'galaxy_cut' : galaxy_cut,
'length_scale' : length_scale,
'halo_data' : self,
}
gal_linker = galaxy_linker.GalaxyLinker( low_memory_mode=False, **galaxy_linker_kwargs )
average_quantity_inside_galaxy = gal_linker.weighted_summed_quantity_inside_galaxy(
s_data.get_data( data_key ),
s_data.get_data( weight_data_key ),
fill_value,
)
return average_quantity_inside_galaxy