Search code examples
pythonnumpyscipyinterpolationsurface

How can I interpolate missing values (undefined areas) of a 3D surface plot (2D array) using Python?


In Python 3.7 using Numpy and matplotlib, I would like to plot a 3D surface for the following equation:

f(x,y)=sin(x)sin(y)/(x*y)

This function is obviously undefined where x=0 or y=0.

To calculate and plot this, I have the following code, which I am currently running in a Jupyter Notebook:

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
%matplotlib notebook

f = lambda x, y: np.sin(x)*np.sin(y)/(x*y)

xs, ys = np.mgrid[-np.pi:np.pi:31j, -np.pi:np.pi:31j]
zs = f(xs, ys)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

ax.plot_surface(X=xs, Y=ys, Z=zs)

Notice the graph, which is missing values: enter image description here

How can I interpolate the missing values so that the graph appears smooth?


Solution

  • Scipy has an interpolation module that can do this. Relying on the above (in the question posting), this code can be run in the next cell:

    from scipy import interpolate
    
    # integer arrays for indexing
    x_indx, y_indx = np.meshgrid(np.arange(0, zs.shape[1]),
                                 np.arange(0, zs.shape[0]))
    
    # mask all invalid values
    zs_masked = np.ma.masked_invalid(zs)
    
    # retrieve the valid, non-Nan, defined values
    valid_xs = x_indx[~zs_masked.mask]
    valid_ys = y_indx[~zs_masked.mask]
    valid_zs = zs_masked[~zs_masked.mask]
    
    # generate interpolated array of z-values
    zs_interp = interpolate.griddata((valid_xs, valid_ys), valid_zs.ravel(),
                                     (x_indx, y_indx), method='cubic')
    
    # finally, plot the data 
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    
    ax.plot_surface(X=xs, Y=ys, Z=zs_interp)
    

    The following graph is returned: enter image description here

    Note that this code has been optimized for readability and comprehensibility rather than memory efficiency. Re-Optimizing this code for memory efficiency is a trivial task that is left up to the reader