Herd Immunity Threshold Graph Plotter

#!/usr/bin/env python3


import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from matplotlib.path import Path
import matplotlib.lines as lines
import matplotlib.patches as patches
import matplotlib.colors as colors
import matplotlib.hatch as mhatch


fathatch_path_vertices = [
  [-0.5, -0.45],
  [-0.5, -0.5],
  [-0.375, -0.5],
  [-0.5, -0.375],
  [-0.5, -0.45],

  [-0.5, 0],
  [-0.5, -0.125],
  [-0.125, -0.5],
  [0.125, -0.5],
  [-0.5, 0.125],
  [-0.5, 0],

  [-0.5, 0.45],
  [-0.5, 0.375],
  [0.375, -0.5],
  [0.5, -0.5],
  [0.5, -0.375],
  [-0.375, 0.5],
  [-0.5, 0.5],
  [-0.5, 0.45],

  [0, 0.5],
  [-0.125, 0.5],
  [0.5, -0.125],
  [0.5, 0.125],
  [0.125, 0.5],
  [0, 0.5],

  [0.45, 0.5],
  [0.375, 0.5],
  [0.5, 0.375],
  [0.5, 0.5],
  [0.45, 0.5],
]

fathatch_path_cmds = [
  Path.MOVETO,
  Path.LINETO,
  Path.LINETO,
  Path.LINETO,
  Path.CLOSEPOLY,

  Path.MOVETO,
  Path.LINETO,
  Path.LINETO,
  Path.LINETO,
  Path.LINETO,
  Path.CLOSEPOLY,

  Path.MOVETO,
  Path.LINETO,
  Path.LINETO,
  Path.LINETO,
  Path.LINETO,
  Path.LINETO,
  Path.LINETO,
  Path.CLOSEPOLY,

  Path.MOVETO,
  Path.LINETO,
  Path.LINETO,
  Path.LINETO,
  Path.LINETO,
  Path.CLOSEPOLY,

  Path.MOVETO,
  Path.LINETO,
  Path.LINETO,
  Path.LINETO,
  Path.CLOSEPOLY,
]

fat_left_hatch_path = patches.Path(
  [[x, y] for x, y in fathatch_path_vertices],
  fathatch_path_cmds,
)

fat_right_hatch_path = patches.Path(
  [[-x, y] for x, y in fathatch_path_vertices],
  fathatch_path_cmds,
)


class FatLHatch(mhatch.Shapes):
    filled = True
    size = 1.0
    path = fat_left_hatch_path
    def __init__(self, hatch, density):
      self.num_rows = 2 * (hatch.count('C')) # * density
      self.shape_vertices = self.path.vertices
      self.shape_codes = self.path.codes
      mhatch.Shapes.__init__(self, hatch, density)


class FatRHatch(mhatch.Shapes):
    filled = True
    size = 1.0
    path = fat_right_hatch_path
    def __init__(self, hatch, density):
      self.num_rows = 2 * (hatch.count('c')) # * density
      self.shape_vertices = self.path.vertices
      self.shape_codes = self.path.codes
      mhatch.Shapes.__init__(self, hatch, density)


def lin_fn(x):
  # Linear reconstruction filter
  return (x + 1 if x < 0 else 1 - x) if -1 < x < 1 else 0


def lanczos(x, a):
  # Lanczos reconstruction filter with size parameter a, a positive integer
  #
  # a = 1 is useless. Larger values increasingly emulate the normalised sinc
  # function. (Lanczos is just the normalised sinc function windowed with the
  # horizontally scaled central lobe of another normalised sinc function.)
  return (np.sinc(x) * np.sinc(x / a)) if -a < x < a else 0


def lanczos_fn(a):
  # Return a 1-argt lanczos function tailored with a specific size parameter.
  return lambda x: lanczos(x, a)


def resample(values, new_x_posns, filter_fn, support):
  result = []
  n = len(values)
  a = max(0, -support[0])
  b = max(0, support[1])
  for x1 in new_x_posns:
    x = n * x1
    rng = range(
      int(np.floor(x)) + support[0] + 1,
      int(np.floor(x)) + support[1] + 1
    )
    s = sum(values[np.clip(i, 0, n - 1)] * filter_fn(x - i) for i in rng)
    result.append(s)
  return result


def resample_lanczos2(values, new_x_posns):
  result = resample(values, new_x_posns, lanczos_fn(2), (-2, 2))
  return result


