Markov Chain Monte Carlo for fun and profit
🎲 ⛓️ 👉 🧪
import numpy as np
import matplotlib.pyplot as plt
from numba import jit
# This loads some custom styles for matplotlib
import json, matplotlib
with open("../_static/matplotlibrc.json") as f:
matplotlib.rcParams.update(json.load(f))
np.random.seed(
42
) # This makes our random numbers reproducible when the notebook is rerun in order
Adding Functionality¶
The main thing we want to be able to do is to take measurements. The code, as I have written it, doesn’t really allow that because it only returns the final state in the chain. Let’s say we have a measurement called average_color(state) that we want to average over the whole chain. We could just stick that inside our definition of mcmc, but we know that we will likely make other measurements too, and we don’t want to keep writing new versions of our core functionality!
Exercise 1¶
Have a think about how you would implement this and what options you have.
Solution 1¶
So I chatted with my mentors on this project on how to best do this, and we came up with a few ideas:
Option 1: Just save all the states and return them¶
The problem with this is the states are very big, and we don’t want to waste all that memory. For an NxN state that uses 8-bit integers (the smallest we can use in NumPy) 1000 samples would already use 2.5GB (2.5 gigabytes) of memory! We will see later that we’d really like to be able to go a bit bigger than 50x50 and 1000 samples!
Option 2: Pass in a function to make measurements¶
def mcmc(initial_state, steps, T, measurement, energy=energy):
...
current_state = initial_state.copy()
E = energy(current_state)
for i in range(steps):
measurements[i] = measurement(state)
...
return measurements
This could work, but it limits how we can store measurements and what shape and type they can be. What if we want to store our measurements in a NumPy array? Or what if your measurement itself is a vector or an object that can’t easily be stored in a NumPy array? We would have to think carefully about what functionality we want.
Option 3: Use Inheritance¶
# This class would define the basic functionality of performing MCMC
class MCMCSampler(object):
def run(self, initial_state, steps, T):
...
for i in range(steps):
self.measurement(state)
# This class would inherit from it and just implement the measurement
class AverageColorSampler(MCMCSampler):
measurements = np.zeros(10)
index = 0
def measurement(self, state):
self.measurements[self.index] = some_function(state)
self.index += 1
color_sampler = AverageColorSampler(...)
measurements = color_sampler.run(...)
This would definitely work, but I personally am not a huge fan of object-oriented programming, so I’m going to skip this option!
Option 4: Use a generator¶
This is the approach I ended up settling on, we will use python generator function. While you may not have come across generator functions before, you almost certainly will have come across generators, range(n) is a generator, (i for i in [1,2,3]) is a generator. Generator functions are a way to build your own generators, by way of example here is range implemented as a generator function:
def my_range(n):
"Behaves like the builtin range function of one argument"
i = 0
while i < n:
yield i # sends i out to whatever function called us
i += 1
return # let's python know that we have nothing else to give
my_range(10), list(my_range(10))
(<generator object my_range at 0x7fcaff25da50>, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
This requires only a very small change to our mcmc function, and suddenly we can do whatever we like with the states! While we’re at it, I’m going to add an argument stepsize that allows us to only sample the state every stepsize MCMC steps. You’ll see why we would want to set this to a value greater than 1 in a moment.
from MCFF.ising_model import energy, all_up_state
from MCFF.mcmc import mcmc_original
@jit(nopython=True, nogil=True)
def mcmc_generator(initial_state, steps, T, stepsize=1000, energy=energy):
N, M = initial_state.shape
assert N == M
current_state = initial_state.copy()
E = energy(current_state)
for _ in range(steps):
for _ in range(stepsize):
i, j = np.random.randint(N), np.random.randint(N)
# modify the state a little, here we just flip a random pixel
current_state[i, j] *= -1
new_E = energy(current_state)
if (new_E < E) or np.exp(-(new_E - E) / T) > np.random.random():
E = new_E
else:
current_state[i, j] *= -1 # reject the change we made
yield current_state.copy()
return
N_steps = 1000
stepsize = 1
initial_state = all_up_state(20)
without_yield = %timeit -o mcmc_original(initial_state, steps = N_steps, T = 5)
with_yield = %timeit -o [np.mean(s) for s in mcmc_generator(initial_state, T = 5, steps = N_steps, stepsize = 1)]
print(f"{with_yield.best / without_yield.best:.0f}x slowdown!")
18.6 ms ± 645 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
13.3 ms ± 1.74 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1x slowdown!
Fun fact: if you replace yield current_state.copy() with yield current_state your python kernel will crash when you run the code. I believe this is a bug in Numba that related to how pointers to numpy arrays work but let’s not worry too much about it.
We take a factor of two slowdown, but that doesn’t seem so much to pay for the fact we can now sample the state at every single step rather than just the last.
%load_ext watermark
%watermark -n -u -v -iv -w -g -r -b -a "Thomas Hodson" -gu "T_Hodson"
Author: Thomas Hodson
Github username: T_Hodson
Last updated: Mon Jul 18 2022
Python implementation: CPython
Python version : 3.9.12
IPython version : 8.4.0
Git hash: 03657e08835fdf23a808f59baa6c6a9ad684ee55
Git repo: https://github.com/ImperialCollegeLondon/ReCoDE_MCMCFF.git
Git branch: main
json : 2.0.9
numpy : 1.21.5
matplotlib: 3.5.1
Watermark: 2.3.1