67 lines
2.5 KiB
Python
67 lines
2.5 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", "--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 3 / trajectory generation ###
|
|
################################################################################
|
|
# 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
|
|
import matplotlib.pyplot as plt
|
|
from mpl_toolkits import mplot3d
|
|
from molecular_dynamics import iter_load, energy, kinetic
|
|
|
|
if config.animate:
|
|
# new plot with 3D axis
|
|
fig = plt.figure()
|
|
ax = fig.add_subplot(111, projection = "3d")
|
|
|
|
count = 0
|
|
for (position, velocity), box_size in iter_load(config.input):
|
|
count += 1
|
|
e_pot = energy(position, box_size)
|
|
e_kin = kinetic(velocity)
|
|
print(f"{count: >3} - {e_pot:10.5e} {e_kin:10.5e} {(e_pot + e_kin):10.5e}")
|
|
|
|
if config.animate:
|
|
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)
|