# /usr/bin/env python
from numpy.random import default_rng
import matplotlib.pyplot as plt
import matplotlib.animation as ani
import numpy as np
from math import *
from scipy.optimize import curve_fit
[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 generate_empty_lattice(lattice_size=100):
"""
Creates an empty lattice as an :math:`N \\times N` array of zeros.
Args:
lattice size (int): system size :math:`N`
Returns:
array: the lattice
"""
grid = np.zeros( [ lattice_size, lattice_size ] )
return grid
[docs]def generate_lattice(particles, lattice_size, label=1, additive=True):
"""
Creates a lattice with particles as an :math:`N \\times N` array.
The lattice sites with no particles are given the value of zero.
The sites with 1 or more particle are given a non-zero value.
Args:
particles (list): list of :class:`Walker` objects
lattice size (int): system size :math:`N`
label (int): value given to sites with particles
additive (bool): If True, the value given to sites is proportional
to the number of particles on that site. If False, all sites
with however many particles are given the same value.
Returns:
array: the lattice
"""
grid = generate_empty_lattice(lattice_size)
for part in particles:
try:
if additive:
grid[part.x, part.y] += label
else:
grid[part.x, part.y] = label
except:
pass
return grid
[docs]def show_system(particles, lattice_size, scale=5):
"""
Draws the system as a :math:`N \\times N` pixel image.
Args:
particles (list): list of :class:`Walker` objects
lattice size (int): system size :math:`N`
scale (int): values above this will be shown as black
"""
plt.clf()
ax = plt.axes()
ax.set_aspect('equal')
plt.pcolormesh(generate_lattice(particles, lattice_size), cmap='Greys', vmax = scale)
plt.show()
[docs]def draw(frame, history, scale=5):
"""
Draws the system for animation.
Args:
frame (int): index of the frame to draw
history (list): list of systems at different times
scale (int): values above this will be shown as black
"""
plt.clf()
ax = plt.axes()
ax.set_aspect('equal')
plt.pcolormesh( history[frame], cmap='Greys', vmax=scale)
[docs]def animate(history):
"""
Animate the simulation.
Args:
history (list): list of systems at different times
lattice size (int): system size :math:`N`
"""
nframes = len(history)
print("animating "+str(nframes)+" frames")
fig = plt.figure()
motion = ani.FuncAnimation(fig, draw, nframes, fargs=( history, ) )
plt.show()
[docs]class Walker:
"""
A particle moving on a lattice.
Since the particle can only move on a lattice,
its coordinates are always integers.
Args:
x (int): initial x coordinate
y (int): initial y coordinate
"""
def __init__(self, x, y):
self.x = x
self.y = y
self.initial_x = x
self.initial_y = y
[docs] def move_randomly(self):
"""
Makes the particle move.
The particle takes a step of length 1 in one of the
cardinal directions: up, down, left or right.
There is a 20 % chance for each direction.
In addition, there is a 20 % chance the particle
stays in place.
.. note ::
This function is incomplete!
"""
pass # todo
[docs] def get_distance_sq(self):
"""
Calculates the squared distance between the initial
and current coordinates of this particle.
Returns:
float: squared distance :math:`r^2 = (x-x_0)^2 + (y-y_0)^2`
"""
dx = self.x - self.initial_x
dy = self.y - self.initial_y
return ( dx*dx + dy*dy )
[docs]def move(particles):
"""
Makes all particles move.
All particles take a random step using :meth:`Walker.move_randomly`.
Args:
particles (list): list of :class:`Walker` objects
"""
for part in particles:
part.move_randomly()
[docs]def calculate_distance_statistics(particles):
"""
Calculates the average and error of mean of all particles squared displacement.
The squared displacement :math:`r_i^2` is calculated with
:meth:`Walker.get_distance_sq` for every particle :math:`i`.
The function then calculates the mean
.. math::
\\langle r^2 \\rangle = \\frac{1}{N} \\sum_i r_i^2
and error of mean
.. math::
\\Delta (r^2) = \\frac{s_{r^2}}{\\sqrt{N}},
where :math:`s_{r^2}` is the sample standard deviation
.. math::
s_{r^2} = \\sqrt{ \\frac{1}{N-1} \\sum_i ( r_i^2 - \\langle r^2 \\rangle )^2 }.
Args:
particles (list): list of :class:`Walker` objects
Returns:
float, float: average :math:`\\langle r^2 \\rangle`, error :math:`\\Delta (r^2)`
"""
n_particles = len(particles)
sq_average = 0.0
sq_deviation = 0.0
sq_distances = np.zeros(n_particles)
for i in range(n_particles):
sq_distances[i] = particles[i].get_distance_sq()
for dist in sq_distances:
sq_average += dist
sq_average /= n_particles
for dist in sq_distances:
sq_deviation += (dist - sq_average)*(dist - sq_average)
sq_deviation = sqrt( sq_deviation / (n_particles - 1) )
sq_error = sq_deviation / sqrt(n_particles)
return sq_average, sq_error
[docs]def calculate_density_histogram(particles, range_steps = 10, max_range = 50):
"""
Calculates a histogram of particle density as a function of displacement.
The, function only takes into account the particles for which
the displacement :math:`r_i` is not greater than a given maximum :math:`R`.
The range of possible values is split in
:math:`k` separate subranges or *bins*, :math:`[0,R/k), [R/k, 2R/k)` etc.
The midpoints of each bin, :math:`\\frac{R}{2k}, \\frac{3R}{2k}, \\frac{5R}{2k}` etc.
are also saved.
Having defined these bins, the function calculates the displacement :math:`r_i`
for each particle using :meth:`Walker.get_distance_sq` and checks
which bin the displacement belongs to.
The function then counts how many particles are placed in each bin, :math:`n(r)`.
Finally, the function calculates the surface density of particles at each
displacement range. That is, the number of particles in a given region
is divided by the area of that region,
.. math::
\\sigma(r) = \\frac{n(r)}{A(r)}.
As this is a 2D simulation, the particles whose displacement hit the range
:math:`[r, r+\\Delta r]` are inside a circular edge with inner radius
:math:`r` and thickness :math:`\\Delta r`.
The area of this edge is :math:`A(r) = \\pi (r + \\frac{1}{2} \\Delta r) \\Delta r`.
Args:
particles (list): list of :class:`Walker` objects
range_steps (int): the number of bins, :math:`k`
max_range (float): maximum diaplacement, :math:`R`
Returns:
array, array, array: average :math:`r`, particle count :math:`n(r)`, particle density :math:`\\sigma(r)` in each bin
"""
# array of counters for how many particles hit each bin
range_bins = np.zeros(range_steps)
# array of density scaled particle counters
density_bins = np.zeros(range_steps)
# the middle points for each subrange
avr_ranges = (np.array( range(range_steps) ) + 0.5 ) * max_range / (range_steps)
dr = avr_ranges[1]-avr_ranges[0]
for part in particles:
distance = sqrt( part.get_distance_sq() )
#range_index = min( floor( distance / max_range * range_steps ), range_steps - 1 )
range_index = floor( distance / max_range * range_steps )
if range_index < range_steps:
range_bins[range_index] += 1
# range_bins stores the number of particles at a distance [r, r+dr] from the starting point
# to get the density of particles at this slice, divide by the area dA = 2 pi r dr
return avr_ranges, range_bins, range_bins/(2*pi*avr_ranges*dr)
[docs]def linear(x, a):
"""Calculates the linear function :math:`y=ax` and returns the result.
Args:
x (float): the variable
a (float): the slope
Returns:
float: the result
"""
return a*x
[docs]def linear_fit(xdata, ydata, name):
"""Fit a linear function :math:`y = ax` to the given xy data.
Also return the optimal fit values obtained for slope a.
To identify this information, you may also name the operation.
Args:
xdata (array): x values
ydata (array): y values
name (str): identifier for printing the result
Returns:
float: slope
"""
params, covariance = curve_fit(linear,
np.array(xdata),
np.array(ydata))
print( "" )
print( "slope for "+name+" fit: "+str(params[0]) )
print( "" )
return params[0]
[docs]def main(n_particles, n_steps, n_plots):
"""
The main program.
Creates particles at the center of the system and then allows them to diffuse
by performing a random walk on a lattice.
Animates the motion and if the number of particles is large enough,
also calculates statistics for the particle distribution.
Args:
n_particles (int): number of particles
n_steps (int): number of simulation steps
n_plots (int): number of times data is recorded
"""
lattice_size = int( sqrt(20*n_steps) )
# if there are at least this many particles,
# calculate statistics
limit = 1000
particles = []
for i in range(n_particles):
x = lattice_size//2
y = lattice_size//2
particles.append( Walker(x,y) )
averages = []
errors = []
steps = []
bins = []
histograms = []
ranges = []
history = []
for i in range(n_steps+1):
# move particles
move(particles)
print_progress( i, n_steps )
# record the system every now and then
if i%(n_steps//n_plots) == 0:
history.append( generate_lattice(particles, lattice_size) )
# calculate stats if we have a lot of particles
if n_particles > limit:
ranges, bins, density_histogram = calculate_density_histogram(particles, int(sqrt(n_particles)/10), max(50, lattice_size//2) )
histograms.append( density_histogram )
avr, err = calculate_distance_statistics(particles)
averages.append(avr)
errors.append(2.0*err)
steps.append(i)
# show simulation
animate( history )
# show statistics if we have a lot of particles
if n_particles > limit:
steps = np.array(steps)
averages = np.array(averages)
errors = np.array(errors)
# do a linear fit to the data
slope = linear_fit(steps, averages, "diffusion")
plt.errorbar(steps, averages, yerr=errors, fmt='-o')
plt.plot(steps, linear(steps, slope), '-')
plt.xlabel("t")
plt.ylabel("$<\\Delta r^2>$")
plt.show()
for hist in histograms[1:]:
plt.plot(ranges, hist, '-')
plt.xlabel("r")
plt.ylabel("$\\sigma (r)$")
plt.show()
if __name__ == "__main__":
random = default_rng()
# first test with a small simulation, increase size when working
n_particles = 1
n_steps = 100
n_plots = 10
main(n_particles, n_steps, n_plots)