Skip to content

Instantly share code, notes, and snippets.

@barrysmyth
barrysmyth / bbc_graph.py
Last active April 11, 2020 21:42
A short snippet of code to reproduce the BBC's COVID-19 lockdown chart.
fig, ax = plt.subplots(figsize=(30, len(use_countries)*1.25))
# ----------------------------
# The Stringency Index Heatmap
# ----------------------------
# For maximum flexibility we use a separate axis to plot the heatmap's
# colour bar (horizontally) as the SI legend.
cbar_ax = fig.add_axes([.15, .17, .38, 3/len(stringency_for_country_by_date)])
@barrysmyth
barrysmyth / simple_heatmap.py
Created April 11, 2020 16:47
A snippet of code to test a simple heatmap.
# Test a simple heatmap for using the SI dataframe for a sample
# of 10 countries
import pandas as pd
import seaborn as sns
sns.heatmap(
stringency_for_country_by_date.loc[use_countries[:10]],
ax=ax, cbar=True
)
# Load in sample dataset
df = pd.read_csv('../data/segmented_bar_graph_example_data.csv').set_index('country').fillna(0)
# Extract the log(deaths)
# Use the sorted index from weeklies to sort deaths in the same way.
deaths = df.loc[weeklies.index]['log_deaths']
# Separate the weekly mobility drops and sort them in descending order of the sum of the
# the weekly drops
weeklies = df.drop(columns=['log_deaths'])\
fig, ax = plt.subplots(figsize=(15,8))
# Plot the segments as a stacked bar for a given x coordinate.
# Each x coordinate value will correspond to a country in our dataset.
def plot_segments(ax, x, segments, width=0.7, min_y=0):
current_y = min_y # The base of each bar.
# Create a new segment for each weekly value and stack the segments.
for segment in segments:
# Use the coolwarm colour map, which comes with Matplotlib.
cmap = plt.cm.coolwarm
# Normalise the log-deaths into a 0, 1 scale
norm = colors.Normalize(vmin=deaths.min(), vmax=deaths.max())
# Now the ith colour in colours corresponds to the ith country in deaths.
colours = cmap(norm(deaths))
# Build a simple colour dictionary indexed by country name.
sns.set_context('talk')
fig, ax = plt.subplots(figsize=(15,8))
# Plot the segments as a stacked bar at the x coordinate.
def plot_segments(ax, x, segments, width=0.7, min_y=0, colour_dict=colour_dict):
current_y = min_y
for segment in segments:
# Plot a single bar at x with height=segment and its base at current_y
# Create 2-row gridspec to hold the colour bar (cax) and the main graph (ax)
fig, (cax, ax) = plt.subplots(nrows=2,figsize=(15,10), gridspec_kw={"height_ratios":[.05, 1]})
# --- Creating the colour map
# Use the coolwarm colour map, which comes with Matplotlib.
cmap = plt.cm.coolwarm
# Normalise the log-deaths into a 0, 1 scale
norm = colors.Normalize(vmin=deaths.min(), vmax=deaths.max())
@barrysmyth
barrysmyth / oeis.py
Last active June 4, 2023 13:59
Two functions to programmatically search for matching sequences on the OEIS and to lookup a given sequence (by id) on the OEIS.
import requests
from lxml import html
def search(seq, n=10, m=10):
"""Search the OEIS for sequences matching seq and return a list of results.
Args:
seq: list of integers
n: number of results to return.
m: number of terms to return per result.
@barrysmyth
barrysmyth / collatz_harriss.py
Last active April 18, 2022 09:51
The Edmund Harriss visualisation of a Collatz orbit.
import math
import numpy as np
import matplotlib.pylab as plt
def turn_angle(n, angle, twist):
"""Compute the turn angle based on whether n is even or odd.
Args:
@barrysmyth
barrysmyth / collatz.py
Created April 18, 2022 09:14
Creating Collatz sequences/orbits from an integer.
def collatz(n):
"""The Collatz operator: If n is even return n/2 else return 3n+1."""
# If n is even then divide by 2.
if n%2==0:
return int(n/2)
# Otherwise return 3n+1
else:
return (3*n)+1