如何解决复数三次方程根绘制等高线时的根混淆问题？

我实现了一个求解复数三次方程根的函数，针对不同的k₀和k₁值求解后，绘制3组根的实部与虚部等高线，但部分等高线存在不连续现象，明显是不同根的片段错位导致的。尝试过多种排序方法（按实部/虚部排序、条件交换根顺序），均无法彻底解决，请问如何避免根的混淆，得到连续的等高线？

原实现代码 #

import numpy as np
import matplotlib.pyplot as plt

# Constants
Ra = 2e4
Pr = 0.1
Omega = 1e5
zeta = 1e-4
deltaN = 0.05
L = 55

def polynomial(k):
    m = 1
    delta_k = m**2 * np.pi**2 + k[0]**2
    a_3 = delta_k
    a_2 = 1j*(Ra * Pr * delta_k * k[0])/Omega + (Pr + zeta + 1)*delta_k**2
    a_1 = 1j*(Ra * Pr * delta_k**2 * k[0] * (Pr + zeta)/Omega) + k[1] * Pr * zeta * (delta_k**2/L**2 + delta_k) - deltaN * Ra * Pr * k[0]**2 + (Pr * zeta + Pr + zeta) * delta_k**3
    a_0 = 1j*(Pr * zeta * k[0] * (Ra * Pr * delta_k**3/Omega + k[1] * Omega * deltaN * delta_k / L**2)) + Pr * zeta * (k[1] * (Pr * delta_k**3 / L**2 + delta_k**2) - deltaN * Ra * delta_k * k[0]**2 + delta_k**4)
    x_K = np.roots([a_3, a_2, a_1, a_0])
    # x_K = np.sort_complex(x_K)
    x_K = sorted(x_K, key=lambda x: x.imag)
    # x_K = sorted(x_K, key=lambda x: x.real)
    # if x_K[2].imag >= 0: 
    #     x_K[-1], x_K[-2] = x_K[-2], x_K[-1]
    # if x_K[0].imag >= x_K[2].imag:
    #     x_K[0], x_K[-1] = x_K[-1], x_K[0]
    if x_K[0].real >= x_K[1].real:
        x_K[0], x_K[1] = x_K[1], x_K[0]
    # if x_K[1].real >= x_K[2].real:
    #     x_K[1], x_K[2] = x_K[2], x_K[1]
    return x_K


# Create arrays of k[0] and k[1] values for contour plot
k0, k1 = np.linspace(0, 5, 100), np.linspace(0, 5e2, 100)
K0, K1 = np.meshgrid(k0, k1)

# Get roots for each pair of k[0], k[1] value
roots = np.array([polynomial([K0[i, j], K1[i, j]]) for i in range(100) for j in range(100)], dtype=complex)

ky_max = []
Qz_max = []

# Plot real and imaginary parts of roots separately in one figure
fig, axs = plt.subplots(2, 3, figsize=(13.6, 7.6), constrained_layout=True)
axs = axs.ravel()
for i in range(3):
    cnt = axs[i].contourf(K0, K1, roots[:, i].real.reshape(K0.shape), levels=20, cmap='coolwarm')
    axs[i].set_title(f'Real part of root {i+1}')
    axs[i].set_xlabel('$k_y$')
    axs[i].set_ylabel('$Q_z$')
    # axs[i].set_yscale('log')
    fig.colorbar(cnt, ax=axs[i])

    cnt = axs[i+3].contourf(K0, K1, roots[:, i].imag.reshape(K0.shape), levels=20, cmap='coolwarm')
    axs[i+3].set_title(f'Imaginary part of root {i+1}')
    axs[i+3].set_xlabel('$k_y$')
    axs[i+3].set_ylabel('$Q_z$')
    # axs[i+3].set_yscale('log')
    cbar1 = fig.colorbar(cnt, ax=axs[i+3])
    cbar1.formatter.set_powerlimits((0, 0))
    
    max_val = np.max(roots[:, i].real)
    print(f'Maximum value for real part of root {i+1} is: {max_val}')
    
    max_val = np.max(roots[:, i].real)
    max_index = np.argmax(roots[:, i].real)
    k0_max, k1_max = K0.flatten()[max_index], K1.flatten()[max_index]
    
    axs[i].scatter(k0_max, k1_max, s=150, color='yellow', marker='x', label=f'Max value {max_val:.4f}')
    axs[i].legend(loc=0)
    
    ky_max.append(K0.flatten()[max_index])
    Qz_max.append(K1.flatten()[max_index])

    
    print(f'k_y for root {i+1} is: {k0_max}')
    print(f'Q_z for  root {i+1} is: {k1_max}')

for axis in ['top','bottom','left','right']:
    axs[2].spines[axis].set_linewidth(3)
    axs[2].spines[axis].set_color("green")
    axs[5].spines[axis].set_linewidth(3)
    axs[5].spines[axis].set_color("green")


# Create a caption
caption = f'Contour plot showing the real and imaginary components of the roots of the cubic for a range of $k_y$ and $Q_z$ values. Where the other variables are given by: Ra$^* = $ {Ra:.1e}, $\Delta N =$ {deltaN}, Pr = {Pr:.1e}, $\zeta =$ {zeta:.1e}, $\Omega =$ {Omega:.1e}, $L$ = {L}.'

# Create a file name
figure_name = f'decay_contour_Ra={Ra:.1e}_Pr={Pr:.1e}_dN={deltaN}'
pdf_file = f'{figure_name}.pdf'
tex_file = f'{figure_name}.tex'

# save the plot as a PDF
plt.savefig(pdf_file)

