8
\$\begingroup\$

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.)

\$\endgroup\$
  • \$\begingroup\$ The Options class could be replaced with a function (or 2), since all it has are static methods. My guess is it works ok as written, but it's more complicated than it needs to be. And a minor point, you can leave off the required=False parameter; that's the default . \$\endgroup\$ – hpaulj Mar 20 at 4:22
  • \$\begingroup\$ Oh it's default? Learn something new everyday. Thank you for your help! \$\endgroup\$ – Duke0200 Mar 20 at 4:41
5
\$\begingroup\$
  • 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.

\$\endgroup\$
  • \$\begingroup\$ On the ternary operator, wouldn't that require the program to get a new set of arguments each time it's called? Because I want it to keep the same state throughout the program, but it can't do that if it needs new arguments. Thank you for the suggestions nonetheless! \$\endgroup\$ – Duke0200 Mar 18 at 12:46
  • \$\begingroup\$ I'm not entirely sure what you mean. the ternary operator I've written is equivalent to the 3 lines in your get_options() method; It won't change any other logic in your code. \$\endgroup\$ – HoboProber Mar 18 at 13:21
  • \$\begingroup\$ Sorry for not being clear. Wouldn't parse_args() in the generate_options() require there to be another set of arguments each time you called generate_options()? Maybe instead of your suggested code (but still with a ternary operator) I could do Options.options = Options.options if Options.options is not None else generate_options() then return Options.options? \$\endgroup\$ – Duke0200 Mar 18 at 13:32
  • \$\begingroup\$ Oh I understand. Yes you're right, you need to set Options.options first before returning it; I hadn't caught that, sorry. \$\endgroup\$ – HoboProber Mar 18 at 13:34
  • \$\begingroup\$ Don't worry. Thank you for your help! \$\endgroup\$ – Duke0200 Mar 18 at 14:35

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.