3D graphing the complex values of a function in Python

Question

This is the real function I am looking to represent in 3D: y = f(x) = x^2 + 1

The complex function would be as follows: w = f(z) = z^2 + 1

Where z = x + iy and w = u + iv. These are four dimentions (x, y, u, v), but one can use u for 3D graphing.

We get: f(x + iy) = x^2 + 2xyi - y^2 + 1

So: u = x^2 - y^2 + 1 and v = 2xy

This u is what is being used in the code below.

import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(-100, 101, 150)
y = np.linspace(-100, 101, 150)
 
X, Y = np.meshgrid(x,y)
U = (X**2) - (Y**2) + 1

fig = plt.figure(dpi = 300)
ax = plt.axes(projection='3d')
ax.plot_surface(X, Y, Z)
plt.show()

The following images are the side-view of the 3D function and the 2D plot for reference. I do not think they are alike.

Likewise, here is the comparison between the 3 side-view and the 2D plot of w = z^3 + 1. They seem to differ as well.

I have not been able to find too many resources regarding plotting in 3D using complex numbers. Because of this and the possible discrepancies mentioned before, I think the code must be flawed, but I can't figure out why. I would be grateful if you could correct me or advise me on any changes.

The inspiration came from Welch Labs' 'Imaginary Numbers are Real' YouTube series where he shows a jaw-dropping representation of the complex values of the function I have been tinkering with.

I was just wondering if anybody could point out any flaws in my reasoning or the execution of my idea since this code would be helpful in explaining the importance of complex numbers to HS students.

Thank you very much for your time.

Answer 1

The f(z) = z^2 + 1 projection (that is, side-view) looks OK to me. You can use this technique to add the projections; this code:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm

def f(z):
    return z**2 + 1

def freal(x, y):
    return x**2 - y**2 + 1

x = np.linspace(-100, 101, 150)
y = np.linspace(-100, 101, 150)
 
yproj = 0 # value of y for which to project xu axes
xproj = 0 # value of x to project onto yu axes

X, Y = np.meshgrid(x,y)
Z = X + 1j * Y
W = f(Z)
U = W.real

fig = plt.figure()
ax = plt.axes(projection='3d')

## surface
ax.plot_surface(X, Y, U, alpha=0.7)

# xu projection
xuproj = freal(x, yproj)
ax.plot(x, xuproj, zs=101, zdir='y', color='red', lw=5)
ax.plot(x, xuproj, zs=yproj, zdir='y', color='red', lw=5)

# yu projection
yuproj = freal(xproj, y)
ax.plot(y, yuproj, zs=101, zdir='x', color='green', lw=5)
ax.plot(y, yuproj, zs=xproj, zdir='x', color='green', lw=5)

# partially reproduce https://www.youtube.com/watch?v=T647CGsuOVU&t=107s
x = np.linspace(-3, 3, 150)
y = np.linspace(0, 3, 150)
X, Y = np.meshgrid(x,y)
U = f(X + 1j*Y).real

fig = plt.figure()
ax = plt.axes(projection='3d')

## surface
ax.plot_surface(X, Y, U, cmap=cm.jet)
ax.set_box_aspect( (np.diff(ax.get_xlim())[0],
                    np.diff(ax.get_ylim())[0],
                    np.diff(ax.get_zlim())[0]))


#ax.set_aspect('equal')

plt.show()

gives this result:

and

The axis ticks don't look very good: you can investigate plt.xticks or ax.set_xticks (and yticks, zticks) to fix this.

There is a way to visualize complex functions using colour as a fourth dimension; see complex-analysis.com for examples.