1from __future__ import annotations
 2import numpy as np
 3
 4
 5def truncated_fourier_weights(mask: np.ndarray, bandwidth, *, clip_negative=False) -> np.ndarray:
 6    """Fourier-project a 2-D or 3-D periodic grid indicator."""
 7    mask = np.asarray(mask, dtype=float)
 8    if mask.ndim not in (2, 3):
 9        raise ValueError("mask must be 2-D or 3-D")
10    if len(bandwidth) != mask.ndim:
11        raise ValueError("bandwidth must have one cutoff per axis")
12    spectrum = np.fft.fftn(mask)
13    keep = np.ones(mask.shape, dtype=bool)
14    for axis, cutoff in enumerate(bandwidth):
15        freq = np.fft.fftfreq(mask.shape[axis]) * mask.shape[axis]
16        shape = [1] * mask.ndim
17        shape[axis] = mask.shape[axis]
18        keep &= np.abs(freq).reshape(shape) <= int(cutoff)
19    weights = np.fft.ifftn(spectrum * keep).real
20    return np.maximum(weights, 0.0) if clip_negative else weights
21
22
23def pool(features: np.ndarray, weights: np.ndarray, eps=1e-12) -> np.ndarray:
24    """Normalized weighted pooling for C,*grid or B,C,*grid."""
25    x, w = np.asarray(features), np.asarray(weights)
26    if x.shape[-w.ndim:] != w.shape or x.ndim not in (w.ndim + 1, w.ndim + 2):
27        raise ValueError("features must be C,*grid or B,C,*grid")
28    axes = tuple(range(x.ndim - w.ndim, x.ndim))
29    return np.sum(x * w, axis=axes) / (np.sum(w) + eps)