orisys-sdk/tests/test_masked_poisson_capfft.py

from pathlib import Path

import numpy as np
import pytest

from orisys.backend_availability import CUPY_AVAILABLE
from orisys.hhd import PoissonSolver, nHHD

REPO_ROOT = Path(__file__).resolve().parents[1]


def _circle_mask(shape, radius_ratio=0.35):
    h, w = shape
    yy, xx = np.mgrid[0:h, 0:w]
    cy = (h - 1) / 2.0
    cx = (w - 1) / 2.0
    radius = min(h, w) * radius_ratio
    return ((yy - cy) ** 2 + (xx - cx) ** 2) <= radius ** 2


def _smooth_rhs(shape):
    h, w = shape
    yy, xx = np.mgrid[0:h, 0:w]
    field = (
        np.sin(np.pi * xx / max(1, w - 1))
        + 0.5 * np.cos(2.0 * np.pi * yy / max(1, h - 1))
    )
    return field.astype(np.float32)


def _radial_flow(shape, scale=0.05):
    h, w = shape
    yy, xx = np.mgrid[0:h, 0:w]
    cy = (h - 1) / 2.0
    cx = (w - 1) / 2.0
    flow = np.zeros((h, w, 2), dtype=np.float32)
    flow[:, :, 0] = (xx - cx) * scale
    flow[:, :, 1] = (yy - cy) * scale
    return flow


def _masked_rhs_vector(rhs, mask, pinned):
    vec = rhs[mask].astype(np.float64)
    vec -= vec.mean()
    vec[pinned] = 0.0
    return vec


def _masked_residual_norm(solver, field, rhs, mask):
    cache = solver._masked_cache if solver._masked_cache_matches(mask) else solver._build_masked_cache(mask)
    rhs_vec = _masked_rhs_vector(rhs, mask, cache["pinned"])
    residual = cache["lap"].dot(field[mask].astype(np.float64)) - rhs_vec
    return float(np.linalg.norm(residual))


def test_capfft_solver_improves_over_raw_fft_on_synthetic_mask():
    shape = (40, 48)
    mask = _circle_mask(shape)
    rhs = _smooth_rhs(shape)
    rhs_masked = np.where(mask, rhs, 0).astype(np.float32)

    sparse_solver = PoissonSolver(
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="sparse",
    )
    capfft_solver = PoissonSolver(
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="capfft",
        masked_poisson_max_boundary=48,
        masked_poisson_max_iters=3,
    )

    sparse = sparse_solver.solve(rhs_masked, mask=mask)
    capfft = capfft_solver.solve(rhs_masked, mask=mask)
    raw_fft = sparse_solver.solve(rhs_masked, mask=None)

    assert np.allclose(capfft[~mask], 0.0, atol=1e-6)

    err_cap = float(np.linalg.norm((capfft - sparse)[mask]))
    err_raw = float(np.linalg.norm((raw_fft - sparse)[mask]))
    assert err_cap < err_raw

    res_cap = _masked_residual_norm(sparse_solver, capfft, rhs_masked, mask)
    res_raw = _masked_residual_norm(sparse_solver, raw_fft, rhs_masked, mask)
    assert res_cap < res_raw


def test_capfft_hhd_forces_track_sparse_reference():
    shape = (36, 44)
    mask = _circle_mask(shape)
    flow = _radial_flow(shape, scale=0.03)
    flow[~mask] = 0.0

    sparse_hhd = nHHD(
        cuda=False,
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="sparse",
    )
    capfft_hhd = nHHD(
        cuda=False,
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="capfft",
        masked_poisson_max_boundary=40,
        masked_poisson_max_iters=3,
    )

    _, sparse_forces = sparse_hhd.get_force(flow, valid_mask=mask)
    _, capfft_forces = capfft_hhd.get_force(flow, valid_mask=mask)

    denom = max(1.0, float(np.linalg.norm(sparse_forces)))
    rel = float(np.linalg.norm(capfft_forces - sparse_forces) / denom)
    assert rel < 0.6

    vnormal = capfft_hhd.get_normal()
    vshear = capfft_hhd.get_shear()
    assert np.allclose(vnormal[~mask], 0.0, atol=1e-5)
    assert np.allclose(vshear[~mask], 0.0, atol=1e-5)


def test_sparse_hhd_normal_sign_tracks_capfft_on_masked_uv_flow():
    shape = (36, 44)
    mask = _circle_mask(shape)
    flow = _radial_flow(shape, scale=0.03)
    flow[~mask] = 0.0

    sparse_hhd = nHHD(
        cuda=False,
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="sparse",
    )
    capfft_hhd = nHHD(
        cuda=False,
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="capfft",
        masked_poisson_max_boundary=40,
        masked_poisson_max_iters=3,
    )

    sparse_hhd.get_force(flow, valid_mask=mask)
    capfft_hhd.get_force(flow, valid_mask=mask)

    sparse_vnormal = sparse_hhd.get_normal()
    capfft_vnormal = capfft_hhd.get_normal()

    # The masked sparse solver should use the same sign convention as capfft/FFT.
    sparse_flat = sparse_vnormal[mask].reshape(-1, 2)
    capfft_flat = capfft_vnormal[mask].reshape(-1, 2)
    dot = float(np.sum(sparse_flat * capfft_flat))
    assert dot > 0.0

    denom = max(1.0, float(np.linalg.norm(capfft_flat)))
    rel = float(np.linalg.norm(sparse_flat - capfft_flat) / denom)
    assert rel < 0.35


