init: ex 2, task 2

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.
+    `<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)