How can I annotate a cell grid in matplotlib?

Question

first, I'm sorry if my question seems really simple, but I am a python beginner. I've been writing for a while this code modeling the spread of an epidemic using cellular automata. The code is the following :

import matplotlib.colors
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors
from matplotlib.animation import FuncAnimation
import matplotlib.animation as ani
import random as rd
import random
import copy
from matplotlib.colors import ListedColormap


cmap=ListedColormap(['k','w','r','b'])
dead=0
untouched=1
ill=2
recovered=3

previous_state_matrix = []
current_state_matrix = []


n=50 #number of array (table of 50*50 : 2500 cells)

def init_graph():
    plt.hlines(y=np.arange(0, 50)+0.5, xmin=np.full(50, 0)-0.5,linewidth=0.25,xmax=np.full(50, 50)-0.5, color="grey")
    plt.vlines(x=np.arange(0, 50)+0.5, ymin=np.full(50, 0)-0.5,linewidth=0.25, ymax=np.full(50, 50)-0.5, color="grey")


def init_matrix_array():
    
    global previous_state_matrix
    global current_state_matrix
    global n
    for i in range (n):
        previous_state_matrix.append([])
        current_state_matrix.append([])
        for j in range (n):
            previous_state_matrix[i].append(1)
            current_state_matrix[i].append(1)
    previous_state_matrix[n//2][n//2]=2
    current_state_matrix[n//2][n//2]=2


def next_to_ill_cell(i,j):
    for x in [i-1,i,i+1]:
        for y in [j-1,j,j+1]:
            if not((x==i and y==j) or x==-1 or y==-1 or x==n or y==n):
                if previous_state_matrix[x][y]==ill:
                    return True

#Rules
def process_next_state ():
    global previous_state_matrix
    global current_state_matrix
    previous_state_matrix = copy.deepcopy(current_state_matrix)
    for i in range (n) :
        for j in range (n) :
            if previous_state_matrix[i][j]==untouched
                if next_to_ill_cell(i,j)== True:
                    k=rd.random()#random
                    if k >=0.5:
                        current_state_matrix[i][j]=ill
                    else:
                        current_state_matrix[i][j]=untouched

            if previous_state_matrix[i][j]==ill:
                s=rd.random()
                if s>= 0.02875:
                    current_state_matrix[i][j]=recovered
                else:
                    current_state_matrix[i][j]=dead

init_graph()
init_matrix_array()

print(current_state_matrix)

for k in range (len(current_state_matrix)):
    for l in range (len(current_state_matrix[1])):
        if current_state_matrix[k][l]==2:
            plt.imshow(current_state_matrix, cmap=cmap, vmin=0, vmax=3)
            process_next_state()
            plt.pause(1)
            plt.imshow(current_state_matrix, cmap=cmap, vmin=0, vmax=3)

plt.show()

I'm sure that this code could be largely improved, but as I said, I'm a beginner. Now what I want to do is to add an annotation at the bottom of the plot to indicate the number of death and the number of recoveries. I also want to indicate at the top of the plot the number of days elapsed, which would correspond to the transition from one table of numbers to the next.

Could you please help me?

Answer 1

To annotate the number of death, recoveries and day, you need at least calculate these three variables. You need something like that:

day = 1
while day < 1000:
    n_death = ...
    n_recovery = ...
    ...

    plt.imshow(...)
    plt.text(x, y, f'day = {day}') 
    plt.text(x, y, f'n_death  = {n_death }') 
    plt.text(x, y, f'n_recovery = {n_recovery}') 

I restructured your code since it is difficult to read. Changing global variables inside a function is very dangerous. You should avoid this.

import matplotlib.colors
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors
from matplotlib.animation import FuncAnimation
import matplotlib.animation as ani
import random as rd
#import random
import copy
from matplotlib.colors import ListedColormap


def init_graph():
    plt.hlines(y=np.arange(n)+0.5, xmin=-0.5, xmax=n-0.5, linewidth=0.25, color="grey")
    plt.vlines(x=np.arange(n)+0.5, ymin=-0.5, ymax=n-0.5, linewidth=0.25, color="grey")


def init_matrix_array(n):
    m = np.ones((n, n))
    m[n//2][n//2]=2
    return m.tolist()


def next_to_ill_cell(current_state_matrix, i, j):
    for x in [i-1,i,i+1]:
        for y in [j-1,j,j+1]:
            if not((x==i and y==j) or x==-1 or y==-1 or x==n or y==n):
                if current_state_matrix[x][y]==ill:
                    return True
    return False

#Rules
def process_next_state (current_state_matrix):
    previous_state_matrix = copy.deepcopy(current_state_matrix)
    for i in range (n) :
        for j in range (n) :
            if previous_state_matrix[i][j] == untouched:
                if next_to_ill_cell(previous_state_matrix, i, j)== True:
                    k = rd.random()#random
                    if k >= 0.5:
                        current_state_matrix[i][j] = ill
                    else:
                        current_state_matrix[i][j] = untouched

            if previous_state_matrix[i][j]==ill:    # elif???????
                s = rd.random()
                if s >= 0.02875:
                    current_state_matrix[i][j] = recovered
                else:
                    current_state_matrix[i][j] = dead

    return current_state_matrix

def number_of_death(current_state_matrix):
    n_death = 0
    for i in range(n):
        for j in range(n):
            if current_state_matrix[i][j] == dead:
                n_death += 1
    return n_death

def number_of_recovery(previous_state_matrix, current_state_matrix):
    """Calculate the number of recovery"""

# for k in range (len(current_state_matrix)):
#     for l in range (len(current_state_matrix[1])):
#         if current_state_matrix[k][l]==2:
#             plt.imshow(current_state_matrix, cmap=cmap, vmin=0, vmax=3)
#             process_next_state()
#             plt.pause(1)
#             plt.imshow(current_state_matrix, cmap=cmap, vmin=0, vmax=3)

if __name__ == '__main__':
    cmap = ListedColormap(['k','w','r','b'])
    dead = 0
    untouched = 1
    ill = 2
    recovered = 3
    n = 50 #number of array (table of 50*50 : 2500 cells)

    init_graph()
    current_state_matrix = init_matrix_array(n)
    day = 1
    while day < 10:
        previous_state_matrix = current_state_matrix

        # Number of death
        n_death = number_of_death(current_state_matrix)
        plt.imshow(current_state_matrix, cmap=cmap, vmin=0, vmax=3)
        plt.text(25, 5, f'day = {day}', horizontalalignment='center')
        plt.text(25, 45, f'number of death = {n_death}', horizontalalignment='center')
        current_state_matrix = process_next_state(current_state_matrix)
        day += 1
        plt.pause(1)
   
    plt.show()

Answer 2

What I meant was that the numbers are overlapping state after state, the 2nd appears over the first ect.. I've hope i'de made myself clear.. I get the following output : unreadable numbers. So I would like the new numbers to replace the old ones and not just appear over them, but I have no idea how to do it.