def compute_intraspecies_tau(distribution, dHdt):
''' compute the single species tau_kk
Parameters
----------
distribution : 3D distribution object
the distribution function for this species
dHdt : float
average dHdt for this species acting on itself
Returns
-------
tau : float
single species relaxation parameter tau_kk
f_eq : distribution
equilibrium distribution relaxing towards
'''
# extract information
m = distribution.mass
n = distribution.density
u = distribution.momentum / m
KE = distribution.kinetic_energy
vx = distribution._x
vy = distribution._y
vz = distribution._z
f = distribution.distribution
f_eq = distributions.discrete_maxwellian3D(vx, vy, vz, m, n, u, KE)
# integrate and compute tau
tau = triple_integral((f_eq - f) * np.log(f), vx, vy, vz) / dHdt
return tau, f_eq
def initialize(self):
''' initialize all the distributions and macro properties that we
need to track throughout the hmm
'''
# initialize the distributions
logging.debug('initializing discrete distributions')
self.distribution = np.empty((self.n_cells_bgk, self.n_species),
dtype=object)
for cell in range(self.n_cells_bgk):
# compute mixture temperature for sufficiently wide velocity grid
mixture_temperature = (np.sum(self.initial_temperature[cell, :] *
self.initial_density[cell, :]) /
self.initial_density[cell, :].sum())
for sp in range(self.n_species):
# set grid based on maximum of mixture and species temperature
T_use = max(mixture_temperature, self.initial_temperature[cell,
sp])
v_thermal = np.sqrt(T_use / self.mass[sp])
vx = vy = vz = np.linspace(-8. * v_thermal, 8. * v_thermal,
self.n_vel)
n = self.initial_density[cell, sp]
m = self.mass[sp]
u = np.array([self.initial_velocity[cell, sp], 0., 0.])
ke = self.initial_temperature[cell, sp] * 3. / 2.
f = distributions.discrete_maxwellian3D(
vx, vy, vz, m, n, u, ke)
self.distribution[cell,sp] = \
distributions.linear_interpolated_rv_3D(vx, vy, vz, f, m)
# initialize taus and errors
logging.debug('intializing other data variables')
if self.only_md:
n_bgks = 0
else:
n_bgks = int(
np.rint(self.final_time /
(self.timestep_bgk * self.tau_update_rate)))
self.tau_hist = np.empty(
(n_bgks, self.n_cells_bgk, self.n_species, self.n_species))
self.tau_error_hist = np.empty(
(n_bgks, self.n_cells_bgk, self.n_species, self.n_species))
self.taus = np.empty(
(self.n_cells_bgk, self.n_species, self.n_species))
# initialize macro properties
self.density = np.empty(self.initial_density.shape)
self.velocity = np.empty(self.initial_velocity.shape)
self.kinetic_energy = np.empty(self.initial_temperature.shape)
for cell in range(self.n_cells_bgk):
for sp in range(self.n_species):
self.density[cell, sp] = self.distribution[cell, sp].density
self.velocity[cell,sp] = \
self.distribution[cell,sp].momentum[0] / self.mass[sp]
self.kinetic_energy[cell,sp] = \
self.distribution[cell,sp].kinetic_energy
# time tracking stuff
self.current_time = 0.0
self.bgk_counter = 0
def simulate_md(params,
distribution,
md,
print_rate=10,
resample=True,
atol=[1e-2, 1.5e-2, 2e-2],
rtol=[1e-2, 1.5e-2, 2e-2, 2e-2],
refresh_rate=0,
save_rate=1000,
write_cross_section=False,
write_distribution=False,
resume=False,
last_step=0,
current_sim=0):
''' run the specified number of md simulations for the desired number of
time steps and collect data on energy, dHdt, and moments
Parameters
----------
params : md_parameters object
parameters of the simulation
distribution : array of distribution objects
the distributions to sample from
md : md module
md simulation object
print_rate : int
how often to print the timestep
resample : boolean
whether to do velocity resampling before each trial
atol : array of 3 floats
absolute tolerance for resampling
rtol : array of 3 floats
relative tolerance for resampling
refresh_rate : int
how often to refresh the distribution based on average temperature
0 -> never
save_rate : int
how often to save the phasespace and data for restart
0 -> never
write_cross_section : boolean
whether to write distribution cross-section every time step
write_distribution : int
whether to write distribution (and how often)
resume : boolean
whether resuming from a previous simulation
last_step : int
step to resume from
0 -> not resuming
current_sim : int
current simulation (for adding sims)
Returns
-------
energy : n_timesteps+1 x 3 array of floats
the total, potential, and kinetic energies throughout equilibration
data : md_data object
all the moments and dHdt from throughout the simulation
'''
# setup
energy = np.empty((params.n_sims, params.n_timesteps + 1, 3))
data = md_data(params, write_cross_section, distribution)
md.parameters_mod.thermostaton = False
pos0, vel0 = get_md_phasespace(md)
prev_r = np.empty((3, params.n_particles))
prev_v = np.empty((3, params.n_particles))
prev_a = np.empty((3, params.n_particles))
if refresh_rate > 0:
refresh_rate = int(refresh_rate)
distribution_log = np.empty(
(params.n_sims, params.n_timesteps / refresh_rate + 1,
params.n_species),
dtype=object)
distribution_log[:,0,:] = \
np.array(distribution)[np.newaxis,np.newaxis,:]
log_counter = 1
if write_distribution > 0:
distribution_file = np.empty(params.n_species, dtype=file)
for sp in range(params.n_species):
distribution_file[sp] = open('distribution' + str(sp), 'wb')
np.save('distribution' + str(sp) + '_grid', distribution[sp]._x)
distribution_counter = 0
v_hist = np.empty((write_distribution, params.n_particles, 3))
# loop over simulations
for sim in range(current_sim, params.n_sims):
# logging.info('---------------------------------------------------')
# logging.info('SIMULATION %d of %d' % (sim+1, params.n_sims))
# logging.info('---------------------------------------------------')
# velocity resample if needed
if resample and not resume:
vel = velocity_resample(distribution, params, atol, rtol)
set_md_phasespace(pos0, vel, md)
# intialize simulation
md.forces()
md.conservation()
md.movie()
md.naughtypair()
if current_sim > 0:
logging.info('reloading simulation data')
energy[:-1] = np.load('energy.npy')
data.dHdt[:-1] = np.load('data.dHdt.npy')
data.momentum[:-1] = np.load('data.momentum.npy')
data.stress[:-1] = np.load('data.stress.npy')
data.kinetic_energy[:-1] = np.load('data.kinetic_energy.npy')
data.heat[:-1] = np.load('data.heat.npy')
data.m4[:-1] = np.load('data.m4.npy')
if not resume:
energy[sim, 0, :] = get_md_conservation(md)
data.get_md_data(sim, 0, params, md, distribution)
else:
logging.info("reinitializing previous simulation data")
energy[sim,:last_step+1,:] = \
np.load('energy.npy')[sim,:last_step+1,:]
data.dHdt[sim,:last_step+1,:,:] = \
np.load('data.dHdt.npy')[sim,:last_step+1,:,:]
data.momentum[sim,:last_step+1,:,:] = \
np.load('data.momentum.npy')[sim,:last_step+1,:,:]
data.stress[sim,:last_step+1,:,:,:] = \
np.load('data.stress.npy')[sim,:last_step+1,:,:,:]
data.kinetic_energy[sim,:last_step+1,:] = \
np.load('data.kinetic_energy.npy')[sim,:last_step+1,:]
data.heat[sim,:last_step+1,:,:] = \
np.load('data.heat.npy')[sim,:last_step+1,:,:]
data.m4[sim,:last_step+1,:] = \
np.load('data.m4.npy')[sim,:last_step+1,:]
# loop over time
for step in range(last_step + 1, params.n_timesteps + 1):
if step % print_rate == 0:
logging.info(
'timestep %d of %d for simulation %d of %d' %
(step, params.n_timesteps, sim + 1, params.n_sims))
logging.info('total energy jump so far: %f',
(energy[sim, step - 1, 0] - energy[sim, 0, 0]) /
energy[sim, 0, 0])
prev_r[0, :] = md.particles_mod.rx[:]
prev_r[1, :] = md.particles_mod.ry[:]
prev_r[2, :] = md.particles_mod.rz[:]
prev_v[0, :] = md.particles_mod.vx[:]
prev_v[1, :] = md.particles_mod.vy[:]
prev_v[2, :] = md.particles_mod.vz[:]
prev_a[0, :] = md.particles_mod.ax[:]
prev_a[1, :] = md.particles_mod.ay[:]
prev_a[2, :] = md.particles_mod.az[:]
# distribution output if needed
if write_distribution > 0:
v_hist[distribution_counter, :, 0] = md.particles_mod.vx[:]
v_hist[distribution_counter, :, 1] = md.particles_mod.vy[:]
v_hist[distribution_counter, :, 2] = md.particles_mod.vz[:]
distribution_counter += 1
if distribution_counter == write_distribution:
distribution_counter = 0
dx = distribution[sp]._x[1] - distribution[sp]._x[0]
dy = distribution[sp]._y[1] - distribution[sp]._y[0]
dz = distribution[sp]._z[1] - distribution[sp]._z[0]
x_grid = np.append(distribution[sp]._x,
distribution[sp]._x[-1] + dx)
y_grid = np.append(distribution[sp]._y,
distribution[sp]._y[-1] + dy)
z_grid = np.append(distribution[sp]._z,
distribution[sp]._z[-1] + dz)
sp_end = np.cumsum(params.particles)
sp_start = sp_end - params.particles
logging.debug(
'saving distribution from particles at step %d' % step)
for sp in range(params.n_species):
sp_vel = v_hist[:, sp_start[sp]:sp_end[sp], :].reshape(
params.particles[sp] * write_distribution, 3)
sp_distribution, _ = np.histogramdd(sp_vel,
bins=(x_grid,
y_grid,
z_grid))
integral = np.trapz(
np.trapz(
np.trapz(sp_distribution, distribution[sp]._x),
distribution[sp]._y), distribution[sp]._z)
sp_distribution *= params.density[sp] / integral
np.save(distribution_file[sp], sp_distribution)
md.evl()
if md.parameters_mod.outofbounds:
raise ValueError('particle flew out of the domain')
md.conservation()
md.naughtypair()
energy[sim, step, :] = get_md_conservation(md)
# MTS
delta = (energy[sim, step, 0] -
energy[sim, step - 1, 0]) / energy[sim, 0, 0]
if abs(delta) > params.mts_threshold:
logging.debug('taking MTS step at step %d' % step)
logging.debug('energy jump is %f' % delta)
totjump = ((energy[sim, step - 1, 0] - energy[sim, 0, 0]) /
energy[sim, 0, 0])
logging.debug('total energy change before this step was %f' %
totjump)
logging.debug('max acceleration is: %f' %
md.particles_mod.maxaij)
md.particles_mod.rx = prev_r[0, :]
md.particles_mod.ry = prev_r[1, :]
md.particles_mod.rz = prev_r[2, :]
md.particles_mod.vx = prev_v[0, :]
md.particles_mod.vy = prev_v[1, :]
md.particles_mod.vz = prev_v[2, :]
md.particles_mod.ax = prev_a[0, :]
md.particles_mod.ay = prev_a[1, :]
md.particles_mod.az = prev_a[2, :]
md.evlmts()
if md.parameters_mod.outofbounds:
raise ValueError('particle flew out of the domain')
logging.debug('Number of MTS particles: %d' %
np.sum(md.particles_mod.ismts))
md.conservation()
energy[sim, step, :] = get_md_conservation(md)
if step % params.movie_rate == 0:
md.movie()
data.get_md_data(sim, step, params, md, distribution)
# save data every save_rate timesteps
if (save_rate > 0
and step % save_rate == 0) or step == params.n_timesteps:
logging.debug('saving simulation data at time step %d' % step)
pos, vel = get_md_phasespace(md)
# write output to files
np.save('energy', energy)
np.save('data.momentum', data.momentum)
np.save('data.stress', data.stress)
np.save('data.kinetic_energy', data.kinetic_energy)
np.save('data.heat', data.heat)
np.save('data.m4', data.m4)
np.save('data.dHdt', data.dHdt)
np.save('data.mass', data.mass)
np.save('data.time', data.time)
np.save('end_pos', pos)
np.save('end_vel', vel)
if write_cross_section:
np.save('md_distribution', data.vx_histo)
np.save('histo_grid',
np.stack([dist._x for dist in distribution]))
# update distributions if needed
if refresh_rate > 0 and step % refresh_rate == 0:
logging.debug('refreshing the distributions')
average_ke = data.kinetic_energy[sim,
step - refresh_rate:step, :]
average_ke = np.squeeze(average_ke).mean(axis=0)
for sp in range(params.n_species):
vx = distribution[sp]._x
vy = distribution[sp]._y
vz = distribution[sp]._z
f = distributions.discrete_maxwellian3D(
vx, vy, vz, params.mass[sp], params.density[sp],
np.zeros((3, )), average_ke[sp])
distribution[sp] = distributions.linear_interpolated_rv_3D(
vx, vy, vz, f, params.mass[sp])
distribution_log[sim, log_counter, :] = distribution
log_counter += 1
# do another equilibration bewteen simulations
if sim < params.n_sims - 1:
md.closefiles()
md.openfiles()
logging.debug('equilibrating between MD simulations')
set_md_phasespace(pos0, vel0, md)
equilibrate_md(params, md, print_rate=200, save_rate=save_rate)
md.parameters_mod.thermostaton = False
md.closefiles()
md.openfiles()
if write_distribution > 0:
for sp in range(params.n_species):
distribution_file[sp].close()
if refresh_rate > 0:
return energy, data, distribution_log
else:
return energy, data, pos0, vel0
def simulate_md(params,
distribution,
md,
print_rate=10,
resample=True,
atol=[1e-2, 1.5e-2, 2e-2],
rtol=[1e-2, 1e-2, 2e-2, 2e-2],
refresh_rate=0,
save_rate=1000):
''' run the specified number of md simulations for the desired number of
time steps and collect data on energy, dHdt, and moments
Parameters
----------
params : md_parameters object
parameters of the simulation
distribution : array of distribution objects
the distributions to sample from
md : md module
md simulation object
print_rate : int
how often to print the timestep
resample : boolean
whether to do velocity resampling before each trial
atol : array of 3 floats
absolute tolerance for resampling
rtol : array of 3 floats
relative tolerance for resampling
refresh_rate : int
how often to refresh the distribution based on average temperature
0 -> never
save_rate : int
how often to save the phasespace
0 -> never
Returns
-------
energy : n_timesteps+1 x 3 array of floats
the total, potential, and kinetic energies throughout equilibration
data : md_data object
all the moments and dHdt from throughout the simulation
'''
# setup
energy = np.empty((params.n_sims, params.n_timesteps + 1, 3))
data = md_data(params)
md.parameters_mod.thermostaton = False
prev_r = np.empty((3, params.n_particles))
prev_v = np.empty((3, params.n_particles))
prev_a = np.empty((3, params.n_particles))
pos0, vel0 = get_md_phasespace(md)
if refresh_rate > 0:
refresh_rate = int(refresh_rate)
distribution_log = np.empty(
(params.n_sims, params.n_timesteps / refresh_rate + 1,
params.n_species),
dtype=object)
distribution_log[:,0,:] = \
np.array(distribution)[np.newaxis,np.newaxis,:]
log_counter = 1
# loop over simulations
for sim in range(params.n_sims):
logging.info('---------------------------------------------------')
logging.info('SIMULATION %d of %d' % (sim + 1, params.n_sims))
logging.info('---------------------------------------------------')
# velocity resample if needed
if resample:
vel = velocity_resample(distribution, params, atol, rtol)
set_md_phasespace(pos0, vel, md)
# intialize simulation
md.forces()
md.conservation()
md.movie()
md.naughtypair()
energy[sim, 0, :] = get_md_conservation(md)
data.get_md_data(sim, 0, params, md, distribution)
# loop over time
for step in range(1, params.n_timesteps + 1):
if step % print_rate == 0:
logging.info(
'timestep %d of %d for simulation %d of %d' %
(step, params.n_timesteps, sim + 1, params.n_sims))
logging.info('total energy jump so far: %f',
(energy[sim, step - 1, 0] - energy[sim, 0, 0]) /
energy[sim, 0, 0])
prev_r[0, :] = md.particles_mod.rx[:]
prev_r[1, :] = md.particles_mod.ry[:]
prev_r[2, :] = md.particles_mod.rz[:]
prev_v[0, :] = md.particles_mod.vx[:]
prev_v[1, :] = md.particles_mod.vy[:]
prev_v[2, :] = md.particles_mod.vz[:]
prev_a[0, :] = md.particles_mod.ax[:]
prev_a[1, :] = md.particles_mod.ay[:]
prev_a[2, :] = md.particles_mod.az[:]
md.evl()
if md.parameters_mod.outofbounds:
raise ValueError('particle flew out of the domain')
md.conservation()
md.naughtypair()
energy[sim, step, :] = get_md_conservation(md)
# MTS
delta = (energy[sim, step, 0] -
energy[sim, step - 1, 0]) / energy[sim, 0, 0]
if abs(delta) > params.mts_threshold:
logging.debug('taking MTS step at step %d' % step)
logging.debug('energy jump is %f' % delta)
totjump = ((energy[sim, step - 1, 0] - energy[sim, 0, 0]) /
energy[sim, 0, 0])
logging.debug('total energy change before this step was %f' %
totjump)
logging.debug('max acceleration is: %f' %
md.particles_mod.maxaij)
md.particles_mod.rx = prev_r[0, :]
md.particles_mod.ry = prev_r[1, :]
md.particles_mod.rz = prev_r[2, :]
md.particles_mod.vx = prev_v[0, :]
md.particles_mod.vy = prev_v[1, :]
md.particles_mod.vz = prev_v[2, :]
md.particles_mod.ax = prev_a[0, :]
md.particles_mod.ay = prev_a[1, :]
md.particles_mod.az = prev_a[2, :]
md.evlmts()
if md.parameters_mod.outofbounds:
raise ValueError('particle flew out of the domain')
logging.debug('Number of MTS particles: %d' %
np.sum(md.particles_mod.ismts))
md.conservation()
energy[sim, step, :] = get_md_conservation(md)
if step % params.movie_rate == 0:
md.movie()
data.get_md_data(sim, step, params, md, distribution)
# save phasespace every 1000 timesteps
if save_rate > 0 and step % save_rate == 0:
logging.debug('saving simulation phase space at time step %d' %
step)
pos, vel = get_md_phasespace(md)
species = np.array(md.particles_mod.sp[:])
species = species.reshape((params.n_particles, 1))
specposvel = np.concatenate((species, pos, vel), axis=1)
filename = 'phasespace_sim%d_step%d.dat' % (sim + 1, step)
specposvel.tofile(filename)
# update distributions if needed
if refresh_rate > 0 and step % refresh_rate == 0:
logging.debug('refreshing the distributions')
average_ke = data.kinetic_energy[sim,
step - refresh_rate:step, :]
average_ke = np.squeeze(average_ke).mean(axis=0)
for sp in range(params.n_species):
vx = distribution[sp]._x
vy = distribution[sp]._y
vz = distribution[sp]._z
f = distributions.discrete_maxwellian3D(
vx, vy, vz, params.mass[sp], params.density[sp],
np.zeros((3, )), average_ke[sp])
distribution[sp] = distributions.linear_interpolated_rv_3D(
vx, vy, vz, f, params.mass[sp])
distribution_log[sim, log_counter, :] = distribution
log_counter += 1
if refresh_rate > 0:
return energy, data, distribution_log
else:
return energy, data