def test_has_attributes(self):
self.fn()
f = h5py.File('tests/data/tracking_output/ptracks_test.hdf5', 'r')
# Load one of the original snapshots to compare
P = readsnap.readsnap('tests/data/test_data_with_new_id_scheme',
600,
0,
True,
cosmological=True)
compare_keys = ['omega_matter', 'omega_lambda', 'hubble']
for key in compare_keys:
npt.assert_allclose(P[key], f.attrs[key])
for key in default_data_p.keys():
if key == 'check_same_sdir' or key == 'custom_fns':
continue
elif key == 'custom_fns_str':
for i, fn in enumerate(default_data_p[key]):
assert inspect.getsource(
fn) == f['parameters'].attrs[key][i]
else:
assert default_data_p[key] == f['parameters'].attrs[key]
def test_basic( self ):
'''Basically, does it work?.'''
# Input
sdir = './tests/data/test_data_with_new_id_scheme'
snum = 600
p_types = [0, ]
target_ids = np.array([ 36091289, 36091289, 3211791, 10952235 ])
target_child_ids = np.array([ 893109954, 1945060136, 0, 0 ])
# My knowledge, by hand
target_inds = [4, 0, 1, 5]
P = readsnap.readsnap( sdir, snum, 0, True, cosmological=True )
expected = {
'ID': target_ids,
'ChildID': target_child_ids,
'Den': np.array([ P['rho'][ind]*constants.UNITDENSITY_IN_NUMDEN for ind in target_inds ]),
}
dfid, redshift, attrs = self.fn( sdir, snum, p_types, target_ids, \
target_child_ids=target_child_ids)
for key in expected.keys():
npt.assert_allclose( dfid[key], expected[key] )
# Make sure the redshift's right too
npt.assert_allclose( redshift, 0. )
def test_reads_potential(self):
sdir = './tests/data/sdir/output'
snum = 600
data = read_snapshot.readsnap(sdir=sdir, snum=snum, ptype=0)
expected_potential = 146365.1875
npt.assert_allclose(expected_potential, data['potential'][-3])
def test_gets_new_ids(self):
self.kwargs['load_additional_ids'] = True
self.p_data = self.test_class(**self.kwargs)
P = read_snapshot.readsnap(self.kwargs['sdir'], self.kwargs['snum'],
self.kwargs['ptype'], True)
npt.assert_allclose(P['id'], self.p_data.data['ID'])
npt.assert_allclose(P['child_id'], self.p_data.data['ChildID'])
npt.assert_allclose(P['id_gen'], self.p_data.data['IDGen'])
def concatenate_particle_data(self, verbose=False):
'''Get all the particle data for the snapshot in one big array, to allow
searching through it.
Args:
verbose (bool): If True, print additional information.
Modifies:
self.full_snap_data (dict):
A dictionary of the concatenated particle data.
'''
if verbose:
print('Reading data...')
sys.stdout.flush()
full_snap_data = {
'ID': [],
'PType': [],
'Den': [],
'SFR': [],
'T': [],
'Z': [],
'M': [],
'P0': [],
'P1': [],
'P2': [],
'V0': [],
'V1': [],
'V2': [],
}
time_start = time.time()
if hasattr(self, 'target_child_ids'):
load_additional_ids = True
full_snap_data['ChildID'] = []
else:
load_additional_ids = False
# Flag for saving header info
saved_header_info = False
for i, p_type in enumerate(self.p_types):
P = read_snapshot.readsnap(self.sdir,
self.snum,
p_type,
load_additional_ids=load_additional_ids,
cosmological=1,
skip_bh=1,
header_only=0)
if P['k'] < 0:
continue
pnum = P['id'].size
# On the first time through, save global information
if not saved_header_info:
# Save the redshift
self.redshift = P['redshift']
# Store the attributes to be used in the final data file.
attrs_keys = ['omega_matter', 'omega_lambda', 'hubble']
self.attrs = {}
for key in attrs_keys:
self.attrs[key] = P[key]
saved_header_info = True
if verbose:
print(' ... ', pnum, ' type', p_type, ' particles')
if 'rho' in P:
Den = P['rho'] * constants.UNITDENSITY_IN_NUMDEN
else:
Den = [
linefinder_config.FLOAT_FILL_VALUE,
] * pnum
if 'sfr' in P:
sfr = P['sfr']
else:
sfr = [
linefinder_config.FLOAT_FILL_VALUE,
] * pnum
if 'u' in P:
T = read_snapshot.gas_temperature(P['u'], P['ne'])
else:
T = [
linefinder_config.FLOAT_FILL_VALUE,
] * pnum
thistype = np.zeros(pnum, dtype='int8')
thistype.fill(p_type)
full_snap_data['ID'].append(P['id'])
full_snap_data['PType'].append(thistype)
full_snap_data['Den'].append(Den)
full_snap_data['SFR'].append(sfr)
full_snap_data['T'].append(T)
full_snap_data['Z'].append(P['z'][:, 0] / constants.Z_MASSFRAC_SUN)
full_snap_data['M'].append(P['m'] * constants.UNITMASS_IN_MSUN)
full_snap_data['P0'].append(P['p'][:, 0])
full_snap_data['P1'].append(P['p'][:, 1])
full_snap_data['P2'].append(P['p'][:, 2])
full_snap_data['V0'].append(P['v'][:, 0])
full_snap_data['V1'].append(P['v'][:, 1])
full_snap_data['V2'].append(P['v'][:, 2])
# Save the potential, if it exists
if 'potential' in P.keys():
try:
full_snap_data['Potential'].append(P['potential'])
except KeyError:
full_snap_data['Potential'] = [
P['potential'],
]
if hasattr(self, 'target_child_ids'):
full_snap_data['ChildID'].append(P['child_id'])
time_end = time.time()
if verbose:
print('readsnap done in ... {:.3g} seconds'.format(time_end -
time_start))
# Convert to numpy arrays
for key in full_snap_data.keys():
full_snap_data[key] = np.concatenate(full_snap_data[key])
self.full_snap_data = full_snap_data
def retrieve_data(self):
# Assume we convert from cosmological units
P = read_snapshot.readsnap(
self.sdir,
self.snum,
self.ptype,
load_additional_ids=self.load_additional_ids,
cosmological=True,
)
# Parse the keys and put in a more general format
# All units should be the standard GIZMO output units
self.data = {}
self.data_attrs = {}
for key in P.keys():
# Get the attributes
attrs_keys = [
'redshift', 'omega_lambda', 'flag_metals', 'flag_cooling',
'omega_matter', 'flag_feedbacktp', 'time', 'boxsize', 'hubble',
'flag_sfr', 'flag_stellarage', 'k'
]
if key in attrs_keys:
self.data_attrs[key] = P[key]
# Get the data
# Gas Density
elif key == 'rho':
self.data['Den'] = P[key]
# Gas Neutral Hydrogen fraction
elif key == 'nh':
self.data['nHI'] = P[key]
# Should be the mean free electron number per proton
elif key == 'ne':
self.data['ne'] = P[key]
# Star Formation Rate
elif key == 'sfr':
self.data['SFR'] = P[key]
# Position
elif key == 'p':
self.data['P'] = P[key].transpose()
# Velocity
elif key == 'v':
self.data['V'] = P[key].transpose()
# Metal mass fraction
elif key == 'z':
self.data['Z'] = P[
key][:, 0] # Total metallicity (everything not H or He)
self.data['Z_Species'] = P[
key][:,
1:] # Details per species, [He, C, N, O, Ne, Mg, Si, S, Ca, Fe], in order
# Particle IDs
elif key == 'id':
self.data['ID'] = P[key]
elif key == 'child_id':
self.data['ChildID'] = P[key]
elif key == 'id_gen':
self.data['IDGen'] = P[key]
# Particle Masses
elif key == 'm':
self.data['M'] = P[key]
# Internal energy
elif key == 'u':
self.data['U'] = P[key]
# Smoothing lengths
elif key == 'h':
self.data['h'] = P[key]
elif key == 'age':
self.data['Age'] = P[key]
elif key == 'potential':
self.data['Potential'] = P[key]
else:
raise Exception('NULL key, key={}'.format(key))