def test_capfft_uv_mask_fixture_improves_over_raw_fft():
    uv_path = REPO_ROOT / "unroll_map_uv.npz"
    if not uv_path.exists():
        return

    full_mask = np.load(uv_path)["valid_mask"].astype(bool)
    mask = full_mask[::24, ::24]
    if mask.sum() < 16:
        mask = full_mask[::16, ::16]

    shape = mask.shape
    rhs = _smooth_rhs(shape)
    rhs_masked = np.where(mask, rhs, 0).astype(np.float32)

    sparse_solver = PoissonSolver(
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="sparse",
    )
    capfft_solver = PoissonSolver(
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="capfft",
        masked_poisson_max_boundary=24,
        masked_poisson_max_iters=3,
    )

    sparse = sparse_solver.solve(rhs_masked, mask=mask)
    capfft = capfft_solver.solve(rhs_masked, mask=mask)
    raw_fft = sparse_solver.solve(rhs_masked, mask=None)

    err_cap = float(np.linalg.norm((capfft - sparse)[mask]))
    err_raw = float(np.linalg.norm((raw_fft - sparse)[mask]))
    assert err_cap < err_raw


@pytest.mark.skipif(not CUPY_AVAILABLE, reason="CuPy not available")
def test_capfft_solver_cupy_backend_uses_gpu_bulk_solve_and_tracks_cpu():
    shape = (32, 40)
    mask = _circle_mask(shape)
    rhs = _smooth_rhs(shape)
    rhs_masked = np.where(mask, rhs, 0).astype(np.float32)

    cpu_solver = PoissonSolver(
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="capfft",
        masked_poisson_max_boundary=40,
        masked_poisson_max_iters=3,
    )
    cupy_solver = PoissonSolver(
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cupy",
        masked_poisson_solver="capfft",
        masked_poisson_max_boundary=40,
        masked_poisson_max_iters=3,
    )

    cpu_capfft = cpu_solver.solve(rhs_masked, mask=mask)
    raw_fft = cpu_solver.solve(rhs_masked, mask=None)

    original_gpu_solve = cupy_solver._rect_solver._solve_gpu_internal
    calls = {"count": 0}

    def counted_gpu_solve(f_gpu):
        calls["count"] += 1
        return original_gpu_solve(f_gpu)

    cupy_solver._rect_solver._solve_gpu_internal = counted_gpu_solve
    cupy_capfft = cupy_solver.solve(rhs_masked, mask=mask)

    assert calls["count"] > 0
    assert np.allclose(cupy_capfft[~mask], 0.0, atol=1e-6)

    err_gpu = float(np.linalg.norm((cupy_capfft - cpu_capfft)[mask]))
    err_raw = float(np.linalg.norm((raw_fft - cpu_capfft)[mask]))
    assert err_gpu < err_raw

    denom = max(1.0, float(np.linalg.norm(cpu_capfft[mask])))
    rel = err_gpu / denom
    assert rel < 1e-2


@pytest.mark.skipif(not CUPY_AVAILABLE, reason="CuPy not available")
def test_capfft_hhd_cupy_backend_tracks_cpu_capfft_reference():
    shape = (36, 44)
    mask = _circle_mask(shape)
    flow = _radial_flow(shape, scale=0.03)
    flow[~mask] = 0.0

    cpu_hhd = nHHD(
        cuda=False,
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cpu",
        masked_poisson_solver="capfft",
        masked_poisson_max_boundary=40,
        masked_poisson_max_iters=3,
    )
    cupy_hhd = nHHD(
        cuda=False,
        grid=shape,
        spacings=(1.0, 1.0),
        poisson_backend="cupy",
        masked_poisson_solver="capfft",
        masked_poisson_max_boundary=40,
        masked_poisson_max_iters=3,
    )

    _, cpu_forces = cpu_hhd.get_force(flow, valid_mask=mask)
    _, cupy_forces = cupy_hhd.get_force(flow, valid_mask=mask)

    denom = max(1.0, float(np.linalg.norm(cpu_forces)))
    rel = float(np.linalg.norm(cupy_forces - cpu_forces) / denom)
    assert rel < 1e-2

    vnormal = cupy_hhd.get_normal()
    vshear = cupy_hhd.get_shear()
    assert np.allclose(vnormal[~mask], 0.0, atol=1e-5)
    assert np.allclose(vshear[~mask], 0.0, atol=1e-5)