test_resampler.py

#!/usr/bin/env python3
# -*- fill-column: 100; -*-
"""Tests for package resampler.

c:/windows/sysnative/wsl -e bash -lc 'flake8 --indent-size 2 --max-line-length=1000 --extend-ignore E302,E741,E131,E305,E402 test_resampler.py && python3 test_resampler.py'
"""
from __future__ import annotations

from collections.abc import Callable
import functools
import itertools
import math
import os
from typing import Any
import unittest
import warnings

import numpy as np
import scipy

import resampler

_ArrayLike = Any
_NDArray = Any
_TensorflowTensor = Any

# pylint: disable=protected-access, missing-function-docstring, too-many-public-methods

# Silence "WARNING:absl:No GPU/TPU found, falling back to CPU".
os.environ['JAX_PLATFORM_NAME'] = 'cpu'


def enable_jax_float64() -> None:
  """Enable use of double-precision float in Jax; this only works at startup."""
  import jax  # ("import jax.config" is disallowed.)

  jax.config.update('jax_enable_x64', True)


def _check_eq(a: Any, b: Any) -> None:
  """If the two values or arrays are not equal, raise an exception with a useful message."""
  are_equal = np.all(a == b) if isinstance(a, np.ndarray) else a == b
  if not are_equal:
    raise AssertionError(f'{a!r} == {b!r}')


class TestResampler(unittest.TestCase):
  """Test class for resampler package."""

  @classmethod
  def setUpClass(cls: type) -> None:
    if 'jax' in resampler.ARRAYLIBS:
      enable_jax_float64()
    # Silence the warning in package flatbuffers.
    warnings.filterwarnings('ignore', message='.*the imp module is deprecated')

  def test_resize_on_list(self) -> None:
    array = [3.0, 5.0, 8.0, 7.0]
    expected = np.array([2.84536097, 3.6902174, 5.58573019, 7.77282572, 7.8097826, 6.79608312])
    np.testing.assert_allclose(resampler.resize(array, (6,)), expected)
    np.testing.assert_allclose(resampler.resize(np.array(array), (6,)), expected)

  def test_precision(self) -> None:
    _check_eq(resampler._real_precision(np.dtype(np.float32)), np.float32)
    _check_eq(resampler._real_precision(np.dtype(np.float64)), np.float64)
    _check_eq(resampler._real_precision(np.dtype(np.complex64)), np.float32)
    _check_eq(resampler._real_precision(np.dtype(np.complex128)), np.float64)

  def test_get_precision(self) -> None:
    _check_eq(
        resampler._get_precision(None, [np.dtype(np.complex64)], [np.dtype(np.float64)]),
        np.complex128,
    )

  def test_cached_sampling(self) -> None:
    radius = 2.0

    def create_scipy_interpolant(
        func: Callable[[_ArrayLike], _NDArray], xmin: float, xmax: float, num_samples: int = 3_600
    ) -> Callable[[_NDArray], _NDArray]:
      samples_x = np.linspace(xmin, xmax, num_samples + 1, dtype=np.float32)
      samples_func = func(samples_x)
      assert np.all(samples_func[[0, -1]] == 0.0)
      interpolator: Callable[[_NDArray], _NDArray] = scipy.interpolate.interp1d(
          samples_x, samples_func, kind='linear', bounds_error=False, fill_value=0
      )
      return interpolator

    def func(x: _ArrayLike) -> _NDArray:  # Lanczos kernel
      x = np.abs(x)
      return np.where(x < radius, resampler._sinc(x) * resampler._sinc(x / radius), 0.0)

    @resampler._cache_sampled_1d_function(xmin=-radius, xmax=radius)
    def func2(x: _ArrayLike) -> _NDArray:
      return func(x)

    scipy_interp = create_scipy_interpolant(func, -radius, radius)

    shape = 2, 8_000
    rng = np.random.default_rng(0)
    array = rng.random(shape, np.float32) * 2 * radius - radius
    result = {'expected': func(array), 'scipy': scipy_interp(array), 'obtained': func2(array)}

    assert all(a.dtype == np.float32 for a in result.values())
    assert all(a.shape == shape for a in result.values())
    assert np.allclose(result['scipy'], result['expected'], rtol=0, atol=1e-6)
    assert np.allclose(result['obtained'], result['expected'], rtol=0, atol=1e-6)

  def test_downsample_in_2d_using_box_filter(self) -> None:
    for shape in [(6, 6), (4, 4)]:
      for ch in [1, 2, 3, 4]:
        array = np.ones((*shape, ch), np.float32)
        new = resampler._downsample_in_2d_using_box_filter(array, (2, 2))
        _check_eq(new.shape, (2, 2, ch))
        assert np.allclose(new, 1.0)

    for shape in [(6, 6), (4, 4)]:
      array = np.ones(shape, np.float32)
      new = resampler._downsample_in_2d_using_box_filter(array, (2, 2))
      _check_eq(new.shape, (2, 2))
      assert np.allclose(new, 1.0)

  def test_block_shape_with_min_size(self) -> None:
    for compact in [True, False]:
      with self.subTest(compact=compact):
        shape = 2, 3, 4
        for min_size in range(1, math.prod(shape) + 1):
          block_shape = resampler._block_shape_with_min_size(shape, min_size, compact=compact)
          assert np.all(np.array(block_shape) >= 1)
          assert np.all(block_shape <= shape)
          assert min_size <= math.prod(block_shape) <= math.prod(shape)

  def test_split_2d(self) -> None:
    numpy_array = np.random.default_rng(0).choice([1, 2, 3, 4], (5, 8))
    for arraylib in resampler.ARRAYLIBS:
      array = resampler._make_array(numpy_array, arraylib)
      blocks = resampler._split_array_into_blocks(array, [2, 3])
      blocks = resampler._map_function_over_blocks(blocks, lambda x: 2 * x)
      new = resampler._merge_array_from_blocks(blocks)
      _check_eq(resampler._arr_arraylib(new), arraylib)
      _check_eq(np.sum(resampler._map_function_over_blocks(blocks, lambda _: 1)), 9)
      _check_eq(resampler._arr_numpy(new), 2 * numpy_array)

  def test_split_3d(self) -> None:
    shape = 4, 3, 2
    numpy_array = np.random.default_rng(0).choice([1, 2, 3, 4], shape)

    for arraylib in resampler.ARRAYLIBS:
      array = resampler._make_array(numpy_array, arraylib)
      for min_size in range(1, math.prod(shape) + 1):
        block_shape = resampler._block_shape_with_min_size(shape, min_size)
        blocks = resampler._split_array_into_blocks(array, block_shape)
        blocks = resampler._map_function_over_blocks(blocks, lambda x: x**2)
        new = resampler._merge_array_from_blocks(blocks)
        _check_eq(resampler._arr_arraylib(new), arraylib)
        _check_eq(resampler._arr_numpy(new), numpy_array**2)

        def check_block_shape(block: _NDArray) -> None:
          assert np.all(np.array(block.shape) >= 1)
          assert np.all(np.array(block.shape) <= shape)

        resampler._map_function_over_blocks(blocks, check_block_shape)

  def test_split_prefix_dims(self) -> None:
    shape = 2, 3, 2
    array = np.arange(math.prod(shape)).reshape(shape)

    for min_size in range(1, math.prod(shape[:2]) + 1):
      block_shape = resampler._block_shape_with_min_size(shape[:2], min_size)
      blocks = resampler._split_array_into_blocks(array, block_shape)

      new_blocks = resampler._map_function_over_blocks(blocks, lambda x: x**2)
      new = resampler._merge_array_from_blocks(new_blocks)
      _check_eq(new, array**2)

      new_blocks = resampler._map_function_over_blocks(blocks, lambda x: x.sum(axis=-1))
      new = resampler._merge_array_from_blocks(new_blocks)
      _check_eq(new, array.sum(axis=-1))

  def test_linear_boundary(self) -> None:
    index = np.array([[-3], [-2], [-1], [0], [1], [2], [3], [2], [3]])
    weight = np.array([[1.0], [1.0], [1.0], [1.0], [1.0], [1.0], [1.0], [0.0], [0.0]])
    size = 2
    index, weight = resampler.LinearExtendSamples()(index, weight, size, resampler.DualGridtype())
    expected_weight = [
        [0, 4, -3, 0, 0],
        [0, 3, -2, 0, 0],
        [0, 2, -1, 0, 0],
        [1, 0, 0, 0, 0],
        [1, 0, 0, 0, 0],
        [0, 0, 0, -1, 2],
        [0, 0, 0, -2, 3],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
    ]
    assert np.allclose(weight, expected_weight)
    expected_index = [
        [0, 0, 1, 0, 0],
        [0, 0, 1, 0, 0],
        [0, 0, 1, 0, 0],
        [0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 0, 1],
        [1, 1, 1, 0, 1],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1],
    ]
    assert np.all(index == expected_index)

  def test_gamma_roundtrip_uint(self) -> None:
    dtypes = 'uint8 uint16 uint32'.split()
    for config in itertools.product(resampler.ARRAYLIBS, resampler.GAMMAS, dtypes):
      arraylib, gamma_name, dtype = config
      gamma = resampler._get_gamma(gamma_name)
      if arraylib == 'torch' and dtype in ['uint16', 'uint32']:
        continue  # "The only supported types are: ..., int64, int32, int16, int8, uint8, and bool."
      with self.subTest(config=config):
        int_max = np.iinfo(dtype).max
        precision = 'float32' if np.iinfo(dtype).bits < 32 else 'float64'
        values = list(range(256)) + list(range(int_max - 255, int_max)) + [int_max]
        array_numpy = np.array(values, dtype)
        array = resampler._make_array(array_numpy, arraylib)
        decoded = gamma.decode(array, np.dtype(precision))
        _check_eq(resampler._arr_dtype(decoded), precision)
        decoded_numpy = resampler._arr_numpy(decoded)
        assert decoded_numpy.min() >= 0.0 and decoded_numpy.max() <= 1.0
        encoded = gamma.encode(decoded, dtype)
        _check_eq(resampler._arr_dtype(encoded), dtype)
        encoded_numpy = resampler._arr_numpy(encoded)
        _check_eq(encoded_numpy, array_numpy)

  def test_gamma_roundtrip_float(self) -> None:
    dtypes = 'float32 float64'.split()
    precisions = 'float32 float64'.split()
    for config in itertools.product(resampler.ARRAYLIBS, resampler.GAMMAS, dtypes, precisions):
      arraylib, gamma_name, dtype, precision = config
      with self.subTest(config=config):
        gamma = resampler._get_gamma(gamma_name)
        array_numpy = np.linspace(0.0, 1.0, 100, dtype=dtype)
        array = resampler._make_array(array_numpy, arraylib)
        decoded = gamma.decode(array, np.dtype(precision))
        _check_eq(resampler._arr_dtype(decoded), precision)
        encoded = gamma.encode(decoded, dtype)
        _check_eq(resampler._arr_dtype(encoded), dtype)
        assert np.allclose(resampler._arr_numpy(encoded), array_numpy)

  def test_create_resize_matrix_for_trapezoid_filter(self) -> None:
    filter = resampler.TrapezoidFilter()
    for src_size, dst_size in [(6, 2), (7, 3), (7, 6), (14, 13), (3, 6), (3, 12), (3, 11), (3, 16)]:
      with self.subTest(src_size=src_size, dst_size=dst_size):
        resize_matrix, unused_cval_weight = resampler._create_resize_matrix(
            src_size,
            dst_size,
            src_gridtype=resampler.DualGridtype(),
            dst_gridtype=resampler.DualGridtype(),
            boundary=resampler._get_boundary('reflect'),
            filter=filter,
        )
        resize_matrix = resize_matrix.toarray()
        assert resize_matrix.sum(axis=0).var() < 1e-10
        assert resize_matrix.sum(axis=1).var() < 1e-10

  def test_sparse_csr_matrix_duplicate_entries_are_summed(self) -> None:
    indptr = np.array([0, 2, 3, 6])
    indices = np.array([0, 2, 2, 0, 1, 0])
    data = np.array([1, 2, 3, 4, 5, 3])
    new = scipy.sparse.csr_matrix((data, indices, indptr), shape=(3, 3)).toarray()
    _check_eq(new, [[1, 0, 2], [0, 0, 3], [7, 5, 0]])

  def test_that_resize_matrices_are_equal_across_arraylib(self) -> None:
    src_sizes = range(1, 6)
    dst_sizes = range(1, 6)
    for config in itertools.product(src_sizes, dst_sizes):
      src_size, dst_size = config
      with self.subTest(config=config):

        def resize_matrix(arraylib: str) -> Any:
          return resampler._create_resize_matrix(
              src_size,
              dst_size,
              src_gridtype=resampler.DualGridtype(),
              dst_gridtype=resampler.DualGridtype(),
              boundary=resampler._get_boundary('reflect'),
              filter=resampler._get_filter('lanczos3'),
              translate=0.8,
              dtype=np.float32,
              arraylib=arraylib,
          )[0]

        results = {}
        for arraylib in resampler.ARRAYLIBS:
          sparse_matrix = resize_matrix(arraylib)
          match arraylib:
            case 'numpy':
              result = sparse_matrix.toarray()
            case 'tensorflow':
              import tensorflow as tf

              result = tf.sparse.to_dense(sparse_matrix).numpy()
            case 'torch':
              result = sparse_matrix.to_dense().numpy()
            case 'jax':
              result = np.array(sparse_matrix.todense())
            case _:
              raise AssertionError
          results[arraylib] = result
        for arraylib in resampler.ARRAYLIBS:
          assert np.allclose(results[arraylib], results['numpy'])

  def test_that_resize_combinations_are_affine(self) -> None:
    dst_sizes = 1, 2, 3, 4, 9, 20, 21, 22, 31
    for config in itertools.product(resampler.BOUNDARIES, dst_sizes):
      boundary, dst_size = config
      with self.subTest(config=config):
        resize_matrix, cval_weight = resampler._create_resize_matrix(
            21,
            dst_size,
            src_gridtype=resampler.DualGridtype(),
            dst_gridtype=resampler.DualGridtype(),
            boundary=resampler._get_boundary(boundary),
            filter=resampler.TriangleFilter(),
            scale=0.5,
            translate=0.3,
        )
        if cval_weight is None:
          row_sum = np.asarray(resize_matrix.sum(axis=1)).ravel()
          assert np.allclose(row_sum, 1.0, rtol=0, atol=1e-6), (resize_matrix.todense(), row_sum)

  def test_linear_precision_of_1d_primal_upsampling(self) -> None:
    array = np.arange(7.0)
    new = resampler.resize(array, (13,), gridtype='primal', filter='triangle')
    with np.printoptions(linewidth=300):
      _check_eq(new, np.arange(13) / 2)

  def test_linear_precision_of_2d_primal_upsampling(self) -> None:
    shape = 3, 5
    new_shape = 5, 9
    array = np.moveaxis(np.indices(shape, np.float32), 0, -1) @ [10, 1]
    new = resampler.resize(array, new_shape, gridtype='primal', filter='triangle')
    with np.printoptions(linewidth=300):
      expected = np.moveaxis(np.indices(new_shape, np.float32), 0, -1) @ [10, 1] / 2
      _check_eq(new, expected)

  def test_resize_of_complex_value_type(self) -> None:
    for arraylib in resampler.ARRAYLIBS:
      array = resampler._make_array([1 + 2j, 3 + 6j], arraylib)
      new = resampler._original_resize(array, (4,), filter='triangle')
      assert np.allclose(new, [1 + 2j, 1.5 + 3j, 2.5 + 5j, 3 + 6j])

  def test_resize_of_integer_type(self) -> None:
    array = np.array([1, 6])
    new = resampler.resize(array, (4,), filter='triangle')
    assert np.allclose(new, [1, 2, 5, 6])

  def test_apply_digital_filter_1d(self) -> None:
    cval = -10.0
    shape = 7, 8
    original = np.arange(math.prod(shape), dtype=np.float32).reshape(shape) + 10
    array1 = original.copy()
    filters = 'cardinal3 cardinal5'.split()
    for config in itertools.product(resampler.GRIDTYPES, resampler.BOUNDARIES, filters):
      gridtype, boundary, filter = config
      with self.subTest(config=config):
        if gridtype == 'primal' and boundary in ('wrap', 'tile'):
          continue  # Last value on each dimension is ignored and so will not match.
        array2 = array1
        for dim in range(array2.ndim):
          array2 = resampler._apply_digital_filter_1d(
              array2,
              resampler._get_gridtype(gridtype),
              resampler._get_boundary(boundary),
              cval,
              resampler._get_filter(filter),
              axis=dim,
          )
        bspline = resampler.BsplineFilter(degree=int(filter[-1:]))
        array3 = resampler.resize(
            array2, array2.shape, gridtype=gridtype, boundary=boundary, cval=cval, filter=bspline
        )
        assert np.allclose(array3, original)

  def test_resample_small_arrays(self) -> None:
    shape = 2, 3
    new_shape = 3, 4
    for arraylib in resampler.ARRAYLIBS:
      with self.subTest(arraylib=arraylib):
        array = np.arange(math.prod(shape) * 3, dtype=np.float32).reshape(shape + (3,))
        coords = np.moveaxis(np.indices(new_shape) + 0.5, 0, -1) / new_shape
        array = resampler._make_array(array, arraylib)
        upsampled = resampler.resample(array, coords)
        _check_eq(upsampled.shape, (*new_shape, 3))
        coords = np.moveaxis(np.indices(shape) + 0.5, 0, -1) / shape
        downsampled = resampler.resample(upsampled, coords)
        rms = np.sqrt(np.mean(np.square(np.array(array) - np.array(downsampled)))).item()
        assert 0.07 <= rms <= 0.08, rms

  def test_identity_resampling_with_many_boundary_rules(self) -> None:
    filter = resampler.LanczosFilter(radius=5, sampled=False)
    for boundary in resampler.BOUNDARIES:
      with self.subTest(boundary=boundary):
        array = np.arange(6, dtype=np.float32).reshape(2, 3)
        coords = (np.moveaxis(np.indices(array.shape), 0, -1) + 0.5) / array.shape
        new_array = resampler.resample(array, coords, boundary=boundary, cval=10000, filter=filter)
        assert np.allclose(new_array, array), boundary

  def test_identity_resampling(self) -> None:
    shape = 3, 2, 5
    array = np.random.default_rng(0).random(shape)
    coords = (np.moveaxis(np.indices(array.shape), 0, -1) + 0.5) / array.shape
    new = resampler.resample(array, coords)
    assert np.allclose(new, array, rtol=0, atol=1e-6)
    new = resampler.resample(array, coords, filter=resampler.LanczosFilter(radius=3, sampled=False))
    assert np.allclose(new, array)

  def test_resample_of_complex_value_type(self) -> None:
    array = np.array([1 + 2j, 3 + 6j])
    new = resampler.resample(array, (0.125, 0.375, 0.625, 0.875), filter='triangle')
    assert np.allclose(new, [1 + 2j, 1.5 + 3j, 2.5 + 5j, 3 + 6j])

  def test_resample_of_integer_type(self) -> None:
    array = np.array([1, 6])
    new = resampler.resample(array, (0.125, 0.375, 0.625, 0.875), filter='triangle')
    assert np.allclose(new, [1, 2, 5, 6])

  def test_resample_using_coords_of_various_shapes(self) -> None:
    for array in [
        8,
        [7],
        [0, 1, 6, 6],
        [[0, 1], [10, 16]],
        [[0], [1], [6], [6]],
    ]:
      with self.subTest(array=array):
        array = np.array(array, np.float64)
        for shape in [(), (1,), (2,), (1, 1), (1, 2), (3, 1), (2, 2)]:
          coords = np.full(shape, 0.4)
          try:
            new = resampler.resample(array, coords, filter='triangle', dtype=np.float32).tolist()
          except ValueError:
            new = None
          # print(f'{array.tolist()!s:30} {coords.shape!s:8} {new!s}')
          _check_eq(new is None, coords.ndim >= 2 and coords.shape[-1] > max(array.ndim, 1))

  def test_resize_using_resample(self) -> None:
    shape = 3, 2, 5
    new_shape = 4, 2, 7
    step = 37
    assert np.all(np.array(shape) <= new_shape)
    array = np.random.default_rng(0).random(shape)
    scale = 1.1
    translate = -0.4, -0.03, 0.4
    gammas = 'identity power2'.split()  # Sublist of resampler.GAMMAS.
    sequences = [resampler.GRIDTYPES, resampler.BOUNDARIES, resampler.FILTERS, gammas]
    assert step == 1 or all(len(sequence) % step != 0 for sequence in sequences)
    configs = itertools.product(*sequences)  # len(configs) = math.prod([2, 12, 19, 2]) = 912.
    for config in itertools.islice(configs, 0, None, step):
      gridtype, boundary, filter, gamma = config
      with self.subTest(config=config):
        kwargs: Any = dict(gridtype=gridtype, boundary=boundary, filter=filter)
        kwargs |= dict(gamma=gamma, scale=scale, translate=translate)
        expected = resampler._original_resize(array, new_shape, **kwargs)
        new_array = resampler._resize_using_resample(array, new_shape, **kwargs)
        assert np.allclose(new_array, expected, rtol=0, atol=1e-7)

  def test_resize_using_resample_of_complex_value_type(self) -> None:
    array = np.array([1 + 2j, 3 + 6j])
    new = resampler._resize_using_resample(array, (4,), filter='triangle')
    assert np.allclose(new, [1 + 2j, 1.5 + 3j, 2.5 + 5j, 3 + 6j])

  def test_resizers_produce_correct_shape(self) -> None:
    configs: list[tuple[Callable[..., Any], str]] = [(resampler.resize, 'lanczos3')]
    for arraylib in resampler.ARRAYLIBS:
      resizer0 = functools.partial(resampler.resize_in_arraylib, arraylib=arraylib)
      configs.append((resizer0, 'lanczos3'))
    configs.append((resampler.pil_image_resize, 'lanczos3'))
    configs.append((resampler.cv_resize, 'lanczos4'))
    configs.append((resampler.scipy_ndimage_resize, 'cardinal3'))
    configs.append((resampler.skimage_transform_resize, 'cardinal3'))
    configs.append((resampler.tf_image_resize, 'lanczos3'))
    configs.append((resampler.torch_nn_resize, 'sharpcubic'))
    configs.append((resampler.jax_image_resize, 'lanczos3'))
    for config in configs:
      resizer, filter = config
      if resizer not in resampler._RESIZERS.values():  # Skip if the package is not installed.
        continue
      with self.subTest(config=config):
        tol: Any = dict(rtol=0, atol=1e-7)
        np.allclose(resizer(np.ones((11,)), (13,), filter=filter), np.ones((13,)), **tol)
        np.allclose(resizer(np.ones((8, 8)), (5, 20), filter=filter), np.ones((5, 20)), **tol)
        np.allclose(resizer(np.ones((9, 8, 3)), (13, 7), filter=filter), np.ones((13, 7, 3)), **tol)


if __name__ == '__main__':
  unittest.main()