Generate new matrix from two square matrices of the same shape based on the diagonal and a schema

Question

I have two matrices A and B which are square and have the same shape. I would like to generate a third square matrix C of the same shape such that for example,

C[7,4] = A[0,0]*B[7,4] + A[1,1]*B[6,3] + A[2,2]*B[5,2] + A[3,3]*B[4,1] + A[4,4]*B[3,0]

Another eg:

C[5,9] = A[0,0]*B[5,9] + A[1,1]*B[4,8] + A[2,2]*B[3,7] + A[3,3]*B[2,6] + A[4,4]*B[1,5] + A[5,5]*B[0,4]

Both indices of B decrease by 1 for each term in the sum, until either index is 0. The indices of A and B must sum to the original index we are trying to populate in C. The indices of A are restricted to that of its diagonal.

I would ideally like to do this vectorized using numpy and in the fastest way possible.

Answer 1

One way to have a fully vectorized version (but at the cost of doing sometimes some useless multiplication with 0) is to create a matrix of the subdiagonals

import numpy as np
# Example
n=10
A=np.arange(n*n).reshape(n,n)
B=np.random.uniform(-1,1,(n,n))

# C computation
Adiag=np.diag(A)[None,None,:]
Bpadded=np.pad(B, ((n-1,0),(n-1,0)))
s1,s2=Bpadded.strides
Bdiags = np.lib.stride_tricks.as_strided(Bpadded[n-1:,n-1:], shape=(n,n,n), strides=(s1,s2,-s1-s2))
C = (Bdiags*Adiag).sum(axis=2)

# Verification
# it works for your examples
C[7,4],  A[0,0]*B[7,4] + A[1,1]*B[6,3] + A[2,2]*B[5,2] + A[3,3]*B[4,1] + A[4,4]*B[3,0]
# (-47.11186336152282, -47.11186336152282)
C[5,9], A[0,0]*B[5,9] + A[1,1]*B[4,8] + A[2,2]*B[3,7] + A[3,3]*B[2,6] + A[4,4]*B[1,5] + A[5,5]*B[0,4]
# (85.27133136528813, 85.27133136528813)

Some explanation: Adiags is easy. It is just the diagonal of A, but reshaped in fake 3d, with only axis 2 being size n (just for later broadcasting)
Bpadded is just a copy of B, but ensuring that there are enough 0 around

And the tricky part is Bdiags. It is a 3d matrix, whose data are the one of Bpadded (no memory copy. It is the same memory address). But with different strides. Since in an array, arr[i1,i2,i3] is the data at address addressof(arr[0,0,0]) + strides1*i1 + strides2*i2 + strides3*i3, in my case, Bdiags[i1,i2,i3] is therefore data at address addressOf(Bpadded[n-1,n-1])+i1*strides1+i2*strides2+i3*(-strides1-strides2) aka adressOf(Bpadded[n-1,n-1])+(i1-i3)*strides1 + (i2-i3)*strides2. So it is Bpadded[n-1+i1-i3, n-1+i2-i3]. Which is B[i1-i3,i2-i3] but for the 0 when i1-i3 or i2-i3 are negative.

So Bdiags is just what you want to multiply with A[k,k]

Rest is simple broadcasting/multiplication/sum.

Answer 2

The easy pass is vectorization over the innermost column with a dot product. You can vectorise one more dimension by using a rolled matrix product. The original shown with the other two methods:

import numpy as np
from numpy.random import default_rng


def diagdot_slow(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    C = np.zeros_like(A)
    for i in range(C.shape[0]):
        for j in range(C.shape[1]):
            trace_length = 1 + min(i, j)
            for t in range(trace_length):
                C[i, j] += A[t, t] * B[i - t, j - t]
    return C


# https://stackoverflow.com/a/56175538/313768
def indep_roll(arr: np.ndarray, shifts: np.ndarray, axis: int=1) -> np.ndarray:
    arr = np.swapaxes(arr,axis,-1)
    all_idcs = np.ogrid[[slice(0,n) for n in arr.shape]]

    # Convert to a positive shift
    shifts[shifts < 0] += arr.shape[-1]
    all_idcs[-1] = all_idcs[-1] - shifts[:, np.newaxis]

    result = arr[tuple(all_idcs)]
    arr = np.swapaxes(result,-1,axis)
    return arr


def diagdot_inner_dot(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    Adiag = A.diagonal()
    C = np.empty_like(A)
    for i in range(C.shape[0]):
        for j in range(C.shape[1]):
            trace_length = 1 + min(i, j)
            Bdiag = B.diagonal(offset=j-i)[trace_length-1::-1]
            C[i, j] = Adiag[:trace_length].dot(Bdiag)
    return C


def diagdot_roll(A: np.ndarray, B: np.ndarray) -> np.ndarray:
    Adiag = A.diagonal()
    n = A.shape[0]
    C = np.empty_like(A)
    shifts = np.arange(0, -n, -1)
    Brolly = indep_roll(arr=B, shifts=shifts, axis=0)
    Brollx = indep_roll(arr=B, shifts=shifts, axis=1)

    for trace_len in range(1, n+1):
        Afactor = Adiag[:trace_len]
        short_dim = slice(trace_len-n-1, -n-1, -1)
        long_dim = n - trace_len + 1

        # Vertical segment; include the primary diagonal
        Bfactory = Brolly[:long_dim, short_dim]
        C[trace_len - 1:, trace_len - 1] = Bfactory @ Afactor

        # Horizontal segment; exclude the primary diagonal
        Bfactorx = Brollx[short_dim, 1: long_dim]
        C[trace_len - 1, trace_len:] = Afactor @ Bfactorx

    return C


def test() -> None:
    rand = default_rng(seed=0)
    n = 20
    A = rand.random((n, n))
    B = rand.random((n, n))
    Cref = diagdot_slow(A, B)

    for method in (diagdot_inner_dot, diagdot_roll):
        C = method(A, B)
        assert np.allclose(C, Cref)


if __name__ == '__main__':
    test()

Vectorising the last dimension is non-trivial because it's not rectangular.