NSSC/Exercise_02/molecular_dynamics.py

138 lines
5.0 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/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) * (velocity**2).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 Newtons 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