How can I create a 4D complex surface plot? [closed]

问题

I have the following Matlab code that I would like to be converted to a Python 3 one.

r = (0:1:15)';                           % create a matrix of complex inputs
theta = pi*(-2:0.05:2);
z = r*exp(1i*theta);
%w = z.^(1/2)  ;                          % calculate the complex outputs
w = sqrt(r)*exp(1i*theta/2);

figure('Name','Graphique complexe','units','normalized','outerposition',[ 0.08 0.1 0.8 0.55]);
subplot(121)

surf(real(z),imag(z),real(w),imag(w))    % visualize the complex function using surf
xlabel('Real(z)')
ylabel('Imag(z)')
zlabel('Real(u)')
cb = colorbar;
colormap jet;                            % gradient from blue to red
cb.Label.String = 'Imag(v)';

subplot(122)
surf(real(z),imag(z),imag(w),real(w))    % visualize the complex function using surf
xlabel('Real(z)')
ylabel('Imag(z)')
zlabel('Imag(v)')
cb = colorbar;
colormap jet;                            % gradient from blue to red
cb.Label.String = 'Real(u)';

The results and original discussions can be found here. There's also a discussion available on this SO page. However, I failed to run and reproduce those codes. What can I try next?

回答1:

This is perfectly straightforward if you spend the time learning how matplotlib (and 3d axes in particular) work:

import numpy as np  
import matplotlib.pyplot as plt  
import matplotlib.cm as cm 
from mpl_toolkits.mplot3d import Axes3D 
 
# compute data to plot 
r, theta = np.mgrid[1:16, -2*np.pi:2*np.pi:50j] 
z = r * np.exp(1j*theta)  
w = np.sqrt(r) * np.exp(1j*theta/2)  
 
# plot data  
fig = plt.figure()  
for plot_index in [1, 2]: 
    if plot_index == 1: 
        z_data, c_data = w.real, w.imag 
        z_comp, c_comp = 'Re', 'Im' 
    else: 
        z_data, c_data = w.imag, w.real 
        z_comp, c_comp = 'Im', 'Re' 
    c_data = (c_data - c_data.min()) / c_data.ptp() 
    colors = cm.viridis(c_data) 
 
    ax = fig.add_subplot(f'12{plot_index}', projection='3d') 
    surf = ax.plot_surface(z.real, z.imag, z_data, facecolors=colors,
                           clim=[z_data.min(), z_data.max()])
    ax.set_xlabel('$Re z$')  
    ax.set_ylabel('$Im z$')   
    ax.set_zlabel(f'${z_comp} w$')  
    cb = plt.colorbar(surf, ax=ax)  
    cb.set_label(f'${c_comp} w$')  
 
plt.show()

The result:

Some things that should be noted:

Viridis colormap is good, jet is bad.
In general there could be rendering issues with complex (interlocking) 3d geometries, because matplotlib has a 2d renderer. Fortunately, in this case the dataset is tightly coupled enough that this doesn't seem to happen, even if you rotate around the figure interactively. (But if you were to plot two intersecting surfaces together, things would probably be different.)
One might want to enable latex rendering of labels to make the result extra crispy.
The shading looks a lot better if you use the default option of colouring according to the z component of the data.

If we also want to port the second part of my MATLAB answer you will have to use a trick to stitch together the two branches of the function (which, as I said, is necessary to render interlocking surfaces properly). For the specific example in the above code this will not give you perfect results, since both branches themselves contain discontinuities in the imaginary part, so regardless of our efforts in rendering the two surfaces nicely, the result will look a bit bad:

import numpy as np 
import matplotlib.pyplot as plt 
import matplotlib.cm as cm 
from mpl_toolkits.mplot3d import Axes3D 
 
# compute data to plot 
r0 = 15 
re, im = np.mgrid[-r0:r0:31j, -r0:r0:31j] 
z = re + 1j*im 
r, theta = abs(z), np.angle(z) 
w1 = np.sqrt(r) * np.exp(1j*theta/2)  # first branch 
w2 = np.sqrt(r) * np.exp(1j*(theta + 2*np.pi)/2)  # second branch 
 
# plot data 
fig = plt.figure() 
for plot_index in [1, 2]: 
    # construct transparent bridge 
    re_bridge = np.vstack([re[-1, :], re[0, :]]) 
    im_bridge = np.vstack([im[-1, :], im[0, :]]) 
    c_bridge = np.full(re_bridge.shape + (4,), [1, 1, 1, 0])  # 0% opacity
 
    re_surf = np.vstack([re, re_bridge, re]) 
    im_surf = np.vstack([im, im_bridge, im]) 
    w12 = np.array([w1, w2]) 
    if plot_index == 1: 
        z_comp, c_comp = 'Re', 'Im' 
        z12, c12 = w12.real, w12.imag 
    else: 
        z_comp, c_comp = 'Im', 'Re' 
        z12, c12 = w12.imag, w12.real 
         
    color_arrays = cm.viridis((c12 - c12.min()) / c12.ptp()) 
    z1,z2 = z12 
    c1,c2 = color_arrays 
     
    z_bridge = np.vstack([z1[-1, :], z2[0, :]]) 
    z_surf = np.vstack([z1, z_bridge, z2]) 
    c_surf = np.vstack([c1, c_bridge, c2]) 
     
    ax = fig.add_subplot(f'12{plot_index}', projection='3d') 
    surf = ax.plot_surface(re_surf, im_surf, z_surf, facecolors=c_surf, 
                           clim=[c12.min(), c12.max()], 
                           rstride=1, cstride=1) 
    ax.set_xlabel('$Re z$') 
    ax.set_ylabel('$Im z$') 
    ax.set_zlabel(f'${z_comp} w$') 
    cb = plt.colorbar(surf, ax=ax) 
    cb.set_label(f'${c_comp} w$') 
  
plt.show()

The ugly jump in the right figure might be fixed with a lot of work, but it won't be easy: it's an actual discontinuity in both surface datasets occuring at negative real arguments. Since your actual problem is probably more like this, you will probably not need to face this issue, and you can use the above stitching (bridging) trick to combine your surfaces.

来源：https://stackoverflow.com/questions/63144394/how-can-i-create-a-4d-complex-surface-plot

标签

python

matlab

matplotlib