2022-04-11 10:14:40 +00:00
|
|
|
|
#!/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()
|
2022-04-11 10:14:40 +00:00
|
|
|
|
|
|
|
|
|
@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
|