# /usr/bin/env python
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 = (10*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, bounds=[[-1,3],[-2,2]] ):
"""
Animates the simulation.
Args:
particles (list): list of :class:`Planet` objects
bounds (array): list of lower and upper bounds for the plot as [[xmin, xmax], [ymin, ymax]]
"""
nframes = len(particles[0].trajectory)
print("animating "+str(nframes)+" frames")
xtraj = []
ytraj = []
ztraj = []
bounds = np.array(bounds)
for i in range(nframes):
xtraj.append([])
ytraj.append([])
ztraj.append([])
for p in particles:
xtraj[-1].append(p.trajectory[i][0])
ytraj[-1].append(p.trajectory[i][1])
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,tail=10,skip=10):
"""
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
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
"""
N = len(particles)
plt.clf()
ax = plt.axes()
ax.set_aspect('equal','box')
for part in particles:
xcoordinates = []
ycoordinates = []
size = []
for pos in reversed(part.trajectory[-tail*skip::skip]):
xcoordinates.append( pos[0] )
ycoordinates.append( pos[1] )
size.append( (10*part.mass+1)*2 )
plt.scatter(xcoordinates, ycoordinates, s=size )
plt.show()
[docs]class Planet:
"""
A planet (or some other point-like object).
A planet 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):
"""
Set a new position for the particle as
.. math::
\\vec{r}(t+\\Delta t) = \\vec{r}(t) + \\vec{v} \Delta t + \\frac{1}{2m}\\vec{F} (\\Delta t)^2
Args:
force (array): force acting on the planet :math:`[F_x, F_y, F_z]`
dt (float): time step :math:`\\Delta t`
"""
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:`Planet` objects
"""
f = open(filename)
lines = f.readlines()
f.close()
particles = []
for line in lines:
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( Planet( position, velocity, mass ) )
return particles
[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:`Planet` 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:`Planet` 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 distance_squared(particle1, particle2):
"""
Calculates the squared distance between two atoms,
.. math ::
r^2_{ij} = (x_i-x_j)^2 + (y_i-y_j)^2 + (z_i-z_j)^2.
Args:
particle1 (Planet): the first body
particle2 (Planet): the second body
Returns:
float: the squared distance :math:`r^2_{ij}`
"""
pos1 = particle1.position
pos2 = particle2.position
dx = pos1[0] - pos2[0]
dy = pos1[1] - pos2[1]
dz = pos1[2] - pos2[2]
return dx*dx + dy*dy + dz*dz
[docs]def vector(particle1, particle2):
"""
The vector pointing from one planet, :math:`i`, to another, :math:`j`,
.. math ::
\\vec{r}_{i \\to j} = \\vec{r}_j - \\vec{r}_i
Args:
particle1 (Planet): the first body
particle2 (Planet): 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}]`
"""
pos1 = particle1.position
pos2 = particle2.position
return pos2 - pos1
[docs]def calculate_forces(particles, g = 0.1):
"""
Calculates the total force applied on each body.
All bodies interact via an inverse square force
(Newton's gravitation). Specifically, the force
on body :math:`j` due to body :math:`i` is
.. math ::
\\vec{F}_{i \\to j} = - \\frac{g m_i m_j}{r^2_{ij}} \\hat{r}_{i \\to j}.
The forces are returned as a numpy array where
each row contains the force on a body
and the columns contain the x, y and z components.
.. code-block::
[ [ fx0, fy0, fz0 ],
[ fx1, fy1, fz1 ],
[ fx2, fy2, fz2 ],
... ]
Args:
particles (list): a list of :class:`Planet` objects
g (float): gravitational constant
Returns:
array: forces
"""
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]
m_i = part_i.mass
m_j = part_j.mass
dist_sq = distance_squared(part_i, part_j)
dist = sqrt(dist_sq)
force_j_to_i = g*m_i*m_j/dist_sq * \
vector(part_i, part_j) / dist
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:`Planet.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:`Planet` 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:`Planet.move` according to
.. math::
\\vec{r}(t+\\Delta t) = \\vec{r}(t) + \\vec{v} \Delta t
Args:
particles (list): a list of :class:`Planet` 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:`Planet.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:`Planet` 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 euler(particles, dt, time, trajectory_dt = 1.0):
"""
Euler algorithm for integrating the equations of motion,
i.e., advancing time.
The algorithm works as follows:
* Forces are calculated at the new positions, :math:`\\vec{F}(t)`.
* Positions are updated using velocities and forces,
.. math ::
\\vec{r}(t+\\Delta t) = \\vec{r}(t) + \\vec{v}(t) \Delta t + \\frac{1}{2m}\\vec{F}(t) (\\Delta t)^2.
* Velocities are updated using the forces
.. math ::
\\vec{v}(t+\\Delta t) = \\vec{v}(t) + \\frac{1}{m}\\vec{F}(t) \Delta t
These operations are done using the methods
:meth:`calculate_forces`,
:meth:`update_velocities` and
:meth:`update_positions`.
Unfortunately, the Euler algorithm is not stable and should not be used
for proper simulations!
Args:
particles (list): a list of :class:`Planet` objects
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
Returns:
array, array: kinetic energy time series, potential energy time series
"""
# the number of steps in the simulation
steps = int(time/dt)
# the number of steps between recording the positions
trajectory_wait = int(trajectory_dt / dt)
# energies
k = []
u = []
for i in range(steps):
forces = calculate_forces(particles)
update_positions(particles, forces, dt)
update_velocities(particles, forces, dt)
# record positions every now and then
if i%trajectory_wait == 0:
for part in particles:
part.save_position()
print_progress(i, steps)
k.append( calculate_kinetic_energy(particles) )
u.append( calculate_potential_energy(particles) )
print_progress(steps, steps)
k = np.array(k)
u = np.array(u)
return k, u
[docs]def velocity_verlet(particles, 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.
.. note ::
This function is incomplete!
Args:
particles (list): a list of :class:`Planet` objects
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
Returns:
array, array: kinetic energy time series, potential energy time series
"""
# the number of steps in the simulation
steps = int(time/dt)
# the number of steps between recording the positions
trajectory_wait = int(trajectory_dt / dt)
# energies
k = []
u = []
# calculate forces in initial configuration
forces = calculate_forces(particles)
for i in range(steps):
#############################
# Implement the algorithm: #
# * update positions #
# * calculate forces #
# * update velocities #
# Note that there are #
# functions for all these. #
#############################
# todo
# record positions every now and then
if i%trajectory_wait == 0:
for part in particles:
part.save_position()
print_progress(i, steps)
k.append( calculate_kinetic_energy(particles) )
u.append( calculate_potential_energy(particles) )
print_progress(steps, steps)
# to calculate proper kinetic energy, return to proper timestep
update_velocities(particles, forces, -0.5*dt)
k = np.array(k)
u = np.array(u)
return k, u
[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:`Planet` 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:`Planet` 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, g=0.1):
"""
Calculates the total potential energy of the system.
All bodies interact via an inverse square force
(Newton's gravitation), which has an inverse distance
potential energy.
.. math ::
U = \\sum_{i < j} - \\frac{g m_i m_j}{r_{ij}}.
Args:
particles (list): a list of :class:`Planet` objects
g (float): gravitational constant
Returns:
float: potential energy :math:`U`
"""
u = 0.0
for i in range(len(particles)):
for j in range(0,i):
part_i = particles[i]
part_j = particles[j]
m_i = part_i.mass
m_j = part_j.mass
dist_sq = distance_squared(part_i, part_j)
dist = sqrt(dist_sq)
u += -g*m_i*m_j/dist
return u
[docs]def main(filename, algo, dt, simulation_time, recording_dt):
"""
The main program.
Read the initial configuration from a file,
run the simulation, and print results.
Args:
filename (str): the file with system information
algo (str): either 'euler' or 'verlet'
dt (float): time step :math:`\\Delta t`
time (float): the total system time to be simulated
recording_dt (float): the positions of particles are
saved at these these time intervals - does not affect
the dynamics in any way
"""
# Read particles from a file.
particles = read_particles_from_file(filename)
# Calculate the initial momentum and energy
p_start = calculate_momentum(particles)
ene_start = calculate_kinetic_energy(particles) + \
calculate_potential_energy(particles)
# Run the simulation.
if algo == 'euler':
# Euler algorithm
k,u = euler(particles, dt, simulation_time, recording_dt)
elif algo == 'verlet':
# velocity-Verlet
k,u = velocity_verlet(particles, dt, simulation_time, recording_dt)
else:
print("Unsupported algorithm: "+str(algo))
# Plot the simulation as a trail.
show_trajectories(particles,1000,1)
# create an animation
animate(particles)
# Calculate the final momentum and energy
p_end = calculate_momentum(particles)
ene_end = calculate_kinetic_energy(particles) + \
calculate_potential_energy(particles)
print( "p at start: ", p_start )
print( "p at end: ", p_end )
print( "change: ", p_end-p_start )
print()
print( "E at start: ", ene_start )
print( "E at end: ", ene_end )
print( "change: ", ene_end-ene_start )
print()
# Plot energies
plt.plot(k, label="kinetic energy")
plt.plot(u, label="potential energy")
plt.plot(u+k, label="total energy")
plt.legend()
plt.xlabel("time")
plt.ylabel("energy")
plt.show()
if __name__ == "__main__":
# choose the algorithm
algo = 'euler'
#algo = 'verlet'
# timestep and total run time
dt = 0.1
simulation_time = 100.0
recording_dt = 0.2
main('orbit.txt', algo, dt, simulation_time, recording_dt)