138 lines
4.9 KiB
Python
138 lines
4.9 KiB
Python
#!/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((-1, 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, delta_t):
|
||
"""
|
||
Performs a single Newton time step with `delta_t` given system state
|
||
through the current particle `position`, `velocity` and `acceleration`.
|
||
"""
|
||
# update position with a second order Taylor expantion
|
||
position += delta_t * velocity + (0.5 * delta_t**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 * delta_t) * (acceleration + acceleration_next)
|
||
# updated state
|
||
return position, velocity, acceleration_next
|