add: file dump subsampling (made live better),

fix: task 04 loop body speedup (JIT)
This commit is contained in:
Daniel Kapla 2022-04-19 16:11:49 +02:00
vanhempi afe7224dc6
commit f47c3609e7
2 muutettua tiedostoa jossa 32 lisäystä ja 15 poistoa

Näytä tiedosto

@ -19,6 +19,9 @@ arg_parser.add_option("-N", "--iterations", action = "store", type = int,
arg_parser.add_option("-o", "--output", action = "store", type = str,
dest = "output", default = "task03.xyz",
help = "output file path (default: 'task03.xyz')")
arg_parser.add_option("-s", "--subset", action = "store", type = int,
dest = "subset", default = 100,
help = "only stores every subset state to output file (default: 100)")
arg_parser.add_option("-v", action = "store_true",
dest = "verbose", default = False,
help = "turn verbosity mode on (default: off a.k.a. silent)")
@ -61,17 +64,25 @@ acceleration = force(position, box_size) / mass
# iterate time steps for system time evolution
with open(config.output, "w") as output:
# store initial state to file
dump(output, position, velocity, box_size)
dump(output, position, velocity, box_size,
comment = f"Initial State with step size {config.step_size}")
# perform iterations many time steps
for iteration in range(1, config.iterations):
# perform a single time step
position, velocity, acceleration = step(position, velocity, acceleration,
box_size, config.step_size)
# dump current state to file
dump(output, position, velocity, box_size)
if (iteration % config.subset) == 0:
dump(output, position, velocity, box_size,
comment = f"Iteration: {iteration}")
# if verbosity turned on, print a progress bar
if config.verbose:
print_prog_bar(iteration, config.iterations)
# In case of final iteration NOT writen to file yet, write also final state
if (iteration % config.subset) != 0:
dump(output, position, velocity, box_size,
comment = f"Iteration: {iteration}")
# report output file path
print(f"Done: trajectories saved to '{config.output}'")

Näytä tiedosto

@ -17,8 +17,8 @@ arg_parser.add_option("-B", "--outputB", action = "store", type = str,
dest = "outputB", default = "task04B.txt",
help = "second output file path (default: 'task04B.txt')")
arg_parser.add_option("-s", "--cdf_sample_count", action = "store", type = int,
dest = "cdf_sample_count", default = 5,
help = "Radius interval sample count for CDF estimate (default: 5)")
dest = "cdf_sample_count", default = 100,
help = "Radius interval sample count for CDF estimate (default: 100)")
arg_parser.add_option("-a", "--animate", action = "store_true",
dest = "animate", default = False,
help = "creates an trajectory animation of the states in the input file")
@ -45,6 +45,7 @@ else:
# note, load module _after_ processing script parameters (no need to load all
# of the heavy numeric modules if only usage or alike is needed)
import numpy as np
from jax import jit, numpy as jnp
from molecular_dynamics import iter_load, energy, kinetic
@ -73,6 +74,20 @@ lower_tri = np.tril_indices(position.shape[0], k = -1)
# maximum distance two particles (in folded space) can be apart
max_dist = box_size / 2.0
@jit
def loop_body(position, box_size):
# magic for all particle pair distances
diff = position[:, jnp.newaxis, :] - position
diff = diff[lower_tri[0], lower_tri[1], :]
# fold space in all directions. The `max_dist` is the maximum distance
# in the folded space untill the distance through the folded space
# border is smaller than inside the box. 3-axis parallel fold
diff = jnp.mod(diff + max_dist, box_size) - max_dist
# count nr particles pairs with distance smaller than r
dist = jnp.linalg.norm(diff, axis = 1)
# accumulate per timepoint partial mean estimates of the CDF
return jnp.mean(dist[:, jnp.newaxis] < radii, axis = 0)
# second iteration to estimate pairwise distance CDF
nr_time_points = 0
for i, ((position, velocity), box_size) in enumerate(iter_load(config.input)):
@ -81,17 +96,8 @@ for i, ((position, velocity), box_size) in enumerate(iter_load(config.input)):
continue
# Count number of actually considured states / time points
nr_time_points += 1
# magic for all particle pair distances
diff = position[:, np.newaxis, :] - position
diff = diff[lower_tri[0], lower_tri[1], :]
# fold space in all directions. The `max_dist` is the maximum distance
# in the folded space untill the distance through the folded space
# border is smaller than inside the box. 3-axis parallel fold
diff = np.mod(diff + max_dist, box_size) - max_dist
# count nr particles pairs with distance smaller than r
dist = np.linalg.norm(diff, axis = 1)
# accumulate per timepoint partial mean estimates of the CDF
CDF += np.mean(dist[:, np.newaxis] < radii, axis = 0)
# aggregate CDF sum
CDF += loop_body(position, box_size)
# final division by timesteps as mean over pairs and time.
CDF /= nr_time_points