import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
plt.rcParams['svg.fonttype'] = 'none'

x = lambda r, t: np.outer(r, np.cos(t))
y = lambda r, t: np.outer(r, np.sin(t))
z = lambda r, t: t

fig = plt.figure(figsize=[8,7])
ax = fig.gca(projection='3d')

t = np.linspace(-np.pi, np.pi, 20)
r = np.linspace(1, 10, 9)
for rr in r:
    ax.plot(x(rr,t)[0], y(rr,t)[0], z(rr,t), color='r', linewidth=.5, alpha=.5)
ax.plot(x([0,10],-np.pi)[:,0], y([0,10],-np.pi)[:,0], z([0,10],-np.pi), color='r', linewidth=1.5)
ax.plot(x([0,10],np.pi)[:,0], y([0,10],np.pi)[:,0], z([0,10],np.pi), color='r', linewidth=1.5)

t = np.linspace(-2*np.pi, 2*np.pi, 50)
r = np.linspace(0, 10, 5)
ax.plot_surface(x(r,t), y(r,t), z(r,t)
                   , antialiased=True, shade=False
                   , rstride=1, cstride=1
                   , cmap=cm.rainbow_r
                   , alpha=1
                  )
ax.view_init(elev=30, azim=120)
ax.set_xticks([-10,-5,0,5,10])
ax.set_yticks([-10,-5,0,5,10])
ax.set_zticks([-2*np.pi,-np.pi,0,np.pi,2*np.pi])
ax.set_zticklabels(['$-2\\pi$','$-\\pi$','0','$\\pi$','$2\\pi$'])
ax.set_xlabel('R')
ax.set_ylabel('l')
ax.set_zlabel('$\\mathrm{arg}(z)$')
fig.savefig('test_3d.svg', bbox_inches='tight', transparent=True)