Recently, I became very interested in a probability practice problem in my textbook for my class. I decided to implement it in code and I think I got most of it implemented. Right now, I'm hoping to see if there is any possible way that I can improve upon it.
The question reads as follows:
A particle starts at (0, 0) and moves in one unit independent
steps with equal probabilities of 1/4 in each of the four
directions: north, south, east, and west. Let S equal
the east-west position and T the north-south position
after n steps.
The code (and more information) can be found here in the GitHub repository.
The code is right here:
simulation.py
"""
A particle starts at (0, 0) and moves in one unit independent
steps with equal probabilities of 1/4 in each of the four
directions: north, south, east, and west. Let S equal
the east-west position and T the north-south position
after n steps.
"""
from random import choice
import numpy as np
from options import Options
# Directions (K -> V is initial of direction -> (dx, dy)
directions = {
'S': (0, -1),
'N': (0, 1),
'E': (1, 0),
'W': (-1, 0)
}
def get_direction():
"""
Get a random direction. Each direction has a 25% chance of occurring.
:returns: the chosen directions changes in x and y
"""
dirs = "NSEW"
return directions[choice(dirs)]
def change_position(curr_pos, change_in_pos):
"""
Updates the current location based on the change in position.
:returns: the update position (x, y)
"""
return curr_pos[0] + change_in_pos[0], curr_pos[1] + change_in_pos[1]
def increment_counter(counter, end_pos):
"""
Increments the provided counter at the given location.
:param counter: an numpy ndarray with the number of all ending locations in the simulation.
:param end_pos: the ending position of the last round of the simulation.
:returns: the updated counter.
"""
counter[end_pos[1], end_pos[0]] += 1
return counter
def get_chance_of_positions(n):
"""
Gets the approximated chance the particle ends at a given location.
Starting location is in the center of the output.
:param n: The number of steps the simulation is to take.
:returns: the number of iterations and an ndarray with the approximated chance the particle would be at each location.
"""
# The starting position starts at n, n so that we can pass in the location
# into the counter without worrying about negative numbers.
starting_pos = (n, n)
options = Options.get_options()
total_num_of_sims = options.num_of_rounds
counter = np.zeros(shape=(2 * n + 1, 2 * n + 1))
for j in range(total_num_of_sims):
curr_pos = starting_pos
for i in range(n):
change_in_pos = get_direction()
curr_pos = change_position(curr_pos, change_in_pos)
counter = increment_counter(counter, curr_pos)
chances = np.round(counter / total_num_of_sims, decimals=n + 1)
return total_num_of_sims, chances
plot.py
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from simulation import get_chance_of_positions as get_chance
from options import Options
def plot(n):
"""
Saves the plots of the chance of positions from simulation.py
:param n: the number of steps the simulation will take.
"""
num_of_iterations, counter = get_chance(n)
fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_title("Position Color Map (n = {})".format(n))
ax.set_xlabel("Number of iterations: {}".format(int(num_of_iterations)))
plt.imshow(counter, cmap=plt.get_cmap('Blues_r'))
ax.set_aspect('equal')
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.05)
plt.colorbar(orientation='vertical', cax=cax)
fig.savefig('plots/plot-{:03}.png'.format(n), dpi=fig.dpi)
def main():
"""
The main function for this file. Makes max_n plots of the simulation.
"""
options = Options.get_options()
for n in range(1, options.num_of_sims + 1):
plot(n)
if __name__ == "__main__":
main()
options.py
import argparse
class Options:
options = None
@staticmethod
def generate_options():
arg_parser = argparse.ArgumentParser()
arg_parser.add_argument('-N', '--num-of-rounds',
type = int,
required = False,
default = 10**5,
help = "The number of rounds to run in each simulation. Should be a big number. Default is 1E5")
arg_parser.add_argument('-n', '--num-of-sims',
type = int,
required = False,
default = 10,
help = "The number of simulations (and plots) to run. Default is 10.")
return arg_parser.parse_args()
@staticmethod
def get_options():
if Options.options is None:
Options.options = Options.generate_options()
return Options.options
I would love for some recommendations on PEP 8, design, and "Pythonic" code (code that Python expects and will thus be better optimized, such as list comprehension and numpy optimizations, wherever they may be.)