from scipy.io import FortranFile
import os.path
import h5py
from .basefuncs import *
[docs]def read_snapshot(fname, input_type='G1/G2', gadget2_header=True, ntypes=6, multiple_snaps=False, ics=False,
has_pot=False, has_accel=False, has_dtentr=False, has_tstp=False, store_pot=False, store_accel=False,
store_dtentr=False, store_tstp=False):
"""Read snapshot(s) for one simulation time output.
Parameters
----------
fname : str
path to snapshot
input_type : {'G1/G2', 'G2_str', 'G4_hdf5'}
Snapshot data type, 'G1/G2' for usual Gadget-2 format.
gadget2_header : bool, optional
Used if snapshot is of 'G1/G2' or 'G2_str' type. See `~readsnapshot.read_gadget`
ntypes: int, optional
Used if snapshot is of 'G1/G2' or 'G2_str' type. See `~readsnapshot.read_gadget`
multiple_snaps: bool, default = `False`
`True` if several snapshots per simulation time output.
ics : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
has_pot : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
has_accel : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
has_dtentr : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
has_tstp : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
store_pot : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
store_accel : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
store_dtentr : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
store_tstp : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
Returns
-------
dict
Data dictionary.
"""
if multiple_snaps:
data = read_sub_snapshots(fname, input_type=input_type, gadget2_header=gadget2_header, ntypes=ntypes, ics=ics,
has_pot=has_pot, has_accel=has_accel, has_dtentr=has_dtentr, has_tstp=has_tstp,
store_pot=store_pot, store_accel=store_accel, store_dtentr=store_dtentr,
store_tstp=store_tstp)
else:
data = read_single_snapshot(fname, input_type=input_type, gadget2_header=gadget2_header, ntypes=ntypes, ics=ics,
has_pot=has_pot, has_accel=has_accel, has_dtentr=has_dtentr, has_tstp=has_tstp,
store_pot=store_pot, store_accel=store_accel, store_dtentr=store_dtentr,
store_tstp=store_tstp)
return data
[docs]def read_single_snapshot(fname, input_type='G1/G2', gadget2_header=True, ntypes=6, ics=False, has_pot=False,
has_accel=False, has_dtentr=False, has_tstp=False, store_pot=False, store_accel=False,
store_dtentr=False, store_tstp=False):
"""Read a snapshot for one simulation time output.
Parameters
----------
fname : str
path to snapshot
input_type : {'G1/G2', 'G2_str', 'G4_hdf5'}
Snapshot data type, 'G1/G2' for usual Gadget-2 format.
gadget2_header : bool, optional `
Used if snapshot is of 'G1/G2' or 'G2_str' type. See `~readsnapshot.read_gadget`
ntypes: int, optional
Used if snapshot is of 'G1/G2' or 'G2_str' type. See `~readsnapshot.read_gadget_header`
ics : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
has_pot : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
has_accel : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
has_dtentr : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
has_tstp : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
store_pot : bool, optional
Used if snapshot is of 'G1/G2' type or 'G2_str' type. See `~readsnapshot.read_gadget`
store_accel : bool, optional
Used if snapshot is of 'G1/G2' type or 'G2_str' type. See `~readsnapshot.read_gadget`
store_dtentr : bool, optional
Used if snapshot is of 'G1/G2' type or 'G2_str' type. See `~readsnapshot.read_gadget`
store_tstp : bool, optional
Used if snapshot is of 'G1/G2' type or 'G2_str' type. See `~readsnapshot.read_gadget`
Returns
-------
dict
Data dictionary.
"""
if input_type == 'G1/G2':
data = read_gadget(fname, gadget2_header=gadget2_header, ntypes=ntypes, ics=ics, has_pot=has_pot,
has_accel=has_accel, has_dtentr=has_dtentr, has_tstp=has_tstp, store_pot=store_pot,
store_accel=store_accel, store_dtentr=store_dtentr, store_tstp=store_tstp)
elif input_type == 'G2_str':
data = read_gadget_str(fname, gadget2_header=gadget2_header, ntypes=ntypes, store_pot=store_pot,
store_accel=store_accel, store_dtentr=store_dtentr, store_tstp=store_tstp)
elif input_type == 'G4_hdf5':
data = read_hdf5(fname)
else:
raise ValueError('Unknown input type. "G1/G2", "G2_str" and "G4_hdf5" can be tried.')
return data
[docs]def read_sub_snapshots(fname, input_type='G1/G2', gadget2_header=True, ntypes=6, ics=False, has_pot=False,
has_accel=False, has_dtentr=False, has_tstp=False, store_pot=False, store_accel=False,
store_dtentr=False, store_tstp=False):
"""Read multiple snapshots corresponding to one simulation time output.
Parameters
----------
fname : str
path to snapshot
input_type : {'G1/G2', 'G2_str', 'G4_hdf5'}
Snapshot data type, 'G1/G2' for usual Gadget-2 format.
gadget2_header : bool, optional `
Used if snapshot is of 'G1/G2' or 'G2_str' type. See `~readsnapshot.read_gadget`
ntypes: int, optional
Used if snapshot is of 'G1/G2' or 'G2_str' type. See `~readsnapshot.read_gadget`
ics : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
has_pot : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
has_accel : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
has_dtentr : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
has_tstp : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
store_pot : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
store_accel : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
store_dtentr : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
store_tstp : bool, optional
Used if snapshot is of 'G1/G2' type. See `~readsnapshot.read_gadget`
Returns
-------
dict
Data dictionary.
"""
if os.path.isdir(fname):
isnap = int(fname.split('_')[-1])
fname += '/snapshot_' + "{0:03d}".format(isnap)
data = read_single_snapshot(fname + '.0', input_type=input_type, gadget2_header=gadget2_header, ntypes=ntypes,
ics=ics, has_pot=has_pot, has_accel=has_accel, has_dtentr=has_dtentr, has_tstp=has_tstp,
store_pot=store_pot, store_accel=store_accel, store_dtentr=store_dtentr,
store_tstp=store_tstp)
nfiles = data['head']['numfiles']
if nfiles > 1:
nparttotal = data['head']['npart_total']
npartinsnap = np.zeros((nfiles, 6))
npartinsnap[0, :] = data['npart']
for i in range(1, nfiles):
noftypeinprevsnaps = np.sum(npartinsnap[:i, :], axis=0).astype(int)
inds0prev, inds1prev, inds2prev, inds3prev, inds4prev, inds5prev = inds_of_p_types(noftypeinprevsnaps)
datasnap = read_single_snapshot(fname + '.' + str(i), input_type=input_type, gadget2_header=gadget2_header,
ntypes=ntypes, ics=ics, has_pot=has_pot, has_accel=has_accel,
has_dtentr=has_dtentr, has_tstp=has_tstp, store_pot=store_pot,
store_accel=store_accel, store_dtentr=store_dtentr, store_tstp=store_tstp)
npartinsnap[i, :] = datasnap['npart']
inds0, inds1, inds2, inds3, inds4, inds5 = inds_of_p_types(npartinsnap[i, :])
for key in data.keys():
if key in ['pos', 'vel', 'acc']:
data[key] = np.vstack((data[key][inds0prev, :], datasnap[key][inds0, :],
data[key][inds1prev, :], datasnap[key][inds1, :],
data[key][inds2prev, :], datasnap[key][inds2, :],
data[key][inds3prev, :], datasnap[key][inds3, :],
data[key][inds4prev, :], datasnap[key][inds4, :],
data[key][inds5prev, :], datasnap[key][inds5, :]))
elif key in ['head', 'time', 'npart']:
continue
else:
data[key] = np.concatenate((data[key][inds0prev], datasnap[key][inds0],
data[key][inds1prev], datasnap[key][inds1],
data[key][inds2prev], datasnap[key][inds2],
data[key][inds3prev], datasnap[key][inds3],
data[key][inds4prev], datasnap[key][inds4],
data[key][inds5prev], datasnap[key][inds5]))
for itype in range(0, 6):
print(itype, np.sum(npartinsnap, axis=0)[itype], nparttotal[itype])
if np.sum(npartinsnap, axis=0)[itype] != nparttotal[itype]:
raise ValueError(
'The sums of particles numbers of present types from sub snapshots are not the theoretical '
'total numbers.')
data['npart'] = data['head']['npart_total']
return data
[docs]def read_hdf5(fname):
"""Reads a HDF5 snapshot file.
Parameters
----------
fname : str
Notes
-----
Consitency with Gadget-2 output hdf5 files has not been checked because of problems with running Gadget-2 with the
right HDF5 library.
Returns
-------
dict
Data dictionary.
"""
if fname[-5:] != '.hdf5':
fname += '.hdf5'
fhdf5 = h5py.File(fname)
list_keys = list(fhdf5.keys())
hdf5_header = fhdf5['Header']
redshift = hdf5_header.attrs['Redshift']
time = hdf5_header.attrs['Time']
numfiles = hdf5_header.attrs['NumFilesPerSnapshot']
hdf5_npart = hdf5_header.attrs['NumPart_ThisFile']
hdf5_nparttotal = hdf5_header.attrs['NumPart_Total']
hdf5_massarr = hdf5_header.attrs['MassTable']
npart = np.zeros((6), dtype=int)
nparttotal = np.zeros((6), dtype=int)
massarr = np.zeros((6))
part_types_pres = np.zeros((6))
for i in range(0, 6):
if 'PartType' + str(i) in list_keys:
part_types_pres[i] = 1
for i in range(0, 6):
if part_types_pres[i]:
npart[i] = hdf5_npart[i]
nparttotal[i] = hdf5_nparttotal[i]
massarr[i] = hdf5_massarr[i]
head_dic = {'npart': npart, 'npart_total': nparttotal, 'massarr': massarr, 'time': time, 'redshift': redshift,
'numfiles': numfiles}
ntot = np.sum(npart)
pos = np.zeros((ntot, 3))
vel = np.zeros((ntot, 3))
ids = np.zeros((ntot))
mass = np.zeros((ntot))
skip = 0
skipblock = 0
for i in range(0, 6):
if part_types_pres[i]:
ptype_hdf5_group = fhdf5['PartType' + str(i)]
pos[skip:skip + npart[i], :] = ptype_hdf5_group['Coordinates']
vel[skip:skip + npart[i], :] = ptype_hdf5_group['Velocities']
ids[skip:skip + npart[i]] = ptype_hdf5_group['ParticleIDs']
if (massarr[i]):
mass[skip:skip + npart[i]] = massarr[i]
else:
mass[skip:skip + npart[i]] = ptype_hdf5_group['Masses']
skipblock += npart[i]
skip += npart[i]
data = {'head': head_dic, 'time': time, 'pos': pos, 'vel': vel, 'ids': ids, 'mass': mass, 'npart': npart}
# we consider here that 'PartType0 ' corresponds to gas particles. Not fully generic.
if 'PartType0' in list_keys:
gas_hdf5_group = fhdf5['PartType0']
list_gas_hdf5 = ['InternalEnergy', 'Density', 'SmoothingLength']
list_gas_data = ['uint', 'rho', 'hsml']
if 'ElectronAbundance' in gas_hdf5_group.keys():
electron_ab = np.zeros((npart[0]))
electron_ab[:] = gas_hdf5_group['ElectronAbundance']
data['electron_ab'] = electron_ab
if 'Pressure' in gas_hdf5_group.keys():
pressure = np.zeros((npart[0]))
pressure[:] = gas_hdf5_group['Pressure']
data['pressure'] = pressure
if 'Entropy' in gas_hdf5_group.keys():
entropy = np.zeros((npart[0]))
entropy[:] = gas_hdf5_group['Entropy']
data['entropy'] = entropy
for str_hdf5, str_data in zip(list_gas_hdf5, list_gas_data):
if str_hdf5 in gas_hdf5_group.keys():
arr = np.zeros((npart[0]))
arr[:] = gas_hdf5_group[str_hdf5]
data[str_data] = arr
if 'PartType4' in list_keys:
newstars_hdf5_group = fhdf5['PartType4']
if 'StellarFormationTime' in newstars_hdf5_group.keys():
stellage = np.zeros((npart[4]))
stellage[:] = newstars_hdf5_group['StellarFormationTime']
data['stellage'] = stellage
has_pot = False
for i in range(0, 6):
if part_types_pres[i]:
ptype_hdf5_group = fhdf5['PartType' + str(i)]
if 'Potential' in ptype_hdf5_group.keys():
has_pot = True
if has_pot:
skip = 0
pot = np.zeros((ntot))
for i in range(0, 6):
if part_types_pres[i]:
ptype_hdf5_group = fhdf5['PartType' + str(i)]
pot[skip:skip + npart[i]] = ptype_hdf5_group['Potential']
skip += npart[i]
data['pot'] = pot
return data
[docs]def read_gadget(fname, gadget2_header=True, ntypes=6, ics=False, has_pot=False, has_accel=False, has_dtentr=False,
has_tstp=False, store_pot=False, store_accel=False, store_dtentr=False, store_tstp=False):
"""Reads a Gadget-1/Gadget-2 file and fills a dictionary.
Parameters
----------
fname : str
path to snapshot
gadget2_header : bool, default = `True`
`True` if classic Gdget-2 header (default in Gadget-2, option in Gadget-4).
ntypes : int, optional
See `~readsnapshot.read_gadget_header`
ics : bool, default = `False`
`True` if initial conditions file, with limited number of blocks.
has_pot : bool, default = `False`
`True` if snapshot contains gravitational potential block.
has_accel : bool, default = `False`
`True` if snapshot contains accelerations block.
has_dtentr : bool, default = `False`
`True` if snapshot contains variation of entropy block.
has_tstp : bool, default = `False`
`True` if snapshot contains time steps block.
store_pot : bool, default = `False`
`True` if one wants to store gravitational potential block (if existing).
store_accel : bool, default = `False`
`True` if one wants to store accelerations block (if existing).
store_dtentr : bool, default = `False`
`True` if one wants to store variation of entropy block (if existing).
store_tstp : bool, default = `False`
(Default value = False)
Returns
-------
dict
Data dictionary.
"""
f = FortranFile(fname)
head_dic = read_gadget_header(f, gadget2_header=gadget2_header, ntypes=ntypes)
npart = head_dic['npart']
massarr = head_dic['massarr']
pos = f.read_reals(dtype='f4')
pos = np.reshape(pos, (-1, 3), 'C')
vel = f.read_reals(dtype='f4')
vel = np.reshape(vel, (-1, 3), 'C')
ids = f.read_ints(dtype='i4')
ntot = np.sum(npart)
mass = np.zeros((ntot))
massblock = None
if np.size(np.where((npart > 0) & (massarr == 0))):
massblock = f.read_reals(dtype='f4')
skip = 0
skipblock = 0
for i in range(0, ntypes):
if npart[i]:
if massarr[i]:
mass[skip:skip + npart[i]] = massarr[i]
else:
mass[skip:skip + npart[i]] = massblock[skipblock:skipblock + npart[i]]
skipblock += npart[i]
skip += npart[i]
time = head_dic['time']
data = {'head': head_dic, 'time': time, 'pos': pos, 'vel': vel, 'ids': ids, 'mass': mass, 'npart': npart}
if (npart[0]):
uint = f.read_reals(dtype='f4')
data['uint'] = uint
if ics is False:
if npart[0]:
rho = f.read_reals(dtype='f4')
hsml = f.read_reals(dtype='f4')
data['rho'] = rho
data['hsml'] = hsml
if has_pot:
pot = f.read_reals(dtype='f4')
if store_pot:
data['pot'] = pot
if has_accel:
acc = f.read_reals(dtype='f4')
if store_accel:
data['accel'] = acc
if npart[0]:
if has_dtentr:
dtentr = f.read_reals(dtype='f4')
if store_dtentr:
data['endt'] = dtentr
if has_tstp:
tstp = f.read_reals(dtype='f4')
if store_tstp:
data['tstp'] = tstp
f.close()
return data
[docs]def read_gadget_str(fname, gadget2_header=True, ntypes=6, store_pot=False, store_accel=False, store_dtentr=False,
store_tstp=False):
"""Reads Gadget-2 file of type legacyformat2,, with blocks preceded by header strrings, and fills a dictionary.
Parameters
----------
fname : str
path to snapshot
gadget2_header : bool, default = `True`
`True` if classic Gdget-2 header (default in Gadget-2, option in Gadget-4).
ntypes : int, optional
See `~readsnapshot.read_gadget_header`
store_pot : bool, default = `False`
`True` if one wants to store gravitational potential block (if existing).
store_accel : bool, default = `False`
`True` if one wants to store accelerations block (if existing).
store_dtentr : bool, default = `False`
`True` if one wants to store variation of entropy block (if existing).
store_tstp : bool, default = `False`
(Default value = False)
Returns
-------
dict
Data dictionary.
"""
f = FortranFile(fname)
list_block_names = ['HEAD', 'POS', 'VEL', 'ID', 'MASS', 'U', 'RHO', 'HSML', 'POT', 'ACCE', 'ENDT', 'TSTP']
data = {}
massblock = None
for iblock in range(0, len(list_block_names)):
try:
block_name = f.read_record('S8').tobytes()[0:4].decode().strip()
except TypeError:
break
if block_name == 'HEAD':
head_dic = read_gadget_header(f, gadget2_header=gadget2_header, ntypes=ntypes)
data['head'] = head_dic
data['npart'] = head_dic['npart']
data['massarr'] = head_dic['massarr']
data['time'] = head_dic['time']
if block_name == 'POS':
pos = f.read_reals(dtype='f4')
pos = np.reshape(pos, (-1, 3), 'C')
data['pos'] = pos
if block_name == 'VEL':
vel = f.read_reals(dtype='f4')
vel = np.reshape(vel, (-1, 3), 'C')
data['vel'] = vel
if block_name == 'ID':
data['ids'] = f.read_ints(dtype='i4')
if block_name == 'MASS':
massblock = f.read_reals(dtype='f4')
if block_name == 'U':
data['uint'] = f.read_reals(dtype='f4')
if block_name == 'RHO':
data['rho'] = f.read_reals(dtype='f4')
if block_name == 'HSML':
data['hsml'] = f.read_reals(dtype='f4')
if block_name == 'POT':
if store_pot:
data['pot'] = f.read_reals(dtype='f4')
if block_name == 'ACC':
if store_accel:
data['acc'] = f.read_reals(dtype='f4')
if block_name == 'ENDT':
if store_dtentr:
data['endt'] = f.read_reals(dtype='f4')
if block_name == 'TSTP':
if store_tstp:
data['tstp'] = f.read_reals(dtype='f4')
f.close()
npart = data['npart']
massarr = data['massarr']
ntot = np.sum(npart)
mass = np.zeros((ntot))
skip = 0
skipblock = 0
for i in range(0, 6):
if npart[i]:
if massarr[i]:
mass[skip:skip + npart[i]] = massarr[i]
else:
if massblock is not None:
mass[skip:skip + npart[i]] = massblock[skipblock:skipblock + npart[i]]
else:
raise ValueError('One particle type has no given mass in mass array and no mass block was read.')
skipblock += npart[i]
skip += npart[i]
data['mass'] = mass
return data
[docs]def inds_of_p_types(nbytype):
"""Returns indexes of particles corresponding to the 6 Gadget-2 types.
Parameters
----------
nbytype: ndarray
Number of particles by type.
Returns
-------
(ndarray,ndarray,ndarray,ndarray,ndarray,ndarray)
indexes of different types
"""
nbytype = nbytype.astype(int)
inds0 = np.zeros(np.sum(nbytype), dtype='bool')
inds0[np.arange(0, np.sum(nbytype[0:1]))] = True
inds1 = np.zeros(np.sum(nbytype), dtype='bool')
inds1[np.arange(int(np.sum(nbytype[0:1])), np.sum(nbytype[0:2]))] = True
inds2 = np.zeros(np.sum(nbytype), dtype='bool')
inds2[np.arange(int(np.sum(nbytype[0:2])), np.sum(nbytype[0:3]))] = True
inds3 = np.zeros(np.sum(nbytype), dtype='bool')
inds3[np.arange(int(np.sum(nbytype[0:3])), np.sum(nbytype[0:4]))] = True
inds4 = np.zeros(np.sum(nbytype), dtype='bool')
inds4[np.arange(int(np.sum(nbytype[0:4])), np.sum(nbytype[0:5]))] = True
inds5 = np.zeros(np.sum(nbytype), dtype='bool')
inds5[np.arange(int(np.sum(nbytype[0:5])), np.sum(nbytype[0:6]))] = True
return inds0, inds1, inds2, inds3, inds4, inds5