Jacobian product

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 = f

$$
\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.

StackExchange Code Review Q#116417, answer score: 3