diff --git a/Exercise_02/task02.py b/Exercise_02/task02.py new file mode 100644 index 0000000..08625a3 --- /dev/null +++ b/Exercise_02/task02.py @@ -0,0 +1,130 @@ +#!/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. + `` + `` + `` + 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)