|
|
|
@ -88,7 +88,7 @@ def energy(position, box_size): |
|
|
|
|
# enforce expected shape |
|
|
|
|
# (`scipy.optimize.minimize` drops shape information) |
|
|
|
|
if len(position.shape) == 1: |
|
|
|
|
position = position.reshape((position.shape[0] // 3, 3)) |
|
|
|
|
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 |
|
|
|
@ -122,16 +122,16 @@ def kinetic(velocity): |
|
|
|
|
return (mass / 2.0) * (velocity**2).sum() |
|
|
|
|
|
|
|
|
|
@jit |
|
|
|
|
def step(position, velocity, acceleration, box_size, step_size): |
|
|
|
|
def step(position, velocity, acceleration, box_size, delta_t): |
|
|
|
|
""" |
|
|
|
|
Performs a single Newton time step with `step_size` given system state |
|
|
|
|
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 += step_size * velocity + (0.5 * step_size**2) * acceleration |
|
|
|
|
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 * step_size) * (acceleration + acceleration_next) |
|
|
|
|
velocity += (0.5 * delta_t) * (acceleration + acceleration_next) |
|
|
|
|
# updated state |
|
|
|
|
return position, velocity, acceleration_next |
|
|
|
|