A quick demonstration of how to plot multivalued complex functions in Python.

simple.py

import numpy
import sympy

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

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)

# color by argument
arguments = numpy.angle(w)
norm = colors.Normalize(arguments.min(), arguments.max())
color = cm.jet(norm(arguments))

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, alpha=0.5, facecolors=color)

# 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, alpha=0.5, facecolors=color)

simple_notebook.ipynb

{
 "metadata": {
  "name": "",
  "signature": "sha256:e64d7262a527a7b1d4f96247f34f5a60ddf1879cb8a56485959d028e3168dab3"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%matplotlib inline"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import numpy\n",
      "import sympy\n",
      "\n",
      "from mpl_toolkits.mplot3d import Axes3D\n",
      "import matplotlib\n",
      "from matplotlib import cm\n",
      "from matplotlib import pyplot as plt"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "A Naive Approach\n",
      "================\n",
      "\n",
      "...consisting of creating a complex grid and simply taking the square root of the grid."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "a,b = -1,1\n",
      "n = 16\n",
      "x,y = numpy.mgrid[a:b:(1j*n),a:b:(1j*n)]\n",
      "z = x + 1j*y\n",
      "w = z**2"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "fig = plt.figure(figsize=(16,8))\n",
      "\n",
      "# plot the real part\n",
      "ax_real = fig.add_subplot(1,2,1,projection='3d')\n",
      "ax_real.plot_surface(z.real, z.imag, w.real,\n",
      "                    rstride=1, cstride=1, cmap=cm.jet)\n",
      "\n",
      "# plot the imaginary part\n",
      "ax_imag = fig.add_subplot(1,2,2,projection='3d')\n",
      "ax_imag.plot_surface(z.real, z.imag, w.imag,\n",
      "                    rstride=1, cstride=1, cmap=cm.jet)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "a,b = -1,1\n",
      "n = 16\n",
      "x,y = numpy.mgrid[a:b:(1j*n),a:b:(1j*n)]\n",
      "z = x + 1j*y\n",
      "w = numpy.sqrt(z)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "fig = plt.figure(figsize=(16,8))\n",
      "\n",
      "# plot the real part\n",
      "ax_real = fig.add_subplot(1,2,1,projection='3d')\n",
      "ax_real.plot_surface(z.real, z.imag, w.real,\n",
      "                    rstride=1, cstride=1, cmap=cm.jet)\n",
      "\n",
      "# plot the imaginary part\n",
      "ax_imag = fig.add_subplot(1,2,2,projection='3d')\n",
      "ax_imag.plot_surface(z.real, z.imag, w.imag,\n",
      "                    rstride=1, cstride=1, cmap=cm.jet)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The Issue\n",
      "=========\n",
      "\n",
      "In many (all) complex square-root functions a branch cut is chosen (usually the negative real axis) and only one branch is computed in order to make square root single-valued.\n",
      "\n",
      "From complex analysis we learn that we can use polar coordinates to easily verify that $w = \\sqrt{z}$, which I will read as $w^2 = z$ since the latter makes it clear that I won't be chosing a branch of $w$ / making a branch cut in the $z$ plane. Because the branch point $z=0$ is 2-ramified / has branching number two I will need to make two rotations to capture all of the behavior near the branch point."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "branching_number = 2\n",
      "\n",
      "Nr = 16\n",
      "Ntheta = 32\n",
      "\n",
      "# compute the theta,R domain\n",
      "theta = numpy.linspace(0,2*numpy.pi*branching_number, Ntheta)\n",
      "r = numpy.linspace(0,1,Nr)\n",
      "Theta, R = numpy.meshgrid(theta,r)\n",
      "\n",
      "z = R*numpy.exp(1j*Theta)\n",
      "\n",
      "# compute w^2 = z. THE KEY IDEA is to pass the exponentiation by 1/2 into exp().\n",
      "w = numpy.sqrt(R)*numpy.exp(1j*Theta/2)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "fig = plt.figure(figsize=(16,8))\n",
      "\n",
      "# plot the real part\n",
      "ax_real = fig.add_subplot(1,2,1,projection='3d')\n",
      "ax_real.plot_surface(z.real, z.imag, w.real,\n",
      "                    rstride=1, cstride=1, cmap=cm.jet, alpha=0.5)\n",
      "\n",
      "# plot the imaginary part\n",
      "ax_imag = fig.add_subplot(1,2,2,projection='3d')\n",
      "ax_imag.plot_surface(z.real, z.imag, w.imag,\n",
      "                    rstride=1, cstride=1, cmap=cm.jet, alpha=0.5)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}

Thanks for this.