Generalized to any dimension. Tested.

SimPEG/utils/interputils.py

@@ -3,35 +3,97 @@ import scipy.sparse as sp
 from sputils import spzeros
 from matutils import mkvc, sub2ind
 
-def interpmat(x,y,z,xr,yr,zr):
+def _interp_point_1D(x, xr_i):
+    im = np.argmin(abs(x-xr_i))
+    if  xr_i - x[im] >= 0:  # Point on the left
+        ind_x1 = im
+        ind_x2 = im+1
+    elif  xr_i - x[im] < 0:  # Point on the right
+        ind_x1 = im-1
+        ind_x2 = im
+    dx1 = xr_i - x[ind_x1]
+    dx2 = x[ind_x2] - xr_i
+    return ind_x1, ind_x2, dx1, dx2
 
+
+def interpmat(locs, x, y=None, z=None):
     """ Local interpolation computed for each receiver point in turn """
+    if y is None and z is None:
+        return interpmat1D(locs, x)
+    elif z is None:
+        return interpmat2D(locs, x, y)
+    else:
+        return interpmat3D(locs, x, y, z)
 
+
+def interpmat1D(locs, x):
+    nx = x.size
+    locs = mkvc(locs)
+    npts = locs.shape[0]
+
+    Q = sp.lil_matrix((npts, nx))
+
+    for i in range(npts):
+        ind_x1, ind_x2, dx1, dx2 = _interp_point_1D(x, locs[i])
+        dv = (x[ind_x2] - x[ind_x1])
+        Dx =  x[ind_x2] - x[ind_x1]
+        # Get the row in the matrix
+        inds = [ind_x1, ind_x2]
+        vals = [(1-dx1/Dx),(1-dx2/Dx)]
+        Q[i, inds] = vals
+    return Q.tocsr()
+
+
+
+def interpmat2D(locs, x, y):
+    nx = x.size
+    ny = y.size
+    npts = locs.shape[0]
+
+    Q = sp.lil_matrix((npts, nx*ny))
+
+
+    for i in range(npts):
+        ind_x1, ind_x2, dx1, dx2 = _interp_point_1D(x, locs[i, 0])
+        ind_y1, ind_y2, dy1, dy2 = _interp_point_1D(y, locs[i, 1])
+
+        dv = (x[ind_x2] - x[ind_x1]) * (y[ind_y2] - y[ind_y1])
+
+        Dx =  x[ind_x2] - x[ind_x1]
+        Dy =  y[ind_y2] - y[ind_y1]
+
+        # Get the row in the matrix
+
+        inds = sub2ind((nx,ny),[
+            ( ind_x1,  ind_y2),
+            ( ind_x1,  ind_y1),
+            ( ind_x2,  ind_y1),
+            ( ind_x2,  ind_y2)])
+
+        vals = [(1-dx1/Dx)*(1-dy2/Dy),
+                (1-dx1/Dx)*(1-dy1/Dy),
+                (1-dx2/Dx)*(1-dy1/Dy),
+                (1-dx2/Dx)*(1-dy2/Dy)]
+
+        Q[i, mkvc(inds)] = vals
+
+    return Q.tocsr()
+
+
+
+def interpmat3D(locs, x, y, z):
     nx = x.size
     ny = y.size
     nz = z.size
-    npts = xr.shape[0]
+    npts = locs.shape[0]
 
     Q = sp.lil_matrix((npts, nx*ny*nz))
 
 
-    def inter1D(x, xr_i):
-        im = np.argmin(abs(x-xr_i))
-        if  xr_i - x[im] >= 0:  # Point on the left
-            ind_x1 = im
-            ind_x2 = im+1
-        elif  xr_i - x[im] < 0:  # Point on the right
-            ind_x1 = im-1
-            ind_x2 = im
-        dx1 = xr_i - x[ind_x1]
-        dx2 = x[ind_x2] - xr_i
-        return ind_x1, ind_x2, dx1, dx2
-
     for i in range(npts):
-        # in x-direction
-        ind_x1, ind_x2, dx1, dx2 = inter1D(x, xr[i])
-        ind_y1, ind_y2, dy1, dy2 = inter1D(y, yr[i])
-        ind_z1, ind_z2, dz1, dz2 = inter1D(z, zr[i])
+        ind_x1, ind_x2, dx1, dx2 = _interp_point_1D(x, locs[i, 0])
+        ind_y1, ind_y2, dy1, dy2 = _interp_point_1D(y, locs[i, 1])
+        ind_z1, ind_z2, dz1, dz2 = _interp_point_1D(z, locs[i, 2])
 
         dv = (x[ind_x2] - x[ind_x1]) * (y[ind_y2] - y[ind_y1]) *(z[ind_z2] - z[ind_z1])
 
@@ -61,5 +123,5 @@ def interpmat(x,y,z,xr,yr,zr):
                 (1-dx2/Dx)*(1-dy2/Dy)*(1-dz2/Dz)]
 
         Q[i, mkvc(inds)] = vals
-    Q = Q.tocsr()
-    return Q
+
+    return Q.tocsr()