With Python, how to apply vector operations to a neighborhood in an n-D image?

I have a 3D image with vector components (i.e., a mapping from R3 to R3). My goal is to replace each vector with the vector of maximum norm within its 3x3x3 neighborhood.

This task is proving to be unexpectedly challenging. I attempted to use scipy.ndimage.generic_filter, but despite its name, this filter only handles scalar inputs and outputs. I also briefly explored skimage and numpy's sliding_window_view, but neither seemed to provide a straightforward solution.

What would be the correct way to implement this?

Here's what I ended up writing. It's not very elegant and pretty slow, but should help understand what I'm trying to do.

import numpy as np
import matplotlib.pyplot as plt

def max_norm_vector(data):
    """Return the vector with the maximum norm."""
    data = data.reshape(-1, 3)
    norms = np.linalg.norm(data, axis=-1)
    idx_max = np.argmax(norms)
    return data[idx_max]

if __name__ == '__main__':
    # Load the image
    range_ = np.linspace(-5, 5, 30)
    x, y, z = np.meshgrid(range_, range_, range_, indexing='ij')
    data = 1 - (x ** 2)

    # Compute the gradient
    grad = np.gradient(data)
    grad = np.stack(grad, axis=-1)  # Stack gradients along a new last axis


    # grad = grad[:5, :5, :5, :]  # Crop the gradient for testing
    max_grad = np.zeros_like(grad)
    for i in range(1,grad.shape[0]-1):
        for j in range(1,grad.shape[1]-1):
            for k in range(2,grad.shape[2]-1):
                max_grad[i, j, k] = max_norm_vector(grad[i-1:i+2, j-1:j+2, k-1:k+2,:])

    # Visualization
    fig = plt.figure(figsize=(12, 6))

    # Plot original data
    ax1 = fig.add_subplot(121, projection='3d')
    ax1.scatter(x.ravel(), y.ravel(), z.ravel(), c=data.ravel(), cmap='viridis', alpha=0.5)
    ax1.set_title('Original Data')

    # Plot maximum gradient vectors
    ax2 = fig.add_subplot(122, projection='3d')

    # Downsample for better performance
    step = 3
    x_down = x[::step, ::step, ::step]
    y_down = y[::step, ::step, ::step]
    z_down = z[::step, ::step, ::step]
    max_grad_down = max_grad[::step, ::step, ::step]

    ax2.quiver(x_down.ravel(), y_down.ravel(), z_down.ravel(),
               max_grad_down[:, :, :, 0].ravel(), max_grad_down[:, :, :, 1].ravel(), max_grad_down[:, :, :, 2].ravel(),
               length=0.1, color='red', alpha=0.7)
    ax2.set_title('Maximum Gradient Vectors')

    plt.tight_layout()
    plt.show()

DIPlib has this function implemented: dip.SelectionFilter().

This is how you'd use it:

grad = ...  # OP's grad array
norm = dip.Norm(grad)
out = dip.SelectionFilter(grad, norm, dip.Kernel(3, "rectangular"), mode="maximum")

You can cast the dip.Image object out to a NumPy array with np.asarray(out) (no copy of the data will be made). NumPy functions will accept the dip.Image object as input, but many functions in scikit-image expect the input array to have a .shape method or similar, which will fail if you don't do the cast explicitly.

Install the package with pip install diplib.

Disclaimer: I'm an author of DIPlib, but I didn't implement this function.