"""Private helper: edge-padding for 2D arrays prior to filtering."""
import numpy as np
from torx.decorators import autodoc_function
_REFLECT_MODES = frozenset({"reflect", "antisymmetric"})
[docs]
@autodoc_function
def pad_array(
array: np.ndarray,
left: int = 0,
right: int = 0,
bottom: int = 0,
top: int = 0,
mode: str = "reflect",
reflect_type: str = "odd",
**pad_kwargs,
) -> np.ndarray:
"""
Pad a 1D or 2D array beyond its limits using anti-reflection by default.
For 1D arrays, either left/right or bottom/top are required, for 2D arrays
all four of these inputs.
Anti-reflection (odd reflect) avoids discontinuities at the boundary,
which would otherwise introduce ringing artefacts in Fourier filters.
Parameters
----------
array:
1D or 2D input array.
left, right, bottom, top:
Number of points to add on each side.
mode:
Padding mode passed to numpy.pad.
reflect_type:
Reflection type passed to numpy.pad when mode='reflect'.
"""
kwargs = {"mode": mode, **pad_kwargs}
if mode in _REFLECT_MODES:
kwargs["reflect_type"] = reflect_type
if array.ndim == 1:
if (left or right) and (bottom or top):
raise ValueError(
"For 1D arrays, only left/right or bottom/top may be set, not both.")
pad_width = (left, right) if (left or right) else (bottom, top)
elif array.ndim == 2:
pad_width = ((bottom, top), (left, right))
else:
raise ValueError(\
f"pad_array only supports 1D and 2D arrays, got {array.ndim}D.")
return np.pad(array, pad_width, **kwargs)
[docs]
@autodoc_function
def pad_array_with_coords(
x_coords: np.ndarray,
y_coords: np.ndarray,
array: np.ndarray,
left: int = 0,
right: int = 0,
bottom: int = 0,
top: int = 0,
mode: str = "reflect",
reflect_type: str = "odd",
**pad_kwargs,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Pad a 2D array and its coordinate vectors beyond their limits.
Returns the padded (x_coords, y_coords, array) tuple.
"""
x_padded = pad_array(
x_coords,
left=left,
right=right,
mode=mode,
reflect_type=reflect_type,
**pad_kwargs,
)
y_padded = pad_array(
y_coords,
bottom=bottom,
top=top,
mode=mode,
reflect_type=reflect_type,
**pad_kwargs,
)
array_padded = pad_array(
array,
left=left,
right=right,
bottom=bottom,
top=top,
mode=mode,
reflect_type=reflect_type,
**pad_kwargs,
)
return x_padded, y_padded, array_padded
