{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Antenna-pattern-aware image formation\n",
    "\n",
    "A radar with a narrow azimuth beam (a scanning / stripmap-style acquisition) only illuminates each ground point for\n",
    "the short part of the aperture when the beam sweeps past it. During the rest of the\n",
    "aperture that point receives almost no signal, just noise.\n",
    "\n",
    "If the image is backprojected with every pulse weighted equally, those many low-gain pulses add noise into each pixel. With antenna pattern and pointing informtion (`g`, `g_extent`, `att`) it's possible to weight each pulse according to the gain toward that pixel.\n",
    "\n",
    "Both `torchbp.ops.backprojection_polar_2d` and `torchbp.ops.ffbp` accept the antenna pattern. This example uses `ffbp` and compares image formation with and without the antenna pattern on a noisy scanning acquisition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from numpy import hamming\n",
    "\n",
    "import torchbp\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(\"Device:\", device)\n",
    "torch.manual_seed(0);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Acquisition geometry\n",
    "\n",
    "A small FMCW radar flies along $y$ at 30 m altitude, looking toward $+x$. The azimuth\n",
    "is sampled at the Nyquist rate (element spacing $\\lambda/4$). Few point targets are spread across the swath so we can watch the beam sweep over each in turn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "fc = 6e9            # RF center frequency (Hz)\n",
    "bw = 200e6          # Sweep bandwidth (Hz)\n",
    "tsweep = 100e-6     # Sweep length (s)\n",
    "fs = 4e6            # Sampling frequency (Hz)\n",
    "nsamples = int(fs * tsweep)\n",
    "nsweeps = 512       # Pulses in the aperture\n",
    "oversample = 2      # Range FFT oversampling\n",
    "wl = 3e8 / fc\n",
    "element_spacing = 0.25 * wl   # Nyquist azimuth sampling\n",
    "altitude = 30.0\n",
    "\n",
    "# Point targets across the swath, all at 120 m ground range.\n",
    "r_mid = 120.0\n",
    "target_y = torch.tensor([-45.0, 0.0, 45.0], device=device)\n",
    "target_pos = torch.stack([torch.full_like(target_y, r_mid), target_y,\n",
    "                          torch.zeros_like(target_y)], dim=1)\n",
    "target_rcs = torch.ones((target_y.numel(), 1), dtype=torch.float32, device=device)\n",
    "\n",
    "# Straight track along y.\n",
    "pos = torch.zeros((nsweeps, 3), dtype=torch.float32, device=device)\n",
    "pos[:, 1] = torch.linspace(-nsweeps / 2, nsweeps / 2, nsweeps, device=device) * element_spacing\n",
    "pos[:, 2] = altitude\n",
    "\n",
    "grid = {\"r\": (95.0, 145.0), \"theta\": (-0.55, 0.55), \"nr\": 256, \"ntheta\": 512}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Antenna pattern and beam steering\n",
    "\n",
    "A tight pencil beam: narrow in azimuth, broad in elevation. The pattern `g`\n",
    "is the square root of the two-way gain on an elevation x azimuth grid, with extent\n",
    "`g_extent = [el0, az0, el1, az1]` in radians. During the pass the antenna yaw is swept\n",
    "linearly across ±25°, scanning the beam over the whole swath. A fixed roll points the elevation beam down at mid-swath."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "naz, nel = 256, 64\n",
    "az = np.linspace(-np.pi / 2, np.pi / 2, naz)\n",
    "el = np.linspace(-np.pi / 2, np.pi / 2, nel)\n",
    "az_bw = np.deg2rad(2)     # azimuth beamwidth\n",
    "el_bw = np.deg2rad(40)    # elevation beamwidth\n",
    "gain = np.exp(-(el[:, None] / el_bw) ** 2) * np.exp(-(az[None, :] / az_bw) ** 2)\n",
    "g = torch.tensor(gain, dtype=torch.float32, device=device)\n",
    "g_extent = [el[0], az[0], el[-1], az[-1]]\n",
    "\n",
    "roll = -np.arctan2(altitude, r_mid) # point elevation beam to mid-swath\n",
    "att = torch.zeros_like(pos)\n",
    "att[:, 0] = roll\n",
    "att[:, 2] = torch.linspace(-np.deg2rad(25), np.deg2rad(25), nsweeps, device=device)  # yaw sweep\n",
    "\n",
    "plt.figure(figsize=(5, 4))\n",
    "plt.imshow(20 * np.log10(gain + 1e-6), origin=\"lower\", aspect=\"auto\",\n",
    "           extent=[np.rad2deg(az[0]), np.rad2deg(az[-1]), np.rad2deg(el[0]), np.rad2deg(el[-1])],\n",
    "           vmin=-40, vmax=0)\n",
    "plt.colorbar(label=\"Gain (dB)\")\n",
    "plt.xlabel(\"Azimuth (deg)\"); plt.ylabel(\"Elevation (deg)\"); plt.title(\"Antenna pattern\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "### Each target is illuminated only briefly\n",
    "\n",
    "As the beam yaw sweeps, the gain toward a fixed target rises and falls. Outside it the target contributes essentially no signal, so those\n",
    "pulses carry only noise into the backprojection sum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "pos_y = pos[:, 1].cpu().numpy()\n",
    "yaw = att[:, 2].cpu().numpy()\n",
    "plt.figure(figsize=(7, 3.5))\n",
    "for ty in target_y.cpu().numpy():\n",
    "    az_to_target = np.arctan2(ty - pos_y, r_mid) - yaw\n",
    "    gk = np.exp(-(az_to_target / az_bw) ** 2)\n",
    "    plt.plot(np.arange(nsweeps), gk, label=f\"target y={ty:+.0f} m\")\n",
    "plt.xlabel(\"Pulse index\"); plt.ylabel(\"Relative gain toward target\")\n",
    "plt.title(\"Per-pulse illumination as the beam sweeps\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## Simulate the data and add noise\n",
    "\n",
    "The beat signals are generated with the antenna pattern and attitude, range-compressed\n",
    "(windowed, oversampled IFFT, spectrum centered with `data_fmod`), then thermal noise is\n",
    "added in the range-compressed domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw = torchbp.util.generate_fmcw_data(target_pos, target_rcs, pos, fc, bw, tsweep, fs,\n",
    "                                      g=g, g_extent=g_extent, att=att)\n",
    "w = torch.tensor(hamming(raw.shape[-1])[None, :], dtype=torch.float32, device=device)\n",
    "clean = torch.fft.ifft(raw * w, dim=-1, n=nsamples * oversample)\n",
    "data_fmod = -math.pi * (1 - (oversample - 1) / oversample)\n",
    "clean = clean * torch.exp(1j * data_fmod * torch.arange(clean.shape[-1], device=device))[None, :]\n",
    "r_res = 3e8 / (2 * bw * oversample)\n",
    "\n",
    "# Add noise\n",
    "sigma = clean.abs().pow(2).mean().sqrt()\n",
    "torch.manual_seed(7)\n",
    "noise = 30 * sigma / math.sqrt(2) * (torch.randn_like(clean.real) + 1j * torch.randn_like(clean.real))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "## Form the image, with and without the antenna pattern\n",
    "\n",
    "Backprojection is linear, so `ffbp(clean + noise) = ffbp(clean) + ffbp(noise)`. We form\n",
    "the signal-only and noise-only images for each method. Their sum is the noisy image to\n",
    "display, and they also give a clean per-target SNR. The uniform is given no\n",
    "antenna information and the antenna-aware gets `att`, `g`, `g_extent`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "def form(data, weighted):\n",
    "    kw = dict(data_fmod=data_fmod, dealias=True)\n",
    "    if weighted:\n",
    "        kw.update(att=att, g=g, g_extent=g_extent)\n",
    "    return torchbp.ops.ffbp(data, grid, fc, r_res, pos, stages=3, divisions=2, **kw).squeeze()\n",
    "\n",
    "S_uni, N_uni = form(clean, False), form(noise, False)   # uniform: signal, noise\n",
    "S_ant, N_ant = form(clean, True),  form(noise, True)    # antenna-aware: signal, noise\n",
    "img_uni = S_uni + N_uni\n",
    "img_ant = S_ant + N_ant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "extent = [*grid[\"r\"], *grid[\"theta\"]]\n",
    "fig, axes = plt.subplots(1, 2, figsize=(13, 4.5))\n",
    "vmax = 20 * torch.log10(img_ant.abs() + 1e-12).max().item()\n",
    "for ax, im, title in zip(axes, [img_uni, img_ant],\n",
    "                         [\"Uniform (no antenna pattern)\", \"Antenna-aware\"]):\n",
    "    db = 20 * torch.log10(im.abs() + 1e-12).cpu().numpy()\n",
    "    h = ax.imshow(db.T, origin=\"lower\", vmin=vmax - 30, vmax=vmax, extent=extent, aspect=\"auto\")\n",
    "    ax.set_title(title); ax.set_xlabel(\"Range (m)\")\n",
    "    plt.colorbar(h, ax=ax, label=\"dB\")\n",
    "axes[0].set_ylabel(\"Angle (sin)\")\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "The targets stand out more cleanly against the noise in the antenna-aware image."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "## Per-target SNR\n",
    "\n",
    "Using the separate signal and noise images, SNR at each target is the signal peak over\n",
    "the local RMS of the noise image at that pixel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_peaks(I, n):\n",
    "    m = I.abs().cpu().numpy().copy(); pts = []\n",
    "    for _ in range(n):\n",
    "        ri, ci = np.unravel_index(m.argmax(), m.shape); pts.append((ri, ci))\n",
    "        m[max(0, ri-12):ri+12, max(0, ci-12):ci+12] = 0\n",
    "    return sorted(pts, key=lambda p: p[1])\n",
    "\n",
    "print(\"target  |  uniform  | antenna-aware |  gain\")\n",
    "for ri, ci in find_peaks(S_uni, target_y.numel()):\n",
    "    nu = N_uni[ri-8:ri+8, ci-8:ci+8].abs().std().item()\n",
    "    na = N_ant[ri-8:ri+8, ci-8:ci+8].abs().std().item()\n",
    "    snr_u = 20 * math.log10(S_uni[ri, ci].abs().item() / nu)\n",
    "    snr_a = 20 * math.log10(S_ant[ri, ci].abs().item() / na)\n",
    "    print(f\"  c{ci:>4} | {snr_u:6.1f} dB | {snr_a:10.1f} dB | {snr_a - snr_u:+5.1f} dB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "Every target gains several dB. The improvement is the matched-filter gain\n",
    "$N\\,\\Sigma g_k^2 / (\\Sigma g_k)^2$. The antenna-aware image formation effectively integrates only\n",
    "the pulses that actually saw the target, instead of diluting the signal with noise-only pulses."
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}