plot-complex/exp.py at master · betanaught/plot-complex · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
# Code demonstrating the plotting of complex functions
import numpy
# import sympy

import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

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 z = w^2
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)
plt.ion()  # plot interactive mode: on
plt.show()
plt.close()