绘制ICMP热力图时Numpy reshape方法报ValueError问题

错误提示 #

Traceback (most recent call last):
  File "/NetworkAutomation/Python_Dev/icmp_heatmap/create_plot.py", line 107, in <module>
    main()
  File "/NetworkAutomation/Python_Dev/icmp_heatmap/create_plot.py", line 72, in main
    grid_data = np.reshape(grid_data, (num_devices, num_devices))
  File "<__array_function__ internals>", line 6, in reshape
  File "/NetworkAutomation/Python_Dev/icmp_heatmap/pyVENV/lib/python3.7/site-packages/numpy/core/fromnumeric.py", line 298, in reshape
    return _wrapfunc(a, 'reshape', newshape, order=order)
  File "/NetworkAutomation/Python_Dev/icmp_heatmap/pyVENV/lib/python3.7/site-packages/numpy/core/fromnumeric.py", line 57, in _wrapfunc
    return bound(*args, **kwds)
ValueError: cannot reshape array of size 2 into shape (2,2)

Process finished with exit code 1

问题背景 #

这段Python代码通过Napalm连接网络设备执行ICMP ping测试，收集RTT数据后用Numpy的reshape处理生成热力图，但运行时触发上述错误。原代码如下：

import sys
import napalm
import yaml
import seaborn as sns
import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
from napalm.base.exceptions import ConnectionClosedException, ConnectionException


def main():
    with open(r"device_list.yml", encoding="utf-8") as file:
        device_list = yaml.safe_load(file)
        driver = napalm.get_network_driver("ios")

        for device in device_list:
            device_drv = driver(
                hostname=device,
                username="abcdce",
                password="asdfasdf",
                timeout=60,
                optional_args={
                    "secret":"cisco"
                },
            )
            try:
                print(f"Verifying connection to {device}")
                device_drv.open()
            except ConnectionException:
                print(f"Unable to connect to {device}. Removing from the list.")

        device_list.remove(device)
        print(device_list)

        grid_data = np.array([])
        for device in device_list:
            device_drv = driver(
                hostname=device,
                username="abcdef",
                password="asdfasdf",
                timeout=60,
                optional_args={
                    "secret":"cisco",
                    "global_delay_factor":2
                },
            )

            print(f"Initiating tests {device}")

            device_drv.open()
            for device2 in device_list:
                try:
                    ping_results = device_drv.ping(device2, timeout=1,  count=10, source_interface="Loopback0")
                    if "success" in ping_results:
                        avg_response = ping_results["success"]["rtt_avg"]
                        print(avg_response)
                        if avg_response == 0.0:
                            grid_data = np.append(grid_data, 9999)
                        else:
                            grid_data = np.append(grid_data, avg_response)
                    else:
                        print(f"Cannot ping {device2} from {device}")
                        grid_data = np.append(grid_data, np.nan)
                except ConnectionClosedException as e:
                    print(f"Error: {e}")
                    print("Ping command took too long ... Need tp exit program to prevent data alignment error.")
                    sys.exit()
            print(grid_data)

            num_devices = len(device_list)

            grid_data = np.reshape(grid_data, (num_devices, num_devices))
            print(grid_data)
            sns.set_theme()
            df = pd.DataFrame(grid_data,
                              index=device_list,
                              columns=device_list)
            print(df)
            df.to_csv("output.csv")

            plt.figure(figsize=(35,30),
                       edgecolor="black")
            plt.rcParams["xtick.bottom"] = plt.rcParams["xtick.labelbottom"] = False
            plt.rcParams["xtick.top"] = plt.rcParams["xtick.labeltop"] = True

            sns.set_context("paper", font_scale=4.0)
            ax = sns.heatmap(
                df,
                annot=True,
                cmap="Blues",
                linewidths=0,
                vmin=0,
                vmax=300,
                fmt="g",
                cbar=True,
                square=True,
                annot_kws={"fontsize":30,
                           "color":"black",
                           "visible":True},
            )
            ax.set_xticklabels(ax.get_xticklabels(), rotation=90)
            ax.set_yticklabels(ax.get_xticklabels(), rotation=0)
            plt.savefig("output.png")
            plt.show()

