Getting a subset of 2D array given indices of center point

Question

Given a 2D array and a specific element with indices (x,y), how to get a subset square 2D array (n x n) centered at this element?

I was able to implement it only if the size of subset array is completely within the bounds of the original array. I'm having problems if the specific element is near the edges or corners of the original array. For this case, the subset array must have nan values for elements outside the original array.

Answer 1

Here's how I would do this:

def fixed_size_subset(a, x, y, size):
    '''
    Gets a subset of 2D array given a x and y coordinates
    and an output size. If the slices exceed the bounds 
    of the input array, the non overlapping values
    are filled with NaNs
    ----
    a: np.array
        2D array from which to take a subset
    x, y: int. Coordinates of the center of the subset
    size: int. Size of the output array
    ----       
    Returns:
        np.array
        Subset of the input array
    '''
    o, r = np.divmod(size, 2)
    l = (x-(o+r-1)).clip(0)
    u = (y-(o+r-1)).clip(0)
    a_ = a[l: x+o+1, u:y+o+1]
    out = np.full((size, size), np.nan, dtype=a.dtype)
    out[:a_.shape[0], :a_.shape[1]] = a_
    return out

---

Samples runs:

# random 2D array
a = np.random.randint(1,5,(6,6))

array([[1, 3, 2, 2, 4, 1],
       [1, 3, 1, 3, 3, 2],
       [1, 1, 4, 4, 2, 4],
       [1, 2, 3, 4, 1, 1],
       [4, 1, 4, 2, 3, 4],
       [3, 3, 2, 3, 2, 1]])

---

fixed_size_subset(a, 3, 3, 5)

array([[3., 1., 3., 3., 2.],
       [1., 4., 4., 2., 4.],
       [2., 3., 4., 1., 1.],
       [1., 4., 2., 3., 4.],
       [3., 2., 3., 2., 1.]])

Let's try with some examples in which the sliced array is smaller than the expected output size:

fixed_size_subset(a, 4, 1, 4)

array([[ 1.,  2.,  3.,  4.],
       [ 4.,  1.,  4.,  2.],
       [ 3.,  3.,  2.,  3.],
       [nan, nan, nan, nan]])

fixed_size_subset(a, 5, 5, 3)

array([[ 3.,  4., nan],
       [ 2.,  1., nan],
       [nan, nan, nan]])

And the following would also work:

fixed_size_subset(a, -1, 0, 3)

array([[ 1.,  3., nan],
       [nan, nan, nan],
       [nan, nan, nan]])