CODEWITHSUNDEEP
pythonmatplotlib
Njha
Njha

Reputation: 113

Plot surface with binary colormap

I would like to make a 3d plot of a surface parametrised by a function, and I would like the surface to be of one color (say white) where it is above some value a, and of another color (say black) where it is below a.

Here is the code to generate and plot the surface (the way the surface is generated is not important, it could be a much simpler function):

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()

The output is:

enter image description here

I would therefore like the surface to be white above the plane and black below.

Many thanks !

Upvotes: 1

Views: 319

Answers (1)

Quang Hoang
Quang Hoang

Reputation: 150755

You can define a custom color map and pass to 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()

Output:

enter image description here

Upvotes: 2

Related Questions