Accumulate sliding windows relative to origin

As I mentioned in the comment, a vectorized solution doesn't always guarantee a better running time. If your matrix is large, you might prefer more efficient methods. And there are several reasons why matrix rotation is slow (though, intuitive), see comment.

Performance comparison:

Solution: Wall time: 61.6 ms
Rotation: Wall time: 3.32 s

Code (tested in jupyter notebook)

import numpy as np

def rotate45_and_sum(A):
    n = len(A) 
    x, y = np.meshgrid(np.arange(n), np.arange(n))  # at least doubled the running time
    xn, yn = x + y, n - x + y - 1   # generating xn and yn at least doubled the running time
    M = np.zeros((2*n -1, 2*n -1))  # at least slows down running time by a factor of 4
    M[xn,yn] = A[x,y] # very inefficient indexing strategy
    return M.sum(1)

def solution(A):
    n = A.shape[0]
    retval = np.zeros(2*n-1)
    for i in range(n):
        retval[i:(i+n)] += A[i, :]
    return retval

A = np.random.randn(10000, 10000)

%time solution(A)

%time rotate45_and_sum(A)

In multidimensional situation:

def solution(A):
    h,w,x,y = A.shape                # change here
    retval = np.zeros((2*x-w,2*y-h)) # change here
    indices = np.ndindex(w, h)       # change here
    for index in indices:
        slices = tuple()
        for i in range(len(index)):
            slices = slices + (slice(index[i], index[i]+x),) # I assume x = y = ..., you need to change here also if the assumption is not correct
        retval[slices] += A[index] # slices is roughly equal `i:(i+x), j:(j+y)` in your code
    return retval

Actually I don't know the how the dimensions (or shapes) are calculated based on your description :(. But i think it could be generalized. The idea is to construct slices as you go. So you need to specify which dimensions correspond to h, w, which correspond to x, y. I think it's not difficult to do that.

Reference: Numpy index array of unknown dimensions?

Regarding https://stackoverflow.com/a/67341994/14923227

def fast(A):
    n = A.shape[0]
    retval = np.zeros(2*n-1)
    for i in range(n):
        retval[i:(i+n)] += A[i, :]
    print(retval.sum())
    return retval

##########################
import threading

class sumThread(threading.Thread):
    def __init__(self, A, mat, threadID, ngroups, size):
        threading.Thread.__init__(self)
        self.threadID = threadID
        self.size = size
        self.ngroups = ngroups
        self.mat = mat
        self.A = A
    def run(self):
        begin = (self.size + self.ngroups) // self.ngroups * self.threadID
        end   = min(self.size, (self.size+self.ngroups)//self.ngroups*(self.threadID+1))
        for i in range(begin, end):
            self.mat[self.threadID, i:(i+self.size)] += self.A[i, :]

def faster(A):
    
    num_threads = max(1, A.shape[0] // 4000) 
    mat = np.zeros((num_threads, 2*A.shape[0]-1))
    threads = []
    for i in range(num_threads):
        t = sumThread(A, mat, i, num_threads, A.shape[0])
        t.start()
        threads.append(t)

    # Wait for all threads to complete
    for t in threads:
        t.join()
    return np.sum(mat, axis=0)
    

Performance for large array:

A = np.random.randn(20000,20000)
%timeit fast(A)   # 263 ms ± 5.21 ms per loop 
%timeit faster(A) # 155 ms ± 3.14 ms per loop

It's trivial to parallelize the for loop in fast. But fast is actually the most cache efficient (even for GPU cache and memory banks) and thus the fastest way to compute it. Ideally, you can parallelize the code with CUDA/OpenCL since there are way more cores in a GPU. If you do it correctly, the running time will be reduced to log(original_fast_time) with base k, where k is the number of cores you have.

However, there are only a few computations in the function. So the transportation of data between memory and GRAM might dominate. (I didn't test it)

IIUC the problem proposed here is equivalent to rotate matrix A by -45 degrees and sum row-wise (at least for the 2D version). For a better understanding of what I mean by rotating the matrix, see this post.

def rotate45_and_sum(A):
    n = len(A) 
    x, y = np.meshgrid(np.arange(n), np.arange(n)) 
    xn, yn = x + y, n - x + y - 1
    M = np.zeros((2*n -1, 2*n -1)) 
    M[xn,yn] = A[x,y] 
    return M.sum(1)

A = np.array([[0, 0, 1],
              [0, 0, 1],
              [0, 0, 1]], dtype=np.float32)

print(rotate45_and_sum(A))
#[0. 0. 1. 1. 1.]

A = np.array([[1, 2, 3],
              [2, 3, 4],
              [3, 4, 5]], dtype=np.float32)

print(rotate45_and_sum(A))
#[1. 4. 9. 8. 5.]

M is the rotated matrix.

Disclaimer: I don't know if this can be generalized for multiple dimensions

Expanding on the fast solution given by @Shihao Xu, I tried translating it to compilable c-code by adding a np.fast_compiled function inside numpy/core/src/multiarray:

NPY_NO_EXPORT PyObject *
arr_fast_compiled(PyObject *NPY_UNUSED(self), PyObject *args)
{
    PyObject *list_obj = NULL;
    PyArrayObject *list_arr = NULL, *ans = NULL;

    npy_intp len, ans_size;
    npy_intp i, j, k;
    double *dans, *numbers;

    if (!PyArg_ParseTuple(args, "O", &list_obj)) {
            goto fail;
    }

    list_arr = (PyArrayObject *)PyArray_ContiguousFromAny(list_obj, NPY_DOUBLE, 2, 2);
    if (list_arr == NULL) {
        goto fail;
    }

    len = PyArray_DIM(list_arr, 0);
    numbers = (double *)PyArray_DATA(list_arr);
    ans_size = 2*len-1;

    ans = (PyArrayObject *)PyArray_ZEROS(1, &ans_size, NPY_DOUBLE, 0);
    if (ans == NULL) {
        goto fail;
    }
    dans = (double *)PyArray_DATA(ans);
    NPY_BEGIN_ALLOW_THREADS;
    for (i = 0; i < len; ++i) {
        k = i * len;
        for (j = i; j < i + len; ++j, ++k) {
            dans[j] += numbers[k];
        }
    }
    NPY_END_ALLOW_THREADS;
    Py_DECREF(list_arr);
    return (PyObject *)ans;

fail:
    Py_XDECREF(list_arr);
    Py_XDECREF(ans);
    return NULL;
}

The for-loop is the most important:

for (i = 0; i < len; ++i) {
    k = i * len;
    for (j = i; j < i + len; ++j, ++k) {   
        dans[j] += numbers[k];
    }
}

numbers is the input argument (A), and we access the elements in numbers and dans in a strided fashion. In the 3x3 example I had we have the following j and k values:

j = [0, 1, 2, 1, 2, 3, 2, 3, 4]
k = [0, 1, 2, 3, 4, 5, 6, 7, 8]

The NPY_BEGIN_ALLOW_THREADS is something I saw frequently used for other numpy functions, but it seems to make no performance difference when I tested it without.

The performance is similar to just_sum_0