Exercise 4: random walks

This exercise requires the use of many of the elements we've discussed (and a few extra ones too, remember the documentation for matplotlib is comprehensive!). We'll start by defining a random walk and some statistical population data for us to plot:


In [1]:
import matplotlib.pyplot as plt
import numpy as np

np.random.seed(1234)

n_steps = 500
t = np.arange(n_steps)

# Probability distribution:
mu = 0.002    # Mean
sigma = 0.01  # Standard deviation

# Generate a random walk, with position X as a function of time:
S = mu + sigma * np.random.randn(n_steps)
X = S.cumsum()

# Calculate the 1 sigma upper and lower analytic population bounds:
lower_bound = mu * t - sigma * np.sqrt(t)
upper_bound = mu * t + sigma * np.sqrt(t)

1. Plot the walker position X against time (t) using a solid blue line of width 2 and give it a label so that it will appear in a legend as "walker position".

2. Plot the population mean (mu*t) against time (t) using a black dashed line of width 1 and give it a label so that it will appear in a legend as "population mean".

3. Fill the space between the variables upper_bound and lower_bound using yellow with alpha (transparency) of 0.5, label this so that it will appear in a legend as "1 sigma range" (hint: see the fill_between method of an axes or pyplot.fill_between).

4. Draw a legend in the upper left corner of the axes (hint: you should have already set the labels for each line when you created them).

5. Label the x-axis "num steps" and the y-axis "position", and draw gridlines on the axes (hint: ax.grid toggles the state of the grid).

6. (harder) Fill the area under the walker position curve that is above the upper bound of the population mean using blue with alpha 0.5 (hint: fill_between can take a keyword argument called where that allows you to limit where filling is drawn).


In [2]:
# 1: Create an axes and plot the random walk in blue
ax = plt.axes()
ax.plot(t, X, linewidth=2, label='walker position', color='blue')

# 2: Add a dashed line indicating the population mean
ax.plot(t, mu * t, linewidth=1, label='population mean', color='black', ls='--')

# 3: Shade the area between the upper and lower bounds for population mean
ax.fill_between(t, lower_bound, upper_bound, facecolor='yellow', alpha=0.5,
                label='1 sigma range')

# 4: Draw a legend in the upper left
ax.legend(loc='upper left')

# 5: Label the axes and draw some grid lines
ax.set_xlabel('num steps')
ax.set_ylabel('position')
ax.grid()

# 6: Shade the area under the random walk that is outside the upper bound
ax.fill_between(t, upper_bound, X, where=X>upper_bound, facecolor='blue', alpha=0.5)

plt.show()