if __name__ == "__main__":
    main()

问题根源 #

1. 绘图代码缩进错误：reshape、热力图生成等代码被嵌套在for device in device_list循环内部，导致每完成一个设备的ping测试就尝试将仅包含num_devices个元素的grid_data重塑为(num_devices, num_devices)的二维数组（需要num_devices²个元素），元素数量不匹配触发错误。
2. 设备列表维护逻辑错误：device_list.remove(device)语句缩进错误，不在except ConnectionException块内且位于循环外部，导致无论设备连接是否成功，都会移除循环的最后一个设备，破坏了设备列表的完整性和长度准确性。

修复方案 #

1. 修正设备列表验证逻辑 #

仅移除连接失败的设备，同时迭代列表副本避免索引混乱，并且及时关闭连接：

# 替换原设备验证部分代码
with open(r"device_list.yml", encoding="utf-8") as file:
    device_list = yaml.safe_load(file)
    driver = napalm.get_network_driver("ios")

valid_devices = []
for device in device_list:
    device_drv = driver(
        hostname=device,
        username="abcdce",
        password="asdfasdf",
        timeout=60,
        optional_args={"secret":"cisco"},
    )
    try:
        print(f"Verifying connection to {device}")
        device_drv.open()
        valid_devices.append(device)
    except ConnectionException:
        print(f"Unable to connect to {device}. Removing from the list.")
    finally:
        device_drv.close()  # 确保连接关闭

device_list = valid_devices
print(device_list)

2. 调整绘图代码缩进 #

将reshape、热力图生成等代码移到所有ping测试完成之后，确保grid_data收集到全部num_devices²个元素：

# 移到外层循环之外，确保所有ping数据收集完成后再处理
grid_data = np.array([])
for device in device_list:
    device_drv = driver(
        hostname=device,
        username="abcdef",
        password="asdfasdf",
        timeout=60,
        optional_args={
            "secret":"cisco",
            "global_delay_factor":2
        },
    )

    print(f"Initiating tests from {device}")
    try:
        device_drv.open()
        for device2 in device_list:
            try:
                ping_results = device_drv.ping(device2, timeout=1, count=10, source_interface="Loopback0")
                if "success" in ping_results:
                    avg_response = ping_results["success"]["rtt_avg"]
                    print(f"{device} -> {device2}: {avg_response}")
                    grid_data = np.append(grid_data, 9999 if avg_response == 0.0 else avg_response)
                else:
                    print(f"Cannot ping {device2} from {device}")
                    grid_data = np.append(grid_data, np.nan)
            except ConnectionClosedException as e:
                print(f"Error: {e}")
                print("Ping command took too long ... Exiting to prevent data alignment error.")
                sys.exit()
    finally:
        device_drv.close()  # 及时关闭连接

print(grid_data)

num_devices = len(device_list)
# 此时grid_data长度为num_devices*num_devices，可正常重塑
grid_data = np.reshape(grid_data, (num_devices, num_devices))
print(grid_data)

# 后续DataFrame、绘图代码保持不变
sns.set_theme()
df = pd.DataFrame(grid_data,
                  index=device_list,
                  columns=device_list)
print(df)
df.to_csv("output.csv")

plt.figure(figsize=(35,30),
           edgecolor="black")
plt.rcParams["xtick.bottom"] = plt.rcParams["xtick.labelbottom"] = False
plt.rcParams["xtick.top"] = plt.rcParams["xtick.labeltop"] = True

