add: task 3 and 4,
cleanup: moved non task specific code to seperate file
This commit is contained in:
parent
cdf9e3104f
commit
1ff5c11866
|
|
@ -0,0 +1,137 @@
|
|||
#!/usr/bin/env python
|
||||
|
||||
import numpy as np
|
||||
from jax import numpy as jnp, jit, grad
|
||||
|
||||
# toy moleculat dynamic system parameters (constants)
|
||||
De = 1.6 # [eV (electronvolt)]
|
||||
alpha = 3.028 # [A^-1]
|
||||
re = 1.411 # [A]
|
||||
mass = 18.998403 # atomic mass of a single particle
|
||||
|
||||
def dump(file, position, velocity, box_size, mode = "w", comment = ""):
|
||||
"""
|
||||
Dumps `position` and `velocity` to `file` with meta data.
|
||||
|
||||
First writes a header to `file` in `mode` file write mode defines as three
|
||||
lines.
|
||||
`<nr of particles [int]>`
|
||||
`<comment (ignored)>`
|
||||
`<box side lenght [float]>`
|
||||
then each line (nr of particles many) have the form `x y z vx vy vx`.
|
||||
"""
|
||||
header = [str(position.shape[0]), comment, str(box_size)]
|
||||
data = np.hstack((position, velocity))
|
||||
np.savetxt(file, data, comments = "", header = "\n".join(header))
|
||||
|
||||
def load(file):
|
||||
"""
|
||||
Loads a state from `file` and returns the `position` and `velocity` system
|
||||
states as two `numpy` arrays.
|
||||
|
||||
As a side effect the config object as agumented with `file` header info.
|
||||
"""
|
||||
if isinstance(file, str):
|
||||
with open(file) as handle:
|
||||
return load(handle)
|
||||
# read header
|
||||
nr_particles = int(file.readline().strip())
|
||||
file.readline() # ignore comment line
|
||||
box_size = float(file.readline().strip())
|
||||
# read the state data, each line constains `x y z vx vy vz`
|
||||
data = np.loadtxt(file, max_rows = nr_particles)
|
||||
return np.hsplit(data, 2), box_size
|
||||
|
||||
def iter_load(file):
|
||||
"""
|
||||
Iterator yielding similar to `load` entries from `file` in the same format.
|
||||
"""
|
||||
if isinstance(file, str):
|
||||
with open(file) as handle:
|
||||
yield from iter_load(handle)
|
||||
else:
|
||||
while True:
|
||||
# read header
|
||||
line = file.readline().strip()
|
||||
if not line:
|
||||
return
|
||||
nr_particles = int(line)
|
||||
file.readline() # ignore comment line
|
||||
line = file.readline().strip()
|
||||
if not line:
|
||||
return
|
||||
box_size = float(line)
|
||||
# read the state data, each line constains `x y z vx vy vz`
|
||||
data = np.loadtxt(file, max_rows = nr_particles)
|
||||
if not data.shape[0] == nr_particles:
|
||||
return
|
||||
yield np.hsplit(data, 2), box_size
|
||||
|
||||
def print_prog_bar(index, max_index):
|
||||
"""
|
||||
Simple progress bar. Just for convenience.
|
||||
"""
|
||||
if (index % 100) and (max_index != index + 1):
|
||||
return
|
||||
progress = (30 * (index + 1)) // max_index
|
||||
print(f"\r{(index + 1)}/{max_index} - " \
|
||||
f"[{'':=>{(progress)}}{'': >{(30 - progress)}}]",
|
||||
end = "\n" if max_index == index + 1 else "",
|
||||
flush = True)
|
||||
|
||||
@jit
|
||||
def energy(position, box_size):
|
||||
"""
|
||||
Computes the potential energy of a system of particles with `position` using
|
||||
Morses potential.
|
||||
"""
|
||||
# enforce expected shape
|
||||
# (`scipy.optimize.minimize` drops shape information)
|
||||
if len(position.shape) == 1:
|
||||
position = position.reshape((position.shape[0] // 3, 3))
|
||||
# compute all pairwise position differences (all!)
|
||||
diff = position[:, jnp.newaxis, :] - position
|
||||
# extract only one of two distance combinations of non-equal particles
|
||||
lower_tri = jnp.tril_indices(position.shape[0], k = -1)
|
||||
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
|
||||
max_dist = box_size / 2.0
|
||||
diff = jnp.mod(diff + max_dist, box_size) - max_dist
|
||||
# Compute distances
|
||||
dist = jnp.linalg.norm(diff, axis = 1)
|
||||
# calc inbetween exp(.) expression
|
||||
ex = jnp.exp(-alpha * (dist - re))
|
||||
# evaluate Morse potential
|
||||
return De * jnp.sum(ex * (ex - 2.0))
|
||||
|
||||
@jit
|
||||
def force(position, box_size):
|
||||
"""
|
||||
Computes the forces acting on each particle in `position` as the negative
|
||||
gradient of the potential energy.
|
||||
"""
|
||||
return -grad(energy)(position, box_size)
|
||||
|
||||
@jit
|
||||
def kinetic(velocity):
|
||||
"""
|
||||
Computes the kenetic energy given a systems `velocity` state.
|
||||
"""
|
||||
return (mass / 2.0) * jnp.linalg.norm(velocity, axis = 1).sum()
|
||||
|
||||
@jit
|
||||
def step(position, velocity, acceleration, box_size, step_size):
|
||||
"""
|
||||
Performs a single Newton time step with `step_size` given system state
|
||||
through the current particle `position`, `velocity` and `acceleration`.
|
||||
"""
|
||||
# update position with a second order Taylor expantion
|
||||
position += step_size * velocity + (0.5 * step_size**2) * acceleration
|
||||
# compute new particle acceleration through Newton’s second law of motion
|
||||
acceleration_next = force(position, box_size) / mass
|
||||
# update velocity with a finite mean approximation
|
||||
velocity += (0.5 * step_size) * (acceleration + acceleration_next)
|
||||
# updated state
|
||||
return position, velocity, acceleration_next
|
||||
|
|
@ -1,130 +1,113 @@
|
|||
#!/usr/bin/env python
|
||||
|
||||
from optparse import OptionParser # to parse script parameters
|
||||
import scipy
|
||||
import numpy as np # static objects like constants (and outside functions)
|
||||
from jax import numpy as jnp # tracable objects like particle positions
|
||||
from jax import jit, grad, vmap
|
||||
################################################################################
|
||||
### parse script parameters ###
|
||||
################################################################################
|
||||
from optparse import OptionParser
|
||||
|
||||
# Parte script parameters
|
||||
arg_parser = 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", help = "Nr. of particles", default = 1000)
|
||||
dest = "nr_particles", default = -1,
|
||||
help = "Nr. of particles (required!)")
|
||||
arg_parser.add_option("-L", "--box_size", action = "store", type = float,
|
||||
dest = "box_size", help = "side box_size [A (angstrom)] of the simulation cube",
|
||||
default = 15.0)
|
||||
arg_parser.add_option("-T", "--temp", action = "store", type = float,
|
||||
dest = "temp", help = "temperature [K] to generate initial conditions",
|
||||
default = 300.0)
|
||||
arg_parser.add_option("-f", "--file", action = "store", type = str,
|
||||
dest = "path", help = "output file path",
|
||||
default = "task02.xyz")
|
||||
arg_parser.add_option("-p", "--plot", action = "store", type = str,
|
||||
dest = "plot", help = "output file path for a plot of the configuration " \
|
||||
"(default: no plot generated)",
|
||||
default = None)
|
||||
arg_parser.add_option("-v", action = "store_true", dest = "verbose",
|
||||
default = False, help = "turn verbosity mode on (default: off a.k.a. silent)")
|
||||
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)
|
||||
|
||||
# toy moleculat dynamic system parameters
|
||||
De = 1.6 # [eV (electronvolt)]
|
||||
alpha = 3.028 # [A^-1]
|
||||
re = 1.411 # [A]
|
||||
mass = 18.998403 # atomic mass of a single particle
|
||||
|
||||
def dump(position, velocity, box_size = config.box_size, \
|
||||
path = config.path, mode = "w", comment = ""):
|
||||
"""
|
||||
Dumps `position` and `velocity` to `path` with meta data.
|
||||
|
||||
First writes a header to `path` in `mode` file write mode defines as three
|
||||
lines.
|
||||
`<nr of particles [int]>`
|
||||
`<comment (ignored)>`
|
||||
`<box side lenght [float]>`
|
||||
then each line (nr of particles many) have the form `x y z vx vy vx`.
|
||||
"""
|
||||
with open(path, mode) as file:
|
||||
# write headers
|
||||
print(position.shape[0], comment, box_size, sep = "\n", file = file)
|
||||
# store position and velocity for with `x y z vx vy vz` entries per line`
|
||||
for (x, y, z), (vx, vy, vz) in zip(position, velocity):
|
||||
print(x, y, z, vx, vy, vz, file = file)
|
||||
|
||||
@jit
|
||||
def energy(position):
|
||||
"""
|
||||
Computes the potential energy of a system of particles with `position` using
|
||||
Morses potential.
|
||||
"""
|
||||
# enforce expected shape
|
||||
# (`scipy.optimize.minimize` drops shape information)
|
||||
position = position.reshape((config.nr_particles, 3))
|
||||
# compute all pairwise position differences (all!)
|
||||
diff = position[:, jnp.newaxis, :] - position
|
||||
# extract only one of two distance combinations of non-equal particles
|
||||
lower_tri = jnp.tril_indices(config.nr_particles, k = -1)
|
||||
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
|
||||
max_dist = config.box_size / 2
|
||||
diff = jnp.mod(diff + max_dist, 2 * max_dist) - max_dist
|
||||
# Compute distances
|
||||
dist = jnp.linalg.norm(diff, axis = 1)
|
||||
# calc inbetween exp(.) expression
|
||||
ex = jnp.exp(-alpha * (dist - re))
|
||||
# evaluate Morse potential
|
||||
return De * jnp.sum(ex * (ex - 2.0))
|
||||
|
||||
@jit
|
||||
def force(position):
|
||||
"""
|
||||
Computes the forces acting on each particle in `position` as the gradient of
|
||||
the potential energy. (a.k.a. the derivative (gradient) of `energy`)
|
||||
"""
|
||||
return grad(energy)(position)
|
||||
################################################################################
|
||||
### 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
|
||||
sd = np.sqrt(scipy.constants.Boltzmann * config.temp / mass)
|
||||
sd = np.sqrt(scipy.constants.Boltzmann * 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)
|
||||
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, x0 = position, jac = force, method = "CG")
|
||||
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)
|
||||
mean_forces = force(position).mean(axis = 0)
|
||||
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(position, velocity)
|
||||
dump(config.output, position, velocity, config.box_size)
|
||||
|
||||
# report stats (if requested by `-v` script argument)
|
||||
if config.verbose:
|
||||
print("Initial Energy: {:.4e}".format(initial_energy))
|
||||
print("Final Energy: {:.4e}".format(final_energy))
|
||||
print("Mean Forces: {:.4e} {:.4e} {:.4e}".format(*mean_forces))
|
||||
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}'")
|
||||
|
||||
# if a plot path is provided
|
||||
if not config.plot is None:
|
||||
from mpl_toolkits import mplot3d
|
||||
import matplotlib.pyplot as plt
|
||||
# new plot with 3D axis
|
||||
plt.figure()
|
||||
ax = plt.axes(projection = "3d")
|
||||
# create 3D position scatter plot
|
||||
ax.scatter(position[:, 0], position[:, 1], position[:, 2])
|
||||
# and save to file
|
||||
plt.savefig(config.plot)
|
||||
# # if a plot path is provided
|
||||
# if not config.plot is None:
|
||||
# from mpl_toolkits import mplot3d
|
||||
# import matplotlib.pyplot as plt
|
||||
# # new plot with 3D axis
|
||||
# plt.figure()
|
||||
# ax = plt.axes(projection = "3d")
|
||||
# # create 3D position scatter plot
|
||||
# ax.scatter(position[:, 0], position[:, 1], position[:, 2])
|
||||
# # and save to file
|
||||
# plt.savefig(config.plot)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,77 @@
|
|||
#!/usr/bin/env python
|
||||
|
||||
################################################################################
|
||||
### parse script parameters ###
|
||||
################################################################################
|
||||
from optparse import OptionParser
|
||||
|
||||
usage = "usage: %prog [options] [<INPUT>] [<STEP_SIZE>] [<ITERATIONS>]"
|
||||
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("-t", "--step_size", action = "store", type = float,
|
||||
dest = "step_size", default = -1.0,
|
||||
help = "time step size (Delta t) for integration (required!)")
|
||||
arg_parser.add_option("-N", "--iterations", action = "store", type = int,
|
||||
dest = "iterations", default = -1,
|
||||
help = "Number of time steps used for Newton's equation integration (required!)")
|
||||
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("-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.input = args[0]
|
||||
if len(args) > 1:
|
||||
config.step_size = float(args[1])
|
||||
if len(args) > 2:
|
||||
config.iterations = int(args[2])
|
||||
except ValueError as expression:
|
||||
arg_parser.print_help()
|
||||
print(f"Error: {expression}")
|
||||
exit(-1)
|
||||
else:
|
||||
# quick and dirty validation
|
||||
if not config.step_size > 0.0 \
|
||||
or not config.iterations > 0 \
|
||||
or 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)
|
||||
from molecular_dynamics import load, dump, force, mass, step, print_prog_bar
|
||||
|
||||
# load initial state from file
|
||||
(position, velocity), box_size = load(config.input)
|
||||
|
||||
# compute initial acceleration using Newton’s second law of motion
|
||||
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)
|
||||
# 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 verbosity turned on, print a progress bar
|
||||
if config.verbose:
|
||||
print_prog_bar(iteration, config.iterations)
|
||||
|
||||
# report output file path
|
||||
print(f"Done: trajectories saved to '{config.output}'")
|
||||
|
|
@ -0,0 +1,66 @@
|
|||
#!/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)
|
||||
Loading…
Reference in New Issue