add: task 3 and 4,

cleanup: moved non task specific code to seperate file

Exercise_02/molecular_dynamics.py

@@ -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

Exercise_02/task02.py

@@ -1,130 +1,113 @@
 #!/usr/bin/env python
 
-from optparse import OptionParser   # to parse script parameters
+################################################################################
-import scipy
+###                         parse script parameters                          ###
-import numpy as np                  # static objects like constants (and outside functions)
+################################################################################
-from jax import numpy as jnp        # tracable objects like particle positions
+from optparse import OptionParser
-from jax import jit, grad, vmap
 
-# Parte script parameters
+usage = "usage: %prog [options] [<NR_PARTICLES>] [<BOX_SIZE>] [<TEMPERATURE>]"
-arg_parser = OptionParser()
+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",
+    dest = "box_size", default = -1.0,
-    default = 15.0)
+    help = "side box_size [A (angstrom)] of the simulation cube (required!)")
-arg_parser.add_option("-T", "--temp", action = "store", type = float,
+arg_parser.add_option("-T", "--temperature", action = "store", type = float,
-    dest = "temp", help = "temperature [K] to generate initial conditions",
+    dest = "temperature", default = -1.0,
-    default = 300.0)
+    help = "temperature [K] to generate initial conditions (required!)")
-arg_parser.add_option("-f", "--file", action = "store", type = str,
+arg_parser.add_option("-o", "--output", action = "store", type = str,
-    dest = "path", help = "output file path",
+    dest = "output", default = "task02.xyz",
-    default = "task02.xyz")
+    help = "output file path (default: 'task02.xyz')")
-arg_parser.add_option("-p", "--plot", action = "store", type = str,
+arg_parser.add_option("-v", action = "store_true",
-    dest = "plot", help = "output file path for a plot of the configuration " \
+    dest = "verbose", default = False,
-                          "(default: no plot generated)",
+    help = "turn verbosity mode on (default: off a.k.a. silent)")
-    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()
+# 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)]
+###                task 2 / generation of initial conditions                 ###
-alpha = 3.028                       # [A^-1]
+################################################################################
-re = 1.411                          # [A]
+# note, load module _after_ processing script parameters (no need to load all
-mass = 18.998403                    # atomic mass of a single particle
+# of the heavy numeric modules if only usage or alike is needed)
-
+import numpy as np
-def dump(position, velocity, box_size = config.box_size, \
+import scipy
-         path = config.path, mode = "w", comment = ""):
+from jax import jit, grad
-    """
+from molecular_dynamics import dump, energy, force, mass
-    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)
+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)
+final_energy = energy(position, config.box_size)
-mean_forces = force(position).mean(axis = 0)
+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(f"Initial Energy:          {initial_energy:.4e}")
-    print("Final Energy:   {:.4e}".format(final_energy))
+    print( "Initial Mean Forces:     {:.4e} {:.4e} {:.4e}".format(*initial_mean_forces))
-    print("Mean Forces:    {:.4e} {:.4e} {:.4e}".format(*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 a plot path is provided
-if not config.plot is None:
+# if not config.plot is None:
-    from mpl_toolkits import mplot3d
+#     from mpl_toolkits import mplot3d
-    import matplotlib.pyplot as plt
+#     import matplotlib.pyplot as plt
-    # new plot with 3D axis
+#     # new plot with 3D axis
-    plt.figure()
+#     plt.figure()
-    ax = plt.axes(projection = "3d")
+#     ax = plt.axes(projection = "3d")
-    # create 3D position scatter plot
+#     # create 3D position scatter plot
-    ax.scatter(position[:, 0], position[:, 1], position[:, 2])
+#     ax.scatter(position[:, 0], position[:, 1], position[:, 2])
-    # and save to file
+#     # and save to file
-    plt.savefig(config.plot)
+#     plt.savefig(config.plot)

Exercise_02/task03.py

@@ -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}'")

Exercise_02/task04.py

@@ -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)