{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c03a600f",
   "metadata": {},
   "source": [
    "# HealPixFFTConv: large-kernel FFT convolution\n",
    "\n",
    "This notebook checks the zero-padded FFT convolution against direct `conv2d`, verifies the identity-kernel projection round trip, tests gradients/CUDA, measures performance, and repeats the geometry test at the North Pole."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3fd76e2a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "device: cuda\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import time\n",
    "import healpix_geo\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from healpix_analyse import HealPixFFTConv\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "torch.manual_seed(0)\n",
    "np.random.seed(0)\n",
    "print(\"device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c5fd140",
   "metadata": {},
   "source": [
    "## Local NESTED patch\n",
    "\n",
    "The patch stays below the default ten-degree gnomonic limit. The same ordered `cell_ids` array defines the last dimension of every input and output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "db0f8d03",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cells: 1599\n"
     ]
    }
   ],
   "source": [
    "LEVEL = 8\n",
    "CENTRE = (30.0, 20.0)\n",
    "RADIUS_DEG = 5.0\n",
    "cell_ids, _, _ = healpix_geo.nested.cone_coverage(\n",
    "    CENTRE, RADIUS_DEG, LEVEL, ellipsoid=\"sphere\"\n",
    ")\n",
    "cell_ids = np.asarray(cell_ids, dtype=np.int64)\n",
    "print(\"cells:\", len(cell_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ed279c72",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HealPixFFTConv(\n",
      "  level=8, nside=256, in=2, out=3, kernel_sz=33, cells=1599, grid=64, fft_shape=(128, 128), radius=5.175 deg\n",
      "  (transform): LocalFFT(level=8, n_cells=1599, grid_shape=(64, 64), radius=5.175 deg, ellipsoid='sphere', dtype=torch.float64)\n",
      ")\n",
      "projected grid: (64, 64)\n",
      "padded FFT: (128, 128)\n"
     ]
    }
   ],
   "source": [
    "KERNEL_SZ = 33\n",
    "layer = HealPixFFTConv(\n",
    "    level=LEVEL,\n",
    "    in_channels=2,\n",
    "    out_channels=3,\n",
    "    kernel_sz=KERNEL_SZ,\n",
    "    cell_ids=cell_ids,\n",
    "    ellipsoid=\"sphere\",\n",
    "    device=device,\n",
    "    dtype=torch.float64,  # strict correctness check; benchmark is converted to float32 later\n",
    ")\n",
    "print(layer)\n",
    "print(\"projected grid:\", layer.transform.grid_shape)\n",
    "print(\"padded FFT:\", layer.fft_shape)\n",
    "assert layer.fft_shape[0] >= layer.grid_size + KERNEL_SZ - 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "484c3d7c",
   "metadata": {},
   "source": [
    "## FFT result versus direct convolution\n",
    "\n",
    "Both paths use the same gnomonic projection and HEALPix back-projection. They differ only in the planar convolution implementation. This strict comparison uses float64 because CUDA `conv2d` may otherwise use cuDNN/TF32 while the FFT path uses complex64 arithmetic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8856f52d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'max_abs': 3.6637359812630166e-15, 'relative_rms': 1.2503260022905908e-15}\n"
     ]
    }
   ],
   "source": [
    "weight = torch.randn(2, 3, KERNEL_SZ, KERNEL_SZ, device=device, dtype=torch.float64) / (2 * KERNEL_SZ**2) ** 0.5\n",
    "bias = torch.randn(3, device=device, dtype=torch.float64) * 0.01\n",
    "layer.set_kernel(weight, bias=bias, requires_grad=True)\n",
    "x = torch.randn(2, 2, len(cell_ids), device=device, dtype=torch.float64)\n",
    "\n",
    "y_fft = layer(x)\n",
    "grid = layer.transform._project_tensor(x)\n",
    "grid_direct = F.conv2d(\n",
    "    grid, layer.weight.permute(1, 0, 2, 3), padding=KERNEL_SZ // 2\n",
    ")\n",
    "y_direct = layer.transform._unproject_tensor(grid_direct) + layer.bias.view(1, -1, 1)\n",
    "\n",
    "max_abs = (y_fft - y_direct).abs().max().item()\n",
    "rel_rms = ((y_fft - y_direct).square().mean().sqrt() / y_direct.square().mean().sqrt().clamp_min(1e-12)).item()\n",
    "print({\"max_abs\": max_abs, \"relative_rms\": rel_rms})\n",
    "assert rel_rms < 2e-10, f\"FFT/direct relative RMS mismatch: {rel_rms:.3e}\"\n",
    "assert max_abs < 2e-9, f\"FFT/direct maximum mismatch: {max_abs:.3e}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc9dcdf1",
   "metadata": {},
   "source": [
    "## Identity kernel\n",
    "\n",
    "A centred delta must reproduce exactly the `project -> unproject` result. The remaining difference from the original HEALPix field is the interpolation error documented for `LocalFFT`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c972496c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "identity versus project/unproject max abs: 5.960464477539062e-07\n"
     ]
    }
   ],
   "source": [
    "identity_layer = HealPixFFTConv(\n",
    "    LEVEL, 1, 1, KERNEL_SZ, cell_ids=cell_ids, ellipsoid=\"sphere\", device=device\n",
    ")\n",
    "delta = torch.zeros(1, 1, KERNEL_SZ, KERNEL_SZ, device=device)\n",
    "delta[0, 0, KERNEL_SZ // 2, KERNEL_SZ // 2] = 1\n",
    "identity_layer.set_kernel(delta)\n",
    "signal = torch.randn(len(cell_ids), device=device)\n",
    "identity_output = identity_layer(signal)\n",
    "reference = identity_layer.transform.unproject(identity_layer.transform.project(signal))\n",
    "identity_max_abs = (identity_output - reference).abs().max().item()\n",
    "print(\"identity versus project/unproject max abs:\", identity_max_abs)\n",
    "assert torch.allclose(identity_output, reference, rtol=3e-5, atol=3e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72560e17",
   "metadata": {},
   "source": [
    "## Autograd\n",
    "\n",
    "The spatial kernel is transformed inside the graph during training, so gradients must reach both the HEALPix input and the kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c21b45e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.17717390947697245, 'input_grad_norm': 0.00801227055490017, 'kernel_grad_norm': 0.38059604792559915}\n"
     ]
    }
   ],
   "source": [
    "x_grad = torch.randn(2, 2, len(cell_ids), device=device, requires_grad=True)\n",
    "loss = layer(x_grad).square().mean()\n",
    "loss.backward()\n",
    "print({\n",
    "    \"loss\": loss.item(),\n",
    "    \"input_grad_norm\": x_grad.grad.norm().item(),\n",
    "    \"kernel_grad_norm\": layer.weight.grad.norm().item(),\n",
    "})\n",
    "assert torch.isfinite(x_grad.grad).all()\n",
    "assert torch.isfinite(layer.weight.grad).all()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c8b7e82",
   "metadata": {},
   "source": [
    "## Timing against direct conv2d\n",
    "\n",
    "This compares complete HEALPix-to-HEALPix paths using the same projection. Results depend strongly on device, batch size, grid and kernel size. Increase `KERNEL_SZ` and batch size to locate the crossover on the target GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a8ae4a4e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'fft_ms': 0.2798809204250574, 'direct_ms': 1.5176092041656375, 'speedup': 5.422338907063876}\n"
     ]
    }
   ],
   "source": [
    "def synchronize():\n",
    "    if device.type == \"cuda\":\n",
    "        torch.cuda.synchronize(device)\n",
    "\n",
    "def benchmark(function, repetitions=20):\n",
    "    for _ in range(3):\n",
    "        function()\n",
    "    synchronize()\n",
    "    start = time.perf_counter()\n",
    "    for _ in range(repetitions):\n",
    "        function()\n",
    "    synchronize()\n",
    "    return 1e3 * (time.perf_counter() - start) / repetitions\n",
    "\n",
    "layer = layer.float()\n",
    "x = x.float()\n",
    "layer.eval()\n",
    "with torch.inference_mode():\n",
    "    fft_ms = benchmark(lambda: layer(x))\n",
    "\n",
    "    def direct_path():\n",
    "        projected = layer.transform._project_tensor(x)\n",
    "        convolved = F.conv2d(\n",
    "            projected, layer.weight.permute(1, 0, 2, 3), padding=KERNEL_SZ // 2\n",
    "        )\n",
    "        return layer.transform._unproject_tensor(convolved)\n",
    "\n",
    "    direct_ms = benchmark(direct_path)\n",
    "print({\"fft_ms\": fft_ms, \"direct_ms\": direct_ms, \"speedup\": direct_ms / fft_ms})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a10ff6c",
   "metadata": {},
   "source": [
    "## North Pole geometry\n",
    "\n",
    "A final identity test verifies that constructing the tangent frame does not depend on a longitude direction at the pole."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "302b5ca1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'centre': (92.96285521726368, 90.0), 'radius_deg': 3.1069851092541816, 'max_abs': 4.76837158203125e-07}\n"
     ]
    }
   ],
   "source": [
    "polar_ids, _, _ = healpix_geo.nested.cone_coverage(\n",
    "    (0.0, 90.0), 3.0, LEVEL, ellipsoid=\"sphere\"\n",
    ")\n",
    "polar_ids = np.asarray(polar_ids, dtype=np.int64)\n",
    "polar = HealPixFFTConv(\n",
    "    LEVEL, 1, 1, 17, cell_ids=polar_ids, ellipsoid=\"sphere\", device=device\n",
    ")\n",
    "polar_delta = torch.zeros(1, 1, 17, 17, device=device)\n",
    "polar_delta[0, 0, 8, 8] = 1\n",
    "polar.set_kernel(polar_delta)\n",
    "polar_x = torch.randn(len(polar_ids), device=device)\n",
    "polar_y = polar(polar_x)\n",
    "polar_reference = polar.transform.unproject(polar.transform.project(polar_x))\n",
    "print({\n",
    "    \"centre\": (polar.transform.centre_lon_deg, polar.transform.centre_lat_deg),\n",
    "    \"radius_deg\": polar.transform.patch_radius_deg,\n",
    "    \"max_abs\": (polar_y - polar_reference).abs().max().item(),\n",
    "})\n",
    "assert torch.allclose(polar_y, polar_reference, rtol=3e-5, atol=3e-6)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}