def drift(self, Δt):
"""Note that the time step size
Δt is really ∫_t^(t + Δt) dt/a**2.
"""
masterprint('Drifting', self.type, '...')
# Extracting variables
posx = self.posx
posy = self.posy
posz = self.posz
momx = self.momx
momy = self.momy
momz = self.momz
# The factor 1/mass*∫_t^(t + Δt) dt/a**2
fac = Δt/self.mass
# Update positions
for i in range(self.N_local):
posx[i] += momx[i]*fac
posy[i] += momy[i]*fac
posz[i] += momz[i]*fac
# Toroidal boundaries
posx[i] = mod(posx[i], boxsize)
posy[i] = mod(posy[i], boxsize)
posz[i] = mod(posz[i], boxsize)
# Some partiles may have drifted out of the local domain.
# Exchange particles to the correct processes.
masterprint('done')
exchange(self)
def load(self, filename, compare_params=True, only_params=False):
# Load all particles
with h5py.File(filename,
mode='r',
driver='mpio',
comm=comm,
) as hdf5_file:
# Load used base units
snapshot_unit_length = eval(hdf5_file.attrs['unit length'], units_dict)
snapshot_unit_time = eval(hdf5_file.attrs['unit time'], units_dict)
snapshot_unit_mass = eval(hdf5_file.attrs['unit mass'], units_dict)
# Load global attributes
self.params['a'] = hdf5_file.attrs['a']
self.params['boxsize'] = hdf5_file.attrs['boxsize']*snapshot_unit_length
self.params['H0'] = hdf5_file.attrs['H0']*(1/snapshot_unit_time)
self.params['Ωm'] = hdf5_file.attrs[unicode('Ω') + 'm']
self.params['ΩΛ'] = hdf5_file.attrs[unicode('Ω') + unicode('Λ')]
# Check if the parameters of the snapshot
# matches those of the current simulation run.
# Display a warning if they do not.
if compare_params:
tol = 1e-4
msg = ''
if np.abs(self.params['a']/a_begin - 1) > tol:
msg += ('\n' + ' '*8 + 'a_begin: {} vs {}'
).format(a_begin, self.params['a'])
if np.abs(self.params['boxsize']/boxsize - 1) > tol:
msg += ('\n' + ' '*8 + 'boxsize: {} vs {} ({})').format(boxsize,
self.params['boxsize'],
units.length)
if np.abs(self.params['H0']/H0 - 1) > tol:
unit = units.km/(units.s*units.Mpc)
msg += ('\n' + ' '*8 + 'H0: {} vs {} ({})').format(H0/unit,
self.params['H0']/unit,
'km s⁻¹ Mpc⁻¹')
if np.abs(self.params['Ωm']/Ωm - 1) > tol:
msg += ('\n' + ' '*8 + unicode('Ω') + 'm: {} vs {}').format(Ωm,
self.params['Ωm'])
if np.abs(self.params['ΩΛ']/ΩΛ - 1) > tol:
msg += ('\n' + ' '*8 + unicode('Ω')
+ unicode('Λ') + ': {} vs {}').format(ΩΛ,
self.params['ΩΛ'])
if msg:
msg = ('Mismatch between current parameters and those in'
+ 'the snapshot "{}":{}').format(filename, msg)
masterwarn(msg, indent=4)
# Initialize the particle_attributes dict
self.params['particle_attributes'] = {}
# Load particle data
all_particles = hdf5_file['particles']
for particle_type in all_particles:
# Load particle attributes
particles_h5 = all_particles[particle_type]
particle_N = particles_h5['posx'].size
particle_mass = particles_h5.attrs['mass']*snapshot_unit_mass
particle_species = particles_h5.attrs['species']
# The keys in the particle_attributes dict are the
# particle types, and the values are new dicts,
# containing the information for each type.
self.params['particle_attributes'][particle_type] = {}
particle_attribute = self.params['particle_attributes'][particle_type]
particle_attribute['N'] = particle_N
particle_attribute['mass'] = particle_mass
particle_attribute['species'] = particle_species
# Done loading particle attributes
if only_params:
continue
# Write out progress message
masterprint('Loading', particle_N, particle_type,
'({}) ...'.format(particle_species),
indent=4)
# Extract HDF5 datasets
posx_h5 = particles_h5['posx']
posy_h5 = particles_h5['posy']
posz_h5 = particles_h5['posz']
momx_h5 = particles_h5['momx']
momy_h5 = particles_h5['momy']
momz_h5 = particles_h5['momz']
# Compute a fair distribution of
# particle data to the processes.
N_locals = ((particle_N//nprocs, )
*(nprocs - (particle_N % nprocs))
+ (particle_N//nprocs + 1, )*(particle_N % nprocs))
N_local = N_locals[rank]
start_local = np.sum(N_locals[:rank], dtype=C2np['Py_ssize_t'])
end_local = start_local + N_local
# Construct a Particles instance
self.particles = construct(particle_type,
particle_species,
mass=particle_mass,
N=particle_N)
# Populate the Particles instance with data from the file
self.particles.populate(posx_h5[start_local:end_local], 'posx')
self.particles.populate(posy_h5[start_local:end_local], 'posy')
self.particles.populate(posz_h5[start_local:end_local], 'posz')
self.particles.populate(momx_h5[start_local:end_local], 'momx')
self.particles.populate(momy_h5[start_local:end_local], 'momy')
self.particles.populate(momz_h5[start_local:end_local], 'momz')
# If the snapshot and the current run uses different
# systems of units, mulitply the particle positions
# and momenta by the snapshot units.
if snapshot_unit_length != 1:
for i in range(N_local):
self.particles.posx[i] *= snapshot_unit_length
self.particles.posy[i] *= snapshot_unit_length
self.particles.posz[i] *= snapshot_unit_length
unit = snapshot_unit_length/snapshot_unit_time*snapshot_unit_mass
if unit != 1:
for i in range(N_local):
self.particles.momx[i] *= unit
self.particles.momy[i] *= unit
self.particles.momz[i] *= unit
# Finalize progress message
masterprint('done')
# Update the "contains" string
if only_params:
self.contains = 'params'
else:
self.contains = 'params and particles'
# Scatter particles to the correct domain-specific process.
# Setting reset_indices_send == True ensures that buffers
# will be reset afterwards, as this initial exchange is not
# representable for those to come.
if 'particles' in self.contains:
exchange(self.particles, reset_buffers=True)
def load(self, filename, compare_params=True, only_params=False):
""" It is assumed that the snapshot on the disk is a GADGET
snapshot of type 2 and that it uses single precision. The
GadgetSnapshot instance stores the data (positions and
velocities) in double precision. Only GADGET type 1 (halo)
particles, corresponding to dark matter particles,
are supported.
"""
# Only type 1 (halo) particles are supported
particle_species = 'dark matter'
particle_type = 'GADGET halos'
# Read in the snapshot
offset = 0
with open(filename, 'rb') as f:
# Read the HEAD block into a params['header'] dict.
# No unit conversion will be done.
offset = self.new_block(f, offset)
name = self.read(f, '4s').decode('utf8') # "HEAD"
size = self.read(f, 'i') # 264
offset = self.new_block(f, offset)
self.params['header'] = collections.OrderedDict()
header = self.params['header']
header['Npart'] = self.read(f, '6I')
header['Massarr'] = self.read(f, '6d')
header['Time'] = self.read(f, 'd')
header['Redshift'] = self.read(f, 'd')
header['FlagSfr'] = self.read(f, 'i')
header['FlagFeedback'] = self.read(f, 'i')
header['Nall'] = self.read(f, '6i')
header['FlagCooling'] = self.read(f, 'i')
header['NumFiles'] = self.read(f, 'i')
header['BoxSize'] = self.read(f, 'd')
header['Omega0'] = self.read(f, 'd')
header['OmegaLambda'] = self.read(f, 'd')
header['HubbleParam'] = self.read(f, 'd')
header['FlagAge'] = self.read(f, 'i')
header['FlagMetals'] = self.read(f, 'i')
header['NallHW'] = self.read(f, '6i')
header['flag_entr_ics'] = self.read(f, 'i')
# Also include some of the header fields as parameters
# directly in the params dict.
self.params['a'] = header['Time']
unit = units.kpc/header['HubbleParam']
self.params['boxsize'] = header['BoxSize']*unit
unit = 100*units.km/(units.s*units.Mpc)
self.params['H0'] = header['HubbleParam']*unit
self.params['Ωm'] = header['Omega0']
self.params['ΩΛ'] = header['OmegaLambda']
# Check if the parameters of the snapshot matches
# those of the current simulation run. Display a warning
# if they do not.
if compare_params:
tol = 1e-4
msg = ''
if np.abs(self.params['a']/a_begin - 1) > tol:
msg += '\n' + ' '*8 + 'a_begin: {} vs {}'.format(a_begin, self.params['a'])
if np.abs(self.params['boxsize']/boxsize - 1) > tol:
msg += ('\n' + ' '*8 + 'boxsize: {} vs {} ({})').format(boxsize,
self.params['boxsize'],
units.length)
if np.abs(self.params['H0']/H0 - 1) > tol:
unit = units.km/(units.s*units.Mpc)
msg += ('\n' + ' '*8 + 'H0: {} vs {} ({})').format(H0/unit,
self.params['H0']/unit,
'km s⁻¹ Mpc⁻¹')
if np.abs(self.params['Ωm']/Ωm - 1) > tol:
msg += ('\n' + ' '*8 + unicode('Ω') + 'm: {} vs {}').format(Ωm,
self.params['Ωm'])
if np.abs(self.params['ΩΛ']/ΩΛ - 1) > tol:
msg += ('\n' + ' '*8 + unicode('Ω')
+ unicode('Λ') + '{} vs {}').format(ΩΛ,
self.params['ΩΛ'])
if msg:
msg = ('Mismatch between current parameters and those in '
+ 'the GADGET snapshot '
+ '"{}":{}').format(filename, msg)
masterwarn(msg, indent=4)
# Initialize the particle_attributes dict
self.params['particle_attributes'] = {}
# The keys in the particle_attributes dict are the
# particle types, and the values are new dicts,
# containing the information for each type.
self.params['particle_attributes'][particle_type] = {}
particle_attribute = (self.params['particle_attributes']
[particle_type])
N = header['Npart'][1]
particle_attribute['N'] = N
unit = 1e+10*units.m_sun/header['HubbleParam']
mass = header['Massarr'][1]*unit
particle_attribute['mass'] = mass
particle_attribute['species'] = particle_species
# Done loading particle attributes
if only_params:
self.contains = 'params'
return
# Write out progress message
masterprint('Loading', N, particle_type, '({}) ...'.format(particle_species), indent=4)
# Compute a fair distribution
# of particle data to the processes.
N_locals = ((N//nprocs, )*(nprocs - (N % nprocs))
+ (N//nprocs + 1, )*(N % nprocs))
N_local = N_locals[rank]
start_local = np.sum(N_locals[:rank], dtype=C2np['Py_ssize_t'])
# Construct a Particles instance
self.particles = construct(particle_type, particle_species, mass=mass, N=N)
# Read in the POS block. The positions are given in kpc/h.
offset = self.new_block(f, offset)
unit = units.kpc/header['HubbleParam']
name = self.read(f, '4s').decode('utf8') # "POS "
size = self.read(f, 'i')
offset = self.new_block(f, offset)
f.seek(12*start_local, 1) # 12 = sizeof(float)*Ndims
file_position = f.tell()
self.particles.populate(asarray(np.fromfile(f, dtype=C2np['float'], count=3*N_local)
[0::3], dtype=C2np['double'])*unit,
'posx')
f.seek(file_position)
self.particles.populate(asarray(np.fromfile(f, dtype=C2np['float'], count=3*N_local)
[1::3], dtype=C2np['double'])*unit,
'posy')
f.seek(file_position)
self.particles.populate(asarray(np.fromfile(f, dtype=C2np['float'], count=3*N_local)
[2::3], dtype=C2np['double'])*unit,
'posz')
# Read in the VEL block. The velocities are peculiar
# velocities u=a*dx/dt divided by sqrt(a), given in km/s.
offset = self.new_block(f, offset)
unit = units.km/units.s*mass*header['Time']**1.5
name = self.read(f, '4s').decode('utf8') # "VEL "
size = self.read(f, 'i')
offset = self.new_block(f, offset)
f.seek(12*start_local, 1) # 12 = sizeof(float)*Ndims
file_position = f.tell()
self.particles.populate(asarray(np.fromfile(f, dtype=C2np['float'], count=3*N_local)
[0::3], dtype=C2np['double'])*unit,
'momx')
f.seek(file_position)
self.particles.populate(asarray(np.fromfile(f, dtype=C2np['float'], count=3*N_local)
[1::3], dtype=C2np['double'])*unit,
'momy')
f.seek(file_position)
self.particles.populate(asarray(np.fromfile(f, dtype=C2np['float'], count=3*N_local)
[2::3], dtype=C2np['double'])*unit,
'momz')
# Read in the ID block.
# The ID's will be distributed among all processes.
offset = self.new_block(f, offset)
name = self.read(f, '4s').decode('utf8') # "ID "
size = self.read(f, 'i')
offset = self.new_block(f, offset)
f.seek(4*start_local, 1) # 4 = sizeof(unsigned int)
file_position = f.tell()
self.ID = np.fromfile(f, dtype=C2np['unsigned int'], count=N_local)
# Finalize progress message
masterprint('done')
# Possible additional meta data ignored
# Update the "contains" string
self.contains = 'params and particles'
# Scatter particles to the correct domain-specific process.
# Setting reset_indices_send == True ensures that buffers
# will be reset afterwards, as this initial exchange is not
# representable for those to come.
exchange(self.particles, reset_buffers=True)
ids, active, XY = IO.load_grid_of_particles(rank, time = 0, input = True)
if rank == 0:
Nparticles = np.sum(active)
# start at initial time
t = t_0
IO.save_grid_of_particles(ids, active, XY, t, rank)
#plot.plot(rank, XY[:,active], t, dt)
# main loop
print('\nstart time = %s' % t)
# Communicating, to avoid bad load balancing at the beginning
print('This is rank %s, communicating before mainloop' % (rank))
ids, XY, active = communication.exchange(comm, mpi_size, rank, ids, XY[0,:], XY[1,:], active)
while t + dt <= t_max:
# Take Ndt timesteps
# only take active particles into transport function
# non-active particles should be in the end of the arrays,
# so if we want we can use the index as XY[:,:np.sum(active)] instead
print('This is rank %s, transporting %s particles' % (rank, np.sum(active)))
sys.stdout.flush()
XY[:,active], t = transport.transport(XY[:,active], active, t, Ndt, dt)
#plot.plot(rank, XY[:,active], t, dt, name = '_after_transport-before_comm_')
#t += dt # this increment is returned from transport-funcion
############
# Then communicate
'''
