Let's say you have an (A, B) array of zeros (a) and (M, N) array of ones (b). Split the problem by axis. There are A - M + 1 positions that the ones can occupy in the first axis, B - N + 1 in the second, and so on for any additional axes. The total number of possibilities is therefore (A - M + 1) * (B - N + 1).
The most intuitive way to visualize the different positions is to think about placing b into a with some offset. But there is another way. We can look at any given position as a view into an array of ones placed in an array of zeros padded on all sides:
m = np.pad(b, np.subtract(a.shape, b.shape))
Now you can create a view into m for each of the locations you are interested in. For example, if we use A = B = 3 and M = N = 2, the first location would be m[1:4, 1:4], the second would be m[1:4, 0:3], etc. m would look like this:
array([[0, 0, 0, 0],
[0, 1, 1, 0],
[0, 1, 1, 0],
[0, 0, 0, 0]])
If you create m as a C-contiguous array, you can use the arcane magic of np.lib.stride_tricks.as_strided to get all the possible views at once. Starting with m[A - M, B - N], you can add two extra dimension that simply step back one element, using the same strides as m in the first two dimensions. This generalizes easily to more than two dimensions:
def get_positions(a_shape, b_shape):
# Check input shapes
if len(a_shape) != len(b_shape):
raise ValueError('a and b must have the same number of dimensions')
d = np.subtract(a_shape, b_shape)
if (d < 0).any():
raise ValueError('a must be larger than b')
# Make padded buffer
m = np.pad(np.ones(b_shape, np.uint8), d)
# Find initial offset as a tuple of slices
offset = tuple(slice(off, None) for off in d)
# Find new dimensions
shape = tuple(d + 1) + a_shape
# Find new strides
strides = tuple(-s for s in m.strides) + m.strides
# Make a view
view = np.lib.stride_tricks.as_strided(m[offset], shape=shape, strides=strides)
# Return a copy with the leading dims merged
return view.reshape(-1, *a_shape)
The final reshape operation is forced to copy the array because it is otherwise a wildly non-contiguous view into the buffer of m. If you are OK with that, just return the view with twice as many dimensions.
Your original example becomes:
>>> get_positions((3, 3), (2, 2))
array([[[1, 1, 0],
[1, 1, 0],
[0, 0, 0]],
[[0, 1, 1],
[0, 1, 1],
[0, 0, 0]],
[[0, 0, 0],
[1, 1, 0],
[1, 1, 0]],
[[0, 0, 0],
[0, 1, 1],
[0, 1, 1]]], dtype=uint8)
This vectorized solution will work for an arbitrary number of dimensions. You can also do a simplified for loop in 2D that does not not make copies of the data:
a = np.zeros((A, B), dtype=np.uint8)
b = np.ones((M, N), dtype=np.uint8)
m = np.pad(b, (A - M, B - N))
for i in range(A - M, -1, -1):
for j in range(B - N, -1, -1):
print(m[i:i + A, j:j + B])
You can easily generalize to arbitrary dimensions using itertools.product:
from itertools import product
a_shape = (...)
b_shape = (...)
d = np.subtract(a_shape, b_shape)
m = np.pad(np.ones(b_shape, dtype=np.uint8), d)
for offset in product(*[range(off, -1, -1) for off in d]):
index = tuple(slice(off, off + sz) for off, sz in zip(offset, a_shape))
print(m[index])
The advantages to using a loop over the get_positions function are likely threefold:
- I expect the loop to be faster
- The loop uses less memory since it provides real views instead of copying the entire array out.
- The loop allows you to work with the full 32 dimensions that numpy supports, while the function uses up half of that number for auxiliary axes.