使用二进制颜色图绘制曲面

Question

我想制作一个由函数参数化的表面的 3d 图，我希望表面是一种颜色（比如白色），它高于某个值 a 和另一种颜色（比如黑色）它低于a。

这里是生成和绘制曲面的代码（生成曲面的方式并不重要，它可能是一个更简单的函数）：

from __future__ import division
import numpy as np
import time,random
random.seed(-2)

def build_spden(N,M, alpha):
    #computes the spectral density in momentum space
    sp_den = np.zeros((N,M))
    for k1 in prange(-N//2, N//2):
        for k2 in prange(-M//2, M//2):
            sp_den[k1,k2] = np.abs(2*(np.cos(2*np.pi*k1/N)+np.cos(2*np.pi*k2/M)-2))
    sp_den[0,0]=1
    return 1/sp_den**(alpha/2)


def gaussian_field(N,M,alpha):
    '''Builds a correlated gaussian field on a surface NxM'''
    spectral_density = build_spden(N,M, alpha)
    # FFT of gaussian noise:
    noise_real = np.random.normal(0, 1, size = (N,M))
    noise_fourier = np.fft.fft2(noise_real)

    # Add correlations by Fourier Filtering Method:
    convolution = noise_fourier*np.sqrt(spectral_density)

    # Take IFFT and exclude residual complex part
    correlated_noise = np.fft.ifft2(convolution).real

    # Return normalized field
    return correlated_noise * (np.sqrt(N*M)/np.sqrt(np.sum(spectral_density)) )


#PLOT

N = 2**5
alpha = .75
a = -.1985
surf = gaussian_field(N,N,alpha)

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

x = np.outer(np.arange(0, N), np.ones(N))
y = x.copy().T # transpose
z = surf
fig = plt.figure()
ax = plt.axes(projection='3d')

ax.plot_surface(x, y, z,alpha=.4)  #plot the surface
z2 = a*np.ones((N,N))
ax.plot_surface(x, y, z2, alpha=0.9)  #plot a plane z = a.
plt.show()

输出是：

因此，我希望平面上方为白色，下方为黑色。

非常感谢！

Answer 1

您可以定义自定义颜色图并传递给plot_surface：

from matplotlib.colors import ListedColormap, BoundaryNorm
cmap = ListedColormap(['r', 'b'])
norm = BoundaryNorm([z.min(), a, z.max()], cmap.N)
ax.plot_surface(x, y, z, cmap=cmap, norm=norm, alpha=.4)  #plot the surface

z2 = a*np.ones((N,N))
ax.plot_surface(x, y, z2, colalpha=0.9)  #plot a plane z = a.
plt.show()

输出：