103 lines
4.6 KiB
Python
103 lines
4.6 KiB
Python
#!/usr/bin/env python
|
|
|
|
################################################################################
|
|
### parse script parameters ###
|
|
################################################################################
|
|
from optparse import OptionParser
|
|
|
|
usage = "usage: %prog [options] [<NR_PARTICLES>] [<BOX_SIZE>] [<TEMPERATURE>]"
|
|
arg_parser = OptionParser(usage = usage)
|
|
arg_parser.add_option("-M", "--particles", action = "store", type = int,
|
|
dest = "nr_particles", default = -1,
|
|
help = "Nr. of particles (required!)")
|
|
arg_parser.add_option("-L", "--box_size", action = "store", type = float,
|
|
dest = "box_size", default = -1.0,
|
|
help = "side box_size [A (angstrom)] of the simulation cube (required!)")
|
|
arg_parser.add_option("-T", "--temperature", action = "store", type = float,
|
|
dest = "temperature", default = -1.0,
|
|
help = "temperature [K] to generate initial conditions (required!)")
|
|
arg_parser.add_option("-o", "--output", action = "store", type = str,
|
|
dest = "output", default = "task02.xyz",
|
|
help = "output file path (default: 'task02.xyz')")
|
|
arg_parser.add_option("-v", action = "store_true",
|
|
dest = "verbose", default = False,
|
|
help = "turn verbosity mode on (default: off a.k.a. silent)")
|
|
# 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.nr_particles = int(args[0])
|
|
if len(args) > 1:
|
|
config.box_size = float(args[1])
|
|
if len(args) > 2:
|
|
config.temperature = float(args[2])
|
|
except ValueError as expression:
|
|
arg_parser.print_help()
|
|
print(f"Error: {expression}")
|
|
exit(-1)
|
|
else:
|
|
# quick and dirty validation
|
|
if not config.nr_particles > 0 \
|
|
or not config.box_size > 0.0 \
|
|
or not config.temperature > 0.0:
|
|
arg_parser.print_help()
|
|
print("Error: missing or illegal argument")
|
|
exit(-1)
|
|
|
|
################################################################################
|
|
### task 2 / generation of initial conditions ###
|
|
################################################################################
|
|
# 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 scipy
|
|
from jax import jit, grad
|
|
from molecular_dynamics import dump, energy, force, mass
|
|
|
|
# Sample random positions in a 3D cube (TODO: make this not just uniform :-})
|
|
position = np.random.uniform(0.0, config.box_size, (config.nr_particles, 3))
|
|
|
|
# Sample particle velocities
|
|
K_b = scipy.constants.Boltzmann / scipy.constants.eV # [eV / K]
|
|
sd = np.sqrt(K_b * config.temperature / mass)
|
|
velocity = np.random.normal(0.0, sd, (config.nr_particles, 3))
|
|
# center velocities
|
|
velocity -= velocity.mean(axis = 0)
|
|
|
|
# remember energy before optimizing for a low energy state
|
|
initial_energy = energy(position, config.box_size)
|
|
forces = force(position, config.box_size)
|
|
initial_mean_forces = forces.mean(axis = 0)
|
|
initial_mean_fnorm = np.linalg.norm(forces, axis = 1).mean()
|
|
|
|
# optimize energy to find low energy system state using Conjugate-Gradients
|
|
optim = scipy.optimize.minimize(energy, # objective func.
|
|
jac = jit(grad(energy)), # jacobian
|
|
x0 = position, # initial position
|
|
args = (config.box_size, ), # further args
|
|
method = "CG")
|
|
# extract (and reshape) optimization result
|
|
position = optim.x.reshape((config.nr_particles, 3))
|
|
# ensure all particles are in the box
|
|
position = np.mod(position, config.box_size)
|
|
|
|
# recompute stats after optimizing for low energy state
|
|
final_energy = energy(position, config.box_size)
|
|
forces = force(position, config.box_size)
|
|
final_mean_forces = forces.mean(axis = 0)
|
|
final_mean_fnorm = np.linalg.norm(forces, axis = 1).mean()
|
|
|
|
# store state snapshot to file (default target file defined by script args)
|
|
dump(config.output, position, velocity, config.box_size)
|
|
|
|
# report stats (if requested by `-v` script argument)
|
|
if config.verbose:
|
|
print(f"Initial Energy: {initial_energy:.4e}")
|
|
print( "Initial Mean Forces: {:.4e} {:.4e} {:.4e}".format(*initial_mean_forces))
|
|
print(f"Initial Mean ||Forces||: {initial_mean_fnorm:.4e}")
|
|
print(f"Final Energy: {final_energy:.4e}")
|
|
print( "Final Mean Forces: {:.4e} {:.4e} {:.4e}".format(*final_mean_forces))
|
|
print(f"Final Mean ||Forces||: {final_mean_fnorm:.4e}")
|
|
print(f"Done: saved inital state to '{config.output}'")
|