Python中动态扩展数组并保留原有数据、缺失值补NaN的方法

问题背景 #

我有一段Python代码，用于寻找两条曲线的交点——当f(frequency)等于正切或余切函数值时的点。需要保存每个频率下的可用模式值，遍历多个频率来分析模式的传播特性。

在循环中，每个频率（角频率）对应的交点由Intersection工具返回，我只取x值存入X_inter，再将每个X_inter存入更大的arr_KX数组。arr_KX的行对应单个频率，列对应单个模式（第一列是第0阶模式，以此类推）。

核心问题是：循环结束前无法确定最大模式数（即数组所需的列数），而低频对应的模式数少于高频。当前代码在调整arr_KX大小时会丢失之前的数据，理想的结果应该是：

若第一个频率有2个模式，第二个频率有4个模式，则arr_KX应为：
arr_KX = [[m1,m2,nan,nan],[m1,m2,m3,m4]]

目前我只能通过从高到低生成频率（确保初始就有最多模式数）来规避问题，但后续物理研究中这种反转顺序会带来麻烦，需要找到无需反转顺序的解决方案。

原始代码 #

import numpy as np
import matplotlib.pyplot as plt
import Intersection as inter
np.seterr(all='ignore')
plt.clf()

height, width, lenght = 500e-6, 500e-6, 100e-6 
d = height/2
c0 = 3e8 
n_core = 5.2 
n_cladding = 1.
FntSze = 15 

"""
求解以下方程
"""
def cot(angle):
    return np.cos(angle)/np.sin(angle)

def TE_LHS_sym(var1):
    return np.tan(var1*d)

def TE_LHS_antisym(var2):
    return -cot(var2*d)

def TE_RHS(var3,pulsation_):
    return ((pulsation_**2*d**2*
             (n_core**2-n_cladding**2)/(c0**2*var3**2*d**2))-1)**(0.5)

start, end, Nx, NN = 100e9, 1000e9 + 1, 1001 ,21 
##角频率范围的起止值，NN为处理的频率数量，Nx为生成的kx数量

f_range = np.linspace(start,end,NN)
w_range = f_range*2*np.pi 
k = np.linspace(0,4*np.pi,Nx)/d 

tol = 40. ##限制正切/余切函数不连续处的浮点值

y_tan = TE_LHS_sym(k) ##绘制tan(kx*d)
y_tan[y_tan > tol] = np.nan ##限制不连续处的异常值
y_tan[y_tan < -tol] = np.nan

y_cot = TE_LHS_antisym(k) ##绘制cot(kx*d)
y_cot[y_cot > tol] = np.nan
y_cot[y_cot < -tol] = np.nan

init=0
size = 0

##绘图
plt.figure(1)
plt.plot(k*d,y_tan,label='tan($k_{x0}d$)' )
plt.plot(k*d,y_cot)#,label='-cot($k_{x0}d$)')

for init,omega in enumerate(w_range): #角频率循环，从最大值到最小值
        y_RHS = TE_RHS(k,omega) #绘制RHS，与正切/余切的交点即为kx_mode
        y_RHS[y_RHS > tol] = np.nan
        plt.plot(k*d,y_RHS,marker='',label='RHS at $f$ = '+ str(round(omega*1e-9/(2*np.pi),2)) + ' GHz')
       
        ##寻找正切/余切的交点，得到给定w下的kx_mode
        x_inter_tan,y_inter_tan = inter.intersection(k, y_tan, k, y_RHS)
        x_inter_cot,y_inter_cot = inter.intersection(k, y_cot, k, y_RHS) 
        plt.plot(x_inter_cot*d, y_inter_cot, '+k') ##绘制余切与RHS的交点
        plt.plot(x_inter_tan*d, y_inter_tan, '+k') ##绘制正切与RHS的交点
        
        plt.xlabel("$k_{x}d$ (in $rad.m^{-1}$)",fontsize=FntSze)
        plt.ylabel("",fontsize=FntSze)
        
        X_inter = np.sort(np.concatenate([x_inter_cot.copy(),x_inter_tan.copy()])*d) #得到给定w下的kx_m*d
        KX_inter = X_inter/d
        if KX_inter.shape[0] > size :    #确定所需的列数
            size = KX_inter.shape[0]     
            arr_KX = np.zeros((NN,size)) #创建NN*size的矩阵，即w_n*kx_modes
        bb = KX_inter.resize(size)       #高w_m对应更多kx_modes，调整为相同形状
        arr_KX[int(init),:] = KX_inter #将当前w_n的kx_modes存入矩阵对应行

最小可复现示例 #

import numpy as np
from numpy.random import default_rng

NN = [2,2,4,6,6]
size = 0
init = 0
for n in NN:
    KX_inter = default_rng().random((n)) #shape (n,)
    print(KX_inter)
    if KX_inter.shape[0] > size :
        size = KX_inter.shape[0]          #确定所需的列数
        arr_KX = np.zeros((len(NN),size)) #创建Nx*size的矩阵，即w_n*kx_modes
    bb = KX_inter.resize(size)            #高w_m对应更多kx_modes，调整为相同形状
    arr_KX[int(init),:] = KX_inter
    
    init = init + 1
print(arr_KX)

解决方案 #

方法1：先收集所有数据再统一构建数组 #

先把每个频率的模式数据存入列表，循环结束后找到最大长度，再将每个列表填充np.nan到最大长度，最后转换为numpy数组：

import numpy as np
from numpy.random import default_rng

NN = [2,2,4,6,6]
data = []
for n in NN:
    KX_inter = default_rng().random((n))
    data.append(KX_inter)

max_size = max(len(arr) for arr in data)
arr_KX = np.full((len(data), max_size), np.nan)
for i, arr in enumerate(data):
    arr_KX[i, :len(arr)] = arr

print(arr_KX)

这种方法简单直观，完全不需要考虑频率顺序，所有数据先缓存再处理，不会丢失之前的内容。

方法2：动态扩展数组（适合内存敏感场景） #

如果数据量极大，无法一次性缓存所有数据，可以在循环中动态扩展arr_KX：

import numpy as np
from numpy.random import default_rng

NN = [2,2,4,6,6]
arr_KX = np.array([]).reshape(0,0)
init = 0

for n in NN:
    KX_inter = default_rng().random((n))
    current_len = len(KX_inter)
    
    if arr_KX.size == 0:
        # 初始化数组
        arr_KX = np.full((len(NN), current_len), np.nan)
    elif current_len > arr_KX.shape[1]:
        # 扩展列数，填充nan
        new_cols = current_len - arr_KX.shape[1]
        arr_KX = np.hstack([arr_KX, np.full((arr_KX.shape[0], new_cols), np.nan)])
    
    arr_KX[init, :current_len] = KX_inter
    init += 1

print(arr_KX)

每次遇到更长的模式列表时，就给数组添加新的列并填充np.nan，之前的数据会被保留。

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