158 lines
6.0 KiB
Python
158 lines
6.0 KiB
Python
#!/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!)")
|
|
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)
|
|
|
|
################################################################################
|
|
### 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
|
|
|
|
|
|
# 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)
|
|
|
|
# second iteration to estimate pairwise distance CDF
|
|
nr_time_points = 0
|
|
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
|
|
# Count number of actually considured states / time points
|
|
nr_time_points += 1
|
|
# aggregate CDF sum
|
|
CDF += loop_body(position, box_size)
|
|
# final division by timesteps as mean over pairs and time.
|
|
CDF /= nr_time_points
|
|
|
|
# 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])
|