NSSC/Exercise_02/task04.py

158 lines
6.0 KiB
Python
Raw Normal View History

#!/usr/bin/env python
################################################################################
### parse script parameters ###
################################################################################
from optparse import OptionParser
usage = "usage: %prog [options] [<INPUT>]"
arg_parser = OptionParser(usage = usage)
arg_parser.add_option("-i", "--input", action = "store", type = str,
dest = "input", default = None,
help = "input file path (required!)")
2022-04-12 16:04:26 +00:00
arg_parser.add_option("-A", "--outputA", action = "store", type = str,
dest = "outputA", default = "task04A.txt",
help = "first output file path (default: 'task04A.txt')")
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 = 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")
# Parse command line arguments (as def. above) or store defaults to `config`
config, args = arg_parser.parse_args()
# overwrite options with positional arguments if supplied
try:
if len(args) > 0:
config.input = args[0]
except ValueError as expression:
arg_parser.print_help()
print(f"Error: {expression}")
exit(-1)
else:
# quick and dirty validation
if config.input is None:
arg_parser.print_help()
print("Error: missing or illegal argument")
exit(-1)
################################################################################
2022-04-12 16:04:26 +00:00
### task 4 / post processing ###
################################################################################
# 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
2022-04-12 16:04:26 +00:00
# counte the number of states (system snapshots) in the input file
nr_states = 0
# first output file and collect some info used in the next analysis!
with open(config.outputA, "w") as output:
# iterate over all state snapshots stored in the `config.input` file
for (position, velocity), box_size in iter_load(config.input):
nr_states += 1
# compute potneital and kinetic energy for current snapshot
e_pot = energy(position, box_size)
e_kin = kinetic(velocity)
# energy results
print(e_pot, e_kin, file = output)
# report file path
print(f"Saved energy change to '{config.outputA}'")
# CDF radius sample points (`box_size` known from last loop)
radii = np.linspace(0, box_size / 2.0, config.cdf_sample_count)[np.newaxis, :]
# radius Commulative Distribution Function estimate
CDF = np.zeros((config.cdf_sample_count, ))
# extract only one of two distance combinations of non-equal particles
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)
2022-04-12 16:04:26 +00:00
# second iteration to estimate pairwise distance CDF
2022-04-19 12:20:34 +00:00
nr_time_points = 0
2022-04-12 16:04:26 +00:00
for i, ((position, velocity), box_size) in enumerate(iter_load(config.input)):
# ignore the first 25% (just keep it simple)
if 4 * i < nr_states:
continue
2022-04-19 12:20:34 +00:00
# Count number of actually considured states / time points
nr_time_points += 1
# aggregate CDF sum
CDF += loop_body(position, box_size)
2022-04-12 16:04:26 +00:00
# final division by timesteps as mean over pairs and time.
2022-04-19 12:20:34 +00:00
CDF /= nr_time_points
2022-04-12 16:04:26 +00:00
# store the CDF estimate to the second file
with open(config.outputB, "w") as output:
for r, cdf in zip(radii.ravel(), CDF.ravel()):
print(r, cdf, file = output)
# report file path
print(f"Saved radius CDF estimate to '{config.outputB}'")
# import matplotlib.pyplot as plt
# from mpl_toolkits import mplot3d
# if config.animate:
# # new plot with 3D axis
# count = 0
# fig = plt.figure()
# ax = fig.add_subplot(111, projection = "3d")
# for (position, velocity), box_size in iter_load(config.input):
# e_pot = energy(position, box_size)
# e_kin = kinetic(velocity)
# print(f"{e_pot:10.5e} {e_kin:10.5e} {(e_pot + e_kin):10.5e}")
# if config.animate:
# count += 1
# ax.cla()
# plt.title(f"time step {count}")
# plt.xlim([0, box_size])
# plt.ylim([0, box_size])
# ax.set_zlim([0, box_size])
# # create 3D position scatter plot
# position = np.mod(position, box_size)
# ax.scatter(position[:, 0], position[:, 1], position[:, 2])
# # # and save to file
# # plt.savefig(config.plot)
# plt.pause(0.01)
# import numpy as np
# import matplotlib.pyplot as plt
# from mpl_toolkits import mplot3d
# from molecular_dynamics import load, energy, kinetic
# (position, velocity), box_size = load("task02.xyz"):
# fig = plt.figure()
# ax = fig.add_subplot(111, projection = "3d")
# ax.cla()
# plt.xlim([0, box_size])
# plt.ylim([0, box_size])
# ax.set_zlim([0, box_size])
# # create 3D position scatter plot
# position = np.mod(position, box_size)
# ax.scatter(position[:, 0], position[:, 1], position[:, 2])