sns.set_context("paper", font_scale=4.0)
ax = sns.heatmap(
    df,
    annot=True,
    cmap="Blues",
    linewidths=0,
    vmin=0,
    vmax=300,
    fmt="g",
    cbar=True,
    square=True,
    annot_kws={"fontsize":30,
               "color":"black",
               "visible":True},
)
ax.set_xticklabels(ax.get_xticklabels(), rotation=90)
ax.set_yticklabels(ax.get_xticklabels(), rotation=0)
plt.savefig("output.png")
plt.show()

3. 其他优化点 #

• 使用finally块确保网络连接始终被关闭，避免连接泄漏；
• 优化打印信息，明确显示ping测试的源和目标设备，便于排查问题；
• 修复了原代码中"Need tp exit"的拼写错误为"Need to exit"。

完整修复后代码 #

import sys
import napalm
import yaml
import seaborn as sns
import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
from napalm.base.exceptions import ConnectionClosedException, ConnectionException


def main():
    with open(r"device_list.yml", encoding="utf-8") as file:
        device_list = yaml.safe_load(file)
        driver = napalm.get_network_driver("ios")

    # 验证设备连接，筛选可用设备
    valid_devices = []
    for device in device_list:
        device_drv = driver(
            hostname=device,
            username="abcdce",
            password="asdfasdf",
            timeout=60,
            optional_args={"secret":"cisco"},
        )
        try:
            print(f"Verifying connection to {device}")
            device_drv.open()
            valid_devices.append(device)
        except ConnectionException:
            print(f"Unable to connect to {device}. Removing from the list.")
        finally:
            device_drv.close()

    device_list = valid_devices
    if not device_list:
        print("No valid devices left. Exiting.")
        sys.exit()
    print(f"Valid devices: {device_list}")

    # 收集所有设备间的ping数据
    grid_data = np.array([])
    for device in device_list:
        device_drv = driver(
            hostname=device,
            username="abcdef",
            password="asdfasdf",
            timeout=60,
            optional_args={
                "secret":"cisco",
                "global_delay_factor":2
            },
        )

        print(f"Initiating tests from {device}")
        try:
            device_drv.open()
            for device2 in device_list:
                try:
                    ping_results = device_drv.ping(device2, timeout=1, count=10, source_interface="Loopback0")
                    if "success" in ping_results:
                        avg_response = ping_results["success"]["rtt_avg"]
                        print(f"{device} -> {device2}: {avg_response}")
                        grid_data = np.append(grid_data, 9999 if avg_response == 0.0 else avg_response)
                    else:
                        print(f"Cannot ping {device2} from {device}")
                        grid_data = np.append(grid_data, np.nan)
                except ConnectionClosedException as e:
                    print(f"Error: {e}")
                    print("Ping command took too long ... Exiting to prevent data alignment error.")
                    sys.exit()
        finally:
            device_drv.close()

    print(f"Collected data: {grid_data}")

    # 处理数据并生成热力图
    num_devices = len(device_list)
    grid_data = np.reshape(grid_data, (num_devices, num_devices))
    print(f"Reshaped data:\n{grid_data}")

    sns.set_theme()
    df = pd.DataFrame(grid_data,
                      index=device_list,
                      columns=device_list)
    print(f"DataFrame:\n{df}")
    df.to_csv("output.csv")

    plt.figure(figsize=(35,30),
               edgecolor="black")
    plt.rcParams["xtick.bottom"] = plt.rcParams["xtick.labelbottom"] = False
    plt.rcParams["xtick.top"] = plt.rcParams["xtick.labeltop"] = True

    sns.set_context("paper", font_scale=4.0)
    ax = sns.heatmap(
        df,
        annot=True,
        cmap="Blues",
        linewidths=0,
        vmin=0,
        vmax=300,
        fmt="g",
        cbar=True,
        square=True,
        annot_kws={"fontsize":30,
                   "color":"black",
                   "visible":True},
    )
    ax.set_xticklabels(ax.get_xticklabels(), rotation=90)
    ax.set_yticklabels(ax.get_xticklabels(), rotation=0)
    plt.savefig("output.png")
    plt.show()

if __name__ == "__main__":
    main()

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