def hit_fn(r0, ver):
  # Re = R_0 (1 - X.E)
  return (1 - 1 / r0) / ver


def soft_edge_grad_image(colour):
  c = colors.to_rgba(colour)
  ramp = [0.2, 0.6, 0.7, 0.8, 0.9, 0.95]
  x = ramp + [1.0] + list(reversed(ramp))
  n = len(x)
  c3 = np.tile(np.atleast_2d(c[:3]), (n, 1))
  alpha = np.atleast_2d(np.array(x)).T * c[3]
  im = np.expand_dims(np.hstack([c3, alpha]), 0)
  return im


def main():

  r0_max = 20

  mhatch._hatch_types.append(FatLHatch)
  mhatch._hatch_types.append(FatRHatch)

  fig = plt.figure(figsize=(5.5, 7.0))
  ax = plt.axes()
  ax.xaxis.set_major_locator(ticker.MultipleLocator(1))
  ax.yaxis.set_major_locator(ticker.MultipleLocator(10))
  ax.yaxis.set_minor_locator(ticker.MultipleLocator(5))
  plt.xlim(0, r0_max)
  plt.ylim(0, 105)

  # ~ P = [
    # ~ (1, 0.0),
    # ~ (1.25, 0.20),
    # ~ (1.6, 0.375),
    # ~ (2.0, 0.50),
    # ~ (2.5, 0.60),
    # ~ (4.0, 0.75),
    # ~ (5.0, 0.80),
    # ~ (8.0, 0.875),
    # ~ (10.0, 0.90),
    # ~ (12.5, 0.92),
    # ~ (16.0, 0.9375),
    # ~ (20.0, 0.95),
  # ~ ]

  legend_patches = []
  labels = []

  num_subdivs = 10
  x = np.linspace(0, r0_max, r0_max * num_subdivs, endpoint=True)[1:]

  lp100 = dict(
    color = 'black',
    linewidth = 1.5,
  )
  lp10 = dict(
    color = 'black',
    linewidth = 0.6,
  )
  lp5 = dict(
    color = 'black',
    linewidth = 0.3,
    linestyle = "dashed"
  )

  for ver_pc in [65, 70, 75, 80, 85, 90, 95, 100]:
    if ver_pc == 100:
      lparams = lp100
    else:
      lparams = lp10 if ver_pc % 10 == 0 else lp5
    label = "{}%".format(ver_pc)
    y = [100 * hit_fn(x_i, ver_pc/100) for x_i in x]
    ax.plot(x, y, label=label, **lparams)

  legend_patches.append(lines.Line2D([], [], **lp100))
  labels.append("VE = 100%")
  legend_patches.append(lines.Line2D([], [], **lp10))
  labels.append("VE 10% multiple")
  legend_patches.append(lines.Line2D([], [], **lp5))
  labels.append("VE 5% multiple")

  y = [100 * hit_fn(x_i, 1) for x_i in x]
  curve_points = np.vstack([x, y]).T
  upper_clip_path = patches.Polygon(
    np.vstack([curve_points, [[r0_max, 100], [0, 100]]]),
    closed=True, transform=ax.transData,
  )
  lower_clip_path = patches.Polygon(
    np.vstack([curve_points, [[r0_max, 0], [0, 0]]]),
    closed=True, transform=ax.transData,
  )

  D = {
    'influenza': dict(
      label = "Influenza",
      r0_span = [1.0, 2.0],
      colour = "red",
    ),
    'sars': dict(
      label = "SARS",
      r0_span = [2.0, 4.0],
      colour = "#00cc55",
    ),
    'polio': dict(
      label = "Polio",
      r0_span = [5.0, 7.0],
      colour = "#0066ee",
    ),
    'chickenpox': dict(
      label = "Chickenpox",
      r0_span = [10.0, 12.0],
      colour = "orange",
    ),
    'measles': dict(
      label = "Measles",
      r0_span = [12.0, 18.0],
      colour = "#ff00cc",
    ),
    'urcovid19': dict(
      label = "Covid-19 (Ancestral)",
      r0_span = [2.3, 3.5],
      colour = "#ffdd00",
      special_hatching = True,
      #hatch = r"//",
      hatch = r"C",
    ),
    'covid19alpha': dict(
      #label = "Covid-19-\N{GREEK SMALL LETTER ALPHA}",
      label = "Covid-19 (Alpha)",
      r0_span = [4.0, 5.0],
      colour = "#ffdd00",
      special_hatching = True,
      #hatch = r"\\",
      hatch = r"CC",
    ),
    'covid19delta': dict(
      label = "Covid-19 (Delta)",
      r0_span = [5.0, 8.0],
      colour = "#ffdd00",
      special_hatching = True,
      #hatch = r"///",
      hatch = r"c",
    ),
  }

  # ~ D = {
    # ~ 'lurgy': dict(
      # ~ label = "Lurgy",
      # ~ r0_span = [1.5, 8.0],
      # ~ colour = "red",
    # ~ ),
  # ~ }

  for key, drec in D.items():

    ver = 1
    band_alpha = 0.33

    r0_span = drec['r0_span']
    hit_span = [100 * hit_fn(x_i, ver) for x_i in r0_span]
    colour = drec.get('colour', None)
    label = drec.get("label", None)
    hatch=drec.get("hatch", None)

    if drec.get('special_hatching', False):
      # Special Covid-19 hatching

      v_xy = [
        (r0_span[0], 0),
        (r0_span[1], 0),
        (r0_span[1], 200),
        (r0_span[0], 200),
      ]
      vpatch = patches.Polygon(
        v_xy,
        facecolor=colour,
        closed=True,
        fill=False,
        edgecolor=colour,
        hatch=hatch,
        label=label,
      )
      vp = ax.add_patch(vpatch)
      vp.set_clip_path(lower_clip_path)

      h_xy = [
        (0, hit_span[0]),
        (r0_max, hit_span[0]),
        (r0_max, hit_span[1]),
        (0, hit_span[1]),
      ]

      hpatch = patches.Polygon(
        h_xy,
        facecolor=colour,
        closed=True,
        fill=False,
        edgecolor=colour,
        hatch=hatch,
      )
      hp = ax.add_patch(hpatch)
      hp.set_clip_path(upper_clip_path)

      lpatch = patches.Patch(
        facecolor=colour,
        fill=False,
        edgecolor=colour,
        hatch=hatch,
        label=label
      )
      legend_patches.append(lpatch)
      labels.append(label)

    else:
      # soft-edge bands

      image = soft_edge_grad_image(colour)
      vim = ax.imshow(
        image,
        aspect='auto',
        extent=(*r0_span, 0, 100),
        origin='lower',
        interpolation='bicubic',
        alpha=band_alpha,
      )
      vim.set_clip_path(lower_clip_path)

      gdivs = 40
      x = np.linspace(*r0_span, gdivs, endpoint=True)
      y = np.array([100 * hit_fn(x_i, ver) for x_i in x])
      xg = np.linspace(0, 1, gdivs, endpoint=True)
      yg = (y - hit_span[0]) / (hit_span[1] - hit_span[0])
      imx = np.interp(xg, yg, xg)
      row = image[0]
      row1 = np.array(resample_lanczos2(row, imx))
      row1 = np.clip(row1, 0, 1)

      image1 = np.swapaxes(np.expand_dims(row1, 0), 0, 1)

      him = ax.imshow(
        image1,
        aspect='auto',
        extent=(0, r0_max, *hit_span),
        origin='lower',
        interpolation='bicubic',
        alpha=band_alpha,
      )
      him.set_clip_path(upper_clip_path)

      lpatch = patches.Patch(
        facecolor=colour,
        edgecolor=None,
        label=label,
        alpha=band_alpha,
      )
      legend_patches.append(lpatch)
      labels.append(label)

  legend = ax.legend(
    legend_patches,
    labels,
  )

  plt.xlabel("Basic Reproduction Number ($R_0$)")
  plt.ylabel("Herd Immunity Threshold (%)")
  plt.title("Herd Immunity Threshold (of Total Population)\n"
            "for Selected Vaccine Efficacies", fontsize=15)
  major_gl_props = dict(
    color = 'black',
    linestyle = "solid",
    linewidth = 0.2,
    alpha = 0.5
  )
  minor_gl_props = dict(
    color = 'black',
    linestyle = "solid",
    linewidth = 0.2,
    alpha = 0.25
  )
  plt.grid(True, 'major', 'both', **major_gl_props)
  plt.grid(True, 'minor', 'both', **minor_gl_props)
  ax.set_axisbelow(True)

  plt.savefig("hit_covid19.png", dpi=300)
  plt.show()


if __name__ == '__main__':
  main()