Social distancing simulation animation

matplotlib_animation.py

"""
https://youtu.be/rQwnOZFuYsU
In this video I use the animation tools in the matplotlib Python library
to help create a dynamic graph as seen in the Washington Post article
"Why outbreaks like coronavirus spread exponentially, and how to 'flatten the curve'".
Andrej Bauers code:
https://github.com/andrejbauer/social-distancing-simulator
Help translate it into more languages and tell Andrej I sent you.
Andrej Bauers page:
https://social-distancing-simulator.andrej.com/english.html
The Washington Post article:
https://www.washingtonpost.com/graphics/2020/world/corona-simulator/
The MathOverflow thread:
https://stackoverflow.com/questions/9401658/how-to-animate-a-scatter-plot
The color schemes:
https://matplotlib.org/3.1.0/tutorials/colors/colormaps.html
"""

import time
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation


_FPS = 90 # seconds
INTERVAL = 1000 / _FPS
NUM_POINTS = 20 # 5
X_RANGE = 20
Y_RANGE = 20


def _r():
    arr = np.random.random(NUM_POINTS) # each in [0, 1]
    return arr - 1/2


def _raw_data_stream():
    F_dX = 1/20
    F_dY = 1/20
    F_dSIZE = 1/20
    F_dCOL = 1/5

    x = _r() * X_RANGE / 2
    y = _r() * Y_RANGE / 2
    size = _r() + 1/2
    col = _r() + 1/2

    idx = 0
    while True:
        x += F_dX * _r()
        y += F_dY * _r()
        size += F_dSIZE * _r()
        col += F_dCOL * _r()
        data = np.c_[x, y, size, col]
        #time.sleep(0.3); print("x\t\ty\t\ts\t\tc\n{}".format(data))

        yield idx, data
        idx += 1


class PlotStream:
    def __init__(self):
        # Set up figure and axes
        self.__fig, self.__ax = plt.subplots() # plt hereby tied to this class context
        self.__ax.axis([-X_RANGE / 2, X_RANGE / 2, -Y_RANGE / 2, Y_RANGE / 2])

        # Set up stream and first frame
        self.__stream = _raw_data_stream()
        self.__first_plot()

    def run(self):
        _animation = animation.FuncAnimation(self.__fig, self.__next_plot, interval=INTERVAL, blit=True)
        # blit related to event loop on Mac?
        
        plt.show()

    def __first_plot(self):
        COLOR_MAP = "jet" # "Set1" # See https://matplotlib.org/3.1.0/tutorials/colors/colormaps.html
        
        idx, data = next(self.__stream)
        assert idx==0, f"Index assert failed with idx={idx}."
        
        x, y, size, color = data.T
        self.__plot = self.__ax.scatter(x, y, s=size, c=color, vmin=0, vmax=1, cmap=COLOR_MAP, edgecolor="black")
        self.__plot.set_sizes([500 for _ in range(NUM_POINTS)])
        
        return self.__plot,

    def __next_plot(self, i):
        idx, data = next(self.__stream)
        assert i <= 1 or idx == i + 2, f"Index assert failed with idx={idx}, i={i}." # remove if this makes problems
        
        self.__plot.set_offsets(data[:, :2]) # change coordinate
        self.__plot.set_sizes(300 * abs(data[:, 2])**1.5 + 100) # change size
        self.__plot.set_array(data[:, 3]) # change color
        
        return self.__plot,


if __name__ == '__main__':
    PlotStream().run()