utils.py

from functools import partial
import jax
import jax.numpy as jnp
import numpy as np


def repeat_vmap(fun, in_axes=[0]):
    for axes in in_axes:
        fun = jax.vmap(fun, in_axes=axes)
    return fun


def make_grid(patch_size: int | tuple[int, int]):
    """Gera grid de coordenadas com segurança numérica"""
    # Garantir tamanho mínimo de 8x8
    if isinstance(patch_size, int):
        h = w = max(8, patch_size)
    else:
        h, w = (max(8, ps) for ps in patch_size)

    # Espaçamento preciso entre pontos
    y_space = np.linspace(-0.5 + 1 / (2 * h), 0.5 - 1 / (2 * h), h)
    x_space = np.linspace(-0.5 + 1 / (2 * w), 0.5 - 1 / (2 * w), w)

    # Criar grid com dimensões (1, H, W, 2)
    grid = np.stack(np.meshgrid(y_space, x_space, indexing='ij'), axis=-1)
    return grid[np.newaxis, ...]


def interpolate_grid(coords, grid, order=0):
    """Interpolação segura com verificação de dimensões"""
    try:
        # Converter para JAX array e validar formato
        coords = jnp.asarray(coords)
        if coords.ndim != 4 or coords.shape[-1] != 2:
            raise ValueError(
                f"Dimensões inválidas: {coords.shape}. Esperado (B, H, W, 2)"
            )

        # Transformação de coordenadas
        coords = coords.transpose((0, 3, 1, 2))
        coords = coords.at[:, 0].set(
            coords[:, 0] * (grid.shape[-3] - 1) + (grid.shape[-3] - 1) / 2
        )
        coords = coords.at[:, 1].set(
            coords[:, 1] * (grid.shape[-2] - 1) + (grid.shape[-2] - 1) / 2
        )

        # Interpolação vetorizada
        map_fn = jax.vmap(jax.vmap(
            partial(jax.scipy.ndimage.map_coordinates, order=order, mode='nearest'),
            in_axes=(2, None),
            out_axes=2
        ))
        return map_fn(grid, coords)

    except Exception as e:
        raise RuntimeError(f"Erro de interpolação: {str(e)}") from e