%matplotlib inline

import numpy
import sympy

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

a,b = -1,1
n = 16
x,y = numpy.mgrid[a:b:(1j*n),a:b:(1j*n)]
z = x + 1j*y
w = z**2

fig = plt.figure(figsize=(16,8))

# plot the real part
ax_real = fig.add_subplot(1,2,1,projection='3d')
ax_real.plot_surface(z.real, z.imag, w.real,
                    rstride=1, cstride=1, cmap=cm.jet)

# plot the imaginary part
ax_imag = fig.add_subplot(1,2,2,projection='3d')
ax_imag.plot_surface(z.real, z.imag, w.imag,
                    rstride=1, cstride=1, cmap=cm.jet)

a,b = -1,1
n = 16
x,y = numpy.mgrid[a:b:(1j*n),a:b:(1j*n)]
z = x + 1j*y
w = numpy.sqrt(z)

fig = plt.figure(figsize=(16,8))

# plot the real part
ax_real = fig.add_subplot(1,2,1,projection='3d')
ax_real.plot_surface(z.real, z.imag, w.real,
                    rstride=1, cstride=1, cmap=cm.jet)

# plot the imaginary part
ax_imag = fig.add_subplot(1,2,2,projection='3d')
ax_imag.plot_surface(z.real, z.imag, w.imag,
                    rstride=1, cstride=1, cmap=cm.jet)

branching_number = 2

Nr = 16
Ntheta = 32

# compute the theta,R domain
theta = numpy.linspace(0,2*numpy.pi*branching_number, Ntheta)
r = numpy.linspace(0,1,Nr)
Theta, R = numpy.meshgrid(theta,r)

z = R*numpy.exp(1j*Theta)

# compute w^2 = z. THE KEY IDEA is to pass the exponentiation by 1/2 into exp().
w = numpy.sqrt(R)*numpy.exp(1j*Theta/2)

fig = plt.figure(figsize=(16,8))

# plot the real part
ax_real = fig.add_subplot(1,2,1,projection='3d')
ax_real.plot_surface(z.real, z.imag, w.real,
                    rstride=1, cstride=1, cmap=cm.jet, alpha=0.5)

# plot the imaginary part
ax_imag = fig.add_subplot(1,2,2,projection='3d')
ax_imag.plot_surface(z.real, z.imag, w.imag,
                    rstride=1, cstride=1, cmap=cm.jet, alpha=0.5)

# plot abs(w)
fig = plt.figure(figsize=(8,8))
ax_real = fig.add_subplot(1,1,1,projection='3d')

ax_real.plot_surface(z.real, z.imag, abs(w),
                     rstride=1, cstride=1, cmap=cm.jet, alpha=0.5)