NSSC/Exercise_02/molecular_dynamics.py

138 lines
5.0 KiB
Python
Raw Normal View History

#!/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.
"""
2022-04-12 16:04:26 +00:00
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