Complex contour of a double well potential

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

np.seterr(invalid='ignore')



def V(q):
    return (q**2-1)**2

def H(q,p):
    return p**2/2 + V(q)
def P(E,q):
    return np.sqrt(2*(E-V(q)))+0.j

fig = plt.figure(figsize=(10,5))
#fig = plt.figure(figsize=plt.figaspect(0.5))
#fig = plt.figure(figsize=plt.figaspect(0.5))
ax = fig.add_subplot(111,projection='3d')

fig.patch.set_visible(False)
ax.axis('off')

x = np.linspace(-2.0, 2.0, 300)
y = np.linspace(-2.0, 2.0, 300)
x,y = np.meshgrid(x,y)
level = [0.1,0.5, 1.0, 2.0, 3.0, 4.0]
ax.contour(x,y,H(x,y),level,zdir='z',offset=0,colors='k')


sample=5000
E=0.5 + 0.j
q = np.linspace(-0.8, 0.8,sample) + 0.j
p = P(E,q)
ax.plot(q.real,p.real, p.imag, c="r",lw=3)

E=0.5
q = np.linspace(-1.4, 1.4,sample)
p = P(E,q)
ax.plot(q.real,p.real,0, c="r",lw=3)
ax.plot(q.real,-p.real,0, c="r",lw=3)


sample=500
q = np.linspace(-1.4, 1.4,sample) + 0.j
pp1 = np.array([],dtype=np.complex128)
qq1 = np.array([],dtype=np.complex128)
pp2 = np.array([],dtype=np.complex128)
qq2 = np.array([],dtype=np.complex128)
levels = np.linspace(1e-20,0.5,500)

for im in levels:
    p = P(E+1.j*im,q)
    pp1 = np.append(pp1, p)
    qq1 = np.append(qq1, q)

    p = -P(E-1.j*im,q)
    pp2 = np.append(pp2, p)
    qq2 = np.append(qq2, q)


p = pp1.reshape(sample, sample)
q = qq1.reshape(sample, sample)
ax.plot_wireframe(q.real,p.real, p.imag,rstride=2,cstride=2,alpha=0.1,colors='c')

p = pp2.reshape(sample, sample)
q = qq2.reshape(sample, sample)
ax.plot_wireframe(q.real,p.real, p.imag,rstride=2,cstride=2,alpha=0.1,colors='c')


ax.set_xlabel(r"$Re(q)$",fontsize=20)
ax.set_ylabel(r"$Re(p)$",fontsize=20)
ax.set_zlabel(r"$Im(p)$",fontsize=20)
ax.set_xticks([-1,0,1])
ax.set_yticks([-1,0,1])
ax.set_zticks([0,0.5,1])
ax.set_xlim(-2.0,2.0)
ax.set_ylim(-2.0,2.0)
ax.set_zlim(0,1.0)
#fig.savefig("complex_doubel_well.png",transparent=True)
fig.tight_layout()
plt.show()