不同长度数组绘制电力供需堆叠柱状图出现IndexError如何解决

核心思路是统一不同年份发电侧、消费侧的数组长度，预先给进口、出口项预留固定位置，无需根据进出口正负动态增减数组元素，仅对未触发的项补0即可，即可解决数组长度不一致无法构造二维数组的问题。

具体修改逻辑 #

• 发电侧固定为6个维度，顺序为可再生能源、天然气、煤炭、储能、核电、进口，当年进口为负时，进口项填0
• 消费侧固定为3个维度，顺序为出口、终端消费、储能，当年进口为正时，出口项填0

修改后完整可运行代码 #

import matplotlib.pyplot as plt
import numpy as np

# 2020年数据
generation_2020 = [20, 30, 40, 20, 10]
consumption_2020 = [50,50]
import_2020 = 20
# 统一长度处理
if import_2020 > 0:
    generation_2020 = np.append(generation_2020, import_2020)
    consumption_2020 = np.insert(consumption_2020, 0, 0, axis=0) # 出口项补0
else:
    generation_2020 = np.append(generation_2020, 0) # 进口项补0
    consumption_2020 = np.insert(consumption_2020, 0, import_2020*-1, axis=0)

# 2025年数据
generation_2025 = np.array([20, 20, 20, 20, 10])
consumption_2025 = np.array([50,50])
import_2025 = -10
# 统一长度处理
if import_2025 > 0:
    generation_2025 = np.append(generation_2025, import_2025)
    consumption_2025 = np.insert(consumption_2025, 0, 0, axis=0) # 出口项补0
else:
    generation_2025 = np.append(generation_2025, 0) # 进口项补0
    consumption_2025 = np.insert(consumption_2025, 0, import_2025*-1, axis=0)

# 构造二维数组不会再报错
data = np.array([generation_2020, generation_2025])
data2 = np.array([consumption_2020, consumption_2025])
# 补充出口标签和原有标签对应
label_data = ['renewables','gas','cole', 'storage', 'nuclear', 'import', 'export', 'consumption', 'storage']

x2 = ['2020', '2025']
x3 = ['Erzeugung' , '\n\n\n\n\n\n2020', 'Verbrauch', 'Erzeugung' , '\n\n\n\n\n\n2025', 'Verbrauch']
x_pos = np.arange(len(x2))
width = 0.2
x = list()
for i in x_pos:
    x.extend([i-width, i, i+width])  

fig, ax = plt.subplots()
# 绘制发电侧堆叠柱
for i in range(data.shape[1]):
    bottom = np.sum(data[:, 0:i], axis=1)
    ax.bar(x_pos - width, data[:, i], bottom=bottom, width=width, label=f"label {i}")

# 绘制消费侧堆叠柱
for i in range(data2.shape[1]):
    bottom = np.sum(data2[:, 0:i], axis=1)
    ax.bar(x_pos+ width, data2[:, i], bottom=bottom, width=width, label=f"label {i}")

ax.set_ylabel('[TWh]')
ax.set_xticks(x)
ax.set_xticklabels(x3)
ax.tick_params(axis='x', which='both',length=0, labelsize = 10)

for label in ax.get_xmajorticklabels():
    if 'u' in label.get_text(): label.set_rotation(90)
ax.margins(y=0.1)

plt.legend(label_data, bbox_to_anchor=(1, 1))
fig.tight_layout()
plt.show()

额外说明 #

如果后续需要新增更多年份的数据，只需要套用相同的长度统一逻辑，即可保证所有年份的发电、消费数组长度一致，不会再触发数组维度不匹配的报错。

内容的提问来源于stack exchange，提问作者Elias