How to stop the roots of a cubic from becoming mixed up when plotting the 3 roots as contour plots?

Question

I have a function which determines the roots of a complex cubic. I am solving the cubic for a variety of k0 and k1 values and showing the solutions as contour plots. Since the cubic has three roots, I produce 3 contour plots for the real parts and 3 for the imaginary parts. However, sometimes you can clearly see that sections of the contour plots for one root really should be swapped with a different contour plot - all the contours should be continuous. I have tried various "sorting methods" which you can see, but none of them fully fix it. What can I do so that the roots don't get mixed up resulting in non-continuous contours.

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()

I've tried np.sort, np.sorted, flapping the roots using if statements etc, nothing works 100%

Answer 1

For two successive polynomials P and Q, I suggest simply solving the assignment problem to pair each root of P to the closest root of Q.

You can use scipy's linear_sum_assignment along with distance_matrix to find the best assignment of P's roots with Q's roots.

import numpy as np
from scipy.optimize import linear_sum_assignment
from scipy.spatial import distance_matrix
import matplotlib.pyplot as plt

def get_root_sequence(sequence_of_polynomials):
    r0 = np.roots(sequence_of_polynomials[0])
    roots = [r0]
    for P in sequence_of_polynomials[1:]:
        r1 = np.roots(P)
        _, idx = linear_sum_assignment(distance_matrix(r0.reshape(3, 1), r1.reshape(3,1)))
        r1 = r1[idx]
        roots.append(r1)
        r0 = r1
    return np.array(roots)

sequence_of_polynomials = np.linspace((1,0,0,-1), (1,-7-2j,15+9j,-10-10j), 100)
roots = get_root_sequence(sequence_of_polynomials)

plt.axes().set_aspect('equal')

for i in range(3):
    r = roots[:, i]
    ordinal = ('first', 'second', 'third')[i]
    plt.plot(r.real, r.imag, label=f'{ordinal} root')

for triangle, label in zip((roots[0], roots[-1]), ('x³-1', '(x-2)(x-2-i)(x-3-i)')):
    triangle = triangle[[0,1,2,0]]
    plt.plot(triangle.real, triangle.imag, label=label)

plt.legend(loc='best')
plt.xlabel('Real part')
plt.ylabel('Imaginary part')
plt.show()