Jacobian product

Question

Do you have in mind a way to compute the following sum

$$R_{i,j,k} = \sum_{l=0}^{2} J_{i,j,l+3k} v_{i,j,l}$$

obviously in a vectorised fashion? This may sound like an equivalent of:

c = np.einsum('ij(l + 3k),ijl->ijk', j, v)

where of course the notation 'ij(l + 3k),ijl->ijk' is not accepted by einsum.

I tried with the following, but I wish I could use one of the existing NumPy or SciPy functions:

def jacobian_product(j_input, v_input):
    """
    :param j_input: jacobian m x n x (4 or 9) jacobian column major
    :param v_input: matrix m x n x (2 or 3) to be multiplied by the jacobian.
    :return: m x n  x (2 or 3) whose each element is the result of the product of the
     jacobian (i,j,:) multiplied by the corresponding element in the vector v (i,j,:).

    In tensor notation: R_{i,j,k} = \sum_{l=0}^{2} J_{i,j,l+3k} v_{i,j,l}
    """
    # dimensions of the problem:
    d = v_input.shape[-1]
    vol = list(v_input.shape[:-1])

    # repeat v input 3 times, one for each row of the jacobian matrix 3x3 or 2x2 in corresponding position:
    v = np.tile(v_input, [1]*d + [d])

    # element-wise product:
    j_times_v = np.multiply(j_input, v)

    # Sum the three blocks in the third dimension:
    return np.sum(j_times_v.reshape(vol + [d, d]), axis=d+1).reshape(vol + [d])

---

Update after Veedrac comment: here there are two tests for both 2d and 3d vectors.

The jacobians are collected into an array of shape [m,n,4] or [m,n,9] (for 2d or 3d. The first equation proposed in this thread is about the 3d case). Vectors to be multiplied by the Jacobian in the corresponding positions have shape [m,n,2] or [m,n,3] respectively.

Sorry if this was not clear, my bad I didn't state dimensions since the beginning!

from numpy.testing import assert_array_equal

def test_jacobian_product_toy_example_2d():

    def function_v(t, x):
        t = float(t); x = [float(y) for y in x]
        return x[1], -1 * x[0]

    def function_jacobian(t, x):
        t = float(t); x = [float(y) for y in x]
        return 0.5, 0.6, x[0], 0.8

    def function_ground_product(t, x):
        t = float(t); x = [float(y) for y in x]
        return 0.5*x[1] - 0.6*x[0], x[0]*x[1] - 0.8*x[0]

    v   = np.zeros([20, 20, 2])
    jac = np.zeros([20, 20, 4])

    ground_jac_v = np.zeros([20, 20, 2])

    for i in range(0, 20):
        for j in range(0, 20):

            v[i, j, :]   = function_v(1, [i, j])
            jac[i, j, :] = function_jacobian(1, [i, j])

            ground_jac_v[i, j, :] = function_ground_product(1, [i, j])

    jac_v = jacobian_product(jac, v)

    assert_array_equal(jac_v, ground_jac_v)


def test_jacobian_product_toy_example_3d():

    def function_v(t, x):
        t = float(t); x = [float(y) for y in x]
        return 2*x[1], 3*x[0], x[0] - x[2]

    def function_jacobian(t, x):
        t = float(t); x = [float(y) for y in x]
        return 0.5*x[0], 0.5*x[1], 0.5*x[2], \
               0.5,      x[0],     0.3,      \
               0.2*x[2], 3.0,      2.1*x[2]

    def function_ground_product(t, x):
        t = float(t); x = [float(y) for y in x]
        return 0.5*2*x[0]*x[1] + 0.5*3*x[0]*x[1] + 0.5*x[2]*(x[0] - x[2]), \
               0.5*2*x[1] + 3*x[0]*x[0] + 0.3*(x[0] - x[2]), \
               0.2*2*x[2]*x[1] + 3*3*x[0] + 2.1*x[2]*(x[0] - x[2])

    v   = np.zeros([20, 20, 20, 3])
    jac = np.zeros([20, 20, 20, 9])

    ground_jac_v = np.zeros([20, 20, 20, 3])

    for i in range(0, 20):
        for j in range(0, 20):
            for k in range(0, 20):
                v[i, j, k, :]   = function_v(1, [i, j, k])
                jac[i, j, k, :] = function_jacobian(1, [i, j, k])

                ground_jac_v[i, j, k, :] = function_ground_product(1, [i, j, k])

    jac_v = jacobian_product(jac, v)

    assert_array_equal(jac_v, ground_jac_v)

Answer 1

This function replicates your jacobian_product.

def alt_jacobian(j_input, v_input):
    n,m = j_input.shape[:2]
    temp = j_input.reshape(n,m,-1,3)
    return np.einsum('...kl,...l->...k', temp, v_input)

For small arrays, e.g. (4,5,9), this einsum is 5x faster.

---

(in response to question changes in the last 2 days)

To handle the 2 v 3 size dimension, you could pull that from v_input.shape[-1], e.g.

def alt_jacobian(j_input, v_input):
    m,n = j_input.shape[:2]
    p = v_input.shape[-1]
    # or m, n, p = v_input.shape
    temp = j_input.reshape(m,n,-1,p)
    return np.einsum('...kl,...l->...k', temp, v_input)

It looks like the 2 test cases use (20,20,4) and (20,20,20,9) for v_input.

The einsum expression can handle both. The only difference is how the last dimension is reshaped from 4 to 2x2 or 9 to 3x3.

Here's a more general reshaping expression - essentially taking the shape of v_input and inserting a -1 2nd to the end.

temp = j_input.reshape(v_input.shape[:-1]+(-1,v_input.shape[-1]))

Answer 2

$$ \begin{align} R_{i,j,k} &= \sum_{l=0}^{2} J_{i,j,l+3k} v_{i,j,l} \\ &= J_{i,j,3k} v_{i,j,0} + J_{i,j,3k+1} v_{i,j,1} + J_{i,j,3k+2} v_{i,j,2} \end{align} $$

is easily expressible as

R = sum(J[..., l::3] * v[..., l, newaxis] for l in range(3))

or, in long,

R = (
    J[..., 0::3] * v[..., 0, newaxis] +
    J[..., 1::3] * v[..., 1, newaxis] +
    J[..., 2::3] * v[..., 2, newaxis]
)

This would be prettier if your axes were in the other order, as the ...s and newaxis would be implied and you'd just do

R = sum(J[l::3] * v[l] for l in range(3))

You can do this as a vectorized summation, but it's not pretty and will probably be slower. First you change the last dimension in J to group it in threes:

J_threes = J.reshape(J.shape[:-1] + (-1, 3))

Then you multiply over a new axis and sum:

R = (J_threes * v[..., newaxis, :]).sum(axis=3)

This doesn't seem to match your example function, though, so I'm somewhat confused. Hopefully this at least helps.