solver: add interpolate_J

usadelndsoc/solver.py

@@ -53,6 +53,7 @@ __all__ = [
     "S_y",
     "S_z",
     "S_0",
+    "interpolate_J",
 ]
 
 _log = logging.getLogger(__name__)
@@ -1521,3 +1522,136 @@ def tr(M):
     Trace over last two dimensions
     """
     return np.trace(M, axis1=-2, axis2=-1)
+
+
+def interpolate_J(x0, y0, J, x, y, method="bilinear"):
+    r"""Interpolate (matrix) current components at given coordinates.
+
+    The current values are defined on links between cells, and not at
+    the cell centers, so special interpolation is necessary.
+
+    The interpolant inside each cell is linear :math:`J(x,y) = (J_0x + x J_1,
+    J_0y + y J_2)`. The component perpendicular to the cell faces is
+    piecewise linear, but the component parallel to cell faces is
+    piecewise constant.  The interpolant is divergenceless inside
+    cells, if the matrix current values do not indicate divergence.
+
+    Parameters
+    ----------
+    x0 : array, shape (nx,)
+        Cell x-coordinates
+    y0 : array, shape (ny,)
+        Cell y-coordinates
+    J : array, shape (nx, ny, 4, ...)
+        Matrix current, defined on links between cells.
+    x : array
+        x-coordinate to interpolate J at.
+    y : array
+        y-coordinate to interpolate J at.
+    method : {'linear', 'bilinear'}
+        Interpolation method.
+
+    Returns
+    -------
+    J_interp : array, shape broadcast_shape(x, y) + (2, ...)
+        Array containing x and y components of the current,
+        interpolated at the coordinates (x, y).
+
+    """
+
+    dx = x0[1] - x0[0]
+    dy = y0[1] - y0[0]
+    nx = len(x0)
+    ny = len(y0)
+
+    i = np.searchsorted(x0 + dx / 2, x).clip(0, nx - 1)
+    j = np.searchsorted(y0 + dy / 2, y).clip(0, ny - 1)
+
+    xc = x - x0[i]
+    yc = y - y0[j]
+
+    m = (
+        (x < x0[0] - dx / 2)
+        | (x > x0[-1] + dx / 2)
+        | (y < y0[0] - dy / 2)
+        | (y > y0[-1] + dy / 2)
+    )
+
+    Jshape = J.shape[3:]
+    sl = (Ellipsis,) + (None,) * len(Jshape)
+    slx = (Ellipsis, 0) + (None,) * len(Jshape)
+    sly = (Ellipsis, 1) + (None,) * len(Jshape)
+
+    I = np.empty(m.shape + (2,) + Jshape, dtype=J.dtype)
+
+    if method == "linear":
+        J0x = (J[i, j, 0] + J[i, j, 2]) / 2
+        J0y = (J[i, j, 1] + J[i, j, 3]) / 2
+
+        J1x = (J[i, j, 0] - J[i, j, 2]) / dx
+        J2y = (J[i, j, 1] - J[i, j, 3]) / dy
+
+        I[slx] = J0x + xc[sl] * J1x
+        I[sly] = J0y + yc[sl] * J2y
+    elif method == "bilinear":
+        # Bilinear polynomial hat functions at corner sites of cells.
+        # They are nonzero only on neighboring cells.
+        #
+        # On each of the edges, they are linear functions along the
+        # edge, and their average should coincide with the current
+        # perpendicular to the edge.
+
+        def M(n):
+            M = np.zeros((n, n + 1))
+            p = np.arange(n)
+            M[p, p] = 1 / 2
+            M[p, p + 1] = 1 / 2
+            return M
+
+        I0 = J[:, :, 0]
+        I1 = J[:, :, 1]
+        I2 = J[:, :, 2]
+        I3 = J[:, :, 3]
+
+        # Solve corner values Iyc: (Iyc(i,j) + Iyc(i+1,j))/2 == I3(i,j)
+        r = np.empty((nx, ny + 1) + I0.shape[2:], dtype=I0.dtype)
+        r[:, :ny] = I3
+        r[:, ny] = I1[:, -1]
+
+        rp = r.reshape(nx, -1)
+        z = np.linalg.lstsq(M(nx), rp, rcond=None)[0]
+        Iyt = z.reshape(nx + 1, *r.shape[1:])
+
+        # Solve corner values Ixc: (Ixc(i,j) + Ixc(i,j+1))/2 == I2(i,j)
+        r = np.empty((nx + 1, ny) + I0.shape[2:], dtype=I0.dtype)
+        r[:nx, :] = I2
+        r[nx, :] = I0[-1, :]
+
+        rp = np.moveaxis(r, 1, 0).reshape(ny, -1)
+        z = np.linalg.lstsq(M(ny), rp, rcond=None)[0]
+        Ixt = np.moveaxis(z.reshape(ny + 1, *((nx + 1,) + rp.shape[2:])), 0, 1)
+
+        # Interpolate with the bilinear hat functions
+        ax = 0.5 - xc / dx
+        bx = 0.5 + xc / dx
+
+        ay = 0.5 - yc / dy
+        by = 0.5 + yc / dy
+
+        I[slx] = (
+            ax * ay * Ixt[i, j]
+            + bx * ay * Ixt[i + 1, j]
+            + ax * by * Ixt[i, j + 1]
+            + bx * by * Ixt[i + 1, j + 1]
+        )
+        I[sly] = (
+            ax * ay * Iyt[i, j]
+            + bx * ay * Iyt[i + 1, j]
+            + ax * by * Iyt[i, j + 1]
+            + bx * by * Iyt[i + 1, j + 1]
+        )
+    else:
+        raise ValueError(f"Unknown {method=}")
+
+    I[m] = 0
+    return I