Riemann surface plot using Python

Question

I have a complex matrix which looks like

H = [[w1-complex(0,g1),k], [k,w2-complex(0,g2)]

and as discussed in this paper,

I want to find the eigenvalues and plot their real and imaginary part on a Riemann surface. It seems that the paper chooses fixed values of g1 and g2 (denoted with Greek letter gamma originally) and uses a detuning parameter delta=(w1-w2)/2 and plots eigenvalues on the k-delta plane. Again, one may choose w2=0 and vary w1 to change delta.

I am attaching the portion of the paper that discusses the plotting.

I couldn't find any discussion on SO that plots Riemann surface plots using Python. Hope the problem is not a duplicate.

I have the following code.

"""
Riemann plot of a non-Hermitian matrix
"""

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np


XR = np.arange(-1, 1, 0.01)
YR = np.arange(-1, 1, 0.01)

delta, kappa = np.meshgrid(XR, YR)

g1 = 0.2; g2 = 0.1
chi = (g1+g2)/2.0
beta = (g1-g2)/2.0
Gamma = delta + 1j*beta

disc = np.sqrt(kappa**2+Gamma**2)

lambda1 = delta-1j*chi+disc 
lambda2 = delta-1j*chi-disc 


F = plt.figure(1)
A = F.gca(projection='3d')
plt.xlabel('$\delta$')
plt.ylabel('$\kappa$')
A.set_zlabel('Re $\lambda$')
S = A.plot_surface( kappa, delta, lambda1.real, rstride=1, cstride=1, cmap=cm.Reds )
S = A.plot_surface( kappa, delta, lambda2.real, rstride=1, cstride=1, cmap=cm.Greys )

F = plt.figure(2)
A = F.gca(projection='3d')
plt.xlabel('$\delta$')
plt.ylabel('$\kappa$')
A.set_zlabel('Im $\lambda$')

S = A.plot_surface(kappa, delta, lambda1.imag, cmap=cm.Reds)
S = A.plot_surface(kappa, delta, lambda2.imag, cmap=cm.Greys)

plt.show()

which produces the following outputs.

Any way to modify it and get the desired results?