Python implementation of approximating the chance a particle is at a location after n steps in the cardinal directions

Question

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 \$\frac{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.)

Answer 1

• Your directions dictionary is only used by get_direction; why not put it inside the function rather than have it as a global? I'd maybe even put the dict inside get_chance_of_positions() and drop the get_direction() function entirely; it's short, only used once and doesn't require a docstring (its obvious how it works). Put dirs straight into the line as well, rather than creating the dirs variable.

• I think you've created more functions than you need in the first file. Both change_position() and increment_counter() should probably be merged into your worker function (get_chance_of_positions()). At the moment, the reader has to jump around functions and their docstrings when a commented line in the worker function would be better.

• When looping, you often use names like n, i or j. Try to be more explicit, e.g. for sim in range(total_num_of_sims). Sometimes you actually define what the variable does in the docstring; a good name would do this job better.

• You could use fstrings in your plotting rather than .format(). You shouldn't need to cast to int either.

• I like ternary operators so I would change get_options to:

  return Options.options if Options.options is not None else Options.generate_options()
  

• Try to be consistent with how you indent when splitting a line across multiple lines as you have in options.Options.generate_options(). By which I mean:

  Instead of:

  arg_parser.add_argument('-N', '--num-of-rounds',
                      type = int,
                      required = False,
                      default = 10**5 ...)
  

  Use:

  arg_parser.add_argument('-N', '--num-of-rounds',
                          type = int,
                          required = False,
                          default = 10**5 ...)
  

  Or:

  arg_parser.add_argument(
      '-N', '--num-of-rounds',
      type = int,
      required = False,
      default = 10**5 ...
  )
  

Overall very readable code. Not many comments but the code generally speaks for itself. Try to avoid compartmentalising your code so much, it doesn't always make sense to split things up.