Extrapolation

#!/usr/bin/python
import numpy as np
import scipy
from scipy.interpolate import griddata
import matplotlib.pyplot as plt
from matplotlib.mlab import griddata
import matplotlib.mlab as ml

def extrapolate_nans(x, y, v):
    if np.ma.is_masked(v):
        nans = v.mask
    else:
        nans = np.isnan(v)
    notnans = np.logical_not(nans)
    v[nans] = scipy.interpolate.griddata((x[notnans], y[notnans]), v[notnans],
        (x[nans], y[nans]), method='nearest').ravel()
    return v


data = np.genfromtxt('input2d.txt')
x1 = data[:,0]
y1 = data[:,1]
z1 = data[:,2]

numcols, numrows = 500, 500
xi = np.linspace(min(x1), max(x1), numcols)
yi = np.linspace(min(y1), max(y1), numrows)
xi, yi = np.meshgrid(xi, yi)
x, y, z = x1, y1, z1
zi = ml.griddata(x, y, z, xi, yi, interp='nn')
extrapolate_nans(xi,yi,zi)
plt.contour(xi, yi, zi, 20, linewidths = 0.5, colors = 'k')
plt.contourf(xi,yi,zi)
plt.colorbar()
plt.show()