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)