# create a text file containing the LaTeX code to include the figure
with open(tex_file, 'w') as f:
    f.write("\\begin{figure}[h]\n")
    f.write("\\centering\n")
    f.write("\\includegraphics[width=0.85\\linewidth]{"+ pdf_file+"}\n")
    f.write("\\caption{"+ caption +"}\n")
    f.write("\\end{figure}\n")


fig2, axs2 = plt.subplots(2, 3, figsize=(11, 8), constrained_layout=True)


for idx_1 in range(3):
    k1_slice = 0
    indices = np.where(K1.flatten() == k1_slice)
    
    root_slice = roots[indices][:,idx_1].real
    
    k1_slice = K0.flatten()[indices]
    root_slice = roots[indices][:,idx_1].real
    
    axs2[0][idx_1].plot(k1_slice, root_slice, color = 'red')
    
    k1_slice_imag = K0.flatten()[indices]
    root_slice_imag = roots[indices][:,idx_1].imag
    
    axs2[1][idx_1].plot(k1_slice, root_slice_imag, color = 'red')
    axs2[1][idx_1].set_xlabel('$k_y$')
        
        
axs2[0][0].set_ylabel('Re$(s)$')
axs2[1][0].set_ylabel('Im$(s)$')

for idx_1 in range(3):
    axs2[0][idx_1].plot(k0, -zeta*(np.pi**2 + k0**2), 'x', markevery=10, color = 'black')

# Create a caption
caption = f'Profiles at the $k_y$ at $Q_z = 0$ showing the real and imaginary components of the roots of the cubic for a range of $k_y$ and $Q_z$ values. Where the other variables are given by: Ra$^* = $ {Ra:.1e}, $\Delta N =$ {deltaN}, Pr = {Pr:.1e}, $\zeta =$ {zeta:.1e}, $\Omega =$ {Omega:.1e}, $L$ = {L}.'

# Create a file name
figure_name = f'decay_profiles_Ra={Ra:.1e}_Pr={Pr:.1e}_dN={deltaN}'
pdf_file = f'{figure_name}.pdf'
tex_file = f'{figure_name}.tex'

# create a text file containing the LaTeX code to include the figure
with open(tex_file, 'w') as f:
    f.write("\\begin{figure}[h]\n")
    f.write("\\centering\n")
    f.write("\\includegraphics[width=0.99\\linewidth]{"+ pdf_file+"}\n")
    f.write("\\caption{"+ caption +"}\n")
    f.write("\\end{figure}\n")
    
for axis in ['top','bottom','left','right']:
    axs2[0][2].spines[axis].set_linewidth(3)
    axs2[0][2].spines[axis].set_color("green")
    axs2[1][2].spines[axis].set_linewidth(3)
    axs2[1][2].spines[axis].set_color("green")

# save the plot as a PDF
plt.savefig(pdf_file)

plt.show()

解决方案：根的连续跟踪 #

单独对每个点的根排序无法解决问题，因为根是随参数连续变化的，必须保证相邻网格点的对应根是最接近的那个，而非全局排序。具体实现如下：

核心思路 #

1. 初始化：先计算第一个网格点的根作为参考基准
2. 逐点遍历：按行/列顺序遍历网格，对每个点计算根后，与上一个相邻点的根做最小距离匹配，重新排序当前根，确保连续性
3. 匹配算法：用匈牙利算法（全局最优）或贪心匹配（适合少量根的场景）完成根的对应

修改后的根计算代码 #

from scipy.optimize import linear_sum_assignment

# Create arrays of k[0] and k[1] values for contour plot
k0, k1 = np.linspace(0, 5, 100), np.linspace(0, 5e2, 100)
K0, K1 = np.meshgrid(k0, k1)

# 初始化根数组，存储每个网格点的3个根
roots = np.zeros((K0.shape[0], K0.shape[1], 3), dtype=complex)

# 计算第一个点的根作为初始参考
roots[0, 0] = polynomial([K0[0,0], K1[0,0]])

# 处理第一行（k1固定，k0从左到右遍历）
for j in range(1, K0.shape[1]):
    current_roots = polynomial([K0[0,j], K1[0,j]])
    # 计算当前根与左侧点根的距离矩阵
    dist_matrix = np.abs(current_roots[:, None] - roots[0,j-1][None, :])
    # 用匈牙利算法找到最优匹配
    row_ind, col_ind = linear_sum_assignment(dist_matrix)
    # 按匹配结果重新排序当前根
    roots[0,j] = current_roots[col_ind]

# 处理剩余行
for i in range(1, K0.shape[0]):
    # 先处理每行第一个点，与上一行同列点匹配
    current_roots = polynomial([K0[i,0], K1[i,0]])
    dist_matrix = np.abs(current_roots[:, None] - roots[i-1,0][None, :])
    row_ind, col_ind = linear_sum_assignment(dist_matrix)
    roots[i,0] = current_roots[col_ind]
    
    # 处理该行剩余点，与左侧相邻点匹配
    for j in range(1, K0.shape[1]):
        current_roots = polynomial([K0[i,j], K1[i,j]])
        dist_matrix = np.abs(current_roots[:, None] - roots[i,j-1][None, :])
        row_ind, col_ind = linear_sum_assignment(dist_matrix)
        roots[i,j] = current_roots[col_ind]

# 展平成原代码需要的形状
roots = roots.reshape(-1, 3)

说明 #

• 匈牙利算法（linear_sum_assignment）能找到根之间的全局最优匹配，避免局部最小距离导致的错位，适合根的距离差异较小的场景
• 遍历顺序选择按行/列连续遍历，保证每次处理的点与上一个点参数变化极小，根的连续性有保障
• 替换原代码中直接生成roots的列表推导式即可，后续绘图代码无需修改

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