import sys
import copy
from math import *
import numpy as np
import numpy.random as random
import matplotlib.pyplot as plt
import matplotlib.animation as ani
[docs]def print_progress(step, total):
"""
Prints a progress bar.
Args:
step (int): progress counter
total (int): counter at completion
"""
message = "simulation progress ["
total_bar_length = 60
percentage = int(step / total * 100)
bar_fill = int(step / total * total_bar_length)
for i in range(total_bar_length):
if i < bar_fill:
message += "|"
else:
message += " "
message += "] "+str(percentage)+" %"
if step < total:
print(message, end="\r")
else:
print(message)
[docs]def draw(frame, xtraj, ytraj, ztraj, bounds):
"""
Draws a representation of the particle system as a scatter plot.
Used for animation.
Args:
frame (int): index of the frame to be drawn
xtraj (array): x-coordinates of all particles at different animation frames
ytraj (array): y-coordinates at all particles at different animation frames
ztraj (array): z-coordinates at all particles at different animation frames
bounds (array): list of lower and upper bounds for the plot as [[xmin, xmax], [ymin, ymax]]
"""
plt.clf()
ax = plt.axes()
ax.set_xlim(bounds[0])
ax.set_ylim(bounds[1])
ax.set_aspect('equal')
size = (ztraj[frame]+1)*2
#size = np.where( size > 1, size, np.ones( len( ztraj[frame] ) ) )
plt.scatter(xtraj[frame], ytraj[frame], marker='o', s=size )
[docs]def animate( particles, box, multiply = [3,3] ):
"""
Animates the simulation.
Args:
particles (list): list of :class:`pbc.Atom` objects
box (pbc.PeriodicBox): supercell
multiply (array): number of periodic images to draw in x and y directions
"""
nframes = len(particles[0].trajectory)
print("animating "+str(nframes)+" frames")
xtraj = []
ytraj = []
ztraj = []
Lx = box.lattice[0]
Ly = box.lattice[1]
# number of periodic images in x and y directions
multix = multiply[0]
multiy = multiply[1]
margin = 0.1
bounds = np.zeros([3,2])
bounds[0,1] = Lx*multix
bounds[1,1] = Ly*multiy
bounds[:,0] -= margin
bounds[:,1] += margin
for i in range(nframes):
xtraj.append([])
ytraj.append([])
ztraj.append([])
for p in particles:
for ix in range(multix):
for iy in range(multiy):
xtraj[-1].append(p.trajectory[i][0] + ix*Lx)
ytraj[-1].append(p.trajectory[i][1] + iy*Ly)
ztraj[-1].append(p.mass)
xtraj=np.array(xtraj)
ytraj=np.array(ytraj)
ztraj=np.array(ztraj)
fig = plt.figure()
motion = ani.FuncAnimation(fig, draw, nframes, interval=10, fargs=(xtraj, ytraj, ztraj, bounds) )
plt.show()
[docs]def show_trajectories(particles,box,tail=10,skip=10,multiply=[3,3]):
"""
Plot a 2D-projection of the trajectory of the system.
The function creates a plot showing the current and past
positions of particles.
Args:
particles (list): list of :class:`Planet` objects
box (pbc.PeriodicBox): supercell
tail (int): the number of past positions to include in the plot
skip (int): only every nth past position is plotted - skip is the number
n, specifying how far apart the plotted positions are in time
multiply (array): number of periodic images to draw in x and y directions
"""
N = len(particles)
nx = multiply[0]
ny = multiply[1]
plt.clf()
ax = plt.axes()
ax.set_aspect('equal','box')
plt.xlim([ 0, box.lattice[0]*multiply[0] ])
plt.ylim([ 0, box.lattice[1]*multiply[1] ])
xcoordinates = []
ycoordinates = []
size = []
for part in particles:
for pos in reversed(part.trajectory[-tail*skip::skip]):
box.shift_inside_box(pos)
for ix in range(nx):
for iy in range(ny):
xcoordinates.append( pos[0] + ix*box.lattice[0] )
ycoordinates.append( pos[1] + iy*box.lattice[1] )
size.append( (10*part.mass+1)*2/np.max(multiply) )
plt.scatter(xcoordinates, ycoordinates, s=size )
plt.show()
[docs]class PeriodicBox:
"""
Class representing a simulation box with periodic boundaries.
The box is orthogonal, i.e., a rectangular volume. As such,
it is specified by the lengths of its edges (lattice constants).
Args:
lattice (array): lattice constants
"""
def __init__(self, lattice):
self.lattice = lattice
[docs] def shift_inside_box(self, position):
"""
If the given position (3-vector) is outside the
box, it is shifted by multiple of lattice vectors until
the new position is inside the box. That is, the function
transforms the position vector to an equivalent position
inside the box.
Args:
position (array): the position to be shifted
"""
# go over x, y and z coordinates
for i in range(3):
while position[i] < 0:
position[i] += self.lattice[i]
while position[i] > self.lattice[i]:
position[i] -= self.lattice[i]
[docs] def distance_squared(self, particle1, particle2):
"""
Calculates and returns the square of the
distance between two particles.
.. math ::
r^2_{ij} = (x_i-x_j)^2 + (y_i-y_j)^2 + (z_i-z_j)^2.
In a periodic system, each particle has an infinite number of
periodic copies. Therefore the distance between two particles is
not unique. The function returns the shortest such distance,
that is, the distance between the the periodic copies which are
closest to each other.
.. note ::
This function is incomplete!
Args:
particle1 (pbc.Atom): the first body
particle2 (pbc.Atom): the second body
Returns:
float: the squared distance :math:`r^2_{ij}`
"""
# todo
return 1.0
[docs] def vector(self, particle1, particle2):
"""
Returns the vector pointing from the position of
particle1 to the position of particle2.
.. math ::
\\vec{r}_{i \\to j} = \\vec{r}_j - \\vec{r}_i
In a periodic system, each particle has an infinite number of
periodic copies. Therefore the displacement between two particles is
not unique. The function returns the shortest such displacement
vector.
Args:
particle1 (pbc.Atom): the first body
particle2 (pbc.Atom): the second body
Returns:
array: components of :math:`\\vec{r}_{i \\to j}`, :math:`[x_{i \\to j}, y_{i \\to j}, z_{i \\to j}]`
"""
vector_1to2 = particle2.position - particle1.position
# loop over x, y, and z directions
for i in range(3):
# If the absolute value of the separation
# in this direction is larger than half the length
# of the simulation box, there must be another
# periodic image of particle2 closer to particle 1.
#
# We can find it by translating the coordinates by
# an integer multiple of lattice vectors.
#
# Note: there is a more efficient way to calculate this.
while vector_1to2[i] < -self.lattice[i]/2:
vector_1to2[i] += self.lattice[i]
while vector_1to2[i] > self.lattice[i]/2:
vector_1to2[i] -= self.lattice[i]
return vector_1to2
[docs]class Atom:
"""
A point like object.
An atom has a position (a 3-vector), a velocity (3-vector)
and a mass (a scalar).
Args:
position (array): coordinates :math:`[x, y, z]`
velocity (array): velocity components :math:`[v_x, v_y, v_z]`
mass (float): mass :math:`m`
"""
def __init__(self, position, velocity, mass = 1.0):
self.position = np.array(position)
self.velocity = np.array(velocity)
self.mass = mass
self.trajectory = []
[docs] def move(self, force, dt):
"""
Move the atom.
Args:
shift (array): coordinate change :math:`[\Delta x, \Delta y, \Delta z]`
"""
self.position += self.velocity * dt + 0.5 * force/self.mass * dt*dt
[docs] def accelerate(self, force, dt):
"""
Set a new velocity for the particle as
.. math::
\\vec{v}(t+\\Delta t) = \\vec{v}(t) + \\frac{1}{2m}\\vec{F} \Delta t
Args:
force (array): force acting on the planet :math:`[F_x, F_y, F_z]`
dt (float): time step :math:`\\Delta t`
"""
self.velocity += force * dt/self.mass
[docs] def save_position(self):
"""
Save the current position of the particle
in the list 'trajectory'.
Note: in a real large-scale simulation one would
never save trajectories in memory. Instead, these
would be written to a file for later analysis.
"""
self.trajectory.append( [ self.position[0], self.position[1], self.position[2] ] )
[docs]def read_particles_from_file(filename):
"""
Reads the properties of planets from a file.
Each line should define a single :class:`Planet` listing its
position in cartesian coordinates, velocity components and mass,
separated by whitespace:
.. code-block::
x0 y0 z0 vx0 vy0 vz0 m0
x1 y1 z1 vx1 vy1 vz1 m1
x2 y2 z2 vx2 vy2 vz2 m2
x3 y3 z3 vx3 vy3 vz3 m3
...
Args:
filename (str): name of the file to read
Returns:
list: list of :class:`pbc.Atom` objects
"""
f = open(filename)
lines = f.readlines()
f.close()
particles = []
lattice = []
parts = lines[0].split()
for i in range(3):
lattice.append(float(parts[i]))
box = PeriodicBox(np.array(lattice))
for line in lines[1:]:
parts = line.split()
if len(parts) > 0:
position = np.array( [0.0]*3 )
velocity = np.array( [0.0]*3 )
position[0] = float(parts[0])
position[1] = float(parts[1])
position[2] = float(parts[2])
velocity[0] = float(parts[3])
velocity[1] = float(parts[4])
velocity[2] = float(parts[5])
mass = float(parts[6])
particles.append( Atom( position, velocity, mass ) )
return particles, box
[docs]def write_particles_to_file(particles, filename):
"""
Write the configuration of the system in a file.
The format is the same as that specified in :meth:`read_particles_from_file`.
Args:
particles (list): list of :class:`pbc.Atom` objects
filename (str): name of the file to write
"""
writelines = ""
for part in particles:
writelines += str(part.position[0])+" "+str(part.position[1])+" "+str(part.position[2])+" "
writelines += str(part.velocity[0])+" "+str(part.velocity[1])+" "+str(part.velocity[2])+" "
writelines += str(part.mass)+"\n"
f = open(filename,'w')
f.write(writelines)
f.close()
[docs]def write_xyz_file(particles, filename):
"""
Write the configuration of the system in a file.
The information is written in so called xyz format, which many
programs can parse.
Args:
particles (list): list of :class:`pbc.Atom` objects
filename (str): name of the file to write
"""
writelines = ""
N = len(particles)
steps = len(particles[0].trajectory)
for i in range(steps):
writelines += str(N)+"\n \n"
for part in particles:
writelines += "H "+str(part.trajectory[i][0])+" "+str(part.trajectory[i][1])+" "+str(part.trajectory[i][2])+"\n"
f = open(filename,'w')
f.write(writelines)
f.close()
[docs]def calculate_forces(particles, box, parameters=[1.0, 0.1]):
"""
Calculates the total force applied on each atom.
The forces are returned as a numpy array where
each row contains the force on an atom
and the columns contain the x, y and z components.
.. code-block::
[ [ fx0, fy0, fz0 ],
[ fx1, fy1, fz1 ],
[ fx2, fy2, fz2 ],
... ]
Args:
atoms (list): a list of :class:`pbc.Atom` objects
box (pbc.PeriodicBox): supercell
parameters (list): a list of physical parameters
Returns:
array: forces
"""
sigma = parameters[0]
epsilon = parameters[1]
sigma_six = sigma*sigma*sigma*sigma*sigma*sigma
forces = np.array( [[0.0, 0.0, 0.0]]*len(particles) )
for i in range(len(particles)):
for j in range(0,i):
part_i = particles[i]
part_j = particles[j]
dist_sq = box.distance_squared(part_i, part_j)
dist_six = dist_sq * dist_sq * dist_sq
dist_12 = dist_six * dist_six
force_j_to_i = - 4.0 * epsilon * (12.0 * sigma_six*sigma_six / dist_12 - \
6.0 * sigma_six / dist_six ) * \
box.vector(part_i, part_j) / dist_sq
forces[i, :] += force_j_to_i[:]
forces[j, :] -= force_j_to_i[:]
return forces
[docs]def update_positions(particles, forces, dt):
"""
Update the positions of all particles
using :meth:`pbc.Atom.move` according to
.. math::
\\vec{r}(t+\\Delta t) = \\vec{r}(t) + \\vec{v} \Delta t + \\frac{1}{2m}\\vec{F} (\\Delta t)^2
Args:
particles (list): a list of :class:`pbc.Atom` objects
force (array): array of forces on all bodies
dt (float): time step :math:`\\Delta t`
"""
for i in range(len(particles)):
part = particles[i]
force = forces[i,:]
part.move(force, dt)
[docs]def update_positions_no_force(particles, dt):
"""
Update the positions of all particles
using :meth:`pbc.Atom.move` according to
.. math::
\\vec{r}(t+\\Delta t) = \\vec{r}(t) + \\vec{v} \Delta t
Args:
particles (list): a list of :class:`pbc.Atom` objects
dt (float): time step :math:`\\Delta t`
"""
for i in range(len(particles)):
part = particles[i]
part.move(0.0, dt)
[docs]def update_velocities(particles, forces, dt):
"""
Update the positions of all particles
using :meth:`pbc.Atom.accelerate` according to
.. math::
\\vec{v}(t+\\Delta t) = \\vec{v}(t) + \\frac{1}{m}\\vec{F} \Delta t
Args:
particles (list): a list of :class:`pbc.Atom` objects
force (array): array of forces on all bodies
dt (float): time step :math:`\\Delta t`
"""
for i in range(len(particles)):
part = particles[i]
force = forces[i,:]
part.accelerate(force, dt)
[docs]def velocity_verlet(particles, box, dt, time, trajectory_dt = 1.0):
"""
Verlet algorithm for integrating the equations of motion,
i.e., advancing time.
There are a few ways to implement Verlet. The leapfrog
version works as follows: First, forces are calculated
for all particles and velocities are updated by half a
time step,
:math:`\\vec{v}(t+\\frac{1}{2}\\Delta t) = \\vec{v}(t) + \\frac{1}{2m}\\vec{F} \Delta t`.
Then, these steps are repeated:
* Positions are updated by a full time step using velocities but not forces,
.. math ::
\\vec{r}(t+\\Delta t) = \\vec{r}(t) + \\vec{v}(t+\\frac{1}{2}\\Delta t) \Delta t.
* Forces are calculated at the new positions, :math:`\\vec{F}(t + \\Delta t)`.
* Velocities are updated by a full time step using the forces
.. math ::
\\vec{v}(t+\\frac{3}{2}\\Delta t) = \\vec{v}(t+\\frac{1}{2}\\Delta t) + \\frac{1}{m}\\vec{F}(t+\\Delta t) \Delta t
These operations are done using the methods
:meth:`calculate_forces`,
:meth:`update_velocities` and
:meth:`update_positions_no_force`.
Because velocities were updated by half a time step in the beginning of the
simulation, positions and velocities are always offset by half a timestep.
You always use the one that has advanced further to update the other and this
results in a stable algorithm.
Args:
particles (list): a list of :class:`pbc.Atom` objects
box (pbc.PeriodicBox): supercell
dt (float): time step :math:`\\Delta t`
time (float): the total system time to be simulated
trajectory_dt (float): the positions of particles are
saved at these these time intervals - does not affect
the dynamics in any way
"""
steps = int(time/dt)
trajectory_wait = int(trajectory_dt / dt)
forces = calculate_forces(particles, box)
# get velocities at half time-step
update_velocities(particles, forces, 0.5*dt)
for i in range(steps):
update_positions_no_force(particles, dt)
forces = calculate_forces(particles, box)
update_velocities(particles, forces, dt)
if i%trajectory_wait == 0:
for part in particles:
part.save_position()
print_progress(i,steps)
# to calculate proper kinetic energy, return to proper timestep
update_velocities(particles, forces, -0.5*dt)
print_progress(steps,steps)
[docs]def calculate_momentum(particles):
"""
Calculates the total momentum of the system.
.. math ::
\\vec{p}_\\text{total} = \\sum_i \\vec{p}_i = \\sum_i m_i \\vec{v}_i
Args:
particles (list): a list of :class:`pbc.Atom` objects
Returns:
array: momentum components :math:`[p_x, p_y, p_z]`
"""
p = np.array( [0.0, 0.0, 0.0] )
for part in particles:
p += part.mass * part.velocity
return p
[docs]def calculate_kinetic_energy(particles):
"""
Calculates the total kinetic energy of the system.
.. math ::
K_\\text{total} = \\sum_i \\frac{1}{2} m_i v_i^2.
Args:
particles (list): a list of :class:`pbc.Atom` objects
Returns:
float: kinetic energy :math:`K`
"""
k = 0.0
for part in particles:
vx = part.velocity[0]
vy = part.velocity[1]
vz = part.velocity[2]
k += 0.5 * part.mass * (vx*vx + vy*vy + vz*vz)
return k
[docs]def calculate_potential_energy(particles, box, parameters=[1.0, 0.1]):
"""
Calculates the total potential energy of the system.
The potential energy is calculated using the Lennard-Jones model
.. math ::
U = \\sum_{i \\ne j} 4 \\epsilon \\left( \\frac{ \\sigma^{12} }{ r^{12}_{ij} }
- \\frac{ \\sigma^6 }{ r^6_{ij} } \\right).
Args:
particles (list): a list of :class:`pbc.Atom` objects
box (pbc.PeriodicBox): supercell
parameters (list): list of parameters
Returns:
float: potential energy :math:`U`
"""
sigma = parameters[0]
epsilon = parameters[1]
sigma_six = sigma*sigma*sigma*sigma*sigma*sigma
u = 0.0
for i in range(len(particles)):
for j in range(0,i):
part_i = particles[i]
part_j = particles[j]
dist_sq = box.distance_squared(part_i, part_j)
dist_six = dist_sq * dist_sq * dist_sq
dist_12 = dist_six * dist_six
u += 4.0 * epsilon * \
( sigma_six*sigma_six/dist_12 - sigma_six/dist_six)
return u
[docs]def main():
"""
The main program.
The program reads the system from a file, runs the simulation,
and plots the trajectory.
"""
copies = [1,1]
particles, box = read_particles_from_file('supercell.txt')
velocity_verlet(particles, box, 0.01, 100.0, 0.3)
show_trajectories(particles, box, 100, 1, copies)
animate(particles, box, copies)
if __name__ == "__main__":
main()