From a772b43fa7ceb9ec6e27f2bd9911e4cbc9db9fa9 Mon Sep 17 00:00:00 2001
From: Tom Hodson <thomas.c.hodson@gmail.com>
Date: Thu, 16 Jun 2022 14:00:01 +0200
Subject: [PATCH] Add improved versions to the MCFF package

---
 code/src/MCFF/ising_model.py | 24 ++++++++++++++++++++++++
 code/src/MCFF/mcmc.py        | 27 ++++++++++++++++++++++++++-
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/code/src/MCFF/ising_model.py b/code/src/MCFF/ising_model.py
index e846fd2..e1865e6 100644
--- a/code/src/MCFF/ising_model.py
+++ b/code/src/MCFF/ising_model.py
@@ -6,6 +6,15 @@ States are represented by NxM numpy arrays with values +/-1. These functions are
 
 import numpy as np
 from numba import jit
+from matplotlib import pyplot as plt
+
+
+def show_state(state, ax=None):
+    "Plot an Ising state to axis or make a new figure to plot to"
+    if ax is None:
+        f, ax = plt.subplots()
+    ax.matshow(state, cmap="Greys", vmin=-1, vmax=1)
+    ax.set(xticks=[], yticks=[])
 
 
 def all_up_state(N):
@@ -85,3 +94,18 @@ def energy_numpy(state):
     """
     E = -np.sum(state[:-1, :] * state[1:, :]) - np.sum(state[:, :-1] * state[:, 1:])
     return 2 * E / np.product(state.shape)
+
+
+@jit(nopython=True, nogil=True)
+def energy_difference(state, site):
+    "The change in energy if we flipped site of state"
+    # loop over the four neighbours of the site, skipping if the site is near an edge
+    N, M = state.shape
+    i, j = site
+    h = 0
+    for di, dj in [[-1, 0], [1, 0], [0, -1], [0, 1]]:  # loop over N,E,S,W neighbours
+        if (0 <= (i + di) < N) and (
+            0 <= (j + dj) < M
+        ):  # ignore neighbours not in the NxN grid
+            h += state[i + di, j + dj]
+    return 4 * state[i, j] * h / (N * M)
diff --git a/code/src/MCFF/mcmc.py b/code/src/MCFF/mcmc.py
index 682e57d..38a9235 100644
--- a/code/src/MCFF/mcmc.py
+++ b/code/src/MCFF/mcmc.py
@@ -2,6 +2,9 @@
 
 
 """
+import numpy as np
+from numba import jit
+from .ising_model import energy, energy_difference
 
 
 def mcmc(initial_state, steps, T, rng=np.random.default_rng()):
@@ -27,7 +30,7 @@ def mcmc(initial_state, steps, T, rng=np.random.default_rng()):
     assert N == M
 
     current_state = initial_state.copy()
-    E = N**2 * energy(state)
+    E = N**2 * energy(initial_state)
     for i in range(steps):
         i, j = rng.integers(N, size=2)
 
@@ -40,3 +43,25 @@ def mcmc(initial_state, steps, T, rng=np.random.default_rng()):
             E = new_E
 
     return current_state
+
+
+@jit(nopython=True, nogil=True)
+def mcmc_generator(
+    initial_state, steps, T, stepsize=1000, energy_difference=energy_difference
+):
+    N, M = initial_state.shape
+    assert N == M
+
+    current_state = initial_state.copy()
+    for _ in range(steps):
+        for _ in range(stepsize):
+            i, j = np.random.randint(N), np.random.randint(N)
+
+            # calculate the energy change if we were to flip this pixel but don't actually do it
+            change_in_E = N**2 * energy_difference(current_state, (i, j))
+
+            if change_in_E < 0 or np.exp(-change_in_E / T) > np.random.random():
+                current_state[i, j] *= -1  # accept the change!
+
+        yield current_state.copy()
+    return