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
parent afe7224dc6
commit f47c3609e7
2 changed files with 32 additions and 15 deletions

View File

@ -19,6 +19,9 @@ arg_parser.add_option("-N", "--iterations", action = "store", type = int,
arg_parser.add_option("-o", "--output", action = "store", type = str, arg_parser.add_option("-o", "--output", action = "store", type = str,
dest = "output", default = "task03.xyz", dest = "output", default = "task03.xyz",
help = "output file path (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", arg_parser.add_option("-v", action = "store_true",
dest = "verbose", default = False, dest = "verbose", default = False,
help = "turn verbosity mode on (default: off a.k.a. silent)") 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 # iterate time steps for system time evolution
with open(config.output, "w") as output: with open(config.output, "w") as output:
# store initial state to file # 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 # perform iterations many time steps
for iteration in range(1, config.iterations): for iteration in range(1, config.iterations):
# perform a single time step # perform a single time step
position, velocity, acceleration = step(position, velocity, acceleration, position, velocity, acceleration = step(position, velocity, acceleration,
box_size, config.step_size) box_size, config.step_size)
# dump current state to file # 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 verbosity turned on, print a progress bar
if config.verbose: if config.verbose:
print_prog_bar(iteration, config.iterations) 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 # report output file path
print(f"Done: trajectories saved to '{config.output}'") print(f"Done: trajectories saved to '{config.output}'")

View File

@ -17,8 +17,8 @@ arg_parser.add_option("-B", "--outputB", action = "store", type = str,
dest = "outputB", default = "task04B.txt", dest = "outputB", default = "task04B.txt",
help = "second output file path (default: 'task04B.txt')") help = "second output file path (default: 'task04B.txt')")
arg_parser.add_option("-s", "--cdf_sample_count", action = "store", type = int, arg_parser.add_option("-s", "--cdf_sample_count", action = "store", type = int,
dest = "cdf_sample_count", default = 5, dest = "cdf_sample_count", default = 100,
help = "Radius interval sample count for CDF estimate (default: 5)") help = "Radius interval sample count for CDF estimate (default: 100)")
arg_parser.add_option("-a", "--animate", action = "store_true", arg_parser.add_option("-a", "--animate", action = "store_true",
dest = "animate", default = False, dest = "animate", default = False,
help = "creates an trajectory animation of the states in the input file") 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 # note, load module _after_ processing script parameters (no need to load all
# of the heavy numeric modules if only usage or alike is needed) # of the heavy numeric modules if only usage or alike is needed)
import numpy as np import numpy as np
from jax import jit, numpy as jnp
from molecular_dynamics import iter_load, energy, kinetic 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 # maximum distance two particles (in folded space) can be apart
max_dist = box_size / 2.0 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 # second iteration to estimate pairwise distance CDF
nr_time_points = 0 nr_time_points = 0
for i, ((position, velocity), box_size) in enumerate(iter_load(config.input)): 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 continue
# Count number of actually considured states / time points # Count number of actually considured states / time points
nr_time_points += 1 nr_time_points += 1
# magic for all particle pair distances # aggregate CDF sum
diff = position[:, np.newaxis, :] - position CDF += loop_body(position, box_size)
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)
# final division by timesteps as mean over pairs and time. # final division by timesteps as mean over pairs and time.
CDF /= nr_time_points CDF /= nr_time_points