numpy cumsum from a specified center

center_cumsum.py

import numpy as np


def center_cumsum(a:np.ndarray, pivot:tuple) -> np.ndarray:
    """
    Compute the cumsum but pivoting a specified index
    """
    assert isinstance(a, np.ndarray)
    assert a.ndim == 2
    a_cum = np.zeros_like(a, dtype=a.dtype)

    assert len(pivot) == 2
    i, j = pivot
    assert (i >= 0) and (i < a.shape[0])
    assert (j >= 0) and (j < a.shape[1])
    
    quad1 = a[:i+1,j:]
    a_cum[:i+1,j:] = quad1.cumsum(axis=1)[::-1].cumsum(axis=0)[::-1]

    quad2 = a[:i+1,:j+1]
    a_cum[:i+1,:j+1] = quad2[::-1, ::-1].cumsum(axis=1).cumsum(axis=0)[::-1, ::-1]

    quad3 = a[i:,:j+1]
    a_cum[i:,:j+1] = quad3[:, ::-1].cumsum(axis=1).cumsum(axis=0)[:, ::-1]

    quad4 = a[i:,j:]
    a_cum[i:,j:] = quad4.cumsum(axis=1).cumsum(axis=0)

    return a_cum

import matplotlib.pyplot as plt
from center_cumsum import center_cumsum

a = np.random.randint(0, 255, (9,9))
plt.imshow(a, cmap="Greys")
_ = plt.axis("off")

a_cum = center_cumsum(a, [4, 4])
plt.imshow(a_cum, cmap="Greys")
_ = plt.axis("off")