exp_function_complex.py (Source)

#!/usr/bin/env python3
# Code to generate graphs of the exponential function for complex values
# as illustration of the blog post 'The exponential function (exp)' at
# https://prfraanje.github.io/posts/exp-function/
# Rufus Fraanje
# Date: 22/12/2020
# License: GPLv3
from mpl_toolkits.mplot3d import axes3d
import matplotlib
matplotlib.rcParams['text.usetex'] = True
matplotlib.rcParams['font.size'] = 14
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import numpy as np
Nx = 53
Ny = 53
x = np.linspace(-3*np.pi,0*2*np.pi,Nx)
y = np.linspace(-2.2,1,Ny)
X, Y = np.meshgrid(x, y)
Z = np.exp(Y+1j*X)
title_ax = [r'$\displaystyle\mbox{\emph{Re}}(e^{\alpha+i\beta})=e^{\alpha}\cos(\beta)$',
            r'$\displaystyle\mbox{\emph{Im}}(e^{\alpha+i\beta})=e^{\alpha}\sin(\beta)$']
fig = plt.figure(figsize=(14,7))
for i,Zs in enumerate([np.real(Z),np.imag(Z)]):
    ax = fig.add_subplot(1, 2, 1+i, projection='3d')
    ax.plot_surface(X, Y, Zs, alpha=0.3)
    ax.plot_wireframe(X, Y, Zs, rstride=Ny//9, cstride=Nx//20,lw=0.3)
    cset = ax.contour(X, Y, Zs, zdir='x', offset=x[0], cmap=cm.coolwarm)
    cset = ax.contour(X, Y, Zs, zdir='y', offset=y[-1], cmap=cm.coolwarm)
    ax.xaxis.set_rotate_label(False)  # disable automatic rotation
    ax.yaxis.set_rotate_label(False)  # disable automatic rotation
    ax.zaxis.set_rotate_label(False)  # disable automatic rotation
    ax.set_xlabel(r'$\beta$',rotation=0,labelpad=10)
    ax.set_xlim(x[0], x[-1])
    ax.set_ylabel(r'$\alpha$',rotation=0,labelpad=10)
    ax.set_ylim(y[0], y[-1])
    ax.set_zlim(np.min(Zs), np.max(Zs))
    
    ax.set_title(title_ax[i])
    ax.set_xticks([-3*np.pi,-2*np.pi,-1*np.pi,0])
    ax.set_xticklabels([r'$-3\pi$',r'$-2\pi$',r'$-\pi$',r'$0$'])
    ax.set_yticks([-2,-1,0,1])
    ax.set_yticklabels([r'$-2$',r'$-1$',r'$0$',r'$1$'])
    ax.set_zticks([-np.exp(1),-1,0,1,np.exp(1)])
    ax.set_zticklabels([r'$-e$',r'$-1$',r'$0$',r'$1$',r'$e$'])
plt.savefig('exp_function_complex.svg')
plt.show()