Commit ecd3ce0b authored by Mark Hymers's avatar Mark Hymers

Initial commit

Signed-off-by: Mark Hymers's avatarMark Hymers <mark.hymers@hankel.co.uk>
parents
Random Scripts
==============
This repo is a place for me to store random scripts which other people might
find useful. Quality not assured.
* `metropolis_plot.py`: Visualisation of the Metropolis-Hastings algorithm
for a discrete case for teaching purposes.
#!/usr/bin/python3
from time import sleep
import numpy as np
import matplotlib.widgets as widgets
import matplotlib.patches as patches
import matplotlib.pyplot as plt
plt.style.use('ggplot')
fig = plt.figure(figsize=(7, 7))
ax = fig.add_axes((0, 0, 1, 1))
slider_ax = fig.add_axes((0.425, 0.48, 0.15, 0.02))
for a in [ax, slider_ax]:
ax.spines['left'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['bottom'].set_visible(False)
ax.set_xticks([])
ax.set_yticks([])
sfreq = widgets.Slider(slider_ax, 'Delay', 0, 2.0, valinit=1.0, valstep=0.05)
grey_inst = patches.Rectangle((-0.4, 0.4), 0.02, 0.15, color='grey')
blue_inst = patches.Rectangle((-0.4, 0.2), 0.02, 0.15, color='blue')
grey_comment = plt.Text(-0.36, 0.4+0.075, '= True Population Proportion',
color='k', verticalalignment='center', fontsize=8)
blue_comment = plt.Text(-0.36, 0.2+0.075, '= Estimated Population Proportion',
color='k', verticalalignment='center', fontsize=8)
ax.add_patch(grey_inst)
ax.add_patch(blue_inst)
ax.add_artist(grey_comment)
ax.add_artist(blue_comment)
plt.ion()
plt.show()
# Set up number of islands
NUM_ISLANDS = 20
# Set starting island
STARTING_ISLAND = 3
BAR_HEIGHT = 0.16
# Randomise island populations
island_pops = np.random.randint(1000, 25000, (NUM_ISLANDS, ))
# max_island_pop is the maximum population which should map to BAR_HEIGHT
# as a proportion of the total
max_island_pop_prop = island_pops.max() / np.sum(island_pops)
# Set up somewhere to store how many times we visit each island
island_counts = np.zeros_like(island_pops)
# Set up graphics
island_centre_xpos = np.cos(2*np.pi*np.arange(0, NUM_ISLANDS) / NUM_ISLANDS)
island_centre_ypos = np.sin(2*np.pi*np.arange(0, NUM_ISLANDS) / NUM_ISLANDS)
island_circles = []
# Indexing is [from, to]
island_arrows = np.empty((NUM_ISLANDS, NUM_ISLANDS), dtype=np.object)
island_props = np.empty((NUM_ISLANDS, 2), dtype=np.object)
ax.set_xlim(-1.5, 1.5)
ax.set_ylim(-1.5, 1.5)
ax.set_aspect('equal')
notes1 = plt.Text(0, -0.2, 'Use slider to change speed', color='k', horizontalalignment='center')
ax.add_artist(notes1)
graph_iter_num = plt.Text(0, -0.4, 'Iteration: 0', color='k', horizontalalignment='center')
ax.add_artist(graph_iter_num)
notes2 = plt.Text(0, -0.5, 'Red arrows: rejected proposal', color='r', horizontalalignment='center')
notes3 = plt.Text(0, -0.6, 'Green arrows: accepted proposal', color='g', horizontalalignment='center')
ax.add_artist(notes2)
ax.add_artist(notes3)
for island in range(NUM_ISLANDS):
x = island_centre_xpos[island]
y = island_centre_ypos[island]
radius = 0.1
island_circle = patches.Circle((x, y), radius=radius,
edgecolor='black',
facecolor='white')
ax.add_patch(island_circle)
island_circles.append(island_circle)
# Need to create arrows between this island and the next one
this_p = np.array([x, y])
next_island = island + 1 if (island + 1) < NUM_ISLANDS else 0
next_p = np.array([island_centre_xpos[next_island],
island_centre_ypos[next_island]])
# Create arrow from this island to the next one
vect = next_p - this_p
vect = vect / np.sqrt(np.sum(vect**2))
arrow_start = this_p + (vect * (radius + 0.01))
arrow_dist = np.sqrt(np.sum((next_p - this_p)**2)) * 0.3
a1 = patches.Arrow(arrow_start[0], arrow_start[1],
vect[0] * arrow_dist, vect[1] * arrow_dist,
color='k', width=0.02, alpha=0.5)
a1.set_visible(False)
ax.add_patch(a1)
island_arrows[island, next_island] = a1
# and from the next island back to this one
vect = this_p - next_p
vect = vect / np.sqrt(np.sum(vect**2))
arrow_start = next_p + (vect * (radius + 0.01))
arrow_dist = np.sqrt(np.sum((this_p - next_p)**2)) * 0.3
a2 = patches.Arrow(arrow_start[0], arrow_start[1],
vect[0] * arrow_dist, vect[1] * arrow_dist,
color='k', width=0.02, alpha=0.5)
a2.set_visible(False)
ax.add_patch(a2)
island_arrows[next_island, island] = a2
# Place a couple of bars inside the circle
# p1 is the proportion of people on the island
this_island_prop = (island_pops[island] / np.sum(island_pops)) * \
(BAR_HEIGHT / max_island_pop_prop)
p1 = patches.Rectangle((x-0.015-0.02, y-0.08), 0.02,
this_island_prop, color='grey')
island_props[island, 0] = p1
# p2 will be the proportion of the time that we have been on the island
p2 = patches.Rectangle((x+0.015, y-0.08), 0.02,
0, color='blue')
island_props[island, 1] = p2
ax.add_patch(p1)
ax.add_patch(p2)
# Initialise our starting island and previous island
cur_island = STARTING_ISLAND
prev_island = STARTING_ISLAND
# Make life easier
last_arrow_pos = [None, None]
for iteration in range(20000):
print("I: On island {}".format(cur_island))
graph_iter_num.set_text('Iteration {}'.format(iteration))
island_counts[cur_island] += 1
# Update which island we are on
island_circles[prev_island].set_facecolor('white')
island_circles[cur_island].set_facecolor('lightgreen')
island_circles[cur_island].set_alpha(0.5)
# Now move (keep a copy of our current island)
prev_island = cur_island
prob_move = np.random.random()
# Step A: Make a proposal
if prob_move >= 0.5:
# Consider moving east
prop_island = cur_island + 1 if cur_island + 1 < NUM_ISLANDS else 0
print("I: Proposing move east ({} to {})".format(cur_island, prop_island))
else:
# Consider moving west
prop_island = cur_island - 1 if cur_island - 1 >= 0 else NUM_ISLANDS - 1
print("I: Proposing move west ({} to {})".format(cur_island, prop_island))
# Step B: Accept / reject that proposal
# We choose whether to accept the proposal randomly with a p of
# (prop_island / cur_island)
# If prop_island > cur_island, this will give prop_p > 1 so we
# will always accept the proposal
prop_p = island_pops[prop_island] / island_pops[cur_island]
choice_p = np.random.random()
accepted = choice_p < prop_p
# Disable any existing arrow
if last_arrow_pos[0] is not None and last_arrow_pos[1] is not None:
island_arrows[last_arrow_pos[0], last_arrow_pos[1]].set_visible(False)
# Enable any arrow between the current island and the proposed island
if island_arrows[cur_island, prop_island] is None:
# Shouldn't really happen
print("W: Something odd happened at {} {}".format(cur_island, prop_island))
last_arrow_pos = [None, None]
else:
c = 'green' if accepted else 'red'
island_arrows[cur_island, prop_island].set_color(c)
island_arrows[cur_island, prop_island].set_visible(True)
last_arrow_pos = [cur_island, prop_island]
island_txt = "({} to {}) ({:.3f}, {:.3f})".format(cur_island, prop_island, choice_p, prop_p)
if accepted:
# Let's move!
print("I: Accepting move to island {}".format(island_txt))
cur_island = prop_island
else:
print("I: Rejecting move to island {}".format(island_txt))
# Update proportions
max_island_count_prop = np.max(island_counts) / np.sum(island_counts)
for island in range(NUM_ISLANDS):
this_island_prop = (island_counts[island] / np.sum(island_counts)) * (BAR_HEIGHT / max_island_count_prop)
island_props[island, 1].set_height(this_island_prop)
fig.canvas.draw()
fig.canvas.flush_events()
if sfreq.val > 0:
sleep(sfreq.val)
print("Island population proportions")
print(island_counts / np.sum(island_counts))
print("Visited proportions")
print(island_pops / np.sum(island_pops))
[flake8]
max-line-length=120
exclude=*.md,*.yaml,setup.cfg,build/,doc/build
extend-ignore=E741
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment