1"""resampler: fast differentiable resizing and warping of arbitrary grids.
   2
   3.. include:: ../README.md
   4"""
   5
   6from __future__ import annotations
   7
   8# Note that pydoc uses the module docstring in both __init__.py (for section headings) and
   9# __init__.pyi (for the actual content)!
  10
  11__docformat__ = 'google'
  12__version__ = '1.1.0'
  13__version_info__ = tuple(int(num) for num in __version__.split('.'))
  14
  15import abc
  16import dataclasses
  17import functools
  18import importlib
  19import itertools
  20import math
  21import os
  22import sys
  23import types
  24import typing
  25from collections.abc import Callable, Iterable, Sequence
  26from typing import Any, Generic, Literal, TypeAlias, TypeVar
  27
  28import numpy as np
  29import numpy.typing as npt
  30import scipy.interpolate
  31import scipy.linalg
  32import scipy.ndimage
  33import scipy.sparse
  34import scipy.sparse.linalg
  35
  36
  37def _noop_decorator(*args: Any, **kwargs: Any) -> Any:
  38  """Return function decorated with no-operation; invocable with or without args."""
  39  if len(args) != 1 or not callable(args[0]) or kwargs:
  40    return _noop_decorator  # Decorator is invoked with arguments; ignore them.
  41  func: Callable[..., Any] = args[0]
  42  return func
  43
  44
  45try:
  46  import numba
  47except ModuleNotFoundError:
  48  numba = sys.modules['numba'] = types.ModuleType('numba')
  49  numba.njit = _noop_decorator  # type: ignore[attr-defined]
  50_USING_NUMBA = hasattr(numba, 'jit')
  51
  52if TYPE_CHECKING:
  53  import jax.experimental.sparse
  54  import jax.numpy
  55  import torch
  56
  57  _DType: TypeAlias = np.dtype[Any]
  58  _NDArray: TypeAlias = npt.NDArray[Any]
  59  _DTypeLike: TypeAlias = npt.DTypeLike
  60  _ArrayLike: TypeAlias = npt.ArrayLike
  61  _TorchTensor: TypeAlias = torch.Tensor
  62  _JaxArray: TypeAlias = jax.numpy.ndarray
  63
  64else:
  65  # Typically, create named types for use in the `pdoc` documentation.
  66  # But here, these are superseded by the declarations in __init__.pyi!
  67  _DType: TypeAlias = Any
  68  _NDArray: TypeAlias = Any
  69  _DTypeLike: TypeAlias = Any
  70  _ArrayLike: TypeAlias = Any
  71  _TorchTensor: TypeAlias = Any
  72  _JaxArray: TypeAlias = Any
  73
  74_Array = TypeVar('_Array', _NDArray, _TorchTensor, _JaxArray)
  75_AnyArray = _NDArray | _TorchTensor | _JaxArray
  76
  77
  78def _check_eq(a: Any, b: Any, /) -> None:
  79  """If the two values or arrays are not equal, raise an exception with a useful message."""
  80  are_equal = np.all(a == b) if isinstance(a, np.ndarray) else a == b
  81  if not are_equal:
  82    raise AssertionError(f'{a!r} == {b!r}')
  83
  84
  85def _real_precision(dtype: _DTypeLike, /) -> _DType:
  86  """Return the type of the real part of a complex number."""
  87  return np.array([], dtype).real.dtype
  88
  89
  90def _complex_precision(dtype: _DTypeLike, /) -> _DType:
  91  """Return a complex type to represent a non-complex type."""
  92  return np.result_type(dtype, np.complex64)
  93
  94
  95def _get_precision(
  96    precision: _DTypeLike | None, dtypes: list[_DType], weight_dtypes: list[_DType], /
  97) -> _DType:
  98  """Return dtype based on desired precision or on data and weight types."""
  99  precision2 = np.dtype(
 100      precision if precision is not None else np.result_type(np.float32, *dtypes, *weight_dtypes)
 101  )
 102  if not np.issubdtype(precision2, np.inexact):
 103    raise ValueError(f'Precision {precision2} is not floating or complex.')
 104  check_complex = [precision2, *dtypes]
 105  is_complex = [np.issubdtype(dtype, np.complexfloating) for dtype in check_complex]
 106  if len(set(is_complex)) != 1:
 107    s_types = ','.join(str(dtype) for dtype in check_complex)
 108    raise ValueError(f'Types {s_types} must be all real or all complex.')
 109  return precision2
 110
 111
 112def _sinc(x: _ArrayLike, /) -> _NDArray:
 113  """Return the value `np.sinc(x)` but improved to:
 114  (1) ignore underflow that occurs at 0.0 for np.float32, and
 115  (2) output exact zero for integer input values.
 116
 117  >>> _sinc(np.array([-3, -2, -1, 0], np.float32))
 118  array([0., 0., 0., 1.], dtype=float32)
 119
 120  >>> _sinc(np.array([-3, -2, -1, 0]))
 121  array([0., 0., 0., 1.])
 122
 123  >>> _sinc(0)
 124  1.0
 125  """
 126  x = np.asarray(x)
 127  x_is_scalar = x.ndim == 0
 128  with np.errstate(under='ignore'):
 129    result = np.sinc(np.atleast_1d(x))
 130    result[x == np.floor(x)] = 0.0
 131    result[x == 0] = 1.0
 132    return result.item() if x_is_scalar else result
 133
 134
 135def _is_symmetric(matrix: Any, /, tol: float = 1e-10) -> bool:
 136  """Return True if the sparse matrix is symmetric."""
 137  norm = scipy.sparse.linalg.norm(matrix - matrix.transpose(), np.inf)
 138  return bool(norm <= tol)
 139
 140
 141def _cache_sampled_1d_function(
 142    xmin: float,
 143    xmax: float,
 144    *,
 145    num_samples: int = 3_600,
 146    enable: bool = True,
 147) -> Callable[[Callable[[_ArrayLike], _NDArray]], Callable[[_ArrayLike], _NDArray]]:
 148  """Function decorator to linearly interpolate cached function values."""
 149  # Speed unchanged up to num_samples=12_000, then slow decrease until 100_000.
 150
 151  def wrap_it(func: Callable[[_ArrayLike], _NDArray]) -> Callable[[_ArrayLike], _NDArray]:
 152    if not enable:
 153      return func
 154
 155    dx = (xmax - xmin) / num_samples
 156    x = np.linspace(xmin, xmax + dx, num_samples + 2, dtype=np.float32)
 157    samples_func = func(x)
 158    assert np.all(samples_func[[0, -1, -2]] == 0.0)
 159
 160    @functools.wraps(func)
 161    def interpolate_using_cached_samples(x: _ArrayLike) -> _NDArray:
 162      x = np.asarray(x)
 163      index_float = np.clip((x - xmin) / dx, 0.0, num_samples)
 164      index = index_float.astype(np.int64)
 165      frac = np.subtract(index_float, index, dtype=np.float32)
 166      return (1 - frac) * samples_func[index] + frac * samples_func[index + 1]
 167
 168    return interpolate_using_cached_samples
 169
 170  return wrap_it
 171
 172
 173class _DownsampleIn2dUsingBoxFilter:
 174  """Fast 2D box-filter downsampling using cached numba-jitted functions."""
 175
 176  def __init__(self) -> None:
 177    # Downsampling function for params (dtype, block_height, block_width, ch).
 178    self._jitted_function: dict[tuple[_DType, int, int, int], Callable[[_NDArray], _NDArray]] = {}
 179
 180  def __call__(self, array: _NDArray, shape: tuple[int, int]) -> _NDArray:
 181    assert _USING_NUMBA
 182    assert array.ndim in (2, 3), array.ndim
 183    _check_eq(len(shape), 2)
 184    dtype = array.dtype
 185    a = array[..., None] if array.ndim == 2 else array
 186    height, width, ch = a.shape
 187    new_height, new_width = shape
 188    if height % new_height != 0 or width % new_width != 0:
 189      raise ValueError(f'Shape {array.shape} not a multiple of {shape}.')
 190    block_height, block_width = height // new_height, width // new_width
 191
 192    def func(array: _NDArray) -> _NDArray:
 193      new_height = array.shape[0] // block_height
 194      new_width = array.shape[1] // block_width
 195      result = np.empty((new_height, new_width, ch), dtype)
 196      totals = np.empty(ch, dtype)
 197      factor = dtype.type(1.0 / (block_height * block_width))
 198      for y in numba.prange(new_height):  # pylint: disable=not-an-iterable
 199        for x in range(new_width):
 200          # Introducing "y2, x2 = y * block_height, x * block_width" is actually slower.
 201          if ch == 1:  # All the branches involve compile-time constants.
 202            total = dtype.type(0.0)
 203            for yy in range(block_height):
 204              for xx in range(block_width):
 205                total += array[y * block_height + yy, x * block_width + xx, 0]
 206            result[y, x, 0] = total * factor
 207          elif ch == 3:
 208            total0 = total1 = total2 = dtype.type(0.0)
 209            for yy in range(block_height):
 210              for xx in range(block_width):
 211                total0 += array[y * block_height + yy, x * block_width + xx, 0]
 212                total1 += array[y * block_height + yy, x * block_width + xx, 1]
 213                total2 += array[y * block_height + yy, x * block_width + xx, 2]
 214            result[y, x, 0] = total0 * factor
 215            result[y, x, 1] = total1 * factor
 216            result[y, x, 2] = total2 * factor
 217          elif block_height * block_width >= 9:
 218            for c in range(ch):
 219              totals[c] = 0.0
 220            for yy in range(block_height):
 221              for xx in range(block_width):
 222                for c in range(ch):
 223                  totals[c] += array[y * block_height + yy, x * block_width + xx, c]
 224            for c in range(ch):
 225              result[y, x, c] = totals[c] * factor
 226          else:
 227            for c in range(ch):
 228              total = dtype.type(0.0)
 229              for yy in range(block_height):
 230                for xx in range(block_width):
 231                  total += array[y * block_height + yy, x * block_width + xx, c]
 232              result[y, x, c] = total * factor
 233      return result
 234
 235    signature = dtype, block_height, block_width, ch
 236    jitted_function = self._jitted_function.get(signature)
 237    if not jitted_function:
 238      if 0:
 239        print(f'Creating numba jit-wrapper for {signature}.')
 240      jitted_function = numba.njit(func, parallel=True, fastmath=True, cache=True)
 241      self._jitted_function[signature] = jitted_function
 242
 243    try:
 244      result = jitted_function(a)
 245    except RuntimeError:
 246      message = (
 247          'resampler: This runtime error may be due to a corrupt resampler/__pycache__;'
 248          ' try deleting that directory.'
 249      )
 250      print(message, file=sys.stdout, flush=True)
 251      print(message, file=sys.stderr, flush=True)
 252      raise
 253
 254    return result[..., 0] if array.ndim == 2 else result
 255
 256
 257_downsample_in_2d_using_box_filter = _DownsampleIn2dUsingBoxFilter()
 258
 259
 260@numba.njit(nogil=True, fastmath=True, cache=True)  # type: ignore[untyped-decorator]
 261def _numba_serial_csr_dense_mult(
 262    indptr: _NDArray,
 263    indices: _NDArray,
 264    data: _NDArray,
 265    src: _NDArray,
 266    dst: _NDArray,
 267) -> None:
 268  """Faster version of scipy.sparse._sparsetools.csr_matvecs().
 269
 270  The single-threaded numba-jitted code is about 2x faster than the scipy C++.
 271  """
 272  assert indptr.ndim == indices.ndim == data.ndim == 1 and src.ndim == dst.ndim == 2
 273  assert len(indptr) == dst.shape[0] + 1 and src.shape[1] == dst.shape[1]
 274  acc = data[0] * src[0]  # Dummy initialization value, to infer correct shape and dtype.
 275  for i in range(dst.shape[0]):
 276    acc[:] = 0
 277    for jj in range(indptr[i], indptr[i + 1]):
 278      j = indices[jj]
 279      acc += data[jj] * src[j]
 280    dst[i] = acc
 281
 282
 283# I tried using the "minimal" "parallel=" config but this did not result in any jit speedup:
 284#  parallel=dict(comprehension=False, prange=True, numpy=True, reduction=False,
 285#                setitem=False, stencil=False, fusion=False)
 286@numba.njit(parallel=True, fastmath=True, cache=True)  # type: ignore[untyped-decorator]
 287def _numba_parallel_csr_dense_mult(
 288    indptr: _NDArray,
 289    indices: _NDArray,
 290    data: _NDArray,
 291    src: _NDArray,
 292    dst: _NDArray,
 293) -> None:
 294  """Faster version of scipy.sparse._sparsetools.csr_matvecs().
 295
 296  The single-threaded numba-jitted code is about 2x faster than the scipy C++.
 297  The introduction of parallel omp threads provides another 2-4x speedup.
 298  However, "parallel=True" leads to slow jitting (~3 s), so we cache the jitted code on disk.
 299  """
 300  assert indptr.ndim == indices.ndim == data.ndim == 1 and src.ndim == dst.ndim == 2
 301  assert len(indptr) == dst.shape[0] + 1 and src.shape[1] == dst.shape[1]
 302  acc0 = data[0] * src[0]  # Dummy initialization value, to infer correct shape and dtype.
 303  # Default is static scheduling, which is fine.
 304  for i in numba.prange(dst.shape[0]):  # pylint: disable=not-an-iterable
 305    acc = np.zeros_like(acc0)  # Numba automatically hoists the allocation outside the loop.
 306    for jj in range(indptr[i], indptr[i + 1]):
 307      j = indices[jj]
 308      acc += data[jj] * src[j]
 309    dst[i] = acc
 310
 311
 312@dataclasses.dataclass
 313class _Arraylib(abc.ABC, Generic[_Array]):
 314  """Abstract base class for abstraction of array libraries."""
 315
 316  arraylib: str
 317  """Name of array library (e.g., `'numpy'`, `'torch'`, `'jax'`)."""
 318
 319  array: _Array
 320
 321  @staticmethod
 322  @abc.abstractmethod
 323  def recognize(array: Any) -> bool:
 324    """Return True if `array` is recognized by this _Arraylib."""
 325
 326  @abc.abstractmethod
 327  def numpy(self) -> _NDArray:
 328    """Return a `numpy` version of `self.array`."""
 329
 330  @abc.abstractmethod
 331  def dtype(self) -> _DType:
 332    """Return the equivalent of `self.array.dtype` as a `numpy` `dtype`."""
 333
 334  @abc.abstractmethod
 335  def astype(self, dtype: _DTypeLike) -> _Array:
 336    """Return the equivalent of `self.array.astype(dtype, copy=False)` with `numpy` `dtype`."""
 337
 338  def reshape(self, shape: tuple[int, ...]) -> _Array:
 339    """Return the equivalent of `self.array.reshape(shape)`."""
 340    array: Any = self.array
 341    return array.reshape(shape)
 342
 343  def possibly_make_contiguous(self) -> _Array:
 344    """Return a contiguous copy of `self.array` or just `self.array` if already contiguous."""
 345    return self.array
 346
 347  @abc.abstractmethod
 348  def clip(self, low: Any, high: Any, dtype: _DTypeLike | None = None) -> _Array:
 349    """Return the equivalent of `self.array.clip(low, high, dtype=dtype)` with `numpy` `dtype`."""
 350
 351  @abc.abstractmethod
 352  def square(self) -> _Array:
 353    """Return the equivalent of `np.square(self.array)`."""
 354
 355  @abc.abstractmethod
 356  def sqrt(self) -> _Array:
 357    """Return the equivalent of `np.sqrt(self.array)`."""
 358
 359  def getitem(self, indices: Any) -> _Array:
 360    """Return the equivalent of `self.array[indices]` (a "gather" operation)."""
 361    array: Any = self.array
 362    return array[indices]
 363
 364  @abc.abstractmethod
 365  def where(self, if_true: Any, if_false: Any) -> _Array:
 366    """Return the equivalent of `np.where(self.array, if_true, if_false)`."""
 367
 368  @abc.abstractmethod
 369  def transpose(self, axes: Sequence[int]) -> _Array:
 370    """Return the equivalent of `np.transpose(self.array, axes)`."""
 371
 372  @abc.abstractmethod
 373  def best_dims_order_for_resize(self, dst_shape: tuple[int, ...]) -> list[int]:
 374    """Return the best order in which to process dims for resizing `self.array` to `dst_shape`."""
 375
 376  @abc.abstractmethod
 377  def premult_with_sparse(self, sparse: Any, num_threads: int | Literal['auto']) -> _Array:
 378    """Return the multiplication of the `sparse` matrix and `self.array`."""
 379
 380  @staticmethod
 381  @abc.abstractmethod
 382  def concatenate(arrays: Sequence[_Array], axis: int) -> _Array:
 383    """Return the equivalent of `np.concatenate(arrays, axis)`."""
 384
 385  @staticmethod
 386  @abc.abstractmethod
 387  def einsum(subscripts: str, *operands: _Array) -> _Array:
 388    """Return the equivalent of `np.einsum(subscripts, *operands, optimize=True)`."""
 389
 390  @staticmethod
 391  @abc.abstractmethod
 392  def make_sparse_matrix(
 393      data: _NDArray, row_ind: _NDArray, col_ind: _NDArray, shape: tuple[int, int]
 394  ) -> Any:
 395    """Return the equivalent of `scipy.sparse.csr_matrix(data, (row_ind, col_ind), shape=shape)`.
 396    However, the indices must be ordered and unique."""
 397
 398
 399class _NumpyArraylib(_Arraylib[_NDArray]):
 400  """Numpy implementation of the array abstraction."""
 401
 402  # pylint: disable=missing-function-docstring
 403
 404  def __init__(self, array: _NDArray) -> None:
 405    super().__init__(arraylib='numpy', array=np.asarray(array))
 406
 407  @staticmethod
 408  def recognize(array: Any) -> bool:
 409    return isinstance(array, (np.ndarray, np.number))
 410
 411  def numpy(self) -> _NDArray:
 412    return self.array
 413
 414  def dtype(self) -> _DType:
 415    dtype: _DType = self.array.dtype
 416    return dtype
 417
 418  def astype(self, dtype: _DTypeLike) -> _NDArray:
 419    return self.array.astype(dtype, copy=False)
 420
 421  def clip(self, low: Any, high: Any, dtype: _DTypeLike | None = None) -> _NDArray:
 422    return self.array.clip(low, high, dtype=dtype)
 423
 424  def square(self) -> _NDArray:
 425    return np.square(self.array)
 426
 427  def sqrt(self) -> _NDArray:
 428    return np.sqrt(self.array)
 429
 430  def where(self, if_true: Any, if_false: Any) -> _NDArray:
 431    condition = self.array
 432    return np.where(condition, if_true, if_false)
 433
 434  def transpose(self, axes: Sequence[int]) -> _NDArray:
 435    return np.transpose(self.array, tuple(axes))
 436
 437  def best_dims_order_for_resize(self, dst_shape: tuple[int, ...]) -> list[int]:
 438    # Our heuristics: (1) a dimension with small scaling (especially minification) gets priority,
 439    # and (2) timings show preference to resizing dimensions with larger strides first.
 440    # The optimal ordering might be related to the logic in np.einsum_path().  (Unfortunately,
 441    # np.einsum() does not support the sparse multiplications that we require here.)
 442    src_shape: tuple[int, ...] = self.array.shape[: len(dst_shape)]
 443    strides: Sequence[int] = self.array.strides
 444    largest_stride_dim = max(range(len(src_shape)), key=lambda dim: strides[dim])
 445
 446    def priority(dim: int) -> float:
 447      scaling = dst_shape[dim] / src_shape[dim]
 448      return scaling * ((0.49 if scaling < 1.0 else 0.65) if dim == largest_stride_dim else 1.0)
 449
 450    return sorted(range(len(src_shape)), key=priority)
 451
 452  def premult_with_sparse(
 453      self, sparse: scipy.sparse.csr_matrix, num_threads: int | Literal['auto']
 454  ) -> _NDArray:
 455    assert self.array.ndim == sparse.ndim == 2 and sparse.shape[1] == self.array.shape[0]
 456    # Empirically faster than with default numba.config.NUMBA_NUM_THREADS (e.g., 24).
 457    if _USING_NUMBA:
 458      num_threads2 = min(6, os.cpu_count() or 1) if num_threads == 'auto' else num_threads
 459      src = np.ascontiguousarray(self.array)  # Like .ravel() in _mul_multivector().
 460      dtype = np.result_type(sparse.dtype, src.dtype)
 461      dst = np.empty((sparse.shape[0], src.shape[1]), dtype)
 462      num_scalar_multiplies = len(sparse.data) * src.shape[1]
 463      is_small_size = num_scalar_multiplies < 200_000
 464      if is_small_size or num_threads2 == 1:
 465        _numba_serial_csr_dense_mult(sparse.indptr, sparse.indices, sparse.data, src, dst)
 466      else:
 467        numba.set_num_threads(num_threads2)
 468        _numba_parallel_csr_dense_mult(sparse.indptr, sparse.indices, sparse.data, src, dst)
 469      return dst
 470
 471    # Note that sicpy.sparse does not use multithreading.  The "@" operation
 472    # calls _spbase.__matmul__() -> _spbase._mul_dispatch() -> _cs_matrix._mul_multivector() ->
 473    # scipy.sparse._sparsetools.csr_matvecs() in
 474    # https://github.com/scipy/scipy/blob/main/scipy/sparse/sparsetools/csr.h
 475    # which iteratively calls the (in theory, LEVEL 1 BLAS) function axpy() in
 476    # https://github.com/scipy/scipy/blob/main/scipy/sparse/sparsetools/dense.h
 477    return sparse @ self.array
 478
 479  @staticmethod
 480  def concatenate(arrays: Sequence[_NDArray], axis: int) -> _NDArray:
 481    return np.concatenate(arrays, axis)
 482
 483  @staticmethod
 484  def einsum(subscripts: str, *operands: _NDArray) -> _NDArray:
 485    return np.einsum(subscripts, *operands, optimize=True)
 486
 487  @staticmethod
 488  def make_sparse_matrix(
 489      data: _NDArray, row_ind: _NDArray, col_ind: _NDArray, shape: tuple[int, int]
 490  ) -> scipy.sparse.csr_matrix:
 491    return scipy.sparse.csr_matrix((data, (row_ind, col_ind)), shape=shape)
 492
 493
 494class _TorchArraylib(_Arraylib[_TorchTensor]):
 495  """Torch implementation of the array abstraction."""
 496
 497  # pylint: disable=missing-function-docstring
 498
 499  def __init__(self, array: _NDArray) -> None:
 500    import torch
 501
 502    self.torch = torch
 503    super().__init__(arraylib='torch', array=self.torch.as_tensor(array))
 504
 505  @staticmethod
 506  def recognize(array: Any) -> bool:
 507    return type(array).__module__ == 'torch'
 508
 509  def numpy(self) -> _NDArray:
 510    return self.array.numpy()
 511
 512  def dtype(self) -> _DType:
 513    numpy_type = {
 514        self.torch.float32: np.float32,
 515        self.torch.float64: np.float64,
 516        self.torch.complex64: np.complex64,
 517        self.torch.complex128: np.complex128,
 518        self.torch.uint8: np.uint8,  # No uint16, uint32, uint64.
 519        self.torch.int16: np.int16,
 520        self.torch.int32: np.int32,
 521        self.torch.int64: np.int64,
 522    }[self.array.dtype]
 523    return np.dtype(numpy_type)
 524
 525  def astype(self, dtype: _DTypeLike) -> _TorchTensor:
 526    torch_types: dict[Any, Any] = {
 527        np.float32: self.torch.float32,
 528        np.float64: self.torch.float64,
 529        np.complex64: self.torch.complex64,
 530        np.complex128: self.torch.complex128,
 531        np.uint8: self.torch.uint8,  # No uint16, uint32, uint64.
 532        np.int16: self.torch.int16,
 533        np.int32: self.torch.int32,
 534        np.int64: self.torch.int64,
 535    }
 536    return self.array.type(torch_types[np.dtype(dtype).type])
 537
 538  def possibly_make_contiguous(self) -> _TorchTensor:
 539    return self.array.contiguous()
 540
 541  def clip(self, low: Any, high: Any, dtype: _DTypeLike | None = None) -> _TorchTensor:
 542    array = self.array
 543    array = _arr_astype(array, dtype) if dtype is not None else array
 544    return array.clip(low, high)
 545
 546  def square(self) -> _TorchTensor:
 547    return self.array.square()
 548
 549  def sqrt(self) -> _TorchTensor:
 550    return self.array.sqrt()
 551
 552  def getitem(self, indices: Any) -> _TorchTensor:
 553    if not isinstance(indices, tuple):
 554      indices = indices.type(self.torch.int64)
 555    return self.array[indices]  # pylint: disable=unsubscriptable-object
 556
 557  def where(self, if_true: Any, if_false: Any) -> _TorchTensor:
 558    condition = self.array
 559    return if_true.where(condition, if_false)
 560
 561  def transpose(self, axes: Sequence[int]) -> _TorchTensor:
 562    return self.torch.permute(self.array, tuple(axes))
 563
 564  def best_dims_order_for_resize(self, dst_shape: tuple[int, ...]) -> list[int]:
 565    # Similar to `_NumpyArraylib`.  We access `array.stride()` instead of `array.strides`.
 566    src_shape: tuple[int, ...] = self.array.shape[: len(dst_shape)]
 567    strides: Sequence[int] = self.array.stride()
 568    largest_stride_dim = max(range(len(src_shape)), key=lambda dim: strides[dim])
 569
 570    def priority(dim: int) -> float:
 571      scaling = dst_shape[dim] / src_shape[dim]
 572      return scaling * ((0.49 if scaling < 1.0 else 0.65) if dim == largest_stride_dim else 1.0)
 573
 574    return sorted(range(len(src_shape)), key=priority)
 575
 576  def premult_with_sparse(self, sparse: Any, num_threads: int | Literal['auto']) -> _TorchTensor:
 577    del num_threads
 578    if np.issubdtype(_arr_dtype(self.array), np.complexfloating):
 579      sparse = _arr_astype(sparse, _arr_dtype(self.array))
 580    return sparse @ self.array  # Calls torch.sparse.mm().
 581
 582  @staticmethod
 583  def concatenate(arrays: Sequence[_TorchTensor], axis: int) -> _TorchTensor:
 584    import torch
 585
 586    return torch.cat(tuple(arrays), axis)
 587
 588  @staticmethod
 589  def einsum(subscripts: str, *operands: _TorchTensor) -> _TorchTensor:
 590    import torch
 591
 592    operands = tuple(torch.as_tensor(operand) for operand in operands)
 593    if any(np.issubdtype(_arr_dtype(array), np.complexfloating) for array in operands):
 594      operands = tuple(
 595          _arr_astype(array, _complex_precision(_arr_dtype(array))) for array in operands
 596      )
 597    return torch.einsum(subscripts, *operands)
 598
 599  @staticmethod
 600  def make_sparse_matrix(
 601      data: _NDArray, row_ind: _NDArray, col_ind: _NDArray, shape: tuple[int, int]
 602  ) -> _TorchTensor:
 603    import torch
 604
 605    indices = np.vstack((row_ind, col_ind))
 606    with torch.sparse.check_sparse_tensor_invariants(enable=False):
 607      return torch.sparse_coo_tensor(torch.as_tensor(indices), torch.as_tensor(data), shape)
 608    # .coalesce() is unnecessary because indices/data are already merged.
 609
 610
 611class _JaxArraylib(_Arraylib[_JaxArray]):
 612  """Jax implementation of the array abstraction."""
 613
 614  def __init__(self, array: _NDArray) -> None:
 615    import jax.numpy
 616
 617    self.jnp = jax.numpy
 618    super().__init__(arraylib='jax', array=self.jnp.asarray(array))
 619
 620  @staticmethod
 621  def recognize(array: Any) -> bool:
 622    # e.g., jaxlib.xla_extension.DeviceArray, jax.interpreters.ad.JVPTracer
 623    return type(array).__module__.startswith(('jaxlib.', 'jax.'))
 624
 625  def numpy(self) -> _NDArray:
 626    # 2023-01-09: jax 0.3.17 "DeviceArray.to_py() has been deprecated. Use np.asarray(x) instead."
 627    # Whereas array.to_py() and np.asarray(array) may return a non-writable np.ndarray,
 628    # np.array(array) always returns a writable array but the copy may be more costly.
 629    # return self.array.to_py()
 630    return np.asarray(self.array)
 631
 632  def dtype(self) -> _DType:
 633    return np.dtype(self.array.dtype)
 634
 635  def astype(self, dtype: _DTypeLike) -> _JaxArray:
 636    return self.array.astype(np.dtype(dtype))  # (copy=False is unavailable)
 637
 638  def possibly_make_contiguous(self) -> _JaxArray:
 639    return self.array.copy()
 640
 641  def clip(self, low: Any, high: Any, dtype: _DTypeLike | None = None) -> _JaxArray:
 642    array = self.array
 643    if dtype is not None:
 644      array = array.astype(np.dtype(dtype))  # (copy=False is unavailable)
 645    return self.jnp.clip(array, low, high)
 646
 647  def square(self) -> _JaxArray:
 648    return self.jnp.square(self.array)
 649
 650  def sqrt(self) -> _JaxArray:
 651    return self.jnp.sqrt(self.array)
 652
 653  def where(self, if_true: Any, if_false: Any) -> _JaxArray:
 654    condition = self.array
 655    return cast(Any, self.jnp.where(condition, if_true, if_false))
 656
 657  def transpose(self, axes: Sequence[int]) -> _JaxArray:
 658    return self.jnp.transpose(self.array, tuple(axes))
 659
 660  def best_dims_order_for_resize(self, dst_shape: tuple[int, ...]) -> list[int]:
 661    # Jax/XLA does not have strides.  Arrays are contiguous, almost always in C order; see
 662    # https://github.com/google/jax/discussions/7544#discussioncomment-1197038.
 663    # We use a heuristic similar to `_TensorflowArraylib`.
 664    src_shape: tuple[int, ...] = self.array.shape[: len(dst_shape)]
 665    dims = list(range(len(src_shape)))
 666    if len(dims) > 1 and dst_shape[0] / src_shape[0] > 1.0:
 667      dims[:2] = [1, 0]
 668    return dims
 669
 670  def premult_with_sparse(
 671      self, sparse: jax.experimental.sparse.BCOO, num_threads: int | Literal['auto']
 672  ) -> _JaxArray:
 673    del num_threads
 674    return sparse @ self.array  # Calls jax.bcoo_multiply_dense().
 675
 676  @staticmethod
 677  def concatenate(arrays: Sequence[_JaxArray], axis: int) -> _JaxArray:
 678    import jax.numpy as jnp
 679
 680    return jnp.concatenate(arrays, axis)
 681
 682  @staticmethod
 683  def einsum(subscripts: str, *operands: _JaxArray) -> _JaxArray:
 684    import jax.numpy as jnp
 685
 686    return jnp.einsum(subscripts, *operands, optimize='greedy')
 687
 688  @staticmethod
 689  def make_sparse_matrix(
 690      data: _NDArray, row_ind: _NDArray, col_ind: _NDArray, shape: tuple[int, int]
 691  ) -> Any:
 692    # https://jax.readthedocs.io/en/latest/jax.experimental.sparse.html
 693    import jax.experimental.sparse
 694    import jax.numpy as jnp
 695
 696    indices = jnp.asarray(np.vstack((row_ind, col_ind)).T)
 697    return jax.experimental.sparse.BCOO(
 698        (jnp.asarray(data), indices), shape=shape, indices_sorted=True, unique_indices=True
 699    )
 700
 701
 702_CANDIDATE_ARRAYLIBS = {
 703    'numpy': _NumpyArraylib,
 704    'torch': _TorchArraylib,
 705    'jax': _JaxArraylib,
 706}
 707
 708
 709def _is_available(arraylib: str) -> bool:
 710  """Return whether the array library (e.g. 'torch') is available as an installed package."""
 711  # Faster than trying to import it.
 712  return importlib.util.find_spec(arraylib) is not None  # type: ignore[attr-defined]
 713
 714
 715_DICT_ARRAYLIBS: dict[str, Any] = {
 716    arraylib: cls for arraylib, cls in _CANDIDATE_ARRAYLIBS.items() if _is_available(arraylib)
 717}
 718
 719ARRAYLIBS = list(_DICT_ARRAYLIBS)
 720"""Array libraries supported automatically in the resize and resampling operations.
 721
 722- The library is selected automatically based on the type of the `array` function parameter.
 723
 724- The class `_Arraylib` provides library-specific implementations of needed basic functions.
 725
 726- The `_arr_*()` functions dispatch the `_Arraylib` methods based on the array type.
 727"""
 728
 729
 730def _as_arr(array: _AnyArray, /) -> _Arraylib[Any]:
 731  """Return `array` wrapped as an `_Arraylib` for dispatch of functions."""
 732  if isinstance(array, (tuple, list)):
 733    raise ValueError(f'{array} not recognized.  Perhaps convert it using np.asarray().')
 734  for cls in _DICT_ARRAYLIBS.values():
 735    if cls.recognize(array):
 736      return cls(array)
 737  raise ValueError(f'{array} {type(array)} {type(array).__module__} unrecognized by {ARRAYLIBS}.')
 738
 739
 740def _arr_arraylib(array: _AnyArray, /) -> str:
 741  """Return the name of the `Arraylib` representing `array`."""
 742  return _as_arr(array).arraylib
 743
 744
 745def _arr_numpy(array: _AnyArray, /) -> _NDArray:
 746  """Return a `numpy` version of `array`."""
 747  return _as_arr(array).numpy()
 748
 749
 750def _arr_dtype(array: _AnyArray, /) -> _DType:
 751  """Return the equivalent of `array.dtype` as a `numpy` `dtype`."""
 752  return _as_arr(array).dtype()
 753
 754
 755def _arr_shape(array: _AnyArray, /) -> tuple[int, ...]:
 756  """Return `array.shape` as a `tuple` of `int`."""
 757  array2: Any = array
 758  return tuple(array2.shape)
 759
 760
 761def _arr_astype(array: _Array, dtype: _DTypeLike, /) -> _Array:
 762  """Return the equivalent of `array.astype(dtype)` with `numpy` `dtype`."""
 763  return _as_arr(array).astype(dtype)
 764
 765
 766def _arr_reshape(array: _Array, shape: tuple[int, ...], /) -> _Array:
 767  """Return the equivalent of `array.reshape(shape)."""
 768  return _as_arr(array).reshape(shape)
 769
 770
 771def _arr_possibly_make_contiguous(array: _Array, /) -> _Array:
 772  """Return a contiguous copy of `array` or just `array` if already contiguous."""
 773  return _as_arr(array).possibly_make_contiguous()
 774
 775
 776def _arr_clip(array: _Array, low: Any, high: Any, /, dtype: _DTypeLike | None = None) -> _Array:
 777  """Return the equivalent of `array.clip(low, high, dtype)` with `numpy` `dtype`."""
 778  return _as_arr(array).clip(low, high, dtype)
 779
 780
 781def _arr_square(array: _Array, /) -> _Array:
 782  """Return the equivalent of `np.square(array)`."""
 783  return _as_arr(array).square()
 784
 785
 786def _arr_sqrt(array: _Array, /) -> _Array:
 787  """Return the equivalent of `np.sqrt(array)`."""
 788  return _as_arr(array).sqrt()
 789
 790
 791def _arr_getitem(array: _Array, indices: Any, /) -> _Array:
 792  """Return the equivalent of `array[indices]`."""
 793  return _as_arr(array).getitem(indices)
 794
 795
 796def _arr_where(condition: _Array, if_true: Any, if_false: Any, /) -> _Array:
 797  """Return the equivalent of `np.where(condition, if_true, if_false)`."""
 798  return _as_arr(condition).where(if_true, if_false)
 799
 800
 801def _arr_transpose(array: _Array, axes: Sequence[int], /) -> _Array:
 802  """Return the equivalent of `np.transpose(array, axes)`."""
 803  return _as_arr(array).transpose(axes)
 804
 805
 806def _arr_best_dims_order_for_resize(array: _AnyArray, dst_shape: tuple[int, ...], /) -> list[int]:
 807  """Return the best order in which to process dims for resizing `array` to `dst_shape`."""
 808  return _as_arr(array).best_dims_order_for_resize(dst_shape)
 809
 810
 811def _arr_matmul_sparse_dense(
 812    sparse: Any, dense: _Array, /, *, num_threads: int | Literal['auto'] = 'auto'
 813) -> _Array:
 814  """Return the multiplication of the `sparse` and `dense` matrices."""
 815  assert num_threads == 'auto' or num_threads >= 1
 816  return _as_arr(dense).premult_with_sparse(sparse, num_threads)
 817
 818
 819def _arr_concatenate(arrays: Sequence[_Array], axis: int, /) -> _Array:
 820  """Return the equivalent of `np.concatenate(arrays, axis)`."""
 821  arraylib = _arr_arraylib(arrays[0])
 822  return _DICT_ARRAYLIBS[arraylib].concatenate(arrays, axis)
 823
 824
 825def _arr_einsum(subscripts: str, /, *operands: _Array) -> _Array:
 826  """Return the equivalent of `np.einsum(subscripts, *operands, optimize=True)`."""
 827  arraylib = _arr_arraylib(operands[0])
 828  return _DICT_ARRAYLIBS[arraylib].einsum(subscripts, *operands)
 829
 830
 831def _arr_swapaxes(array: _Array, axis1: int, axis2: int, /) -> _Array:
 832  """Return the equivalent of `np.swapaxes(array, axis1, axis2)`."""
 833  ndim = len(_arr_shape(array))
 834  assert 0 <= axis1 < ndim and 0 <= axis2 < ndim, (axis1, axis2, ndim)
 835  axes = list(range(ndim))
 836  axes[axis1] = axis2
 837  axes[axis2] = axis1
 838  return _arr_transpose(array, axes)
 839
 840
 841def _arr_moveaxis(array: _Array, source: int, destination: int, /) -> _Array:
 842  """Return the equivalent of `np.moveaxis(array, source, destination)`."""
 843  ndim = len(_arr_shape(array))
 844  assert 0 <= source < ndim and 0 <= destination < ndim, (source, destination, ndim)
 845  axes = [n for n in range(ndim) if n != source]
 846  axes.insert(destination, source)
 847  return _arr_transpose(array, axes)
 848
 849
 850def _make_sparse_matrix(
 851    data: _NDArray, row_ind: _NDArray, col_ind: _NDArray, shape: tuple[int, int], arraylib: str, /
 852) -> Any:
 853  """Return the equivalent of `scipy.sparse.csr_matrix(data, (row_ind, col_ind), shape=shape)`.
 854  However, indices must be ordered and unique."""
 855  return _DICT_ARRAYLIBS[arraylib].make_sparse_matrix(data, row_ind, col_ind, shape)
 856
 857
 858def _make_array(array: _ArrayLike, arraylib: str, /) -> Any:
 859  """Return an array from the library `arraylib` initialized with the `numpy` `array`."""
 860  return _DICT_ARRAYLIBS[arraylib](np.asarray(array)).array
 861
 862
 863# Because np.ndarray supports strides, np.moveaxis() and np.transpose() are constant-time.
 864# However, ndarray.reshape() creates a copy whenever the new shape cannot be expressed using
 865# strides, e.g. dim=1 in an RGB image.
 866#
 867# torch.Tensor also supports strides, so torch.movedim() is constant-time and
 868# Tensor.reshape() has the same copy-when-necessary behavior as numpy.
 869#
 870# In contrast, tf.Tensor does not support strides, so tf.transpose() returns a new permuted
 871# tensor.  However, tf.reshape() is always efficient.
 872#
 873# For jax.Array, both operations lower to XLA ops; under jit the compiler often fuses away
 874# the transpose and turns the reshape into a bitcast, so neither has a fixed cost.
 875
 876
 877def _block_shape_with_min_size(
 878    shape: tuple[int, ...], min_size: int, compact: bool = True
 879) -> tuple[int, ...]:
 880  """Return shape of block (of size at least `min_size`) to subdivide shape."""
 881  if math.prod(shape) < min_size:
 882    raise ValueError(f'Shape {shape} smaller than min_size {min_size}.')
 883  if compact:
 884    root = math.ceil(min_size ** (1 / len(shape)))
 885    block_shape = np.minimum(shape, root)
 886    for dim in range(len(shape)):
 887      if block_shape[dim] == 2 and block_shape.prod() >= min_size * 2:
 888        block_shape[dim] = 1
 889    for dim in range(len(shape) - 1, -1, -1):
 890      if block_shape.prod() < min_size:
 891        block_shape[dim] = shape[dim]
 892  else:
 893    block_shape = np.ones_like(shape)
 894    for dim in range(len(shape) - 1, -1, -1):
 895      if block_shape.prod() < min_size:
 896        block_shape[dim] = min(shape[dim], math.ceil(min_size / block_shape.prod()))
 897  return tuple(block_shape)
 898
 899
 900def _array_split(array: _Array, axis: int, num_sections: int) -> list[_Array]:
 901  """Split `array` into `num_sections` along `axis`."""
 902  assert 0 <= axis < len(_arr_shape(array))
 903  assert 1 <= num_sections <= _arr_shape(array)[axis]
 904
 905  if 0:
 906    split = np.array_split(array, num_sections, axis=axis)  # Numpy-specific.
 907
 908  else:
 909    # Adapted from https://github.com/numpy/numpy/blob/main/numpy/lib/shape_base.py#L739-L792.
 910    num_total = _arr_shape(array)[axis]
 911    num_each, num_extra = divmod(num_total, num_sections)
 912    section_sizes = [0] + num_extra * [num_each + 1] + (num_sections - num_extra) * [num_each]
 913    div_points = np.array(section_sizes).cumsum()
 914    split = []
 915    tmp: Any = _arr_swapaxes(array, axis, 0)
 916    for i in range(num_sections):
 917      split.append(_arr_swapaxes(tmp[div_points[i] : div_points[i + 1]], axis, 0))
 918
 919  return split
 920
 921
 922def _split_array_into_blocks(array: Any, block_shape: Sequence[int], start_axis: int = 0) -> Any:
 923  """Split `array` into nested lists of blocks of size at most `block_shape`."""
 924  # See https://stackoverflow.com/a/50305924.  (If the block_shape is known to
 925  # exactly partition the array, see https://stackoverflow.com/a/16858283.)
 926  if len(block_shape) > len(_arr_shape(array)):
 927    raise ValueError(f'Block ndim {len(block_shape)} > array ndim {len(_arr_shape(array))}.')
 928  if start_axis == len(block_shape):
 929    return array
 930
 931  num_sections = math.ceil(_arr_shape(array)[start_axis] / block_shape[start_axis])
 932  split = _array_split(array, start_axis, num_sections)
 933  return [_split_array_into_blocks(split_a, block_shape, start_axis + 1) for split_a in split]
 934
 935
 936def _map_function_over_blocks(blocks: Any, func: Callable[[Any], Any]) -> Any:
 937  """Apply `func` to each block in the nested lists of `blocks`."""
 938  if isinstance(blocks, list):
 939    return [_map_function_over_blocks(block, func) for block in blocks]
 940  return func(blocks)
 941
 942
 943def _merge_array_from_blocks(blocks: Any, axis: int = 0) -> Any:
 944  """Merge an array from the nested lists of array blocks in `blocks`."""
 945  # More general than np.block() because the blocks can have additional dims.
 946  if isinstance(blocks, list):
 947    new_blocks = [_merge_array_from_blocks(block, axis + 1) for block in blocks]
 948    return _arr_concatenate(new_blocks, axis)
 949  return blocks
 950
 951
 952@dataclasses.dataclass(frozen=True)
 953class Gridtype(abc.ABC):
 954  """Abstract base class for grid-types such as `'dual'` and `'primal'`.
 955
 956  In resampling operations, the grid-type may be specified separately as `src_gridtype` for the
 957  source domain and `dst_gridtype` for the destination domain.  Moreover, the grid-type may be
 958  specified per domain dimension.
 959
 960  Examples:
 961    `resize(source, shape, gridtype='primal')`  # Sets both src and dst to be `'primal'` grids.
 962
 963    `resize(source, shape, src_gridtype=['dual', 'primal'],
 964            dst_gridtype='dual')`  # Source is `'dual'` in dim0 and `'primal'` in dim1.
 965  """
 966
 967  name: str
 968  """Gridtype name."""
 969
 970  @abc.abstractmethod
 971  def min_size(self) -> int:
 972    """Return the necessary minimum number of grid samples."""
 973
 974  @abc.abstractmethod
 975  def size_in_samples(self, size: int, /) -> int:
 976    """Return the domain size in units of inter-sample spacing."""
 977
 978  @abc.abstractmethod
 979  def point_from_index(self, index: _NDArray, size: int, /) -> _NDArray:
 980    """Return [0.0, 1.0] coordinates given [0, size - 1] indices."""
 981
 982  @abc.abstractmethod
 983  def index_from_point(self, point: _NDArray, size: int, /) -> _NDArray:
 984    """Return location x given coordinates [0.0, 1.0], where x == 0.0 is the first grid sample
 985    and x == size - 1.0 is the last grid sample."""
 986
 987  @abc.abstractmethod
 988  def reflect(self, index: _NDArray, size: int, /) -> _NDArray:
 989    """Map integer sample indices to interior ones using boundary reflection."""
 990
 991  @abc.abstractmethod
 992  def wrap(self, index: _NDArray, size: int, /) -> _NDArray:
 993    """Map integer sample indices to interior ones using wrapping."""
 994
 995  @abc.abstractmethod
 996  def reflect_clamp(self, index: _NDArray, size: int, /) -> _NDArray:
 997    """Map integer sample indices to interior ones using reflect-clamp."""
 998
 999
1000class DualGridtype(Gridtype):
1001  """Samples are at the center of cells in a uniform partition of the domain.
1002
1003  For a unit-domain dimension with N samples, each sample 0 <= i < N has position (i + 0.5) / N,
1004  e.g., [0.125, 0.375, 0.625, 0.875] for N = 4.
1005  """
1006
1007  def __init__(self) -> None:
1008    super().__init__(name='dual')
1009
1010  def min_size(self) -> int:
1011    return 1
1012
1013  def size_in_samples(self, size: int, /) -> int:
1014    return size
1015
1016  def point_from_index(self, index: _NDArray, size: int, /) -> _NDArray:
1017    return (index + 0.5) / size
1018
1019  def index_from_point(self, point: _NDArray, size: int, /) -> _NDArray:
1020    return point * size - 0.5
1021
1022  def reflect(self, index: _NDArray, size: int, /) -> _NDArray:
1023    index = np.mod(index, size * 2)
1024    return np.where(index < size, index, 2 * size - 1 - index)
1025
1026  def wrap(self, index: _NDArray, size: int, /) -> _NDArray:
1027    return np.mod(index, size)
1028
1029  def reflect_clamp(self, index: _NDArray, size: int, /) -> _NDArray:
1030    return np.minimum(np.where(index < 0, -1 - index, index), size - 1)
1031
1032
1033class PrimalGridtype(Gridtype):
1034  """Samples are at the vertices of cells in a uniform partition of the domain.
1035
1036  For a unit-domain dimension with N samples, each sample 0 <= i < N has position i / (N - 1),
1037  e.g., [0, 1/3, 2/3, 1] for N = 4.
1038  """
1039
1040  def __init__(self) -> None:
1041    super().__init__(name='primal')
1042
1043  def min_size(self) -> int:
1044    return 2
1045
1046  def size_in_samples(self, size: int, /) -> int:
1047    return size - 1
1048
1049  def point_from_index(self, index: _NDArray, size: int, /) -> _NDArray:
1050    return index / (size - 1)
1051
1052  def index_from_point(self, point: _NDArray, size: int, /) -> _NDArray:
1053    return point * (size - 1)
1054
1055  def reflect(self, index: _NDArray, size: int, /) -> _NDArray:
1056    index = np.mod(index, size * 2 - 2)
1057    return np.where(index < size, index, 2 * size - 2 - index)
1058
1059  def wrap(self, index: _NDArray, size: int, /) -> _NDArray:
1060    return np.mod(index, size - 1)
1061
1062  def reflect_clamp(self, index: _NDArray, size: int, /) -> _NDArray:
1063    return np.minimum(np.abs(index), size - 1)
1064
1065
1066_DICT_GRIDTYPES = {
1067    'dual': DualGridtype(),
1068    'primal': PrimalGridtype(),
1069}
1070
1071GRIDTYPES = list(_DICT_GRIDTYPES)
1072r"""Shortcut names for the two predefined grid types (specified per dimension):
1073
1074| `gridtype` | `'dual'`<br/>`DualGridtype()`<br/>(default) | `'primal'`<br/>`PrimalGridtype()`<br/>&nbsp; |
1075| --- |:---:|:---:|
1076| Sample positions in 2D<br/>and in 1D at different resolutions | ![Dual](https://github.com/hhoppe/resampler/raw/main/media/dual_grid_small.png) | ![Primal](https://github.com/hhoppe/resampler/raw/main/media/primal_grid_small.png) |
1077| Nesting of samples across resolutions | The samples positions do *not* nest. | The *even* samples remain at coarser scale. |
1078| Number $N_\ell$ of samples (per-dimension) at resolution level $\ell$ | $N_\ell=2^\ell$ | $N_\ell=2^\ell+1$ |
1079| Position of sample index $i$ within domain $[0, 1]$ | $\frac{i + 0.5}{N}$ ("half-integer" coordinates) | $\frac{i}{N-1}$ |
1080| Image resolutions ($N_\ell\times N_\ell$) for dyadic scales | $1\times1, ~~2\times2, ~~4\times4, ~~8\times8, ~\ldots$ | $2\times2, ~~3\times3, ~~5\times5, ~~9\times9, ~\ldots$ |
1081
1082See the source code for extensibility.
1083"""
1084
1085
1086def _get_gridtype(gridtype: str | Gridtype) -> Gridtype:
1087  """Return a `Gridtype`, which can be specified as a name in `GRIDTYPES`."""
1088  return gridtype if isinstance(gridtype, Gridtype) else _DICT_GRIDTYPES[gridtype]
1089
1090
1091def _get_gridtypes(
1092    gridtype: str | Gridtype | None,
1093    src_gridtype: str | Gridtype | Iterable[str | Gridtype] | None,
1094    dst_gridtype: str | Gridtype | Iterable[str | Gridtype] | None,
1095    src_ndim: int,
1096    dst_ndim: int,
1097) -> tuple[list[Gridtype], list[Gridtype]]:
1098  """Return per-dim source and destination grid types given all parameters."""
1099  if gridtype is None and src_gridtype is None and dst_gridtype is None:
1100    gridtype = 'dual'
1101  if gridtype is not None:
1102    if src_gridtype is not None:
1103      raise ValueError('Cannot have both gridtype and src_gridtype.')
1104    if dst_gridtype is not None:
1105      raise ValueError('Cannot have both gridtype and dst_gridtype.')
1106    src_gridtype = dst_gridtype = gridtype
1107  src_gridtype2 = [_get_gridtype(g) for g in np.broadcast_to(np.array(src_gridtype), src_ndim)]
1108  dst_gridtype2 = [_get_gridtype(g) for g in np.broadcast_to(np.array(dst_gridtype), dst_ndim)]
1109  return src_gridtype2, dst_gridtype2
1110
1111
1112@dataclasses.dataclass(frozen=True)
1113class RemapCoordinates(abc.ABC):
1114  """Abstract base class for modifying the specified coordinates prior to evaluating the
1115  reconstruction kernels."""
1116
1117  @abc.abstractmethod
1118  def __call__(self, point: _NDArray, /) -> _NDArray:
1119    ...
1120
1121
1122@dataclasses.dataclass(frozen=True)
1123class NoRemapCoordinates(RemapCoordinates):
1124  """The coordinates are not remapped."""
1125
1126  def __call__(self, point: _NDArray, /) -> _NDArray:
1127    return point
1128
1129
1130@dataclasses.dataclass(frozen=True)
1131class MirrorRemapCoordinates(RemapCoordinates):
1132  """The coordinates are reflected across the domain boundaries so that they lie in the unit
1133  interval.  The resulting function is continuous but not smooth across the boundaries."""
1134
1135  def __call__(self, point: _NDArray, /) -> _NDArray:
1136    point = np.mod(point, 2.0)
1137    return np.where(point >= 1.0, 2.0 - point, point)
1138
1139
1140@dataclasses.dataclass(frozen=True)
1141class TileRemapCoordinates(RemapCoordinates):
1142  """The coordinates are mapped to the unit interval using a "modulo 1.0" operation.  The resulting
1143  function is generally discontinuous across the domain boundaries."""
1144
1145  def __call__(self, point: _NDArray, /) -> _NDArray:
1146    return np.mod(point, 1.0)
1147
1148
1149@dataclasses.dataclass(frozen=True)
1150class ExtendSamples(abc.ABC):
1151  """Abstract base class for replacing references to grid samples exterior to the unit domain by
1152  affine combinations of interior sample(s) and possibly the constant value (`cval`)."""
1153
1154  uses_cval: bool = False
1155  """True if some exterior samples are defined in terms of `cval`, i.e., if the computed weight
1156  is non-affine."""
1157
1158  @abc.abstractmethod
1159  def __call__(
1160      self, index: _NDArray, weight: _NDArray, size: int, gridtype: Gridtype, /
1161  ) -> tuple[_NDArray, _NDArray]:
1162    """Detect references to exterior samples, i.e., entries of `index` that lie outside the
1163    interval [0, size), and update these indices (and possibly their associated weights) to
1164    reference only interior samples.  Return `new_index, new_weight`."""
1165
1166
1167@dataclasses.dataclass(frozen=True)
1168class ReflectExtendSamples(ExtendSamples):
1169  """Find the interior sample by reflecting across domain boundaries."""
1170
1171  def __call__(
1172      self, index: _NDArray, weight: _NDArray, size: int, gridtype: Gridtype, /
1173  ) -> tuple[_NDArray, _NDArray]:
1174    index = gridtype.reflect(index, size)
1175    return index, weight
1176
1177
1178@dataclasses.dataclass(frozen=True)
1179class WrapExtendSamples(ExtendSamples):
1180  """Wrap the interior samples periodically.  For a `'primal'` grid, the last
1181  sample is ignored as its value is replaced by the first sample."""
1182
1183  def __call__(
1184      self, index: _NDArray, weight: _NDArray, size: int, gridtype: Gridtype, /
1185  ) -> tuple[_NDArray, _NDArray]:
1186    index = gridtype.wrap(index, size)
1187    return index, weight
1188
1189
1190@dataclasses.dataclass(frozen=True)
1191class ClampExtendSamples(ExtendSamples):
1192  """Use the nearest interior sample."""
1193
1194  def __call__(
1195      self, index: _NDArray, weight: _NDArray, size: int, gridtype: Gridtype, /
1196  ) -> tuple[_NDArray, _NDArray]:
1197    index = index.clip(0, size - 1)
1198    return index, weight
1199
1200
1201@dataclasses.dataclass(frozen=True)
1202class ReflectClampExtendSamples(ExtendSamples):
1203  """Extend the grid samples from [0, 1] into [-1, 0] using reflection and then define grid
1204  samples outside [-1, 1] as that of the nearest sample."""
1205
1206  def __call__(
1207      self, index: _NDArray, weight: _NDArray, size: int, gridtype: Gridtype, /
1208  ) -> tuple[_NDArray, _NDArray]:
1209    index = gridtype.reflect_clamp(index, size)
1210    return index, weight
1211
1212
1213@dataclasses.dataclass(frozen=True)
1214class BorderExtendSamples(ExtendSamples):
1215  """Let all exterior samples have the constant value (`cval`)."""
1216
1217  def __init__(self) -> None:
1218    super().__init__(uses_cval=True)
1219
1220  def __call__(
1221      self, index: _NDArray, weight: _NDArray, size: int, gridtype: Gridtype, /
1222  ) -> tuple[_NDArray, _NDArray]:
1223    low = index < 0
1224    weight[low] = 0.0
1225    index[low] = 0
1226    high = index >= size
1227    weight[high] = 0.0
1228    index[high] = size - 1
1229    return index, weight
1230
1231
1232@dataclasses.dataclass(frozen=True)
1233class ValidExtendSamples(ExtendSamples):
1234  """Assign all domain samples weight 1 and all outside samples weight 0.
1235  Compute a weighted reconstruction and divide by the reconstructed weight."""
1236
1237  def __init__(self) -> None:
1238    super().__init__(uses_cval=True)
1239
1240  def __call__(
1241      self, index: _NDArray, weight: _NDArray, size: int, gridtype: Gridtype, /
1242  ) -> tuple[_NDArray, _NDArray]:
1243    low = index < 0
1244    weight[low] = 0.0
1245    index[low] = 0
1246    high = index >= size
1247    weight[high] = 0.0
1248    index[high] = size - 1
1249    sum_weight = weight.sum(axis=-1)
1250    nonzero_sum = sum_weight != 0.0
1251    np.divide(weight, sum_weight[..., None], out=weight, where=nonzero_sum[..., None])
1252    return index, weight
1253
1254
1255@dataclasses.dataclass(frozen=True)
1256class LinearExtendSamples(ExtendSamples):
1257  """Linearly extrapolate beyond boundary samples."""
1258
1259  def __call__(
1260      self, index: _NDArray, weight: _NDArray, size: int, gridtype: Gridtype, /
1261  ) -> tuple[_NDArray, _NDArray]:
1262    if size < 2:
1263      index = gridtype.reflect(index, size)
1264      return index, weight
1265    # For each boundary, define new columns in index and weight arrays to represent the last and
1266    # next-to-last samples.  When we later construct the sparse resize matrix, we will sum the
1267    # duplicate index entries.
1268    low = index < 0
1269    high = index >= size
1270    w = np.empty((*weight.shape[:-1], weight.shape[-1] + 4), weight.dtype)
1271    x = index
1272    w[..., -4] = ((1 - x) * weight).sum(where=low, axis=-1)
1273    w[..., -3] = ((x) * weight).sum(where=low, axis=-1)
1274    x = (size - 1) - index
1275    w[..., -2] = ((x) * weight).sum(where=high, axis=-1)
1276    w[..., -1] = ((1 - x) * weight).sum(where=high, axis=-1)
1277    weight[low] = 0.0
1278    index[low] = 0
1279    weight[high] = 0.0
1280    index[high] = size - 1
1281    w[..., :-4] = weight
1282    weight = w
1283    new_index = np.empty(w.shape, index.dtype)
1284    new_index[..., :-4] = index
1285    # Let matrix (including zero values) be banded.
1286    new_index[..., -4:] = np.where(w[..., -4:] != 0.0, [0, 1, size - 2, size - 1], index[..., :1])
1287    index = new_index
1288    return index, weight
1289
1290
1291@dataclasses.dataclass(frozen=True)
1292class QuadraticExtendSamples(ExtendSamples):
1293  """Quadratically extrapolate beyond boundary samples."""
1294
1295  def __call__(
1296      self, index: _NDArray, weight: _NDArray, size: int, gridtype: Gridtype, /
1297  ) -> tuple[_NDArray, _NDArray]:
1298    # [Keys 1981] suggests this as x[-1] = 3*x[0] - 3*x[1] + x[2], calling it "cubic precision",
1299    # but it seems just quadratic.
1300    if size < 3:
1301      index = gridtype.reflect(index, size)
1302      return index, weight
1303    low = index < 0
1304    high = index >= size
1305    w = np.empty((*weight.shape[:-1], weight.shape[-1] + 6), weight.dtype)
1306    x = index
1307    w[..., -6] = (((0.5 * x - 1.5) * x + 1) * weight).sum(where=low, axis=-1)
1308    w[..., -5] = (((-x + 2) * x) * weight).sum(where=low, axis=-1)
1309    w[..., -4] = (((0.5 * x - 0.5) * x) * weight).sum(where=low, axis=-1)
1310    x = (size - 1) - index
1311    w[..., -3] = (((0.5 * x - 0.5) * x) * weight).sum(where=high, axis=-1)
1312    w[..., -2] = (((-x + 2) * x) * weight).sum(where=high, axis=-1)
1313    w[..., -1] = (((0.5 * x - 1.5) * x + 1) * weight).sum(where=high, axis=-1)
1314    weight[low] = 0.0
1315    index[low] = 0
1316    weight[high] = 0.0
1317    index[high] = size - 1
1318    w[..., :-6] = weight
1319    weight = w
1320    new_index = np.empty(w.shape, index.dtype)
1321    new_index[..., :-6] = index
1322    # Let matrix (including zero values) be banded.
1323    new_index[..., -6:] = np.where(
1324        w[..., -6:] != 0.0, [0, 1, 2, size - 3, size - 2, size - 1], index[..., :1]
1325    )
1326    index = new_index
1327    return index, weight
1328
1329
1330@dataclasses.dataclass(frozen=True)
1331class OverrideExteriorValue:
1332  """Abstract base class to set the value outside some domain extent to a
1333  constant value (`cval`)."""
1334
1335  boundary_antialiasing: bool = True
1336  """Antialias the pixel values adjacent to the boundary of the extent."""
1337
1338  uses_cval: bool = False
1339  """Modify some weights to introduce references to the `cval` constant value."""
1340
1341  def __call__(self, weight: _NDArray, point: _NDArray, /) -> None:
1342    """For all `point` outside some extent, modify the weight to be zero."""
1343
1344  def override_using_signed_distance(
1345      self, weight: _NDArray, point: _NDArray, signed_distance: _NDArray, /
1346  ) -> None:
1347    """Reduce sample weights for "outside" values based on the signed distance function,
1348    to effectively assign the constant value `cval`."""
1349    all_points_inside_domain = np.all(signed_distance <= 0.0)
1350    if all_points_inside_domain:
1351      return
1352    if self.boundary_antialiasing and min(point.shape) >= 2:
1353      # For discontinuous coordinate mappings, we may need to somehow ignore
1354      # the large finite differences computed across the map discontinuities.
1355      gradient = np.gradient(point)
1356      gradient_norm = np.linalg.norm(np.atleast_2d(gradient), axis=0)
1357      signed_distance_in_samples = signed_distance / (gradient_norm + 1e-20)
1358      # Opacity is in linear space, which is correct if Gamma is set.
1359      opacity = (0.5 - signed_distance_in_samples).clip(0.0, 1.0)
1360      weight *= opacity[..., None]
1361    else:
1362      is_outside = signed_distance > 0.0
1363      weight[is_outside, :] = 0.0
1364
1365
1366@dataclasses.dataclass(frozen=True)
1367class NoOverrideExteriorValue(OverrideExteriorValue):
1368  """The function value is not overridden."""
1369
1370  def __call__(self, weight: _NDArray, point: _NDArray, /) -> None:
1371    pass
1372
1373
1374@dataclasses.dataclass(frozen=True)
1375class UnitDomainOverrideExteriorValue(OverrideExteriorValue):
1376  """Values outside the unit interval [0, 1] are replaced by the constant `cval`."""
1377
1378  def __init__(self, **kwargs: Any) -> None:
1379    super().__init__(uses_cval=True, **kwargs)
1380
1381  def __call__(self, weight: _NDArray, point: _NDArray, /) -> None:
1382    signed_distance = abs(point - 0.5) - 0.5  # Boundaries at 0.0 and 1.0.
1383    self.override_using_signed_distance(weight, point, signed_distance)
1384
1385
1386@dataclasses.dataclass(frozen=True)
1387class PlusMinusOneOverrideExteriorValue(OverrideExteriorValue):
1388  """Values outside the interval [-1, 1] are replaced by the constant `cval`."""
1389
1390  def __init__(self, **kwargs: Any) -> None:
1391    super().__init__(uses_cval=True, **kwargs)
1392
1393  def __call__(self, weight: _NDArray, point: _NDArray, /) -> None:
1394    signed_distance = abs(point) - 1.0  # Boundaries at -1.0 and 1.0.
1395    self.override_using_signed_distance(weight, point, signed_distance)
1396
1397
1398@dataclasses.dataclass(frozen=True)
1399class Boundary:
1400  """Domain boundary rules.  These define the reconstruction over the source domain near and beyond
1401  the domain boundaries.  The rules may be specified separately for each domain dimension."""
1402
1403  name: str = ''
1404  """Boundary rule name."""
1405
1406  coord_remap: RemapCoordinates = NoRemapCoordinates()
1407  """Modify specified coordinates prior to evaluating the reconstruction kernels."""
1408
1409  extend_samples: ExtendSamples = ReflectExtendSamples()
1410  """Define the value of each grid sample outside the unit domain as an affine combination of
1411  interior sample(s) and possibly the constant value (`cval`)."""
1412
1413  override_value: OverrideExteriorValue = NoOverrideExteriorValue()
1414  """Set the value outside some extent to a constant value (`cval`)."""
1415
1416  @property
1417  def uses_cval(self) -> bool:
1418    """True if weights may be non-affine, involving the constant value (`cval`)."""
1419    return self.extend_samples.uses_cval or self.override_value.uses_cval
1420
1421  def preprocess_coordinates(self, point: _NDArray, /) -> _NDArray:
1422    """Modify coordinates prior to evaluating the filter kernels."""
1423    # Antialiasing across the tile boundaries may be feasible but seems hard.
1424    point = self.coord_remap(point)
1425    return point
1426
1427  def apply(
1428      self, index: _NDArray, weight: _NDArray, point: _NDArray, size: int, gridtype: Gridtype, /
1429  ) -> tuple[_NDArray, _NDArray]:
1430    """Replace exterior samples by combinations of interior samples."""
1431    index, weight = self.extend_samples(index, weight, size, gridtype)
1432    self.override_reconstruction(weight, point)
1433    return index, weight
1434
1435  def override_reconstruction(self, weight: _NDArray, point: _NDArray, /) -> None:
1436    """For points outside an extent, modify weight to zero to assign `cval`."""
1437    self.override_value(weight, point)
1438
1439
1440_DICT_BOUNDARIES = {
1441    'reflect': Boundary('reflect', extend_samples=ReflectExtendSamples()),
1442    'wrap': Boundary('wrap', extend_samples=WrapExtendSamples()),
1443    'tile': Boundary(
1444        'title', coord_remap=TileRemapCoordinates(), extend_samples=ReflectExtendSamples()
1445    ),
1446    'clamp': Boundary('clamp', extend_samples=ClampExtendSamples()),
1447    'border': Boundary('border', extend_samples=BorderExtendSamples()),
1448    'natural': Boundary(
1449        'natural',
1450        extend_samples=ValidExtendSamples(),
1451        override_value=UnitDomainOverrideExteriorValue(),
1452    ),
1453    'linear_constant': Boundary(
1454        'linear_constant',
1455        extend_samples=LinearExtendSamples(),
1456        override_value=UnitDomainOverrideExteriorValue(),
1457    ),
1458    'quadratic_constant': Boundary(
1459        'quadratic_constant',
1460        extend_samples=QuadraticExtendSamples(),
1461        override_value=UnitDomainOverrideExteriorValue(),
1462    ),
1463    'reflect_clamp': Boundary('reflect_clamp', extend_samples=ReflectClampExtendSamples()),
1464    'constant': Boundary(
1465        'constant',
1466        extend_samples=ReflectExtendSamples(),
1467        override_value=UnitDomainOverrideExteriorValue(),
1468    ),
1469    'linear': Boundary('linear', extend_samples=LinearExtendSamples()),
1470    'quadratic': Boundary('quadratic', extend_samples=QuadraticExtendSamples()),
1471}
1472
1473BOUNDARIES = list(_DICT_BOUNDARIES)
1474"""Shortcut names for some predefined boundary rules (as defined by `_DICT_BOUNDARIES`):
1475
1476| name                   | a.k.a. / comments |
1477|------------------------|-------------------|
1478| `'reflect'`            | *reflected*, *symm*, *symmetric*, *mirror*, *grid-mirror* |
1479| `'wrap'`               | *periodic*, *repeat*, *grid-wrap* |
1480| `'tile'`               | like `'reflect'` within unit domain, then tile discontinuously |
1481| `'clamp'`              | *clamped*, *nearest*, *edge*, *clamp-to-edge*, repeat last sample |
1482| `'border'`             | *grid-constant*, use `cval` for samples outside unit domain |
1483| `'natural'`            | *renormalize* using only interior samples, use `cval` outside domain |
1484| `'reflect_clamp'`      | *mirror-clamp-to-edge* |
1485| `'constant'`           | like `'reflect'` but replace by `cval` outside unit domain |
1486| `'linear'`             | extrapolate from 2 last samples |
1487| `'quadratic'`          | extrapolate from 3 last samples |
1488| `'linear_constant'`    | like `'linear'` but replace by `cval` outside unit domain |
1489| `'quadratic_constant'` | like `'quadratic'` but replace by `cval` outside unit domain |
1490
1491These boundary rules may be specified per dimension.  See the source code for extensibility
1492using the classes `RemapCoordinates`, `ExtendSamples`, and `OverrideExteriorValue`.
1493
1494**Boundary rules illustrated in 1D:**
1495
1496<center>
1497<img src="https://github.com/hhoppe/resampler/raw/main/media/boundary_rules_in_1D.png" width="100%"/>
1498</center>
1499
1500**Boundary rules illustrated in 2D:**
1501
1502<center>
1503<img src="https://github.com/hhoppe/resampler/raw/main/media/boundary_rules_in_2D.png" width="100%"/>
1504</center>
1505"""
1506
1507_OFTUSED_BOUNDARIES = (
1508    'reflect wrap tile clamp border natural linear_constant quadratic_constant'.split()
1509)
1510"""A useful subset of `BOUNDARIES` for visualization in figures."""
1511
1512
1513def _get_boundary(boundary: str | Boundary, /) -> Boundary:
1514  """Return a `Boundary`, which can be specified as a name in `BOUNDARIES`."""
1515  return boundary if isinstance(boundary, Boundary) else _DICT_BOUNDARIES[boundary]
1516
1517
1518@dataclasses.dataclass(frozen=True)
1519class Filter(abc.ABC):
1520  """Abstract base class for filter kernel functions.
1521
1522  Each kernel is assumed to be a zero-phase filter, i.e., to be symmetric in a support
1523  interval [-radius, radius].  (Some sites instead define kernels over the interval [0, N]
1524  where N = 2 * radius.)
1525
1526  Portions of this code are adapted from the C++ library in
1527  https://github.com/hhoppe/Mesh-processing-library/blob/main/libHh/Filter.cpp
1528
1529  See also https://hhoppe.com/proj/filtering/.
1530  """
1531
1532  name: str
1533  """Filter kernel name."""
1534
1535  radius: float
1536  """Max absolute value of x for which self(x) is nonzero."""
1537
1538  interpolating: bool = True
1539  """True if self(0) == 1.0 and self(i) == 0.0 for all nonzero integers i."""
1540
1541  continuous: bool = True
1542  """True if the kernel function has $C^0$ continuity."""
1543
1544  partition_of_unity: bool = True
1545  """True if the convolution of the kernel with a Dirac comb reproduces the
1546  unity function."""
1547
1548  unit_integral: bool = True
1549  """True if the integral of the kernel function is 1."""
1550
1551  requires_digital_filter: bool = False
1552  """True if the filter needs a pre/post digital filter for interpolation."""
1553
1554  @abc.abstractmethod
1555  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1556    """Return evaluation of filter kernel at locations x."""
1557
1558
1559class ImpulseFilter(Filter):
1560  """See https://en.wikipedia.org/wiki/Dirac_delta_function."""
1561
1562  def __init__(self) -> None:
1563    super().__init__(name='impulse', radius=1e-20, continuous=False, partition_of_unity=False)
1564
1565  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1566    raise AssertionError('The Impulse is infinitely narrow, so cannot be directly evaluated.')
1567
1568
1569class BoxFilter(Filter):
1570  """See https://en.wikipedia.org/wiki/Box_function.
1571
1572  The kernel function has value 1.0 over the half-open interval [-.5, .5).
1573  """
1574
1575  def __init__(self) -> None:
1576    super().__init__(name='box', radius=0.5, continuous=False)
1577
1578  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1579    use_asymmetric = True
1580    if use_asymmetric:
1581      x = np.asarray(x)
1582      return np.where((-0.5 <= x) & (x < 0.5), 1.0, 0.0)
1583    x = np.abs(x)
1584    return np.where(x < 0.5, 1.0, np.where(x == 0.5, 0.5, 0.0))
1585
1586
1587class TrapezoidFilter(Filter):
1588  """Filter for antialiased "area-based" filtering.
1589
1590  Args:
1591    radius: Specifies the support [-radius, radius] of the filter, where 0.5 < radius <= 1.0.
1592      The special case `radius = None` is a placeholder that indicates that the filter will be
1593      replaced by a trapezoid of the appropriate radius (based on scaling) for correct
1594      antialiasing in both minification and magnification.
1595
1596  This filter is similar to the BoxFilter but with linearly sloped sides.  It has value 1.0
1597  in the interval abs(x) <= 1.0 - radius and decreases linearly to value 0.0 in the interval
1598  1.0 - radius <= abs(x) <= radius, always with value 0.5 at x = 0.5.
1599  """
1600
1601  def __init__(self, *, radius: float | None = None) -> None:
1602    if radius is None:
1603      super().__init__(name='trapezoid', radius=0.0)
1604      return
1605    if not 0.5 < radius <= 1.0:
1606      raise ValueError(f'Radius {radius} is outside the range (0.5, 1.0].')
1607    super().__init__(name=f'trapezoid_{radius}', radius=radius)
1608
1609  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1610    x = np.abs(x)
1611    assert 0.5 < self.radius <= 1.0
1612    return ((0.5 + 0.25 / (self.radius - 0.5)) - (0.5 / (self.radius - 0.5)) * x).clip(0.0, 1.0)
1613
1614
1615class TriangleFilter(Filter):
1616  """See https://en.wikipedia.org/wiki/Triangle_function.
1617
1618  Also known as the hat or tent function.  It is used for piecewise-linear
1619  (or bilinear, or trilinear, ...) interpolation.
1620  """
1621
1622  def __init__(self) -> None:
1623    super().__init__(name='triangle', radius=1.0)
1624
1625  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1626    return (1.0 - np.abs(x)).clip(0.0, 1.0)
1627
1628
1629class CubicFilter(Filter):
1630  """Family of cubic filters parameterized by two scalar parameters.
1631
1632  Args:
1633    b: first scalar parameter.
1634    c: second scalar parameter.
1635
1636  See https://en.wikipedia.org/wiki/Mitchell%E2%80%93Netravali_filters and
1637  https://doi.org/10.1145/378456.378514.
1638
1639  [D. P. Mitchell and A. N. Netravali. Reconstruction filters in computer graphics.
1640  Computer Graphics (Proceedings of ACM SIGGRAPH 1988), 22(4):221-228, 1988.]
1641
1642  - The filter has quadratic precision iff b + 2 * c == 1.
1643  - The filter is interpolating iff b == 0.
1644  - (b=1, c=0) is the (non-interpolating) cubic B-spline basis;
1645  - (b=1/3, c=1/3) is the Mitchell filter;
1646  - (b=0, c=0.5) is the Catmull-Rom spline (which has cubic precision);
1647  - (b=0, c=0.75) is the "sharper cubic" used in Photoshop and OpenCV.
1648  """
1649
1650  def __init__(self, *, b: float, c: float, name: str | None = None) -> None:
1651    name = f'cubic_b{b}_c{c}' if name is None else name
1652    interpolating = b == 0
1653    super().__init__(name=name, radius=2.0, interpolating=interpolating)
1654    self.b, self.c = b, c
1655
1656  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1657    x = np.abs(x)
1658    b, c = self.b, self.c
1659    f3, f2, f0 = 2 - 9 / 6 * b - c, -3 + 2 * b + c, 1 - 1 / 3 * b
1660    g3, g2, g1, g0 = -b / 6 - c, b + 5 * c, -2 * b - 8 * c, 8 / 6 * b + 4 * c
1661    # (np.polynomial.polynomial.polyval(x, [f0, 0, f2, f3]) is almost
1662    # twice as slow; see also https://stackoverflow.com/questions/24065904)
1663    v01 = ((f3 * x + f2) * x) * x + f0
1664    v12 = ((g3 * x + g2) * x + g1) * x + g0
1665    return np.where(x < 1.0, v01, np.where(x < 2.0, v12, 0.0))
1666
1667
1668class CatmullRomFilter(CubicFilter):
1669  """Cubic filter with cubic precision.  Also known as Keys filter.
1670
1671  [E. Catmull, R. Rom.  A class of local interpolating splines.  Computer aided geometric
1672  design, 1974]
1673  [Wikipedia](https://en.wikipedia.org/wiki/Cubic_Hermite_spline#Catmull%E2%80%93Rom_spline)
1674
1675  [R. G. Keys.  Cubic convolution interpolation for digital image processing.
1676  IEEE Trans. on Acoustics, Speech, and Signal Processing, 29(6), 1981.]
1677  https://ieeexplore.ieee.org/document/1163711/.
1678  """
1679
1680  def __init__(self) -> None:
1681    super().__init__(b=0, c=0.5, name='cubic')
1682
1683
1684class MitchellFilter(CubicFilter):
1685  """See https://doi.org/10.1145/378456.378514.
1686
1687  [D. P. Mitchell and A. N. Netravali.  Reconstruction filters in computer graphics.  Computer
1688  Graphics (Proceedings of ACM SIGGRAPH 1988), 22(4):221-228, 1988.]
1689  """
1690
1691  def __init__(self) -> None:
1692    super().__init__(b=1 / 3, c=1 / 3, name='mitchell')
1693
1694
1695class SharpCubicFilter(CubicFilter):
1696  """Cubic filter that is sharper than Catmull-Rom filter.
1697
1698  Used by some tools including OpenCV and Photoshop.
1699
1700  See https://en.wikipedia.org/wiki/Mitchell%E2%80%93Netravali_filters and
1701  https://entropymine.com/resamplescope/notes/photoshop/.
1702  """
1703
1704  def __init__(self) -> None:
1705    super().__init__(b=0, c=0.75, name='sharpcubic')
1706
1707
1708class LanczosFilter(Filter):
1709  """High-quality filter: sinc function modulated by a sinc window.
1710
1711  Args:
1712    radius: Specifies the support window [-radius, radius] over which the filter is nonzero.
1713    sampled: If True, use a discretized approximation for improved speed.
1714
1715  See https://en.wikipedia.org/wiki/Lanczos_kernel.
1716  """
1717
1718  def __init__(self, *, radius: int, sampled: bool = True) -> None:
1719    super().__init__(
1720        name=f'lanczos_{radius}', radius=radius, partition_of_unity=False, unit_integral=False
1721    )
1722
1723    @_cache_sampled_1d_function(xmin=-radius, xmax=radius, enable=sampled)
1724    def _eval(x: _ArrayLike) -> _NDArray:
1725      x = np.abs(x)
1726      # Note that window[n] = sinc(2*n/N - 1), with 0 <= n <= N.
1727      # But, x = n - N/2, or equivalently, n = x + N/2, with -N/2 <= x <= N/2.
1728      window = _sinc(x / radius)  # Zero-phase function w_0(x).
1729      return np.where(x < radius, _sinc(x) * window, 0.0)
1730
1731    self._function = _eval
1732
1733  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1734    return self._function(x)
1735
1736
1737class GeneralizedHammingFilter(Filter):
1738  """Sinc function modulated by a Hamming window.
1739
1740  Args:
1741    radius: Specifies the support window [-radius, radius] over which the filter is nonzero.
1742    a0: Scalar parameter, where 0.0 < a0 < 1.0.  The case of a0=0.5 is the Hann filter.
1743
1744  See https://en.wikipedia.org/wiki/Window_function#Hann_and_Hamming_windows,
1745  and hamming() in https://github.com/scipy/scipy/blob/main/scipy/signal/windows/_windows.py.
1746
1747  Note that `'hamming3'` is `(radius=3, a0=25/46)`, which close to but different from `a0=0.54`.
1748
1749  See also np.hamming() and np.hanning().
1750  """
1751
1752  def __init__(self, *, radius: int, a0: float) -> None:
1753    super().__init__(
1754        name=f'hamming_{radius}',
1755        radius=radius,
1756        partition_of_unity=False,  # 1:1.00242  av=1.00188  sd=0.00052909
1757        unit_integral=False,  # 1.00188
1758    )
1759    assert 0.0 < a0 < 1.0
1760    self.a0 = a0
1761
1762  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1763    x = np.abs(x)
1764    # Note that window[n] = a0 - (1 - a0) * cos(2 * pi * n / N), 0 <= n <= N.
1765    # With n = x + N/2, we get the zero-phase function w_0(x):
1766    window = self.a0 + (1.0 - self.a0) * np.cos(np.pi / self.radius * x)
1767    return np.where(x < self.radius, _sinc(x) * window, 0.0)
1768
1769
1770class KaiserFilter(Filter):
1771  """Sinc function modulated by a Kaiser-Bessel window.
1772
1773  See https://en.wikipedia.org/wiki/Kaiser_window, and example use in:
1774  [Karras et al. 20201.  Alias-free generative adversarial networks.
1775  https://arxiv.org/pdf/2106.12423.pdf].
1776
1777  See also np.kaiser().
1778
1779  Args:
1780    radius: Value L/2 in the definition.  It may be fractional for a (digital) resizing filter
1781      (sample spacing s != 1) with an even number of samples (dual grid), e.g., Eq. (6)
1782      in [Karras et al. 2021] --- this effects the precise shape of the window function.
1783    beta: Determines the trade-off between main-lobe width and side-lobe level.
1784    sampled: If True, use a discretized approximation for improved speed.
1785  """
1786
1787  def __init__(self, *, radius: float, beta: float, sampled: bool = True) -> None:
1788    assert beta >= 0.0
1789    super().__init__(
1790        name=f'kaiser_{radius}_{beta}', radius=radius, partition_of_unity=False, unit_integral=False
1791    )
1792
1793    @_cache_sampled_1d_function(xmin=-math.ceil(radius), xmax=math.ceil(radius), enable=sampled)
1794    def _eval(x: _ArrayLike) -> _NDArray:
1795      x = np.abs(x)
1796      window = np.i0(beta * np.sqrt((1.0 - np.square(x / radius)).clip(0.0, 1.0))) / np.i0(beta)
1797      return np.where(x <= radius + 1e-6, _sinc(x) * window, 0.0)
1798
1799    self._function = _eval
1800
1801  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1802    return self._function(x)
1803
1804
1805class BsplineFilter(Filter):
1806  """B-spline of a non-negative degree.
1807
1808  Args:
1809    degree: The polynomial degree of the B-spline segments.
1810      With `degree=0`, it is like `BoxFilter` except with f(0.5) = f(-0.5) = 0.
1811      With `degree=1`, it is identical to `TriangleFilter`.
1812      With `degree >= 2`, it is no longer interpolating.
1813
1814  See [Carl de Boor.  A practical guide to splines.  Springer, 2001.]
1815  https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.BSpline.html
1816  """
1817
1818  def __init__(self, *, degree: int) -> None:
1819    if degree < 0:
1820      raise ValueError(f'Bspline of degree {degree} is invalid.')
1821    radius = (degree + 1) / 2
1822    interpolating = degree <= 1
1823    super().__init__(name=f'bspline{degree}', radius=radius, interpolating=interpolating)
1824    t = list(range(degree + 2))
1825    self._bspline = scipy.interpolate.BSpline.basis_element(t)
1826
1827  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1828    x = np.abs(x)
1829    return np.where(x < self.radius, self._bspline(x + self.radius), 0.0)
1830
1831
1832class CardinalBsplineFilter(Filter):
1833  """Interpolating B-spline, achieved with aid of digital pre or post filter.
1834
1835  Args:
1836    degree: The polynomial degree of the B-spline segments.
1837    sampled: If True, use a discretized approximation for improved speed.
1838
1839  See [Hou and Andrews.  Cubic splines for image interpolation and digital filtering, 1978] and
1840  [Unser et al.  Fast B-spline transforms for continuous image representation and interpolation,
1841  1991].
1842  """
1843
1844  def __init__(self, *, degree: int, sampled: bool = True) -> None:
1845    self.degree = degree
1846    if degree < 0:
1847      raise ValueError(f'Bspline of degree {degree} is invalid.')
1848    radius = (degree + 1) / 2
1849    super().__init__(
1850        name=f'cardinal{degree}',
1851        radius=radius,
1852        requires_digital_filter=degree >= 2,
1853        continuous=degree >= 1,
1854    )
1855    t = list(range(degree + 2))
1856    bspline = scipy.interpolate.BSpline.basis_element(t)
1857
1858    @_cache_sampled_1d_function(xmin=-radius, xmax=radius, enable=sampled)
1859    def _eval(x: _ArrayLike) -> _NDArray:
1860      x = np.abs(x)
1861      return np.where(x < radius, bspline(x + radius), 0.0)
1862
1863    self._function = _eval
1864
1865  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1866    return self._function(x)
1867
1868
1869class OmomsFilter(Filter):
1870  """OMOMS interpolating filter, with aid of digital pre or post filter.
1871
1872  Args:
1873    degree: The polynomial degree of the filter segments.
1874
1875  Optimal MOMS (maximal-order-minimal-support) function; see [Blu and Thevenaz, MOMS: Maximal-order
1876  interpolation of minimal support, 2001].
1877  https://infoscience.epfl.ch/record/63074/files/blu0101.pdf
1878  """
1879
1880  def __init__(self, *, degree: int) -> None:
1881    if degree not in (3, 5):
1882      raise ValueError(f'Degree {degree} not supported.')
1883    super().__init__(name=f'omoms{degree}', radius=(degree + 1) / 2, requires_digital_filter=True)
1884    self.degree = degree
1885
1886  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1887    x = np.abs(x)
1888    match self.degree:
1889      case 3:
1890        v01 = ((0.5 * x - 1.0) * x + 3 / 42) * x + 26 / 42
1891        v12 = ((-7 / 42 * x + 1.0) * x - 85 / 42) * x + 58 / 42
1892        return np.where(x < 1.0, v01, np.where(x < 2.0, v12, 0.0))
1893      case 5:
1894        v01 = ((((-1 / 12 * x + 1 / 4) * x - 5 / 99) * x - 9 / 22) * x - 1 / 792) * x + 229 / 440
1895        v12 = (
1896            (((1 / 24 * x - 3 / 8) * x + 505 / 396) * x - 83 / 44) * x + 1351 / 1584
1897        ) * x + 839 / 2640
1898        v23 = (
1899            (((-1 / 120 * x + 1 / 8) * x - 299 / 396) * x + 101 / 44) * x - 27811 / 7920
1900        ) * x + 5707 / 2640
1901        return np.where(x < 1.0, v01, np.where(x < 2.0, v12, np.where(x < 3.0, v23, 0.0)))
1902      case _:
1903        raise ValueError(self.degree)
1904
1905
1906class GaussianFilter(Filter):
1907  r"""See https://en.wikipedia.org/wiki/Gaussian_function.
1908
1909  Args:
1910    standard_deviation: Sets the Gaussian $\sigma$.  The default value is 1.25/3.0, which
1911      creates a kernel that is as-close-as-possible to a partition of unity.
1912  """
1913
1914  DEFAULT_STANDARD_DEVIATION = 1.25 / 3.0
1915  """This value creates a kernel that is as-close-as-possible to a partition of unity; see
1916  mesh_processing/test/GridOp_test.cpp: `0.93503:1.06497     av=1           sd=0.0459424`.
1917  Another possibility is 0.5, as suggested on p. 4 of [Ken Turkowski.  Filters for common
1918  resampling tasks, 1990] for kernels with a support of 3 pixels.
1919  https://cadxfem.org/inf/ResamplingFilters.pdf
1920  """
1921
1922  def __init__(self, *, standard_deviation: float = DEFAULT_STANDARD_DEVIATION) -> None:
1923    super().__init__(
1924        name=f'gaussian_{standard_deviation:.3f}',
1925        radius=np.ceil(8.0 * standard_deviation),  # Sufficiently large.
1926        interpolating=False,
1927        partition_of_unity=False,
1928    )
1929    self.standard_deviation = standard_deviation
1930
1931  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1932    x = np.abs(x)
1933    sdv = self.standard_deviation
1934    v0r = np.exp(np.square(x / sdv) / -2.0) / (np.sqrt(math.tau) * sdv)
1935    return np.where(x < self.radius, v0r, 0.0)
1936
1937
1938class NarrowBoxFilter(Filter):
1939  """Compact footprint, used for visualization of grid sample location.
1940
1941  Args:
1942    radius: Specifies the support [-radius, radius] of the narrow box function.  (The default
1943      value 0.199 is an inexact 0.2 to avoid numerical ambiguities.)
1944  """
1945
1946  def __init__(self, *, radius: float = 0.199) -> None:
1947    super().__init__(
1948        name='narrowbox',
1949        radius=radius,
1950        continuous=False,
1951        unit_integral=False,
1952        partition_of_unity=False,
1953    )
1954
1955  def __call__(self, x: _ArrayLike, /) -> _NDArray:
1956    radius = self.radius
1957    magnitude = 1.0
1958    x = np.asarray(x)
1959    return np.where((-radius <= x) & (x < radius), magnitude, 0.0)
1960
1961
1962_DEFAULT_FILTER = 'lanczos3'
1963
1964_DICT_FILTERS = {
1965    'impulse': ImpulseFilter(),
1966    'box': BoxFilter(),
1967    'trapezoid': TrapezoidFilter(),
1968    'triangle': TriangleFilter(),
1969    'cubic': CatmullRomFilter(),
1970    'sharpcubic': SharpCubicFilter(),
1971    'lanczos3': LanczosFilter(radius=3),
1972    'lanczos5': LanczosFilter(radius=5),
1973    'lanczos10': LanczosFilter(radius=10),
1974    'cardinal3': CardinalBsplineFilter(degree=3),
1975    'cardinal5': CardinalBsplineFilter(degree=5),
1976    'omoms3': OmomsFilter(degree=3),
1977    'omoms5': OmomsFilter(degree=5),
1978    'hamming3': GeneralizedHammingFilter(radius=3, a0=25 / 46),
1979    'kaiser3': KaiserFilter(radius=3.0, beta=7.12),
1980    'gaussian': GaussianFilter(),
1981    'bspline3': BsplineFilter(degree=3),
1982    'mitchell': MitchellFilter(),
1983    'narrowbox': NarrowBoxFilter(),
1984    # Not in FILTERS:
1985    'hamming1': GeneralizedHammingFilter(radius=1, a0=0.54),
1986    'hann3': GeneralizedHammingFilter(radius=3, a0=0.5),
1987    'lanczos4': LanczosFilter(radius=4),
1988}
1989
1990FILTERS = list(itertools.takewhile(lambda x: x != 'hamming1', _DICT_FILTERS))
1991r"""Shortcut names for some predefined filter kernels (specified per dimension).
1992The names expand to:
1993
1994| name           | `Filter`                      | a.k.a. / comments |
1995|----------------|-------------------------------|-------------------|
1996| `'impulse'`    | `ImpulseFilter()`             | *nearest* |
1997| `'box'`        | `BoxFilter()`                 | non-antialiased box, e.g. ImageMagick |
1998| `'trapezoid'`  | `TrapezoidFilter()`           | *area* antialiasing, e.g. `cv.INTER_AREA` |
1999| `'triangle'`   | `TriangleFilter()`            | *linear*  (*bilinear* in 2D), spline `order=1` |
2000| `'cubic'`      | `CatmullRomFilter()`          | *catmullrom*, *keys*, *bicubic* |
2001| `'sharpcubic'` | `SharpCubicFilter()`          | `cv.INTER_CUBIC`, `torch 'bicubic'` |
2002| `'lanczos3'`   | `LanczosFilter`(radius=3)     | support window [-3, 3] |
2003| `'lanczos5'`   | `LanczosFilter`(radius=5)     | [-5, 5] |
2004| `'lanczos10'`  | `LanczosFilter`(radius=10)    | [-10, 10] |
2005| `'cardinal3'`  | `CardinalBsplineFilter`(degree=3) | *spline interpolation*, `order=3`, *GF* |
2006| `'cardinal5'`  | `CardinalBsplineFilter`(degree=5) | *spline interpolation*, `order=5`, *GF* |
2007| `'omoms3'`     | `OmomsFilter`(degree=3)       | non-$C^1$, [-3, 3], *GF* |
2008| `'omoms5'`     | `OmomsFilter`(degree=5)       | non-$C^1$, [-5, 5], *GF* |
2009| `'hamming3'`   | `GeneralizedHammingFilter`(...) | (radius=3, a0=25/46) |
2010| `'kaiser3'`    | `KaiserFilter`(radius=3.0, beta=7.12) | |
2011| `'gaussian'`   | `GaussianFilter()`            | non-interpolating, default $\sigma=1.25/3$ |
2012| `'bspline3'`   | `BsplineFilter`(degree=3)     | non-interpolating |
2013| `'mitchell'`   | `MitchellFilter()`            | *mitchellcubic* |
2014| `'narrowbox'`  | `NarrowBoxFilter()`           | for visualization of sample positions |
2015
2016The comment label *GF* denotes a [generalized filter](https://hhoppe.com/proj/filtering/), formed
2017as the composition of a finitely supported kernel and a discrete inverse convolution.
2018
2019**Some example filter kernels:**
2020
2021<center>
2022<img src="https://github.com/hhoppe/resampler/raw/main/media/filter_summary.png" width="100%"/>
2023</center>
2024
2025<br/>A more extensive set of filters is presented [here](#plots_of_filters) in the
2026[notebook](https://colab.research.google.com/github/hhoppe/resampler/blob/main/resampler_notebook.ipynb),
2027together with visualizations and analyses of the filter properties.
2028See the source code for extensibility.
2029"""
2030
2031
2032def _get_filter(filter: str | Filter, /) -> Filter:
2033  """Return a `Filter`, which can be specified as a name string key in `FILTERS`."""
2034  return filter if isinstance(filter, Filter) else _DICT_FILTERS[filter]
2035
2036
2037def _to_float_01(array: _Array, /, dtype: _DTypeLike) -> _Array:
2038  """Scale uint to the range [0.0, 1.0], and clip float to [0.0, 1.0]."""
2039  array_dtype = _arr_dtype(array)
2040  dtype = np.dtype(dtype)
2041  assert np.issubdtype(dtype, np.floating)
2042  match array_dtype.type:
2043    case np.uint8 | np.uint16 | np.uint32:
2044      if _arr_arraylib(array) == 'numpy':
2045        result: Any = np.multiply(array, 1 / np.iinfo(array_dtype).max, dtype=dtype)
2046        return result
2047      return cast(_Array, _arr_astype(array, dtype) / np.iinfo(array_dtype).max)
2048    case _:
2049      assert np.issubdtype(array_dtype, np.floating)
2050      return _arr_clip(array, 0.0, 1.0, dtype)
2051
2052
2053def _from_float(array: _Array, /, dtype: _DTypeLike) -> _Array:
2054  """Convert a float in range [0.0, 1.0] to uint or float type."""
2055  assert np.issubdtype(_arr_dtype(array), np.floating)
2056  dtype = np.dtype(dtype)
2057  match dtype.type:
2058    case np.uint8 | np.uint16:
2059      scale32 = cast(float, np.float32(np.iinfo(dtype.name).max))
2060      return _arr_astype(cast(_Array, array * scale32 + 0.5), dtype)
2061    case np.uint32:
2062      scale64 = cast(float, np.float64(np.iinfo(dtype.name).max))
2063      return _arr_astype(cast(_Array, array * scale64 + 0.5), dtype)
2064    case _:
2065      assert np.issubdtype(dtype, np.floating)
2066      return _arr_astype(array, dtype)
2067
2068
2069@dataclasses.dataclass(frozen=True)
2070class Gamma(abc.ABC):
2071  """Abstract base class for transfer functions on sample values.
2072
2073  Image/video content is often stored using a color component transfer function.
2074  See https://en.wikipedia.org/wiki/Gamma_correction.
2075
2076  Converts between integer types and [0.0, 1.0] internal value range.
2077  """
2078
2079  name: str
2080  """Name of component transfer function."""
2081
2082  @abc.abstractmethod
2083  def decode(self, array: _Array, /, dtype: _DTypeLike = np.float32) -> _Array:
2084    """Decode source sample values into floating-point, possibly nonlinearly.
2085
2086    Uint source values are mapped to the range [0.0, 1.0].
2087    """
2088
2089  @abc.abstractmethod
2090  def encode(self, array: _Array, /, dtype: _DTypeLike) -> _Array:
2091    """Encode float signal into destination samples, possibly nonlinearly.
2092
2093    Uint destination values are mapped from the range [0.0, 1.0].
2094
2095    Note that non-integer destination types are not clipped to the range [0.0, 1.0].
2096    If that is desired, it can be performed as a postprocess using `output.clip(0.0, 1.0)`.
2097    """
2098
2099
2100class IdentityGamma(Gamma):
2101  """Identity component transfer function."""
2102
2103  def __init__(self) -> None:
2104    super().__init__('identity')
2105
2106  def decode(self, array: _Array, /, dtype: _DTypeLike = np.float32) -> _Array:
2107    dtype = np.dtype(dtype)
2108    assert np.issubdtype(dtype, np.inexact)
2109    if np.issubdtype(_arr_dtype(array), np.unsignedinteger):
2110      return _to_float_01(array, dtype)
2111    return _arr_astype(array, dtype)
2112
2113  def encode(self, array: _Array, /, dtype: _DTypeLike) -> _Array:
2114    dtype = np.dtype(dtype)
2115    assert np.issubdtype(dtype, np.number)
2116    if np.issubdtype(dtype, np.unsignedinteger):
2117      return _from_float(_arr_clip(array, 0.0, 1.0), dtype)
2118    if np.issubdtype(dtype, np.integer):
2119      return _arr_astype(cast(_Array, array + 0.5), dtype)
2120    return _arr_astype(array, dtype)
2121
2122
2123class PowerGamma(Gamma):
2124  """Gamma correction using a power function."""
2125
2126  def __init__(self, power: float) -> None:
2127    super().__init__(name=f'power_{power}')
2128    self.power = power
2129
2130  def decode(self, array: _Array, /, dtype: _DTypeLike = np.float32) -> _Array:
2131    dtype = np.dtype(dtype)
2132    assert np.issubdtype(dtype, np.floating)
2133    if _arr_dtype(array) == np.uint8 and self.power != 2:
2134      arraylib = _arr_arraylib(array)
2135      decode_table = _make_array(self.decode(np.arange(256, dtype=dtype) / 255), arraylib)
2136      return _arr_getitem(decode_table, array)
2137
2138    array = _to_float_01(array, dtype)
2139    return _arr_square(array) if self.power == 2 else cast(_Array, array**self.power)
2140
2141  def encode(self, array: _Array, /, dtype: _DTypeLike) -> _Array:
2142    array = _arr_clip(array, 0.0, 1.0)
2143    if self.power == 2:
2144      array = _arr_sqrt(array)
2145    else:
2146      array = cast(_Array, array ** (1.0 / self.power))
2147    return _from_float(array, dtype)
2148
2149
2150class SrgbGamma(Gamma):
2151  """Gamma correction using sRGB; see https://en.wikipedia.org/wiki/SRGB."""
2152
2153  def __init__(self) -> None:
2154    super().__init__(name='srgb')
2155
2156  def decode(self, array: _Array, /, dtype: _DTypeLike = np.float32) -> _Array:
2157    dtype = np.dtype(dtype)
2158    assert np.issubdtype(dtype, np.floating)
2159    if _arr_dtype(array) == np.uint8:
2160      arraylib = _arr_arraylib(array)
2161      decode_table = _make_array(self.decode(np.arange(256, dtype=dtype) / 255), arraylib)
2162      return _arr_getitem(decode_table, array)
2163
2164    x = _to_float_01(array, dtype)
2165    return _arr_where(cast(_Array, x > 0.04045), ((x + 0.055) / 1.055) ** 2.4, x / 12.92)
2166
2167  def encode(self, array: _Array, /, dtype: _DTypeLike) -> _Array:
2168    x = _arr_clip(array, 0.0, 1.0)
2169    # Unfortunately, exponentiation is slow, and np.digitize() is even slower.
2170    x = _arr_where(
2171        cast(_Array, x > 0.0031308), x ** (1.0 / 2.4) * 1.055 - (0.055 - 1e-17), x * 12.92
2172    )
2173    return _from_float(x, dtype)
2174
2175
2176_DICT_GAMMAS = {
2177    'identity': IdentityGamma(),
2178    'power2': PowerGamma(2.0),
2179    'power22': PowerGamma(2.2),
2180    'srgb': SrgbGamma(),
2181}
2182
2183GAMMAS = list(_DICT_GAMMAS)
2184r"""Shortcut names for some predefined gamma-correction schemes:
2185
2186| name | `Gamma` | Decoding function<br/> (linear space from stored value) | Encoding function<br/> (stored value from linear space) |
2187|---|---|:---:|:---:|
2188| `'identity'` | `IdentityGamma()` | $l = e$ | $e = l$ |
2189| `'power2'` | `PowerGamma`(2.0) | $l = e^{2.0}$ | $e = l^{1/2.0}$ |
2190| `'power22'` | `PowerGamma`(2.2) | $l = e^{2.2}$ | $e = l^{1/2.2}$ |
2191| `'srgb'` | `SrgbGamma()` | $l = \left(\left(e + 0.055\right) / 1.055\right)^{2.4}$ | $e = l^{1/2.4} * 1.055 - 0.055$ |
2192
2193See the source code for extensibility.
2194"""
2195
2196
2197def _get_gamma(gamma: str | Gamma, /) -> Gamma:
2198  """Return a `Gamma`, which can be specified as a name in `GAMMAS`."""
2199  return gamma if isinstance(gamma, Gamma) else _DICT_GAMMAS[gamma]
2200
2201
2202def _get_src_dst_gamma(
2203    gamma: str | Gamma | None,
2204    src_gamma: str | Gamma | None,
2205    dst_gamma: str | Gamma | None,
2206    src_dtype: _DType,
2207    dst_dtype: _DType,
2208) -> tuple[Gamma, Gamma]:
2209  if gamma is None and src_gamma is None and dst_gamma is None:
2210    src_uint = np.issubdtype(src_dtype, np.unsignedinteger)
2211    dst_uint = np.issubdtype(dst_dtype, np.unsignedinteger)
2212    if src_uint and dst_uint:
2213      # The default might ideally be 'srgb' but that conversion is costlier.
2214      gamma = 'power2'
2215    elif not src_uint and not dst_uint:
2216      gamma = 'identity'
2217    else:
2218      raise ValueError(f'Gamma must be specified because {src_dtype=} and {dst_dtype=}.')
2219  if gamma is not None:
2220    if src_gamma is not None:
2221      raise ValueError('Cannot specify both gamma and src_gamma.')
2222    if dst_gamma is not None:
2223      raise ValueError('Cannot specify both gamma and dst_gamma.')
2224    src_gamma = dst_gamma = gamma
2225  assert src_gamma and dst_gamma
2226  src_gamma = _get_gamma(src_gamma)
2227  dst_gamma = _get_gamma(dst_gamma)
2228  return src_gamma, dst_gamma
2229
2230
2231def _create_resize_matrix(
2232    src_size: int,
2233    dst_size: int,
2234    src_gridtype: Gridtype,
2235    dst_gridtype: Gridtype,
2236    boundary: Boundary,
2237    filter: Filter,
2238    prefilter: Filter | None = None,
2239    scale: float = 1.0,
2240    translate: float = 0.0,
2241    dtype: _DTypeLike = np.float64,
2242    arraylib: str = 'numpy',
2243) -> tuple[Any, _AnyArray | None]:
2244  """Compute affine weights for 1D resampling from `src_size` to `dst_size`.
2245
2246  Compute a sparse matrix in which each row expresses a destination sample value as a combination
2247  of source sample values depending on the boundary rule.  If the combination is non-affine,
2248  the remainder (returned as `cval_weight`) is the contribution of the special constant value
2249  (cval) defined outside the domain.
2250
2251  Args:
2252    src_size: The number of samples within the source 1D domain.
2253    dst_size: The number of samples within the destination 1D domain.
2254    src_gridtype: Placement of the samples in the source domain grid.
2255    dst_gridtype: Placement of the output samples in the destination domain grid.
2256    boundary: The reconstruction boundary rule.
2257    filter: The reconstruction kernel (used for upsampling/magnification).
2258    prefilter: The prefilter kernel (used for downsampling/minification).  If it is `None`,
2259      `filter` is used.
2260    scale: Scaling factor applied when mapping the source domain onto the destination domain.
2261    translate: Offset applied when mapping the scaled source domain onto the destination domain.
2262    dtype: Precision of computed resize matrix entries.
2263    arraylib: Representation of output.  Must be an element of `ARRAYLIBS`.
2264
2265  Returns:
2266    sparse_matrix: Matrix whose rows express output sample values as affine combinations of the
2267      source sample values.
2268    cval_weight: Optional vector expressing the additional contribution of the constant value
2269      (`cval`) to the combination in each row of `sparse_matrix`.  It equals one minus the sum of
2270      the weights in each matrix row.
2271  """
2272  if src_size < src_gridtype.min_size():
2273    raise ValueError(f'Source size {src_size} is too small for resize.')
2274  if dst_size < dst_gridtype.min_size():
2275    raise ValueError(f'Destination size {dst_size} is too small for resize.')
2276  prefilter = filter if prefilter is None else prefilter
2277  dtype = np.dtype(dtype)
2278  assert np.issubdtype(dtype, np.floating)
2279
2280  scaling = dst_gridtype.size_in_samples(dst_size) / src_gridtype.size_in_samples(src_size) * scale
2281  is_minification = scaling < 1.0
2282  filter = prefilter if is_minification else filter
2283  if filter.name == 'trapezoid':
2284    radius = 0.5 + 0.5 * min(scaling, 1.0 / scaling)
2285    filter = TrapezoidFilter(radius=radius)
2286  radius = filter.radius
2287  num_samples = int(np.ceil(radius * 2 / scaling) if is_minification else np.ceil(radius * 2))
2288
2289  dst_index = np.arange(dst_size, dtype=dtype)
2290  # Destination sample locations in unit domain [0, 1].
2291  dst_position = dst_gridtype.point_from_index(dst_index, dst_size)
2292
2293  src_position = (dst_position - translate) / scale
2294  src_position = boundary.preprocess_coordinates(src_position)
2295
2296  # Sample positions mapped back to source unit domain [0, 1].
2297  src_float_index = src_gridtype.index_from_point(src_position, src_size)
2298  src_first_index = (
2299      np.floor(src_float_index + (0.5 if num_samples % 2 == 1 else 0.0)).astype(np.int32)
2300      - (num_samples - 1) // 2
2301  )
2302
2303  sample_index = np.arange(num_samples, dtype=np.int32)
2304  src_index = src_first_index[:, None] + sample_index  # (dst_size, num_samples)
2305
2306  def get_weight_matrix() -> _NDArray:
2307    if filter.name == 'impulse':
2308      return np.ones(src_index.shape, dtype)
2309    if is_minification:
2310      x = (src_float_index[:, None] - src_index.astype(dtype)) * scaling
2311      return filter(x) * scaling
2312    # Either same size or magnification.
2313    x = src_float_index[:, None] - src_index.astype(dtype)
2314    return filter(x)
2315
2316  weight = get_weight_matrix().astype(dtype, copy=False)
2317
2318  if filter.name != 'narrowbox' and (is_minification or not filter.partition_of_unity):
2319    weight = weight / weight.sum(axis=-1)[..., None]
2320
2321  src_index, weight = boundary.apply(src_index, weight, src_position, src_size, src_gridtype)
2322  shape = dst_size, src_size
2323
2324  def prepare_sparse_resize_matrix() -> tuple[_NDArray, _NDArray, _NDArray]:
2325    linearized = (src_index + np.indices(src_index.shape)[0] * src_size).ravel()
2326    values = weight.ravel()
2327    # Remove the zero weights.
2328    nonzero = values != 0.0
2329    linearized, values = linearized[nonzero], values[nonzero]
2330    # Sort and merge the duplicate indices.
2331    unique, unique_inverse = np.unique(linearized, return_inverse=True)
2332    data2 = np.ones(len(linearized), np.float32)
2333    row_ind2 = unique_inverse
2334    col_ind2 = np.arange(len(linearized))
2335    shape2 = len(unique), len(linearized)
2336    csr = scipy.sparse.csr_matrix((data2, (row_ind2, col_ind2)), shape=shape2)
2337    data = csr * values  # Merged values.
2338    row_ind, col_ind = unique // src_size, unique % src_size  # Merged indices.
2339    return data, row_ind, col_ind
2340
2341  data, row_ind, col_ind = prepare_sparse_resize_matrix()
2342  resize_matrix = _make_sparse_matrix(data, row_ind, col_ind, shape, arraylib)
2343
2344  uses_cval = boundary.uses_cval or filter.name == 'narrowbox'
2345  cval_weight = _make_array(1.0 - weight.sum(axis=-1), arraylib) if uses_cval else None
2346
2347  return resize_matrix, cval_weight
2348
2349
2350def _apply_digital_filter_1d(
2351    array: _Array,
2352    gridtype: Gridtype,
2353    boundary: Boundary,
2354    cval: _ArrayLike,
2355    filter: Filter,
2356    /,
2357    *,
2358    axis: int = 0,
2359) -> _Array:
2360  """Apply inverse convolution to the specified dimension of the array.
2361
2362  Find the array coefficients such that convolution with the (continuous) filter (given
2363  gridtype and boundary) interpolates the original array values.
2364  """
2365  assert filter.requires_digital_filter
2366  arraylib = _arr_arraylib(array)
2367
2368  if arraylib == 'torch':
2369    import torch.autograd
2370
2371    class InverseConvolution(torch.autograd.Function):  # type: ignore[misc] # pylint: disable=abstract-method
2372      """Differentiable wrapper for _apply_digital_filter_1d."""
2373
2374      @staticmethod
2375      # pylint: disable-next=arguments-differ
2376      def forward(ctx: Any, *args: _TorchTensor, **kwargs: Any) -> _TorchTensor:
2377        del ctx
2378        assert not kwargs
2379        (x,) = args
2380        a = _apply_digital_filter_1d_numpy(
2381            x.detach().numpy(), gridtype, boundary, cval, filter, axis, False
2382        )
2383        return torch.as_tensor(a)
2384
2385      @staticmethod
2386      def backward(ctx: Any, *grad_outputs: _TorchTensor) -> _TorchTensor:
2387        del ctx
2388        (grad_output,) = grad_outputs
2389        a = _apply_digital_filter_1d_numpy(
2390            grad_output.detach().numpy(), gridtype, boundary, cval, filter, axis, True
2391        )
2392        return torch.as_tensor(a)
2393
2394    return InverseConvolution.apply(array)
2395
2396  if arraylib == 'jax':
2397    import jax
2398    import jax.numpy as jnp
2399    # It seems rather difficult to implement this digital filter (inverse convolution) in Jax.
2400    # https://jax.readthedocs.io/en/latest/jax.scipy.html sadly omits scipy.signal.filtfilt().
2401    # To include a (non-traceable) numpy function in Jax requires jax.experimental.host_callback
2402    # and/or defining a new jax.core.Primitive (which allows differentiability).  See
2403    # https://github.com/google/jax/issues/1142#issuecomment-544286585
2404    # https://github.com/google/jax/blob/main/docs/notebooks/How_JAX_primitives_work.ipynb  :-(
2405    # https://github.com/google/jax/issues/5934
2406
2407    @jax.custom_gradient  # type: ignore[untyped-decorator]
2408    def jax_inverse_convolution(x: _JaxArray) -> Any:
2409      # This function is not jax-traceable due to the presence of to_py(), so jit and grad fail.
2410      x_py = np.asarray(x)  # to_py() deprecated.
2411      a = _apply_digital_filter_1d_numpy(x_py, gridtype, boundary, cval, filter, axis, False)
2412      y = jnp.asarray(a)
2413
2414      def grad(grad_output: _JaxArray) -> _JaxArray:
2415        grad_output_py = np.asarray(grad_output)  # to_py() deprecated.
2416        a = _apply_digital_filter_1d_numpy(
2417            grad_output_py, gridtype, boundary, cval, filter, axis, True
2418        )
2419        return jnp.asarray(a)
2420
2421      return y, grad
2422
2423    return jax_inverse_convolution(array)
2424
2425  assert arraylib == 'numpy'
2426  array_np: Any = array
2427  result: Any = _apply_digital_filter_1d_numpy(
2428      array_np, gridtype, boundary, cval, filter, axis, False
2429  )
2430  return result
2431
2432
2433def _apply_digital_filter_1d_numpy(
2434    array: _NDArray,
2435    gridtype: Gridtype,
2436    boundary: Boundary,
2437    cval: _ArrayLike,
2438    filter: Filter,
2439    axis: int,
2440    compute_backward: bool,
2441    /,
2442) -> _NDArray:
2443  """Version of _apply_digital_filter_1d` specialized to numpy array."""
2444  assert np.issubdtype(array.dtype, np.inexact)
2445  cval = np.asarray(cval).astype(array.dtype, copy=False)
2446
2447  # Use fast spline_filter1d() if we have a compatible gridtype, boundary, and filter:
2448  mode = {
2449      ('reflect', 'dual'): 'reflect',
2450      ('reflect', 'primal'): 'mirror',
2451      ('wrap', 'dual'): 'grid-wrap',
2452      ('wrap', 'primal'): 'wrap',
2453  }.get((boundary.name, gridtype.name))
2454  filter_is_compatible = isinstance(filter, CardinalBsplineFilter)
2455  use_split_filter1d = filter_is_compatible and mode
2456  if use_split_filter1d:
2457    assert isinstance(filter, CardinalBsplineFilter)  # Help mypy.
2458    assert filter.degree >= 2
2459    # compute_backward=True is same: matrix is symmetric and cval is unused.
2460    return scipy.ndimage.spline_filter1d(
2461        array, axis=axis, order=filter.degree, mode=mode, output=array.dtype
2462    )
2463
2464  array_dim = np.moveaxis(array, axis, 0)
2465  l = original_l = math.ceil(filter.radius) - 1
2466  x = np.arange(-l, l + 1, dtype=array.real.dtype)
2467  values = filter(x)
2468  size = array_dim.shape[0]
2469  src_index = np.arange(size)[:, None] + np.arange(len(values)) - l
2470  weight: _NDArray = np.full((size, len(values)), values)
2471  src_position = np.broadcast_to(0.5, len(values))
2472  src_index, weight = boundary.apply(src_index, weight, src_position, size, gridtype)
2473  if gridtype.name == 'primal' and boundary.name == 'wrap':
2474    # Overwrite redundant last row to preserve unreferenced last sample and thereby make the
2475    # matrix non-singular.
2476    src_index[-1] = [size - 1] + [0] * (src_index.shape[1] - 1)
2477    weight[-1] = [1.0] + [0.0] * (weight.shape[1] - 1)
2478  bandwidth = abs(src_index - np.arange(size)[:, None]).max()
2479  is_banded = bandwidth <= l + 1  # Add one for quadratic boundary and l == 1.
2480  # Currently, matrix is always banded unless boundary.name == 'wrap'.
2481
2482  data = weight.reshape(-1).astype(array.dtype, copy=False)
2483  row_ind = np.arange(size).repeat(src_index.shape[1])
2484  col_ind = src_index.reshape(-1)
2485  matrix = scipy.sparse.csr_matrix((data, (row_ind, col_ind)), shape=(size, size))
2486  if compute_backward:
2487    matrix = matrix.transpose()
2488
2489  if boundary.uses_cval and not compute_backward:
2490    cval_weight = 1.0 - np.asarray(matrix.sum(axis=-1))[:, 0]
2491    if array_dim.ndim == 2:  # Handle the case that we have array_flat.
2492      cval = np.tile(cval.reshape(-1), array_dim[0].size // cval.size)
2493    array_dim = array_dim - cval_weight.reshape(-1, *(1,) * array_dim[0].ndim) * cval
2494
2495  array_flat = array_dim.reshape(array_dim.shape[0], -1)
2496
2497  if is_banded:
2498    matrix = matrix.todia()
2499    assert np.all(np.diff(matrix.offsets) == 1)  # Consecutive, often [-l, l].
2500    l, u = -matrix.offsets[0], matrix.offsets[-1]
2501    assert l <= original_l + 1 and u <= original_l + 1, (l, u, original_l)
2502    options = dict(check_finite=False, overwrite_ab=True, overwrite_b=False)
2503    if _is_symmetric(matrix):
2504      array_flat = scipy.linalg.solveh_banded(matrix.data[-1 : l - 1 : -1], array_flat, **options)
2505    else:
2506      array_flat = scipy.linalg.solve_banded((l, u), matrix.data[::-1], array_flat, **options)
2507
2508  else:
2509    lu = scipy.sparse.linalg.splu(matrix.tocsc(), permc_spec='NATURAL')
2510    assert all(s <= size * len(values) for s in (lu.L.nnz, lu.U.nnz))  # Sparse.
2511    array_flat = lu.solve(array_flat)
2512
2513  array_dim = np.asarray(array_flat).reshape(array_dim.shape)
2514  return np.moveaxis(array_dim, 0, axis)
2515
2516
2517def resize(
2518    array: _Array,
2519    /,
2520    shape: Iterable[int],
2521    *,
2522    gridtype: str | Gridtype | None = None,
2523    src_gridtype: str | Gridtype | Iterable[str | Gridtype] | None = None,
2524    dst_gridtype: str | Gridtype | Iterable[str | Gridtype] | None = None,
2525    boundary: str | Boundary | Iterable[str | Boundary] = 'auto',
2526    cval: _ArrayLike = 0.0,
2527    filter: str | Filter | Iterable[str | Filter] = _DEFAULT_FILTER,
2528    prefilter: str | Filter | Iterable[str | Filter] | None = None,
2529    gamma: str | Gamma | None = None,
2530    src_gamma: str | Gamma | None = None,
2531    dst_gamma: str | Gamma | None = None,
2532    scale: float | Iterable[float] = 1.0,
2533    translate: float | Iterable[float] = 0.0,
2534    precision: _DTypeLike | None = None,
2535    dtype: _DTypeLike | None = None,
2536    dim_order: Iterable[int] | None = None,
2537    num_threads: int | Literal['auto'] = 'auto',
2538) -> _Array:
2539  """Resample `array` (a grid of sample values) onto a grid with resolution `shape`.
2540
2541  The source `array` is any object recognized by `ARRAYLIBS`.  It is interpreted as a grid
2542  with `len(shape)` domain coordinate dimensions, where each grid sample value has shape
2543  `array.shape[len(shape):]`.
2544
2545  Some examples:
2546
2547  - A grayscale image has `array.shape = height, width` and resizing it with `len(shape) == 2`
2548    produces a new image of scalar values.
2549  - An RGB image has `array.shape = height, width, 3` and resizing it with `len(shape) == 2`
2550    produces a new image of RGB values.
2551  - An 3D grid of 3x3 Jacobians has `array.shape = Z, Y, X, 3, 3` and resizing it with
2552    `len(shape) == 3` produces a new 3D grid of Jacobians.
2553
2554  This function also allows scaling and translation from the source domain to the output domain
2555  through the parameters `scale` and `translate`.  For more general transforms, see `resample`.
2556
2557  Args:
2558    array: Regular grid of source sample values, as an array object recognized by `ARRAYLIBS`.
2559      The array must have numeric type.  Its first `len(shape)` dimensions are the domain
2560      coordinate dimensions.  Each grid dimension must be at least 1 for a `'dual'` grid or
2561      at least 2 for a `'primal'` grid.
2562    shape: The number of grid samples in each coordinate dimension of the output array.  The source
2563      `array` must have at least as many dimensions as `len(shape)`.
2564    gridtype: Placement of samples on all dimensions of both the source and output domain grids,
2565      specified as either a name in `GRIDTYPES` or a `Gridtype` instance.  It defaults to `'dual'`
2566      if `gridtype`, `src_gridtype`, and `dst_gridtype` are all kept `None`.
2567    src_gridtype: Placement of the samples in the source domain grid for each dimension.
2568      Parameters `gridtype` and `src_gridtype` cannot both be set.
2569    dst_gridtype: Placement of the samples in the output domain grid for each dimension.
2570      Parameters `gridtype` and `dst_gridtype` cannot both be set.
2571    boundary: The reconstruction boundary rule for each dimension in `shape`, specified as either
2572      a name in `BOUNDARIES` or a `Boundary` instance.  The special value `'auto'` uses `'reflect'`
2573      for upsampling and `'clamp'` for downsampling.
2574    cval: Constant value used beyond the samples by some boundary rules.  It must be broadcastable
2575      onto `array.shape[len(shape):]`.  It is subject to `src_gamma`.
2576    filter: The reconstruction kernel for each dimension in `shape`, specified as either a filter
2577      name in `FILTERS` or a `Filter` instance.  It is used during upsampling (i.e., magnification).
2578    prefilter: The prefilter kernel for each dimension in `shape`, specified as either a filter
2579      name in `FILTERS` or a `Filter` instance.  It is used during downsampling
2580      (i.e., minification).  If `None`, it inherits the value of `filter`.  The default
2581      `'lanczos3'` is good for natural images.  For vector graphics images, `'trapezoid'` is better
2582      because it avoids ringing artifacts.
2583    gamma: Component transfer functions (e.g., gamma correction) applied when reading samples from
2584      `array` and when creating output grid samples.  It is specified as either a name in `GAMMAS`
2585      or a `Gamma` instance.  If both `array.dtype` and `dtype` are `uint`, the default is
2586      `'power2'`.  If both are non-`uint`, the default is `'identity'`.  Otherwise, `gamma` or
2587      `src_gamma`/`dst_gamma` must be set.   Gamma correction assumes that float values are in the
2588      range [0.0, 1.0].
2589    src_gamma: Component transfer function used to "decode" `array` samples.
2590      Parameters `gamma` and `src_gamma` cannot both be set.
2591    dst_gamma: Component transfer function used to "encode" the output samples.
2592      Parameters `gamma` and `dst_gamma` cannot both be set.
2593    scale: Scaling factor applied to each dimension of the source domain when it is mapped onto
2594      the destination domain.
2595    translate: Offset applied to each dimension of the scaled source domain when it is mapped onto
2596      the destination domain.
2597    precision: Inexact precision of intermediate computations.  If `None`, it is determined based
2598      on `array.dtype` and `dtype`.
2599    dtype: Desired data type of the output array.  If `None`, it is taken to be `array.dtype`.
2600      If it is a uint type, the intermediate float values are rescaled from the [0.0, 1.0] range
2601      to the uint range.
2602    dim_order: Override the automatically selected order in which the grid dimensions are resized.
2603      Must contain a permutation of `range(len(shape))`.
2604    num_threads: Used to determine multithread parallelism if `array` is from `numpy`.  If set to
2605      `'auto'`, it is selected automatically.  Otherwise, it must be a positive integer.
2606
2607  Returns:
2608    An array of the same class as the source `array`, with shape `shape + array.shape[len(shape):]`
2609      and data type `dtype`.
2610
2611  **Example of image upsampling:**
2612
2613  >>> array = np.random.default_rng(1).random((4, 6, 3))  # 4x6 RGB image.
2614  >>> upsampled = resize(array, (128, 192))  # To 128x192 resolution.
2615
2616  <center>
2617  <img src="https://github.com/hhoppe/resampler/raw/main/media/example_array_upsampled.png"/>
2618  </center>
2619
2620  **Example of image downsampling:**
2621
2622  >>> yx = (np.moveaxis(np.indices((96, 192)), 0, -1) + (0.5, 0.5)) / 96
2623  >>> radius = np.linalg.norm(yx - (0.75, 0.5), axis=-1)
2624  >>> array = np.cos((radius + 0.1) ** 0.5 * 70.0) * 0.5 + 0.5
2625  >>> downsampled = resize(array, (24, 48))
2626
2627  <center>
2628  <img src="https://github.com/hhoppe/resampler/raw/main/media/example_array_downsampled2.png"/>
2629  </center>
2630
2631  **Unit test:**
2632
2633  >>> result = resize(np.array([1.0, 4.0, 5.0]), shape=(4,))
2634  >>> assert np.allclose(result, [0.74240461, 2.88088827, 4.68647155, 5.02641199])
2635  """
2636  arraylib = _arr_arraylib(array)
2637  array_dtype = _arr_dtype(array)
2638  if not np.issubdtype(array_dtype, np.number):
2639    raise ValueError(f'Type {array_dtype} is not numeric.')
2640  shape2 = tuple(shape)
2641  array_ndim = len(_arr_shape(array))
2642  if not 0 < len(shape2) <= array_ndim:
2643    raise ValueError(f'Shape {_arr_shape(array)} cannot be resized to {shape2}.')
2644  src_shape = _arr_shape(array)[: len(shape2)]
2645  src_gridtype2, dst_gridtype2 = _get_gridtypes(
2646      gridtype, src_gridtype, dst_gridtype, len(shape2), len(shape2)
2647  )
2648  boundary2 = np.broadcast_to(np.array(boundary), len(shape2))
2649  cval = np.broadcast_to(cval, _arr_shape(array)[len(shape2) :])
2650  prefilter = filter if prefilter is None else prefilter
2651  filter2 = [_get_filter(f) for f in np.broadcast_to(np.array(filter), len(shape2))]
2652  prefilter2 = [_get_filter(f) for f in np.broadcast_to(np.array(prefilter), len(shape2))]
2653  dtype = array_dtype if dtype is None else np.dtype(dtype)
2654  src_gamma2, dst_gamma2 = _get_src_dst_gamma(gamma, src_gamma, dst_gamma, array_dtype, dtype)
2655  scale2 = np.broadcast_to(np.array(scale), len(shape2))
2656  translate2 = np.broadcast_to(np.array(translate), len(shape2))
2657  del shape, src_gridtype, dst_gridtype, boundary, filter, prefilter
2658  del src_gamma, dst_gamma, scale, translate
2659  precision = _get_precision(precision, [array_dtype, dtype], [])
2660  weight_precision = _real_precision(precision)
2661
2662  is_noop = (
2663      all(src == dst for src, dst in zip(src_shape, shape2, strict=True))
2664      and all(gt1 == gt2 for gt1, gt2 in zip(src_gridtype2, dst_gridtype2, strict=True))
2665      and all(f.interpolating for f in prefilter2)
2666      and np.all(scale2 == 1.0)
2667      and np.all(translate2 == 0.0)
2668      and src_gamma2 == dst_gamma2
2669  )
2670  if is_noop:
2671    return array
2672
2673  if dim_order is None:
2674    dim_order = _arr_best_dims_order_for_resize(array, shape2)
2675  else:
2676    dim_order = tuple(dim_order)
2677    if sorted(dim_order) != list(range(len(shape2))):
2678      raise ValueError(f'{dim_order} not a permutation of {list(range(len(shape2)))}.')
2679
2680  array = src_gamma2.decode(array, precision)
2681  cval = _arr_numpy(src_gamma2.decode(cval, precision))
2682
2683  can_use_fast_box_downsampling = (
2684      _USING_NUMBA
2685      and arraylib == 'numpy'
2686      and len(shape2) == 2
2687      and array_ndim in (2, 3)
2688      and all(src > dst for src, dst in zip(src_shape, shape2, strict=True))
2689      and all(src % dst == 0 for src, dst in zip(src_shape, shape2, strict=True))
2690      and all(gridtype.name == 'dual' for gridtype in src_gridtype2)
2691      and all(gridtype.name == 'dual' for gridtype in dst_gridtype2)
2692      and all(f.name in ('box', 'trapezoid') for f in prefilter2)
2693      and np.all(scale2 == 1.0)
2694      and np.all(translate2 == 0.0)
2695  )
2696  if can_use_fast_box_downsampling:
2697    array2 = _downsample_in_2d_using_box_filter(cast(_NDArray, array), shape2)
2698    return cast(_Array, dst_gamma2.encode(array2, dtype))
2699
2700  # Multidimensional resize can be expressed using einsum() with multiple per-dim resize matrices,
2701  # e.g., as in jax.image.resize().  A benefit is to seek the optimal order of multiplications.
2702  # However, efficiency often requires sparse resize matrices, which are unsupported in einsum().
2703  # Sparse tensors requested for tf.einsum: https://github.com/tensorflow/tensorflow/issues/43497
2704  # https://github.com/tensor-compiler/taco: C++ library that computes tensor algebra expressions
2705  # on sparse and dense tensors; however it does not interoperate with tensorflow, torch, or jax.
2706
2707  for dim in dim_order:
2708    skip_resize_on_this_dim = (
2709        shape2[dim] == _arr_shape(array)[dim]
2710        and scale2[dim] == 1.0
2711        and translate2[dim] == 0.0
2712        and filter2[dim].interpolating
2713    )
2714    if skip_resize_on_this_dim:
2715      continue
2716
2717    def get_is_minification() -> bool:
2718      src_in_samples = src_gridtype2[dim].size_in_samples(_arr_shape(array)[dim])  # noqa: B023
2719      dst_in_samples = dst_gridtype2[dim].size_in_samples(shape2[dim])  # noqa: B023
2720      return dst_in_samples / src_in_samples * scale2[dim] < 1.0  # noqa: B023
2721
2722    is_minification = get_is_minification()
2723    boundary_dim = boundary2[dim]
2724    if boundary_dim == 'auto':
2725      boundary_dim = 'clamp' if is_minification else 'reflect'
2726    boundary_dim = _get_boundary(boundary_dim)
2727    resize_matrix, cval_weight = _create_resize_matrix(
2728        _arr_shape(array)[dim],
2729        shape2[dim],
2730        src_gridtype=src_gridtype2[dim],
2731        dst_gridtype=dst_gridtype2[dim],
2732        boundary=boundary_dim,
2733        filter=filter2[dim],
2734        prefilter=prefilter2[dim],
2735        scale=scale2[dim],
2736        translate=translate2[dim],
2737        dtype=weight_precision,
2738        arraylib=arraylib,
2739    )
2740
2741    array_dim: _Array = _arr_moveaxis(array, dim, 0)
2742    array_flat: Any = _arr_reshape(array_dim, (_arr_shape(array_dim)[0], -1))
2743    array_flat = _arr_possibly_make_contiguous(array_flat)
2744    if not is_minification and filter2[dim].requires_digital_filter:
2745      array_flat = _apply_digital_filter_1d(
2746          array_flat, src_gridtype2[dim], boundary_dim, cval, filter2[dim]
2747      )
2748
2749    array_flat = _arr_matmul_sparse_dense(resize_matrix, array_flat, num_threads=num_threads)
2750    if cval_weight is not None:
2751      cval_flat = np.broadcast_to(cval, _arr_shape(array_dim)[1:]).reshape(-1)
2752      cval_weight2: Any = cval_weight
2753      if np.issubdtype(array_dtype, np.complexfloating):
2754        cval_weight2 = _arr_astype(cval_weight2, array_dtype)  # (Only necessary for 'tensorflow'.)
2755      array_flat += cval_weight2[:, None] * cval_flat
2756
2757    if is_minification and filter2[dim].requires_digital_filter:  # use prefilter2[dim]?
2758      array_flat = _apply_digital_filter_1d(
2759          array_flat, dst_gridtype2[dim], boundary_dim, cval, filter2[dim]
2760      )
2761    array_dim = _arr_reshape(array_flat, (_arr_shape(array_flat)[0], *_arr_shape(array_dim)[1:]))
2762    array = _arr_moveaxis(array_dim, 0, dim)
2763
2764  array = dst_gamma2.encode(cast(_Array, array), dtype)
2765  return array
2766
2767
2768_original_resize = resize
2769
2770
2771def resize_in_arraylib(array: _NDArray, /, *args: Any, arraylib: str, **kwargs: Any) -> _NDArray:
2772  """Evaluate the `resize()` operation using the specified array library from `ARRAYLIBS`."""
2773  _check_eq(_arr_arraylib(array), 'numpy')
2774  return _arr_numpy(_original_resize(_make_array(array, arraylib), *args, **kwargs))
2775
2776
2777def resize_in_numpy(array: _NDArray, /, *args: Any, **kwargs: Any) -> _NDArray:
2778  """Evaluate the `resize()` operation using the `numpy` library."""
2779  return resize_in_arraylib(array, *args, arraylib='numpy', **kwargs)
2780
2781
2782def resize_in_torch(array: _NDArray, /, *args: Any, **kwargs: Any) -> _NDArray:
2783  """Evaluate the `resize()` operation using the `torch` library."""
2784  return resize_in_arraylib(array, *args, arraylib='torch', **kwargs)
2785
2786
2787def resize_in_jax(array: _NDArray, /, *args: Any, **kwargs: Any) -> _NDArray:
2788  """Evaluate the `resize()` operation using the `jax` library."""
2789  return resize_in_arraylib(array, *args, arraylib='jax', **kwargs)
2790
2791
2792def _resize_possibly_in_arraylib(
2793    array: _AnyArray, /, *args: Any, arraylib: str, **kwargs: Any
2794) -> _AnyArray:
2795  """If `array` is from numpy, evaluate `resize()` using the array library from `ARRAYLIBS`."""
2796  if _arr_arraylib(array) == 'numpy':
2797    return _arr_numpy(
2798        _original_resize(_make_array(cast(_NDArray, array), arraylib), *args, **kwargs)
2799    )
2800  return _original_resize(cast(Any, array), *args, **kwargs)
2801
2802
2803@functools.cache
2804def _create_jaxjit_resize() -> Callable[..., Any]:
2805  """Lazily invoke `jax.jit` on `resize`."""
2806  import jax
2807
2808  jitted: Any = jax.jit(
2809      _original_resize,
2810      static_argnums=(1,),
2811      static_argnames=list(_original_resize.__kwdefaults__ or []),
2812  )
2813  return jitted
2814
2815
2816def jaxjit_resize(array: _Array, /, *args: Any, **kwargs: Any) -> _Array:
2817  """Compute `resize` but with resize function jitted using Jax."""
2818  return _create_jaxjit_resize()(array, *args, **kwargs)  # pylint: disable=not-callable
2819
2820
2821def uniform_resize(
2822    array: _Array,
2823    /,
2824    shape: Iterable[int],
2825    *,
2826    object_fit: Literal['contain', 'cover'] = 'contain',
2827    gridtype: str | Gridtype | None = None,
2828    src_gridtype: str | Gridtype | Iterable[str | Gridtype] | None = None,
2829    dst_gridtype: str | Gridtype | Iterable[str | Gridtype] | None = None,
2830    boundary: str | Boundary | Iterable[str | Boundary] = 'natural',  # Instead of 'auto' default.
2831    scale: float | Iterable[float] = 1.0,
2832    translate: float | Iterable[float] = 0.0,
2833    **kwargs: Any,
2834) -> _Array:
2835  """Resample `array` onto a grid with resolution `shape` but with uniform scaling.
2836
2837  Calls function `resize` with `scale` and `translate` set such that the aspect ratio of `array`
2838  is preserved.  The effect is similar to CSS `object-fit: contain`.
2839  The parameter `boundary` (whose default is changed to `'natural'`) determines the values assigned
2840  outside the source domain.
2841
2842  Args:
2843    array: Regular grid of source sample values.
2844    shape: The number of grid samples in each coordinate dimension of the output array.  The source
2845      `array` must have at least as many dimensions as `len(shape)`.
2846    object_fit: Like CSS `object-fit`.  If `'contain'`, `array` is resized uniformly to fit within
2847      `shape`. If `'cover'`, `array` is resized to fully cover `shape`.
2848    gridtype: Placement of samples on all dimensions of both the source and output domain grids.
2849    src_gridtype: Placement of the samples in the source domain grid for each dimension.
2850    dst_gridtype: Placement of the samples in the output domain grid for each dimension.
2851    boundary: The reconstruction boundary rule for each dimension in `shape`, specified as either
2852      a name in `BOUNDARIES` or a `Boundary` instance.  The default is `'natural'`, which assigns
2853      `cval` to output points that map outside the source unit domain.
2854    scale: Parameter may not be specified.
2855    translate: Parameter may not be specified.
2856    **kwargs: Additional parameters for `resize` function (including `cval`).
2857
2858  Returns:
2859    An array with shape `shape + array.shape[len(shape):]`.
2860
2861  >>> uniform_resize(np.ones((2, 2)), (2, 4), filter='trapezoid')
2862  array([[0., 1., 1., 0.],
2863         [0., 1., 1., 0.]])
2864
2865  >>> uniform_resize(np.ones((4, 8)), (2, 7), filter='trapezoid')
2866  array([[0. , 0.5, 1. , 1. , 1. , 0.5, 0. ],
2867         [0. , 0.5, 1. , 1. , 1. , 0.5, 0. ]])
2868
2869  >>> a = np.arange(6.0).reshape(2, 3)
2870  >>> uniform_resize(a, (2, 2), filter='trapezoid', object_fit='cover')
2871  array([[0.5, 1.5],
2872         [3.5, 4.5]])
2873  """
2874  if scale != 1.0 or translate != 0.0:
2875    raise ValueError('`uniform_resize()` does not accept `scale` or `translate` parameters.')
2876  shape = tuple(shape)
2877  array_ndim = len(_arr_shape(array))
2878  if not 0 < len(shape) <= array_ndim:
2879    raise ValueError(f'Shape {_arr_shape(array)} cannot be resized to {shape}.')
2880  src_gridtype2, dst_gridtype2 = _get_gridtypes(
2881      gridtype, src_gridtype, dst_gridtype, len(shape), len(shape)
2882  )
2883  raw_scales = [
2884      dst_gridtype2[dim].size_in_samples(shape[dim])
2885      / src_gridtype2[dim].size_in_samples(_arr_shape(array)[dim])
2886      for dim in range(len(shape))
2887  ]
2888  scale0 = {'contain': min(raw_scales), 'cover': max(raw_scales)}[object_fit]
2889  scale2 = scale0 / np.array(raw_scales)
2890  translate = (1.0 - scale2) / 2
2891  return resize(array, shape, boundary=boundary, scale=scale2, translate=translate, **kwargs)
2892
2893
2894_MAX_BLOCK_SIZE_RECURSING = -999  # Special value to indicate re-invocation on partitioned blocks.
2895
2896
2897def resample(
2898    array: _Array,
2899    /,
2900    coords: _ArrayLike,
2901    *,
2902    gridtype: str | Gridtype | Iterable[str | Gridtype] = 'dual',
2903    boundary: str | Boundary | Iterable[str | Boundary] = 'auto',
2904    cval: _ArrayLike = 0.0,
2905    filter: str | Filter | Iterable[str | Filter] = _DEFAULT_FILTER,
2906    prefilter: str | Filter | Iterable[str | Filter] | None = None,
2907    gamma: str | Gamma | None = None,
2908    src_gamma: str | Gamma | None = None,
2909    dst_gamma: str | Gamma | None = None,
2910    jacobian: _ArrayLike | None = None,
2911    precision: _DTypeLike | None = None,
2912    dtype: _DTypeLike | None = None,
2913    max_block_size: int = 40_000,
2914    debug: bool = False,
2915) -> _Array:
2916  """Interpolate `array` (a grid of samples) at specified unit-domain coordinates `coords`.
2917
2918  The last dimension of `coords` contains unit-domain coordinates at which to interpolate the
2919  domain grid samples in `array`.
2920
2921  The number of coordinates (`coords.shape[-1]`) determines how to interpret `array`: its first
2922  `coords.shape[-1]` dimensions define the grid, and the remaining dimensions describe each grid
2923  sample (e.g., scalar, vector, tensor).
2924
2925  Concretely, the grid has shape `array.shape[:coords.shape[-1]]` and each grid sample has shape
2926  `array.shape[coords.shape[-1]:]`.
2927
2928  Examples include:
2929
2930  - Resample a grayscale image with `array.shape = height, width` onto a new grayscale image with
2931    `new.shape = height2, width2` by using `coords.shape = height2, width2, 2`.
2932
2933  - Resample an RGB image with `array.shape = height, width, 3` onto a new RGB image with
2934    `new.shape = height2, width2, 3` by using `coords.shape = height2, width2, 2`.
2935
2936  - Sample an RGB image at `num` 2D points along a line segment by using `coords.shape = num, 2`.
2937
2938  - Sample an RGB image at a single 2D point by using `coords.shape = (2,)`.
2939
2940  - Sample a 3D grid of 3x3 Jacobians with `array.shape = nz, ny, nx, 3, 3` along a 2D plane by
2941    using `coords.shape = height, width, 3`.
2942
2943  - Map a grayscale image through a color map by using `array.shape = 256, 3` and
2944    `coords.shape = height, width`.
2945
2946  Args:
2947    array: Regular grid of source sample values, as an array object recognized by `ARRAYLIBS`.
2948      The array must have numeric type.  The coordinate dimensions appear first, and
2949      each grid sample may have an arbitrary shape.  Each grid dimension must be at least 1 for
2950      a `'dual'` grid or at least 2 for a `'primal'` grid.
2951    coords: Grid of points at which to resample `array`.  The point coordinates are in the last
2952      dimension of `coords`.  The domain associated with the source grid is a unit hypercube,
2953      i.e. with a range [0, 1] on each coordinate dimension.  The output grid has shape
2954      `coords.shape[:-1]` and each of its grid samples has shape `array.shape[coords.shape[-1]:]`.
2955    gridtype: Placement of the samples in the source domain grid for each dimension, specified as
2956      either a name in `GRIDTYPES` or a `Gridtype` instance.  It defaults to `'dual'`.
2957    boundary: The reconstruction boundary rule for each dimension in `coords.shape[-1]`, specified
2958      as either a name in `BOUNDARIES` or a `Boundary` instance.  The special value `'auto'` uses
2959      `'reflect'` for upsampling and `'clamp'` for downsampling.
2960    cval: Constant value used beyond the samples by some boundary rules.  It must be broadcastable
2961      onto the shape `array.shape[coords.shape[-1]:]`.  It is subject to `src_gamma`.
2962    filter: The reconstruction kernel for each dimension in `coords.shape[-1]`, specified as either
2963      a filter name in `FILTERS` or a `Filter` instance.
2964    prefilter: The prefilter kernel for each dimension in `coords.shape[:-1]`, specified as either
2965      a filter name in `FILTERS` or a `Filter` instance.  It is used during downsampling
2966      (i.e., minification).  If `None`, it inherits the value of `filter`.
2967    gamma: Component transfer functions (e.g., gamma correction) applied when reading samples
2968      from `array` and when creating output grid samples.  It is specified as either a name in
2969      `GAMMAS` or a `Gamma` instance.  If both `array.dtype` and `dtype` are `uint`, the default
2970      is `'power2'`.  If both are non-`uint`, the default is `'identity'`.  Otherwise, `gamma` or
2971      `src_gamma`/`dst_gamma` must be set.   Gamma correction assumes that float values are in the
2972      range [0.0, 1.0].
2973    src_gamma: Component transfer function used to "decode" `array` samples.
2974      Parameters `gamma` and `src_gamma` cannot both be set.
2975    dst_gamma: Component transfer function used to "encode" the output samples.
2976      Parameters `gamma` and `dst_gamma` cannot both be set.
2977    jacobian: Optional array, which must be broadcastable onto the shape
2978      `coords.shape[:-1] + (coords.shape[-1], coords.shape[-1])`, storing for each point in the
2979      output grid the Jacobian matrix of the map from the unit output domain to the unit source
2980      domain.  If omitted, it is estimated by computing finite differences on `coords`.
2981    precision: Inexact precision of intermediate computations.  If `None`, it is determined based
2982      on `array.dtype`, `coords.dtype`, and `dtype`.
2983    dtype: Desired data type of the output array.  If `None`, it is taken to be `array.dtype`.
2984      If it is a uint type, the intermediate float values are rescaled from the [0.0, 1.0] range
2985      to the uint range.
2986    max_block_size: If nonzero, maximum number of grid points in `coords` before the resampling
2987      evaluation gets partitioned into smaller blocks for reduced memory usage and better caching.
2988    debug: Show internal information.
2989
2990  Returns:
2991    A new sample grid of shape `coords.shape[:-1]`, represented as an array of shape
2992    `coords.shape[:-1] + array.shape[coords.shape[-1]:]`, of the same array library type as
2993    the source array.
2994
2995  **Example of resample operation:**
2996
2997  <center>
2998  <img src="https://github.com/hhoppe/resampler/raw/main/media/example_warp_coords.png"/>
2999  </center>
3000
3001  For reference, the identity resampling for a scalar-valued grid with the default grid-type
3002  `'dual'` is:
3003
3004  >>> array = np.random.default_rng(1).random((5, 7, 3))
3005  >>> coords = (np.moveaxis(np.indices(array.shape), 0, -1) + 0.5) / array.shape
3006  >>> new_array = resample(array, coords)
3007  >>> assert np.allclose(new_array, array)
3008
3009  It is more efficient to use the function `resize` for the special case where the `coords` are
3010  obtained as simple scaling and translation of a new regular grid over the source domain:
3011
3012  >>> scale, translate, new_shape = (1.1, 1.2), (0.1, -0.2), (6, 8)
3013  >>> coords = (np.moveaxis(np.indices(new_shape), 0, -1) + 0.5) / new_shape
3014  >>> coords = (coords - translate) / scale
3015  >>> resampled = resample(array, coords)
3016  >>> resized = resize(array, new_shape, scale=scale, translate=translate)
3017  >>> assert np.allclose(resampled, resized)
3018  """
3019  arraylib = _arr_arraylib(array)
3020  if len(_arr_shape(array)) == 0:
3021    array = _arr_reshape(array, (1,))
3022  coords = np.atleast_1d(coords)
3023  if not np.issubdtype(_arr_dtype(array), np.number):
3024    raise ValueError(f'Type {_arr_dtype(array)} is not numeric.')
3025  if not np.issubdtype(coords.dtype, np.floating):
3026    raise ValueError(f'Type {coords.dtype} is not floating.')
3027  array_ndim = len(_arr_shape(array))
3028  if coords.ndim == 1 and coords.shape[0] > 1 and array_ndim == 1:
3029    coords = coords[:, None]
3030  grid_ndim = coords.shape[-1]
3031  grid_shape = _arr_shape(array)[:grid_ndim]
3032  sample_shape = _arr_shape(array)[grid_ndim:]
3033  resampled_ndim = coords.ndim - 1
3034  resampled_shape = coords.shape[:-1]
3035  if grid_ndim > array_ndim:
3036    raise ValueError(
3037        f'There are more coordinate dimensions ({grid_ndim}) in {coords=}'
3038        f' than in array.shape={_arr_shape(array)}.'
3039    )
3040  gridtype2 = [_get_gridtype(g) for g in np.broadcast_to(np.array(gridtype), grid_ndim)]
3041  boundary2 = np.broadcast_to(np.array(boundary), grid_ndim).tolist()
3042  cval = np.broadcast_to(cval, sample_shape)
3043  prefilter = filter if prefilter is None else prefilter
3044  filter2 = [_get_filter(f) for f in np.broadcast_to(np.array(filter), grid_ndim)]
3045  prefilter2 = [_get_filter(f) for f in np.broadcast_to(np.array(prefilter), resampled_ndim)]
3046  dtype = _arr_dtype(array) if dtype is None else np.dtype(dtype)
3047  src_gamma2, dst_gamma2 = _get_src_dst_gamma(gamma, src_gamma, dst_gamma, _arr_dtype(array), dtype)
3048  del gridtype, boundary, filter, prefilter, src_gamma, dst_gamma
3049  if jacobian is not None:
3050    jacobian = np.broadcast_to(jacobian, resampled_shape + (coords.shape[-1],) * 2)
3051  precision = _get_precision(precision, [_arr_dtype(array), dtype], [coords.dtype])
3052  weight_precision = _real_precision(precision)
3053  coords = coords.astype(weight_precision, copy=False)
3054  is_minification = False  # Current limitation; no prefiltering!
3055  assert max_block_size >= 0 or max_block_size == _MAX_BLOCK_SIZE_RECURSING
3056  for dim in range(grid_ndim):
3057    if boundary2[dim] == 'auto':
3058      boundary2[dim] = 'clamp' if is_minification else 'reflect'
3059    boundary2[dim] = _get_boundary(boundary2[dim])
3060
3061  if max_block_size != _MAX_BLOCK_SIZE_RECURSING:
3062    array = src_gamma2.decode(array, precision)
3063    for dim in range(grid_ndim):
3064      assert not is_minification
3065      if filter2[dim].requires_digital_filter:
3066        array = _apply_digital_filter_1d(
3067            array, gridtype2[dim], boundary2[dim], cval, filter2[dim], axis=dim
3068        )
3069    cval = _arr_numpy(src_gamma2.decode(cval, precision))
3070
3071  if math.prod(resampled_shape) > max_block_size > 0:
3072    block_shape = _block_shape_with_min_size(resampled_shape, max_block_size)
3073    if debug:
3074      print(f'(resample: splitting coords into blocks {block_shape}).')
3075    coord_blocks = _split_array_into_blocks(coords, block_shape)
3076
3077    def process_block(coord_block: _NDArray) -> _Array:
3078      return resample(
3079          cast(Any, array),
3080          coord_block,
3081          gridtype=gridtype2,
3082          boundary=boundary2,
3083          cval=cval,
3084          filter=filter2,
3085          prefilter=prefilter2,
3086          src_gamma='identity',
3087          dst_gamma=dst_gamma2,
3088          jacobian=jacobian,
3089          precision=precision,
3090          dtype=dtype,
3091          max_block_size=_MAX_BLOCK_SIZE_RECURSING,
3092      )
3093
3094    result_blocks = _map_function_over_blocks(coord_blocks, process_block)
3095    array = _merge_array_from_blocks(result_blocks)
3096    return array
3097
3098  # A concrete example of upsampling:
3099  #   array = np.ones((5, 7, 3))  # source RGB image has height=5 width=7
3100  #   coords = np.random.default_rng(1).random((8, 9, 2))  # output RGB image has height=8 width=9
3101  #   resample(array, coords, filter=('cubic', 'lanczos3'))
3102  #   grid_shape = 5, 7  grid_ndim = 2
3103  #   resampled_shape = 8, 9  resampled_ndim = 2
3104  #   sample_shape = (3,)
3105  #   src_float_index.shape = 8, 9
3106  #   src_first_index.shape = 8, 9
3107  #   sample_index.shape = (4,) for dim == 0, then (6,) for dim == 1
3108  #   weight = [shape(8, 9, 4), shape(8, 9, 6)]
3109  #   src_index = [shape(8, 9, 4), shape(8, 9, 6)]
3110
3111  # Both:[shape(8, 9, 4), shape(8, 9, 6)]
3112  weight: list[_NDArray] = [np.array([]) for _ in range(grid_ndim)]
3113  src_index: list[_NDArray] = [np.array([]) for _ in range(grid_ndim)]
3114  uses_cval = False
3115  all_num_samples = []  # will be [4, 6]
3116
3117  for dim in range(grid_ndim):
3118    src_size = grid_shape[dim]  # scalar
3119    coords_dim = coords[..., dim]  # (8, 9)
3120    radius = filter2[dim].radius  # scalar
3121    num_samples = int(np.ceil(radius * 2))  # scalar
3122    all_num_samples.append(num_samples)
3123
3124    boundary_dim = boundary2[dim]
3125    coords_dim = boundary_dim.preprocess_coordinates(coords_dim)
3126
3127    # Sample positions mapped back to source unit domain [0, 1].
3128    src_float_index = gridtype2[dim].index_from_point(coords_dim, src_size)  # (8, 9)
3129    src_first_index = (
3130        np.floor(src_float_index + (0.5 if num_samples % 2 == 1 else 0.0)).astype(np.int32)
3131        - (num_samples - 1) // 2
3132    )  # (8, 9)
3133
3134    sample_index = np.arange(num_samples, dtype=np.int32)  # (4,) then (6,)
3135    src_index[dim] = src_first_index[..., None] + sample_index  # (8, 9, 4) then (8, 9, 6)
3136    if filter2[dim].name == 'trapezoid':
3137      # (It might require changing the filter radius at every sample.)
3138      raise ValueError('resample() cannot use adaptive `trapezoid` filter.')
3139    if filter2[dim].name == 'impulse':
3140      weight[dim] = np.ones_like(src_index[dim], weight_precision)
3141    else:
3142      x = src_float_index[..., None] - src_index[dim].astype(weight_precision)
3143      weight[dim] = filter2[dim](x).astype(weight_precision, copy=False)
3144      if filter2[dim].name != 'narrowbox' and (
3145          is_minification or not filter2[dim].partition_of_unity
3146      ):
3147        weight[dim] = weight[dim] / weight[dim].sum(axis=-1)[..., None]
3148
3149    src_index[dim], weight[dim] = boundary_dim.apply(
3150        src_index[dim], weight[dim], coords_dim, src_size, gridtype2[dim]
3151    )
3152    if boundary_dim.uses_cval or filter2[dim].name == 'narrowbox':
3153      uses_cval = True
3154
3155  # Gather the samples.
3156
3157  # Recall that src_index = [shape(8, 9, 4), shape(8, 9, 6)].
3158  src_index_expanded = []
3159  for dim in range(grid_ndim):
3160    src_index_dim = np.moveaxis(
3161        src_index[dim].reshape(src_index[dim].shape + (1,) * (grid_ndim - 1)),
3162        resampled_ndim,
3163        resampled_ndim + dim,
3164    )
3165    src_index_expanded.append(src_index_dim)
3166  indices = tuple(src_index_expanded)  # (shape(8, 9, 4, 1), shape(8, 9, 1, 6))
3167  samples = _arr_getitem(array, indices)  # (8, 9, 4, 6, 3)
3168
3169  # Indirectly derive samples.ndim (which is unavailable during Tensorflow grad computation).
3170  samples_ndim = resampled_ndim + grid_ndim + len(sample_shape)
3171
3172  # Compute an Einstein summation over the samples and each of the per-dimension weights.
3173
3174  def label(dims: Iterable[int]) -> str:
3175    return ''.join(chr(ord('a') + i) for i in dims)
3176
3177  operands: list[Any] = [samples]  # (8, 9, 4, 6, 3)
3178  assert samples_ndim < 26  # Letters 'a' through 'z'.
3179  labels = [label(range(samples_ndim))]  # ['abcde']
3180  for dim in range(grid_ndim):
3181    operands.append(weight[dim])  # (8, 9, 4), then (8, 9, 6)
3182    labels.append(label(list(range(resampled_ndim)) + [resampled_ndim + dim]))  # 'abc' then 'abd'
3183  output_label = label(
3184      list(range(resampled_ndim)) + list(range(resampled_ndim + grid_ndim, samples_ndim))
3185  )  # 'abe'
3186  subscripts = ','.join(labels) + '->' + output_label  # 'abcde,abc,abd->abe'
3187  # Starting in numpy 2.0, np.einsum() outputs np.float64 even with all np.float32 inputs;
3188  # GPT: "aligns np.einsum with other functions where intermediate calculations use higher
3189  # precision (np.float64) regardless of input type when floating-point arithmetic is involved."
3190  # we could explicitly add the parameter `dtype=precision`.
3191  array = _arr_einsum(subscripts, *operands)  # (8, 9, 3)
3192
3193  # Gathering `samples` is the memory bottleneck.  It would be ideal if the gather() and einsum()
3194  # computations could be fused.  In Jax, https://github.com/google/jax/issues/3206 suggests
3195  # that this may become possible.  In any case, for large outputs it helps to partition the
3196  # evaluation over output tiles (using max_block_size).
3197
3198  if uses_cval:
3199    cval_weight = 1.0 - np.multiply.reduce(
3200        [weight[dim].sum(axis=-1) for dim in range(resampled_ndim)]
3201    )  # (8, 9)
3202    cval_weight_reshaped = cval_weight.reshape(cval_weight.shape + (1,) * len(sample_shape))
3203    array += _make_array((cval_weight_reshaped * cval).astype(precision, copy=False), arraylib)
3204
3205  array = dst_gamma2.encode(array, dtype)
3206  return array
3207
3208
3209def resample_affine(
3210    array: _Array,
3211    /,
3212    shape: Iterable[int],
3213    matrix: _ArrayLike,
3214    *,
3215    gridtype: str | Gridtype | None = None,
3216    src_gridtype: str | Gridtype | Iterable[str | Gridtype] | None = None,
3217    dst_gridtype: str | Gridtype | Iterable[str | Gridtype] | None = None,
3218    filter: str | Filter | Iterable[str | Filter] = _DEFAULT_FILTER,
3219    prefilter: str | Filter | Iterable[str | Filter] | None = None,
3220    precision: _DTypeLike | None = None,
3221    dtype: _DTypeLike | None = None,
3222    **kwargs: Any,
3223) -> _Array:
3224  """Resample a source array using an affinely transformed grid of given shape.
3225
3226  The `matrix` transformation can be linear,
3227    `source_point = matrix @ destination_point`,
3228  or it can be affine where the last matrix column is an offset vector,
3229    `source_point = matrix @ (destination_point, 1.0)`.
3230
3231  Args:
3232    array: Regular grid of source sample values, as an array object recognized by `ARRAYLIBS`.
3233      The array must have numeric type.  The number of grid dimensions is determined from
3234      `matrix.shape[0]`; the remaining dimensions are for each sample value and are all
3235      linearly interpolated.
3236    shape: Dimensions of the desired destination grid.  The number of destination grid dimensions
3237      may be different from that of the source grid.
3238    matrix: 2D array for a linear or affine transform from unit-domain destination points
3239      (in a space with `len(shape)` dimensions) into unit-domain source points (in a space with
3240      `matrix.shape[0]` dimensions).  If the matrix has `len(shape) + 1` columns, the last column
3241      is the affine offset (i.e., translation).
3242    gridtype: Placement of samples on all dimensions of both the source and output domain grids,
3243      specified as either a name in `GRIDTYPES` or a `Gridtype` instance.  It defaults to `'dual'`
3244      if `gridtype`, `src_gridtype`, and `dst_gridtype` are all kept `None`.
3245    src_gridtype: Placement of samples in the source domain grid for each dimension.
3246      Parameters `gridtype` and `src_gridtype` cannot both be set.
3247    dst_gridtype: Placement of samples in the output domain grid for each dimension.
3248      Parameters `gridtype` and `dst_gridtype` cannot both be set.
3249    filter: The reconstruction kernel for each dimension in `matrix.shape[0]`, specified as either
3250      a filter name in `FILTERS` or a `Filter` instance.
3251    prefilter: The prefilter kernel for each dimension in `len(shape)`, specified as either
3252      a filter name in `FILTERS` or a `Filter` instance.  It is used during downsampling
3253      (i.e., minification).  If `None`, it inherits the value of `filter`.
3254    precision: Inexact precision of intermediate computations.  If `None`, it is determined based
3255      on `array.dtype` and `dtype`.
3256    dtype: Desired data type of the output array.  If `None`, it is taken to be `array.dtype`.
3257      If it is a uint type, the intermediate float values are rescaled from the [0.0, 1.0] range
3258      to the uint range.
3259    **kwargs: Additional parameters for `resample` function.
3260
3261  Returns:
3262    An array of the same class as the source `array`, representing a grid with specified `shape`,
3263    where each grid value is resampled from `array`.  Thus the shape of the returned array is
3264    `shape + array.shape[matrix.shape[0]:]`.
3265  """
3266  shape = tuple(shape)
3267  matrix = np.asarray(matrix)
3268  dst_ndim = len(shape)
3269  if matrix.ndim != 2:
3270    raise ValueError(f'Array {matrix} is not 2D matrix.')
3271  src_ndim = matrix.shape[0]
3272  # grid_shape = array.shape[:src_ndim]
3273  is_affine = matrix.shape[1] == dst_ndim + 1
3274  if src_ndim > len(_arr_shape(array)):
3275    raise ValueError(
3276        f'Matrix {matrix} has more rows ({matrix.shape[0]}) than ndim in'
3277        f' array.shape={_arr_shape(array)}.'
3278    )
3279  if matrix.shape[1] != dst_ndim and not is_affine:
3280    raise ValueError(
3281        f'Matrix has {matrix.shape=}, but we expect either {dst_ndim} or {dst_ndim + 1} columns.'
3282    )
3283  src_gridtype2, dst_gridtype2 = _get_gridtypes(
3284      gridtype, src_gridtype, dst_gridtype, src_ndim, dst_ndim
3285  )
3286  prefilter = filter if prefilter is None else prefilter
3287  filter2 = [_get_filter(f) for f in np.broadcast_to(np.array(filter), src_ndim)]
3288  prefilter2 = [_get_filter(f) for f in np.broadcast_to(np.array(prefilter), dst_ndim)]
3289  del src_gridtype, dst_gridtype, filter, prefilter
3290  dtype = _arr_dtype(array) if dtype is None else np.dtype(dtype)
3291  precision = _get_precision(precision, [_arr_dtype(array), dtype], [])
3292  weight_precision = _real_precision(precision)
3293
3294  dst_position_list = []  # per dimension
3295  for dim in range(dst_ndim):
3296    dst_size = shape[dim]
3297    dst_index = np.arange(dst_size, dtype=weight_precision)
3298    dst_position_list.append(dst_gridtype2[dim].point_from_index(dst_index, dst_size))
3299  dst_position = np.meshgrid(*dst_position_list, indexing='ij')
3300
3301  linear_matrix = matrix[:, :-1] if is_affine else matrix
3302  src_position = np.tensordot(linear_matrix, dst_position, 1)
3303  coords = np.moveaxis(src_position, 0, -1)
3304  if is_affine:
3305    coords += matrix[:, -1]
3306
3307  # TODO: Based on grid_shape, shape, linear_matrix, and prefilter, determine a
3308  # convolution prefilter and apply it to bandlimit 'array', using boundary for padding.
3309
3310  return resample(
3311      array,
3312      coords,
3313      gridtype=src_gridtype2,
3314      filter=filter2,
3315      prefilter=prefilter2,
3316      precision=precision,
3317      dtype=dtype,
3318      **kwargs,
3319  )
3320
3321
3322def _resize_using_resample(
3323    array: _Array,
3324    /,
3325    shape: Iterable[int],
3326    *,
3327    scale: _ArrayLike = 1.0,
3328    translate: _ArrayLike = 0.0,
3329    filter: str | Filter | Iterable[str | Filter] = _DEFAULT_FILTER,
3330    fallback: bool = False,
3331    **kwargs: Any,
3332) -> _Array:
3333  """Use the more general `resample` operation for `resize`, as a debug tool."""
3334  shape = tuple(shape)
3335  scale = np.broadcast_to(scale, len(shape))
3336  translate = np.broadcast_to(translate, len(shape))
3337  # TODO: let resample() do prefiltering for proper downsampling.
3338  src_shape = _arr_shape(array)[: len(shape)]
3339  has_minification = np.any(np.array(shape) < src_shape) or np.any(scale < 1.0)
3340  filter2 = [_get_filter(f) for f in np.broadcast_to(np.array(filter), len(shape))]
3341  has_auto_trapezoid = any(f.name == 'trapezoid' for f in filter2)
3342  if fallback and (has_minification or has_auto_trapezoid):
3343    return _original_resize(array, shape, scale=scale, translate=translate, filter=filter, **kwargs)
3344  offset = -translate / scale
3345  matrix = np.concatenate([np.diag(1.0 / scale), offset[:, None]], axis=1)
3346  return resample_affine(array, shape, matrix, filter=filter, **kwargs)
3347
3348
3349def rotation_about_center_in_2d(
3350    src_shape: _ArrayLike,
3351    /,
3352    angle: float,
3353    *,
3354    new_shape: _ArrayLike | None = None,
3355    scale: float = 1.0,
3356) -> _NDArray:
3357  """Return the 3x3 matrix mapping destination into a source unit domain.
3358
3359  The returned matrix accounts for the possibly non-square domain shapes.
3360
3361  Args:
3362    src_shape: Resolution `(ny, nx)` of the source domain grid.
3363    angle: Angle in radians (positive from x to y axis) applied when mapping the source domain
3364      onto the destination domain.
3365    new_shape: Resolution `(ny, nx)` of the destination domain grid; it defaults to `src_shape`.
3366    scale: Scaling factor applied when mapping the source domain onto the destination domain.
3367  """
3368
3369  def translation_matrix(vector: _NDArray) -> _NDArray:
3370    matrix = np.eye(len(vector) + 1)
3371    matrix[:-1, -1] = vector
3372    return matrix
3373
3374  def scaling_matrix(scale: _NDArray) -> _NDArray:
3375    return np.diag(tuple(scale) + (1.0,))
3376
3377  def rotation_matrix_2d(angle: float) -> _NDArray:
3378    cos, sin = np.cos(angle), np.sin(angle)
3379    return np.array([[cos, sin, 0], [-sin, cos, 0], [0, 0, 1]])
3380
3381  src_shape = np.asarray(src_shape)
3382  new_shape = src_shape if new_shape is None else np.asarray(new_shape)
3383  _check_eq(src_shape.shape, (2,))
3384  _check_eq(new_shape.shape, (2,))
3385  half = np.array([0.5, 0.5])
3386  matrix = (
3387      translation_matrix(half)
3388      @ scaling_matrix(min(src_shape) / src_shape)
3389      @ rotation_matrix_2d(angle)
3390      @ scaling_matrix(scale * new_shape / min(new_shape))
3391      @ translation_matrix(-half)
3392  )
3393  assert np.allclose(matrix[-1], [0.0, 0.0, 1.0])
3394  return matrix
3395
3396
3397def rotate_image_about_center(
3398    image: _NDArray,
3399    /,
3400    angle: float,
3401    *,
3402    new_shape: _ArrayLike | None = None,
3403    scale: float = 1.0,
3404    num_rotations: int = 1,
3405    **kwargs: Any,
3406) -> _NDArray:
3407  """Return a copy of `image` rotated about its center.
3408
3409  Args:
3410    image: Source grid samples; the first two dimensions are spatial (ny, nx).
3411    angle: Angle in radians (positive from x to y axis) applied when mapping the source domain
3412      onto the destination domain.
3413    new_shape: Resolution `(ny, nx)` of the output grid; it defaults to `image.shape[:2]`.
3414    scale: Scaling factor applied when mapping the source domain onto the destination domain.
3415    num_rotations: Number of rotations (each by `angle`).  Successive resamplings are useful in
3416      analyzing the filtering quality.
3417    **kwargs: Additional parameters for `resample_affine`.
3418  """
3419  new_shape = image.shape[:2] if new_shape is None else np.asarray(new_shape)
3420  matrix = rotation_about_center_in_2d(image.shape[:2], angle, new_shape=new_shape, scale=scale)
3421  for _ in range(num_rotations):
3422    image = resample_affine(image, new_shape, matrix[:-1], **kwargs)
3423  return image
3424
3425
3426def _pil_image_resize(
3427    array: _ArrayLike,
3428    /,
3429    shape: Iterable[int],
3430    *,
3431    filter: str,
3432    boundary: str = 'natural',
3433    cval: float = 0.0,
3434) -> _NDArray:
3435  """Invoke `PIL.Image.resize` using the same parameters as `resize`."""
3436  import PIL.Image
3437
3438  if boundary != 'natural':
3439    raise ValueError(f"{boundary=} must equal 'natural'.")
3440  del cval
3441  array = np.asarray(array)
3442  assert 1 <= array.ndim <= 3
3443  assert np.issubdtype(array.dtype, np.floating)
3444  shape = tuple(shape)
3445  _check_eq(len(shape), 2 if array.ndim >= 2 else 1)
3446  if array.ndim == 1:
3447    return _pil_image_resize(array[None], (1, *shape), filter=filter)[0]
3448  if not hasattr(PIL.Image, 'Resampling'):  # Pillow<9.0
3449    PIL.Image.Resampling = PIL.Image  # type: ignore
3450  filters = {
3451      'impulse': PIL.Image.Resampling.NEAREST,
3452      'box': PIL.Image.Resampling.BOX,
3453      'triangle': PIL.Image.Resampling.BILINEAR,
3454      'hamming1': PIL.Image.Resampling.HAMMING,
3455      'cubic': PIL.Image.Resampling.BICUBIC,
3456      'lanczos3': PIL.Image.Resampling.LANCZOS,
3457  }
3458  if filter not in filters:
3459    raise ValueError(f'{filter=} not in {filters=}.')
3460  pil_resample = filters[filter]
3461  ny, nx = shape
3462  if array.ndim == 2:
3463    return np.array(PIL.Image.fromarray(array).resize((nx, ny), resample=pil_resample), array.dtype)
3464  stack = []
3465  for channel in np.moveaxis(array, -1, 0):
3466    pil_image = PIL.Image.fromarray(channel).resize((nx, ny), resample=pil_resample)
3467    stack.append(np.array(pil_image, array.dtype))
3468  return np.dstack(stack)
3469
3470
3471def _cv_resize(
3472    array: _ArrayLike,
3473    /,
3474    shape: Iterable[int],
3475    *,
3476    filter: str,
3477    boundary: str = 'clamp',
3478    cval: float = 0.0,
3479) -> _NDArray:
3480  """Invoke `cv.resize` using the same parameters as `resize`."""
3481  import cv2 as cv
3482
3483  if boundary != 'clamp':
3484    raise ValueError(f"{boundary=} must equal 'clamp'.")
3485  del cval
3486  array = np.asarray(array)
3487  assert 1 <= array.ndim <= 3
3488  shape = tuple(shape)
3489  _check_eq(len(shape), 2 if array.ndim >= 2 else 1)
3490  if array.ndim == 1:
3491    return _cv_resize(array[None], (1, *shape), filter=filter)[0]
3492  filters = {
3493      'impulse': cv.INTER_NEAREST,  # Or consider cv.INTER_NEAREST_EXACT.
3494      'triangle': cv.INTER_LINEAR_EXACT,  # Or just cv.INTER_LINEAR.
3495      'trapezoid': cv.INTER_AREA,
3496      'sharpcubic': cv.INTER_CUBIC,
3497      'lanczos4': cv.INTER_LANCZOS4,
3498  }
3499  if filter not in filters:
3500    raise ValueError(f'{filter=} not in {filters=}.')
3501  interpolation = filters[filter]
3502  result = cv.resize(array, shape[::-1], interpolation=interpolation)
3503  if array.ndim == 3 and result.ndim == 2:
3504    assert array.shape[2] == 1
3505    return result[..., None]  # Add back the last dimension dropped by cv.resize().
3506  return result
3507
3508
3509def _scipy_ndimage_resize(
3510    array: _ArrayLike,
3511    /,
3512    shape: Iterable[int],
3513    *,
3514    filter: str,
3515    boundary: str = 'reflect',
3516    cval: float = 0.0,
3517    scale: float | Iterable[float] = 1.0,
3518    translate: float | Iterable[float] = 0.0,
3519) -> _NDArray:
3520  """Invoke `scipy.ndimage.map_coordinates` using the same parameters as `resize`."""
3521  array = np.asarray(array)
3522  shape = tuple(shape)
3523  assert 1 <= len(shape) <= array.ndim
3524  filters = {'box': 0, 'triangle': 1} | {f'cardinal{i}': i for i in range(2, 6)}
3525  if filter not in filters:
3526    raise ValueError(f'{filter=} not in {filters=}.')
3527  order = filters[filter]
3528  boundaries = {'reflect': 'reflect', 'wrap': 'grid-wrap', 'clamp': 'nearest', 'border': 'constant'}
3529  if boundary not in boundaries:
3530    raise ValueError(f'{boundary=} not in {boundaries=}.')
3531  mode = boundaries[boundary]
3532  shape_all = shape + array.shape[len(shape) :]
3533  coords = np.moveaxis(np.indices(shape_all, array.dtype), 0, -1)
3534  coords[..., : len(shape)] = (
3535      (coords[..., : len(shape)] + 0.5) / shape - np.asarray(translate)
3536  ) / np.asarray(scale) * np.array(array.shape)[: len(shape)] - 0.5
3537  coords = np.moveaxis(coords, -1, 0)
3538  return scipy.ndimage.map_coordinates(array, coords, order=order, mode=mode, cval=cval)
3539
3540
3541def _skimage_transform_resize(
3542    array: _ArrayLike,
3543    /,
3544    shape: Iterable[int],
3545    *,
3546    filter: str,
3547    boundary: str = 'reflect',
3548    cval: float = 0.0,
3549) -> _NDArray:
3550  """Invoke `skimage.transform.resize` using the same parameters as `resize`."""
3551  import skimage.transform
3552
3553  array = np.asarray(array)
3554  shape = tuple(shape)
3555  assert 1 <= len(shape) <= array.ndim
3556  filters = {'box': 0, 'triangle': 1} | {f'cardinal{i}': i for i in range(2, 6)}
3557  if filter not in filters:
3558    raise ValueError(f'{filter=} not in {filters=}.')
3559  order = filters[filter]
3560  boundaries = {'reflect': 'symmetric', 'wrap': 'wrap', 'clamp': 'edge', 'border': 'constant'}
3561  if boundary not in boundaries:
3562    raise ValueError(f'{boundary=} not in {boundaries=}.')
3563  mode = boundaries[boundary]
3564  shape_all = shape + array.shape[len(shape) :]
3565  # Default anti_aliasing=None automatically enables (poor) Gaussian prefilter if downsampling.
3566  # clip=False is the default behavior in `resampler` if the output type is non-integer.
3567  return skimage.transform.resize(
3568      array, shape_all, order=order, mode=mode, cval=cval, clip=False
3569  )  # type: ignore[no-untyped-call]
3570
3571
3572_TORCH_INTERPOLATE_MODE_FROM_FILTER = {
3573    'impulse': 'nearest-exact',  # ('nearest' matches buggy OpenCV's INTER_NEAREST)
3574    'trapezoid': 'area',
3575    'triangle': 'bilinear',
3576    'sharpcubic': 'bicubic',
3577}
3578
3579
3580def _torch_nn_resize(
3581    array: _ArrayLike,
3582    /,
3583    shape: Iterable[int],
3584    *,
3585    filter: str,
3586    boundary: str = 'clamp',
3587    cval: float = 0.0,
3588    antialias: bool = False,
3589) -> _TorchTensor:
3590  """Invoke `torch.nn.functional.interpolate` using the same parameters as `resize`."""
3591  import torch
3592
3593  if filter not in _TORCH_INTERPOLATE_MODE_FROM_FILTER:
3594    raise ValueError(f'{filter=} not in {_TORCH_INTERPOLATE_MODE_FROM_FILTER=}.')
3595  if boundary != 'clamp':
3596    raise ValueError(f"{boundary=} must equal 'clamp'.")
3597  del cval
3598  a = torch.as_tensor(array)
3599  del array
3600  assert 1 <= a.ndim <= 3
3601  shape = tuple(shape)
3602  _check_eq(len(shape), 2 if a.ndim >= 2 else 1)
3603  mode = _TORCH_INTERPOLATE_MODE_FROM_FILTER[filter]
3604
3605  def local_resize(a: _TorchTensor) -> _TorchTensor:
3606    # For upsampling, BILINEAR antialias is same PSNR and slower,
3607    #  and BICUBIC antialias is worse PSNR and faster.
3608    # For downsampling, antialias improves PSNR for both BILINEAR and BICUBIC.
3609    # Default align_corners=None corresponds to False which is what we desire.
3610    return torch.nn.functional.interpolate(a, shape, mode=mode, antialias=antialias)
3611
3612  match a.ndim:
3613    case 1:
3614      shape = (1, *shape)
3615      return local_resize(a[None, None, None])[0, 0, 0]
3616    case 2:
3617      return local_resize(a[None, None])[0, 0]
3618    case _:
3619      return local_resize(a.moveaxis(2, 0)[None])[0].moveaxis(0, 2)
3620
3621
3622def _jax_image_resize(
3623    array: _ArrayLike,
3624    /,
3625    shape: Iterable[int],
3626    *,
3627    filter: str,
3628    boundary: str = 'natural',
3629    cval: float = 0.0,
3630    scale: float | Iterable[float] = 1.0,
3631    translate: float | Iterable[float] = 0.0,
3632) -> _JaxArray:
3633  """Invoke `jax.image.scale_and_translate` using the same parameters as `resize`."""
3634  import jax.image
3635  import jax.numpy as jnp
3636
3637  filters = 'triangle cubic lanczos3 lanczos5'.split()
3638  if filter not in filters:
3639    raise ValueError(f'{filter=} not in {filters=}.')
3640  if boundary != 'natural':
3641    raise ValueError(f"{boundary=} must equal 'natural'.")
3642  # When `scale` or `translate` are applied, any region outside the unit domain is assigned value 0.
3643  # To be consistent, the parameter `cval` must be zero.
3644  if scale != 1.0 and cval != 0.0:
3645    raise ValueError(f'Non-unity {scale=} requires that {cval=} be zero.')
3646  if translate != 0.0 and cval != 0.0:
3647    raise ValueError(f'Nonzero {translate=} requires that {cval=} be zero.')
3648  array2 = jnp.asarray(array)
3649  del array
3650  shape = tuple(shape)
3651  assert len(shape) <= array2.ndim
3652  completed_shape = shape + (1,) * (array2.ndim - len(shape))
3653  spatial_dims = list(range(len(shape)))
3654  scale2 = np.broadcast_to(np.array(scale), len(shape))
3655  scale2 = scale2 / np.array(array2.shape[: len(shape)]) * np.array(shape)
3656  translate2 = np.broadcast_to(np.array(translate), len(shape))
3657  translate2 = translate2 * np.array(shape)
3658  return jax.image.scale_and_translate(
3659      array2, completed_shape, spatial_dims, scale2, translate2, filter
3660  )
3661
3662
3663_CANDIDATE_RESIZERS = {
3664    'resampler.resize': resize,
3665    'PIL.Image.resize': _pil_image_resize,
3666    'cv.resize': _cv_resize,
3667    'scipy.ndimage.map_coordinates': _scipy_ndimage_resize,
3668    'skimage.transform.resize': _skimage_transform_resize,
3669    'torch.nn.functional.interpolate': _torch_nn_resize,
3670    'jax.image.scale_and_translate': _jax_image_resize,
3671}
3672
3673
3674def _resizer_is_available(library_function: str) -> bool:
3675  """Return whether the resizer is available as an installed package."""
3676  top_name = library_function.split('.', 1)[0]
3677  module = {'PIL': 'Pillow', 'cv': 'cv2'}.get(top_name, top_name)
3678  return importlib.util.find_spec(module) is not None  # type: ignore[attr-defined]
3679
3680
3681_RESIZERS = {
3682    library_function: resizer
3683    for library_function, resizer in _CANDIDATE_RESIZERS.items()
3684    if _resizer_is_available(library_function)
3685}
3686
3687
3688def _find_closest_filter(filter: str, resizer: Callable[..., Any]) -> str:
3689  """Return the filter supported by `resizer` (i.e., `*_resize`) that is closest to `filter`."""
3690  match filter:
3691    case 'box_like':
3692      return {
3693          _cv_resize: 'trapezoid',
3694          _skimage_transform_resize: 'box',
3695          _torch_nn_resize: 'trapezoid',
3696      }.get(resizer, 'box')
3697    case 'cubic_like':
3698      return {
3699          _cv_resize: 'sharpcubic',
3700          _scipy_ndimage_resize: 'cardinal3',
3701          _skimage_transform_resize: 'cardinal3',
3702          _torch_nn_resize: 'sharpcubic',
3703      }.get(resizer, 'cubic')
3704    case 'high_quality':
3705      return {
3706          _pil_image_resize: 'lanczos3',
3707          _cv_resize: 'lanczos4',
3708          _scipy_ndimage_resize: 'cardinal5',
3709          _skimage_transform_resize: 'cardinal5',
3710          _torch_nn_resize: 'sharpcubic',
3711      }.get(resizer, 'lanczos5')
3712    case _:
3713      return filter
3714
3715
3716# For Emacs:
3717# Local Variables:
3718# fill-column: 100
3719# End: