Sums over overlapping rectangular slices of a 2D array

I need to calculate multiple sums, each sum being over an axis oriented rectangular slice of a 2D array. Many of the slices will overlap, and therefore there will be many shared terms in the sums. To reduce the needed computations, how can I automatically split the sums so that sums over shared terms are calculated separately to be then included in the sums in which they appear?

Figure 1. Toy example data with sums calculated over two overlapping slices (red and blue). The terms 1 4 1 appear in both slices.

For my application, the speed of the algorithm is not very critical because determining the split will be done at compile time ahead of computation on GPU (using the JAX framework). The values in the 2D array are not known at compile time.

Here is a Python function that splits the area covered by the given rectangles into sub-rectangles, each sub-rectangle being fully overlapped by a number of the input rectangles. The algorithm first lists all the x and y coordinates appearing in the rectangle corner coordinates, and sorts each list. Then it can more compactly use the indexes to those lists to encode each x and y. The approach was inspired by this answer. In a 2D array indexed by the x and y encodings, the algorithm lists which areas are covered by which input rectangles, using a 64-bit binary representation. Multiple such arrays are used if there are more than 64 input rectangles. The 2D array is efficiently populated using a 2D cumulative sum. Then the array is scanned for the upper left corner of sub-rectangles. Each sub-rectangle is enlarged, first horizontally and then vertically, until it would no longer only contain the same value throughout.

The algorithm doesn't necessarily provide the minimum number of splits.

Print statements are included to more easily follow the algorithm. They can be safely removed.

import numpy as np

# Split axis oriented input rectangles and their overlaps to subrectangles that are each either fully covered or not at all covered by an input rectangle
# Input should be a list of tuples ((start_y, start_x), (end_y, end_x)) where the values must be of sortable type
# All returned subrectangles are covered by at least one input rectangle
# Returns a list of subrectangle tuples ((start_y, start_x), (end_y, end_x), indexes)) where indexes is a tuple of indexes to the input list indicating the input rectangles that cover the subrectangle
def get_subrectangles(rectangles):
  # Find x and y values that appear in the rectangle coordinates
  ys = sorted(set([rectangle[0][0] for rectangle in rectangles] + [rectangle[1][0] for rectangle in rectangles]))
  xs = sorted(set([rectangle[0][1] for rectangle in rectangles] + [rectangle[1][1] for rectangle in rectangles]))
  y_dict = {}
  x_dict = {}
  for index, x in enumerate(xs):
    x_dict[x] = index
  for index, y in enumerate(ys):
    y_dict[y] = index
  print("\nys:")
  print(ys)
  print("xs:")
  print(xs)
  print("y_dict:")
  print(y_dict)
  print("x_dict:")
  print(x_dict)

  # Use binary representation to mark areas covered by each rectangle
  membership = np.zeros([len(ys), len(xs), ((len(rectangles) - 1) >> 6) + 1], dtype=np.int64)
  for index, rectangle in enumerate(rectangles):
    signature_int64 = 1 << np.int64(index & 0b111111)
    signature_table = index >> 6
    with np.errstate(over='ignore'):
      print(f"membership[{y_dict[rectangle[0][0]]}, {x_dict[rectangle[0][1]]}, {signature_table}] += {signature_int64}")
      print(f"membership[{y_dict[rectangle[1][0]]}, {x_dict[rectangle[0][1]]}, {signature_table}] -= {signature_int64}")
      print(f"membership[{y_dict[rectangle[1][0]]}, {x_dict[rectangle[1][1]]}, {signature_table}] += {signature_int64}")
      print(f"membership[{y_dict[rectangle[0][0]]}, {x_dict[rectangle[1][1]]}, {signature_table}] -= {signature_int64}")
      membership[y_dict[rectangle[0][0]], x_dict[rectangle[0][1]], signature_table] += signature_int64
      membership[y_dict[rectangle[1][0]], x_dict[rectangle[0][1]], signature_table] -= signature_int64
      membership[y_dict[rectangle[1][0]], x_dict[rectangle[1][1]], signature_table] += signature_int64
      membership[y_dict[rectangle[0][0]], x_dict[rectangle[1][1]], signature_table] -= signature_int64
  print("2D cumsum")
  membership = membership.cumsum(axis=0).cumsum(axis=1)

  # Find subrectangles
  print("\nmembership: (table axis moved first for convenience)")
  print(np.moveaxis(membership, -1, 0))
  subrectangles = []
  for start_y in range(membership.shape[0] - 1):
    for start_x in range(membership.shape[1] - 1):
      # Use top left as the membership signature that the rest of the rectangle must match
      signature = membership[start_y, start_x, :]
      if np.any(signature != 0):
        # First enlarge the rectangle horizontally
        end_x = start_x + 1
        while np.all(membership[start_y, end_x, :] == signature):
          end_x += 1
        # Then enlarge the rectangle vertically
        end_y = start_y + 1
        while np.all(membership[end_y, start_x:end_x, :] == signature):
          end_y += 1
        indexes = []
        for index in range(len(rectangles)):
          signature_int64 = 1 << np.int64(index & 0b111111)
          signature_table = index >> 6
          if signature_int64 & signature[signature_table]:
            indexes.append(index)
        if len(indexes):
          subrectangles.append(((ys[start_y], xs[start_x]), (ys[end_y], xs[end_x]), tuple(indexes)))
        membership[start_y:end_y, start_x:end_x, :] = 0
        print("membership: (table axis moved first for convenience)")
        print(np.moveaxis(membership, -1, 0))
  return subrectangles

For example an input list of two rectangles, each element being a tuple of tuples denoting the upper left and the lower right corner coordinates:

[((1, 0), (6, 9)), ((3, 8), (9, 10))]

results in the following list of sub-rectangles:

[((1, 0), (3, 9), (0,)), ((3, 0), (6, 8), (0,)), ((3, 8), (6, 9), (0, 1)), ((3, 9), (9, 10), (1,)), ((6, 8), (9, 9), (1,))]

The third tuple in each gives the indexes of the original rectangles that the sub-rectangle fully overlaps with.

Figure 1. Left: the two input rectangles. Right: The five sub-rectangles found by the algorithm.

The list of sub-rectangles can then be used to calculate the sums as follows. Validation of the results is included:

# Split multiple sums over slices of data to slices with shared terms
# Start (included) and end (excluded) coordinates for rectangles, over which sums should be calculated

from random import randint, seed

def randStartEnd(N):
  a = randint(0, N)
  b = randint(0, N)
  if a > b:
    a, b = b, a
  if a == b:
    if a == 0:
      b = 1
    elif b == N:
      a = N-1
    else: b = a + 1
  return a, b

# Configuration
M = 2  # number of rectangles
N = 10  # width and height
seed(46)
np.random.seed(47)

data = np.random.randint(0, 10, (N, N))
print("\ndata:")
print(data)
rectangles = []
sums_validation = np.zeros((M,))
sums = np.zeros((M,))
orig_work = 0
work = 0
for i in range(M):
  start_y, end_y = randStartEnd(N)
  start_x, end_x = randStartEnd(N)
  rectangles.append(((start_y, start_x), (end_y, end_x)))
  sums_validation[i] = np.sum(data[start_y:end_y, start_x:end_x])
  orig_work += (end_y - start_y)*(end_x - start_x)
print("rectangles:")
print(rectangles)
subrectangles = get_subrectangles(rectangles)
print("\nsubrectangles:")
print(subrectangles)
print("Contributions of subrectangle sums to rectangle sums:")
for subrectangle in subrectangles:
  start_y = subrectangle[0][0]
  start_x = subrectangle[0][1]
  end_y = subrectangle[1][0]
  end_x = subrectangle[1][1]
  work += (end_y - start_y)*(end_x - start_x)
  terms_sum = np.sum(data[start_y:end_y, start_x:end_x])
  terms_sum_str = f"np.sum(data[{start_y}:{end_y}, {start_x}:{end_x}])"
  if len(subrectangle[-1]) > 1:
    print(f"shared_terms_sum = {terms_sum_str}")
  for index in subrectangle[-1]:
    print(f"sums[{index}] += {'shared_terms_sum' if len(subrectangle[-1]) > 1 else terms_sum_str}")
    sums[index] += terms_sum

print(f"Calculated sums over {100*work/orig_work} % of terms of the original sums")
print("sums:")
print(sums)
print("sums_validation:")
print(sums_validation)
print("diff:")
print(sums_validation - sums)

Example printout:

data:
[[7 6 7 8 8 3 0 7 0 7]
 [7 1 7 2 2 1 7 4 8 9]
 [2 9 1 5 0 9 2 0 2 1]
 [4 3 4 5 9 9 6 6 1 5]
 [3 9 8 9 4 3 5 6 4 3]
 [1 4 4 2 2 0 9 0 1 8]
 [4 0 8 9 6 6 9 5 2 2]
 [9 2 8 8 5 8 0 4 4 6]
 [7 4 8 1 7 8 9 8 9 6]
 [1 7 6 0 1 9 2 2 7 8]]
rectangles:
[((1, 0), (6, 9)), ((3, 8), (9, 10))]

ys:
[1, 3, 6, 9]
xs:
[0, 8, 9, 10]
y_dict:
{1: 0, 3: 1, 6: 2, 9: 3}
x_dict:
{0: 0, 8: 1, 9: 2, 10: 3}
membership[0, 0, 0] += 1
membership[2, 0, 0] -= 1
membership[2, 2, 0] += 1
membership[0, 2, 0] -= 1
membership[1, 1, 0] += 2
membership[3, 1, 0] -= 2
membership[3, 3, 0] += 2
membership[1, 3, 0] -= 2
2D cumsum

membership: (table axis moved first for convenience)
[[[1 1 0 0]
  [1 3 2 0]
  [0 2 2 0]
  [0 0 0 0]]]
membership: (table axis moved first for convenience)
[[[0 0 0 0]
  [1 3 2 0]
  [0 2 2 0]
  [0 0 0 0]]]
membership: (table axis moved first for convenience)
[[[0 0 0 0]
  [0 3 2 0]
  [0 2 2 0]
  [0 0 0 0]]]
membership: (table axis moved first for convenience)
[[[0 0 0 0]
  [0 0 2 0]
  [0 2 2 0]
  [0 0 0 0]]]
membership: (table axis moved first for convenience)
[[[0 0 0 0]
  [0 0 0 0]
  [0 2 0 0]
  [0 0 0 0]]]
membership: (table axis moved first for convenience)
[[[0 0 0 0]
  [0 0 0 0]
  [0 0 0 0]
  [0 0 0 0]]]

subrectangles:
[((1, 0), (3, 9), (0,)), ((3, 0), (6, 8), (0,)), ((3, 8), (6, 9), (0, 1)), ((3, 9), (9, 10), (1,)), ((6, 8), (9, 9), (1,))]
Contributions of subrectangle sums to rectangle sums:
sums[0] += np.sum(data[1:3, 0:9])
sums[0] += np.sum(data[3:6, 0:8])
shared_terms_sum = np.sum(data[3:6, 8:9])
sums[0] += shared_terms_sum
sums[1] += shared_terms_sum
sums[1] += np.sum(data[3:9, 9:10])
sums[1] += np.sum(data[6:9, 8:9])
Calculated sums over 94.73684210526316 % of terms of the original sums
sums:
[190.  51.]
sums_validation:
[190.  51.]
diff:
[0. 0.]

The code is copyright my employer HAMK Häme University of Applied Sciences (does not affect licensing).