def div(X, u):
    """compute the divergence for a vector field"""
    div = _np.zeros(u.shape[1:], dtype=u.dtype)

    poly = _interp(X[0], u[0], axis=0)
    div += poly(X[0], nu=1)

    poly = _interp(X[1], u[1], axis=1)
    div += poly(X[1], nu=1)

    poly = _interp(X[2], u[2], axis=2)
    div += poly(X[2], nu=1)

    return div
def scl_grad(X, phi):
    """compute the gradient vector for a scalar field"""
    shape = list(phi.shape)
    shape.insert(0, 3)
    grad = _np.zeros(shape, dtype=phi.dtype)

    polynomial = _interp(X[0], phi, axis=0)
    grad[0, ...] = polynomial(X[0], nu=1)

    polynomial = _interp(X[1], phi, axis=1)
    grad[1, ...] = polynomial(X[1], nu=1)

    polynomial = _interp(X[2], phi, axis=2)
    grad[2, ...] = polynomial(X[2], nu=1)

    return grad
def grad(X, u):
    """compute the gradient tensor (Jacobian) for a vector field"""
    Jshape = list(u.shape)
    Jshape[1] -= 4  #subtract ghost cell count from shape
    Jshape.insert(0, 3)
    J = _np.zeros(Jshape, dtype=u.dtype)

    for i in range(3):
        polynomial = _interp(X[0], u[i, :, :, :], axis=0)
        J[0, i, ...] = polynomial(X[0][2:-2], nu=1)

        polynomial = _interp(X[1], u[i, 2:-2, :, :], axis=1)
        J[1, i, ...] = polynomial(X[1], nu=1)

        polynomial = _interp(X[2], u[i, 2:-2, :, :], axis=2)
        J[2, i, ...] = polynomial(X[2], nu=1)

    return J
def deriv(x, phi, order=1, axis=0):
    """compute the derivative of a scalar field"""
    deriv = _np.zeros_like(phi)
    copy = _np.ascontiguousarray

    if axis == 0:
        for x1 in xrange(phi.shape[1]):
            for x2 in xrange(phi.shape[2]):
                polynomial = _interp(X[0], copy(phi[:, x1, x2]))
                deriv[:, x1, x2] = polynomial(X[0], nu=order)
    elif axis == 1:
        for x0 in xrange(phi.shape[0]):
            for x2 in xrange(phi.shape[2]):
                polynomial = _interp(X[1], copy(phi[x0, :, x2]))
                deriv[x0, :, x2] = polynomial(X[1], nu=order)
    else:
        for x0 in xrange(phi.shape[0]):
            for x1 in xrange(phi.shape[1]):
                polynomial = _interp(X[2], copy(phi[x0, x1, :]))
                deriv[x0, x1, :] = polynomial(X[2], nu=order)

    return deriv
def helmholtz(X, u):
    """Computes the Helmholtz decomposition of a vector field into a dilatational (zero curl) and
    solenoidal (zero div) component by solving the poisson equation for the dilatational component."""
    ud = _np.zeros_like(u)
    copy = _np.ascontiguousarray

    theta = div(X, u)
    gth = scl_grad(X, theta)

    for i in range(3):
        poly1 = _interp(X[0], gth[i, :, :, :], axis=0)
        poly2 = poly1.antiderivative(nu=2)
        ud[i, ...] += poly2(X[0])

        poly1 = _interp(X[1], gth[i, :, :, :], axis=1)
        poly2 = poly1.antiderivative(nu=2)
        ud[i, ...] += poly2(X[1])

        poly1 = _interp(X[2], gth[i, :, :, :], axis=2)
        poly2 = poly1.antiderivative(nu=2)
        ud[i, ...] += poly2(X[2])

    return ud
def deriv_bak(phi, h, axis=0):
    """
    deriv(phi, h, axis=0):

    deriv computes the k'th derivative of a uniform gridded array along the
    prescribed axis using Akima spline approximation.

    Arguments
    ---------
    phi   - input array
    h     - uniform grid spacing
    bc    -
    k     - order of the derivative
    axis  -

    Output
    ------
    f - d^k/dx^k(phi)
    """
    axis = axis % phi.ndim

    if axis != 0:
        phi = _np.swapaxes(phi, axis, 0)

    s = list(phi.shape)
    x = _np.arange(-3, s[0] + 3, dtype=_np.float64)
    xi = _np.arange(-0.5, s[0] + 0.5, 1.0, dtype=_np.float64)
    deriv = _np.empty_like(phi)

    s[0] += 6
    tmp = _np.empty(s, dtype=phi.dtype)

    tmp[3:-3] = phi
    tmp[:3] = phi[-3:]
    tmp[-3:] = phi[:3]

    for j in range(s[-2]):
        for i in range(s[-1]):
            spline = _interp(x, tmp[..., j, i])
            phi2 = spline(xi)
            deriv[..., j, i] = (phi2[1:] - phi2[:-1]) * (1.0 / h)

    if axis != 0:
        deriv = _np.swapaxes(deriv, axis, 0)

    return deriv
def deriv(phi, h, axis=0):
    """
    deriv(phi, h, axis=0):

    deriv computes the k'th derivative of a uniform gridded array along the
    prescribed axis using Akima spline approximation.

    Arguments
    ---------
    phi   - input array
    h     - uniform grid spacing
    axis  -

    Output
    ------
    f - d^k/dx^k(phi)
    """
    if phi.ndim != 3:
        print("ERROR: phi.ndim not equal to 3!")

    axis = axis % phi.ndim
    if axis != 2:
        phi = _np.swapaxes(phi, axis, 2)

    s = list(phi.shape)
    x = _np.arange(-3, s[2] + 3, dtype=phi.dtype)
    xh = _np.arange(-0.5, s[2] + 0.5, 1.0, dtype=phi.dtype)
    deriv = _np.empty_like(phi)

    nx = s[2] + 6
    tmp = _np.empty(nx, dtype=phi.dtype)

    for k in range(s[0]):
        for j in range(s[1]):
            tmp[3:-3] = phi[k, j, :]
            tmp[:3] = phi[k, j, -3:]
            tmp[-3:] = phi[k, j, :3]

            spline = _interp(x, tmp)
            phih = spline(xh)
            deriv[k, j, :] = (1.0 / h) * (phih[1:] - phih[:-1])

    if axis != 2:
        deriv = _np.swapaxes(deriv, axis, 2)

    return deriv