From 1ff5c1186639405287fdcf5713b9242b8007f5cd Mon Sep 17 00:00:00 2001 From: daniel Date: Mon, 11 Apr 2022 12:14:40 +0200 Subject: [PATCH] add: task 3 and 4, cleanup: moved non task specific code to seperate file --- Exercise_02/molecular_dynamics.py | 137 ++++++++++++++++++++++ Exercise_02/task02.py | 183 ++++++++++++++---------------- Exercise_02/task03.py | 77 +++++++++++++ Exercise_02/task04.py | 66 +++++++++++ 4 files changed, 363 insertions(+), 100 deletions(-) create mode 100644 Exercise_02/molecular_dynamics.py create mode 100644 Exercise_02/task03.py create mode 100644 Exercise_02/task04.py diff --git a/Exercise_02/molecular_dynamics.py b/Exercise_02/molecular_dynamics.py new file mode 100644 index 0000000..5e516a8 --- /dev/null +++ b/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. + `` + `` + `` + 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 diff --git a/Exercise_02/task02.py b/Exercise_02/task02.py index 08625a3..ea0dd18 100644 --- a/Exercise_02/task02.py +++ b/Exercise_02/task02.py @@ -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] [] [] []" +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. - `` - `` - `` - 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) diff --git a/Exercise_02/task03.py b/Exercise_02/task03.py new file mode 100644 index 0000000..cf005fb --- /dev/null +++ b/Exercise_02/task03.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python + +################################################################################ +### parse script parameters ### +################################################################################ +from optparse import OptionParser + +usage = "usage: %prog [options] [] [] []" +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}'") diff --git a/Exercise_02/task04.py b/Exercise_02/task04.py new file mode 100644 index 0000000..793e489 --- /dev/null +++ b/Exercise_02/task04.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python + +################################################################################ +### parse script parameters ### +################################################################################ +from optparse import OptionParser + +usage = "usage: %prog [options] []" +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)