'''
Suppose you have a 3-D array of (negative) depths
with indices level, x, and y. You also have a field,
say T, indexed the same way, and you want to interpolate
it to a given depth. Here is a quick example of how to
do it with no loops. This is not carefully tested, may
be buggy or incorrect, needs docstrings, etc.
It is intended to work for arrays of any dimension greater than
1 provided the first dimension is the one indexed by level.
'''
import numpy as np
def brackets(z, z0):
below = (z > z0).astype(np.int16).sum(axis=0)
above = below - 1
badmask = ((below == z.shape[0]) | (above < 0)).astype(np.bool)
above[badmask] = 0
below[badmask] = z.shape[0] - 1
return below, above, badmask
def interp(T, z, z0):
below, above, badmask = brackets(z, z0)
ind = np.indices(z.shape[1:])
ibelow = tuple([below] + list(ind))
iabove = tuple([above] + list(ind))
Tb = T[ibelow]
zb = z[ibelow]
dT = T[iabove] - Tb
dz = z[iabove] - zb
Ti = Tb + (dT/dz) * (z0 - zb)
return Ti, badmask
#Example use:
# Make depth levels varying linearly from the bottom to the surface:
nlevels = 5 # Use larger values to verify convergence.
z0 = np.linspace(0.0, -5000.0, num=nlevels)
z0.shape = nlevels,1,1
z1 = np.linspace(0.5, 1.0, num=4)
z1.shape = 1,4,1
z2 = np.linspace(0.1, 1.0, num=3)
z2.shape = 1,1,3
z = z0*z1*z2
# Temperature increases quadratically from 0 at -5000 m to 20 at the surface:
T = 20.0 * (1 + z/5000.0)**2
T1500, T1500badmask = interp(T, z, -1500)
# Now you might want to use a masked array to keep track of
# the valid values:
T1500m = np.ma.array(T1500, mask=T1500badmask)
print T1500m