python - 如何在 matplotlib 中注释单元格网格？

Question

首先，如果我的问题看起来很简单，我很抱歉，但我是 python 初学者。我已经写了一段时间了，这段代码使用元胞自动机对流行病的传播进行建模。代码如下：

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()

我确信这段代码可以大大改进，但正如我所说，我是一个初学者。现在我要做的是在图的底部添加一个注释，以指示死亡人数和康复人数。我还想在图的顶部指出经过的天数，这将对应于从一个数字表到下一个数字表的转换。

请你帮助我好吗？

标签： pythonmatplotlibannotationscellular-automata

Answer 1

要标注死亡人数、康复人数和天数，至少需要计算这三个变量。你需要这样的东西：

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}') 

我重组了你的代码，因为它很难阅读。在函数内部更改全局变量是非常危险的。你应该避免这种情况。

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()