add: file dump subsampling (made live better),

fix: task 04 loop body speedup (JIT)

Exercise_02/task03.py

@@ -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}'")

Exercise_02/task04.py

@@ -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,
+    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",
     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
+    # 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.
 CDF /= nr_time_points