在python中在图形上创建框

Question

I want this kind of result . I want my code to read elements of a text file and if element=='healthy' it should create a box in a graph and its color should be green ('healthy written in box'). else if element=='unhealthy' it should create a box and its color should be red (with 'unhealthy written in box'). boxes should be horizontally aligned, and if more than 5 then remaining should start from the next row. (every row should contain only 5 boxes or less). The end result should display a graph that contains boxes, red denoting 'unhealthy' and green denoting 'healthy'

I found the following code, but it is not working they way I want it to.

import matplotlib.pyplot as plt



plt.style.use('seaborn-white')
import numpy as np
from matplotlib import colors


#open text file (percen) that contains healthy/unhealthy
with open('percen.txt', 'r') as f:   
    result= [int(line) for line in f]


data = np.random.rand(10,10) * 20

cmap = colors.ListedColormap(['green'])

cmap1 = colors.ListedColormap(['red'])

bounds = [0,10,20]
norm = colors.BoundaryNorm(bounds, cmap.N)

fig, ax = plt.subplots(2,5 , sharex='col', sharey='row')


for i in range(2):

    for j in range(5):
        for element in result:
            if (element=='healthy'):
                ax[i,j].text(1, -3, 'healthy',
                  fontsize=15, ha='center', color='green')
                ax[i,j].imshow(data,cmap=cmap, norm=norm)
            else:
                ax[i,j].text(1, -3, 'unhealthy',
                 fontsize=15, ha='center', color='red')
                ax[i,j].imshow(data,cmap=cmap1,norm=norm)

fig
plt.show()

Answer 1

There are a few different ways you can do this and your code is probably not the best but we can use it as a starting point. Your issue is that you are looping through the plots and then looping through your data again for each plot. Your current code also adds text above the plot. If you want the text above I would recommend adding the label as a title, otherwise when you set your text inside the plot you need to specify the coordinates within the grid.

Below is a modified form of your code, play around with it some more to get what you want.

import matplotlib.pyplot as plt

plt.style.use('seaborn-white')
import numpy as np
from matplotlib import colors


result = ['healthy', 'unhealthy', 'healthy', 'unhealthy', 'healthy', 'unhealthy', 'healthy', 'healthy', 'unhealthy', 'unhealthy']

data = np.random.rand(10,10) * 20

cmap = colors.ListedColormap(['green'])

cmap1 = colors.ListedColormap(['red'])

bounds = [0,10,20]
norm = colors.BoundaryNorm(bounds, cmap.N)

fig, ax = plt.subplots(2,5 , sharex='col', sharey='row',figsize=(15,8)) # Added figsize to better show your plot



element_index = 0
for i in range(2):

    for j in range(5):
        element = result[element_index] #Instead of the for loop, get the corresponding element

        if (element=='healthy'):
            ax[i,j].text(4.5,4.5, 'healthy',fontsize=15, ha='center' ,color='black',zorder=100) #Change zorder so label is over plot
            ax[i,j].imshow(data,cmap=cmap, norm=norm)
            ax[i,j].set_yticklabels('') #To remove arbitrary numbers on y axis
            ax[i,j].set_xticklabels('') #To remove arbitrary numbers on y axis


        elif element == 'unhealthy':
            ax[i,j].text(4.5,4.5, 'unhealthy',fontsize=15, ha='center' ,color='black',zorder=100)
            ax[i,j].imshow(data,cmap=cmap1,norm=norm)
            ax[i,j].set_yticklabels('') #To remove arbitrary numbers on y axis
            ax[i,j].set_xticklabels('') #To remove arbitrary numbers on x axis



        element_index+=1       #Add 1 to the index so we get the next value for the next plot

fig
plt.show()