python 热成像图绘制

我的excel是3个旗县，有“name”列，记录的是乌审旗、杭锦旗和鄂托克前旗，有“time”列记录的是2013、2014、2017年5-9月，记录格式是2013-5这样；有“PA_rank”"SPI1_rank""MI_rank""MCI_rank""RH_rank"列，就是各干旱指数的干旱等级，记录的格式是1、5这样（1是轻旱、5是特旱），那我想做个热成像图，图片分为三行，第一行是乌审旗、第二行是杭锦旗、第三行是鄂托克前旗，第一行的第一张图是2013年5-9月各干旱指数，横轴是时间5-9月，纵轴是各干旱指数由下到上排列是PA、SPI、MI、MCI、RH这样，第二行的第二张图是2014年5-9月各干旱指数，横轴是时间5-9月，纵轴是各干旱指数由下到上排列是PA、SPI、MI、MCI、RH这样，第三行的第三张图是2017年5-9月各干旱指数，横轴是时间5-9月，纵轴是各干旱指数由下到上排列是PA、SPI、MI、MCI、RH这样；第二行的杭锦旗、第三行的鄂托克前旗图仿照第一行做个一样的。热成像图的颜色，无旱是绿色，轻旱是黄色，中旱是橙色，重旱是红色，特旱是深红色。

import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from matplotlib.colors import ListedColormap
matplotlib.use('TkAgg')
matplotlib.rc("font",family='Microsoft YaHei')


# 假设数据已经加载到df中，如果没有，请取消下面的注释并替换为您的文件路径
df = pd.read_excel('G:/装备科/论文/05-鄂尔多斯干旱/热成像图/3旗县干旱.xlsx')


# 定义干旱等级标签和颜色
drought_levels = ['None', 'Mild', 'Moderate', 'Severe', 'Extreme']
colors = ['green', 'yellow', 'orange', 'red', 'darkred']
cmap = ListedColormap(colors)

# 获取唯一的旗县和年份
counties = df['name'].unique()
years = sorted(df['time'].str[:4].unique())

# 创建子图
fig, axes = plt.subplots(len(counties), len(years), figsize=(10, 5),
                         sharex=True, sharey=True, squeeze=False)
# 关键调整：仅修改这一行，其他代码不变！
plt.subplots_adjust(wspace=0.01, hspace=0.01)  # 调整间距

# 为每个旗县和年份创建热图
for i, county in enumerate(counties):
    for j, year in enumerate(years):
        # 筛选数据
        county_year_data = df[(df['name'] == county) & (df['time'].str.startswith(year))]

        # 提取5-9月的数据并排序
        county_year_data = county_year_data.sort_values('time')

        # 提取干旱指数排名数据
        rank_data = county_year_data[['PA_rank', 'SPI1_rank', 'MI_rank', 'MCI_rank', 'RH_rank']].T

        # 创建热图
        sns.heatmap(rank_data, ax=axes[i, j], cmap=cmap, vmin=1, vmax=5,
                    cbar=False, linewidths=0.5, square=False, linecolor='black')

        # 设置标题和轴标签
        if i == 0:
            axes[i, j].set_title(f'{year}',fontsize=15)
        if j == 0:
            axes[i, j].set_ylabel(county,fontsize=15)

        # 设置x轴为月份
        axes[i, j].set_xticks(np.arange(5) + 0.5)
        axes[i, j].set_xticklabels(['5', '6', '7', '8', '9'],fontsize=15)

        # 设置y轴为干旱指数
        axes[i, j].set_yticks(np.arange(5) + 0.5)
        axes[i, j].set_yticklabels(['PA', 'SPI', 'MI', 'MCI', 'Rsm'],fontsize=15)

# 添加颜色条
# 添加横坐标总标题
fig.text(0.5, 0.01, '月份', ha='center', fontsize=15)
cbar_ax = fig.add_axes([0.87, 0.15, 0.02, 0.7])
cbar = plt.colorbar(axes[0, 0].collections[0], cax=cbar_ax)
cbar.set_ticks([1.4, 2.2, 3.0, 3.8, 4.6])
cbar.set_ticklabels(drought_levels,fontsize=15)

# 将标题放在颜色条上方
cbar.ax.set_title('Rank', pad=10, fontsize=15)  # pad调整标题与颜色条的距离
cbar.ax.title.set_position((1.5, 1.05))  # (x, y) 相对于颜色条的坐标

plt.subplots_adjust(
    left=0.1,      # 左边距
    right=0.85,    # 右边距（给颜色条留空间）
    bottom=0.1,    # 底部边距（给"月份"标题留空间）
    top=0.9,       # 顶部边距
    wspace=0.02,   # 年份间距（2013/2014/2017）
    hspace=0.02    # 旗县间距
)
plt.show()