131 lines
5.3 KiB
Python
131 lines
5.3 KiB
Python
#!/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
|
|
|
|
# Parte script parameters
|
|
arg_parser = OptionParser()
|
|
arg_parser.add_option("-M", "--particles", action = "store", type = int,
|
|
dest = "nr_particles", help = "Nr. of particles", default = 1000)
|
|
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)")
|
|
# Parse command line arguments (as def. above) or store defaults to `config`
|
|
config, args = arg_parser.parse_args()
|
|
|
|
# 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)
|
|
|
|
# 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)
|
|
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)
|
|
|
|
# optimize energy to find low energy system state using Conjugate-Gradients
|
|
optim = scipy.optimize.minimize(energy, x0 = position, jac = force, method = "CG")
|
|
# extract (and reshape) optimization result
|
|
position = optim.x.reshape((config.nr_particles, 3))
|
|
|
|
# recompute stats after optimizing for low energy state
|
|
final_energy = energy(position)
|
|
mean_forces = force(position).mean(axis = 0)
|
|
|
|
# store state snapshot to file (default target file defined by script args)
|
|
dump(position, velocity)
|
|
|
|
# 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))
|
|
|
|
# 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)
|