{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Generalized phase gradient autofocus (GPGA)\n",
    "\n",
    "In real Synthetic Aperture Radar (SAR) acquisitions the platform never follows\n",
    "its nominal trajectory exactly. Even small uncompensated motion errors translate\n",
    "into phase errors across the synthetic aperture, which blur and defocus the\n",
    "backprojected image.\n",
    "\n",
    "Generalized phase gradient autofocus (GPGA) with trajectory deviation estimation estimates these errors directly from\n",
    "the data. The image is split into several subimages spread across range and\n",
    "azimuth. The slant-range error to each subimage is estimated from its phase\n",
    "history, and a 3D position error of the platform is then solved from the\n",
    "collection of slant-range errors.\n",
    "\n",
    "This example simulates a moving FMCW radar with a random motion error,\n",
    "images it with an backprojection, and then recovers the\n",
    "trajectory using `torchbp.autofocus.gpga_bp_polar_tde`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import torchbp\n",
    "import matplotlib.pyplot as plt\n",
    "from numpy import hamming\n",
    "from torchbp.util import detrend\n",
    "\n",
    "torch.manual_seed(125);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "Use CUDA if available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = \"cuda\"\n",
    "else:\n",
    "    device = \"cpu\"\n",
    "print(\"Device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Radar and acquisition geometry\n",
    "\n",
    "We model a small FMCW radar flown along the $y$ axis at 20 m altitude, looking in\n",
    "the $+x$ direction. The platform records `nsweeps` measurements that together\n",
    "form the synthetic aperture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ntheta = 512    # Azimuth points in the image\n",
    "nsweeps = 512   # Number of measurements (aperture positions)\n",
    "fc = 6e9        # RF center frequency (Hz)\n",
    "bw = 300e6      # RF bandwidth (Hz)\n",
    "tsweep = 200e-6 # Sweep length (s)\n",
    "fs = 3e6        # Sampling frequency (Hz)\n",
    "nsamples = int(fs * tsweep)  # Time-domain samples per sweep\n",
    "\n",
    "wl = 3e8 / fc                # Wavelength (m)\n",
    "range_res = 3e8 / (2 * bw)   # Range resolution (m)\n",
    "\n",
    "# Imaging grid. Azimuth angle \"theta\" is sin of radians.\n",
    "nr = 2 * int(1 + (70 - 30) / range_res)\n",
    "grid_polar = {\"r\": (30, 70), \"theta\": (-0.7, 0.7), \"nr\": nr, \"ntheta\": ntheta}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## Targets\n",
    "\n",
    "The scene is a 5x5 grid of point targets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "targets = []\n",
    "for x in np.linspace(35, 65, 5):\n",
    "    for y in np.linspace(-20, 20, 5):\n",
    "        targets.append([x, y, 0])\n",
    "target_pos = torch.tensor(targets, dtype=torch.float32, device=device)\n",
    "target_rcs = torch.ones(target_pos.shape[0], dtype=torch.float32, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## Nominal trajectory and motion error\n",
    "\n",
    "`pos` is the true trajectory the platform actually flew and `pos_error` is the ideal trajectory: a straight line along $y$ at a\n",
    "fixed altitude. The difference between them is a smooth, band-limited random motion error added to\n",
    "all three axes. It's generated by low-pass filtering white noise\n",
    "in the frequency domain.\n",
    "\n",
    "Linear trend is removed since it can't be solved since it only shifts the image.\n",
    "\n",
    "The autofocus algorithm only gets to see the nominal `pos_error` and its job is to recover `pos`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "pos = torch.zeros([nsweeps, 3], dtype=torch.float32, device=device)\n",
    "pos[:, 1] = torch.linspace(-nsweeps / 2, nsweeps / 2, nsweeps) * 0.25 * wl\n",
    "pos[:, 2] = 20\n",
    "\n",
    "# Nominal trajectory without motion error\n",
    "pos_error = pos.clone()\n",
    "\n",
    "window_width = nsweeps // 32  # Low-pass cutoff for the error spectrum\n",
    "error_scale = 10\n",
    "\n",
    "def bandlimited_error(scale):\n",
    "    \"Smooth band-limited random error in units of wavelength * scale\"\n",
    "    r = scale * torch.rand(nsweeps, device=device) * wl\n",
    "    fdata = torch.fft.fft(r)\n",
    "    fdata[window_width // 2:-window_width // 2] = 0\n",
    "    return torch.fft.ifft(fdata).real\n",
    "\n",
    "# Range (x), Along-track (y) and vertical (z) motion errors.\n",
    "pos[:, 0] = detrend(bandlimited_error(error_scale))\n",
    "pos[:, 0] -= torch.mean(pos[:, 0])\n",
    "\n",
    "ey = bandlimited_error(0.25 * error_scale)\n",
    "ey -= torch.mean(ey)\n",
    "pos[:, 1] += detrend(ey)\n",
    "\n",
    "ez = bandlimited_error(0.5 * error_scale)\n",
    "ez -= torch.mean(ez)\n",
    "pos[:, 2] += detrend(ez)\n",
    "\n",
    "print(\"Track length:\", (pos[-1, 1] - pos[0, 1]).item(), \"m\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Motion error (true - nominal)\")\n",
    "plt.plot((pos[:, 0] - pos_error[:, 0]).cpu().numpy() / wl, label=\"x error\")\n",
    "plt.plot((pos[:, 1] - pos_error[:, 1]).cpu().numpy() / wl, label=\"y error\")\n",
    "plt.plot((pos[:, 2] - pos_error[:, 2]).cpu().numpy() / wl, label=\"z error\")\n",
    "plt.xlabel(\"Sweep index\")\n",
    "plt.ylabel(\"Position error (wavelengths)\")\n",
    "plt.legend(loc=\"best\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Generate raw radar data\n",
    "\n",
    "`generate_fmcw_data` simulates the range-compressed FMCW echoes seen along the\n",
    "*true* trajectory `pos`. We apply a Hamming window in range and inverse FFT with\n",
    "oversampling to reduce interpolation errors during backprojection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Oversampling input data decreases interpolation errors.\n",
    "oversample = 3\n",
    "\n",
    "with torch.no_grad():\n",
    "    data = torchbp.util.generate_fmcw_data(\n",
    "        target_pos, target_rcs, pos, fc, bw, tsweep, fs, rvp=False)\n",
    "    # Apply windowing function in range direction.\n",
    "    w = torch.tensor(hamming(data.shape[-1])[None, :],\n",
    "                     dtype=torch.float32, device=device)\n",
    "    data = torch.fft.ifft(data * w, dim=-1, n=nsamples * oversample)\n",
    "\n",
    "r_res = 3e8 / (2 * bw * oversample)  # Range bin size in the input data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "## Uncorrected image\n",
    "\n",
    "Backprojecting with the nominal trajectory `pos_error` has motion errors, so the point targets are defocused."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "img = torchbp.ops.backprojection_polar_2d(data, grid_polar, fc, r_res, pos_error)\n",
    "img = img.squeeze()  # Remove singular batch dimension\n",
    "\n",
    "img_db = 20 * torch.log10(torch.abs(img)).detach()\n",
    "m = torch.max(img_db)\n",
    "extent = [*grid_polar[\"r\"], *grid_polar[\"theta\"]]\n",
    "\n",
    "plt.figure()\n",
    "plt.title(\"Uncorrected image\")\n",
    "plt.imshow(img_db.cpu().numpy().T, origin=\"lower\", vmin=m - 30, vmax=m,\n",
    "           extent=extent, aspect=\"auto\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Angle (sin radians)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "## Run GPGA autofocus\n",
    "\n",
    "`gpga_bp_polar_tde` divides the image into `range_divisions` x\n",
    "`azimuth_divisions` subimages, estimates the slant-range error to each, and\n",
    "solves for the 3D platform position. `estimate_z=True` recovers the cross-track\n",
    "$z$ error as well, which is possible here because the look angle variation is large.\n",
    "\n",
    "It returns the refocused image and the recovered trajectory `pos_new`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_pga, pos_new = torchbp.autofocus.gpga_bp_polar_tde(\n",
    "    img, data, pos_error, fc, r_res, grid_polar,\n",
    "    range_divisions=3, azimuth_divisions=3, estimate_z=True)\n",
    "\n",
    "rms = lambda a, b: torch.sqrt(torch.mean(torch.square(a - b))).item()\n",
    "print(\"RMS error total:\", 1e3*rms(pos_new, pos), \"mm\")\n",
    "print(\"RMS error X:\", 1e3*rms(pos_new[:, 0], pos[:, 0]), \"mm\")\n",
    "print(\"RMS error Y:\", 1e3*rms(pos_new[:, 1], pos[:, 1]), \"mm\")\n",
    "print(\"RMS error Z:\", 1e3*rms(pos_new[:, 2], pos[:, 2]), \"mm\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "## Recovered trajectory\n",
    "\n",
    "The residual error between the recovered trajectory and the true one should be small if everything went correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "pos_new_np = pos_new.cpu().numpy().T\n",
    "\n",
    "plt.figure()\n",
    "plt.title(\"Residual error after autofocus in wavelengths\")\n",
    "plt.plot((pos[:, 0].cpu().numpy() - pos_new_np[0]) / wl, label=\"x error\")\n",
    "plt.plot((pos[:, 1].cpu().numpy() - pos_new_np[1]) / wl, label=\"y error\")\n",
    "plt.plot((pos[:, 2].cpu().numpy() - pos_new_np[2]) / wl, label=\"z error\")\n",
    "plt.legend(loc=\"best\")\n",
    "plt.xlabel(\"Sweep index\")\n",
    "plt.ylabel(\"Position error (wavelengths)\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(3, 1, figsize=(10, 8), sharex=True)\n",
    "\n",
    "axs[0].plot(pos[:, 0].cpu().numpy(), label=\"Actual X position\")\n",
    "axs[0].plot(pos_new_np[0], label=\"Solved X position\")\n",
    "axs[0].legend(loc=\"upper right\")\n",
    "axs[0].set_ylabel(\"X\")\n",
    "\n",
    "# Subtract the (large) nominal track from y so the error is visible.\n",
    "axs[1].plot(pos[:, 1].cpu().numpy() - pos_error[:, 1].cpu().numpy(),\n",
    "            label=\"Actual Y position\")\n",
    "axs[1].plot(pos_new_np[1] - pos_error[:, 1].cpu().numpy(),\n",
    "            label=\"Solved Y position\")\n",
    "axs[1].legend(loc=\"upper right\")\n",
    "axs[1].set_ylabel(\"Y\")\n",
    "\n",
    "axs[2].plot(pos[:, 2].cpu().numpy(), label=\"Actual Z position\")\n",
    "axs[2].plot(pos_new_np[2], label=\"Solved Z position\")\n",
    "axs[2].legend(loc=\"upper right\")\n",
    "axs[2].set_ylabel(\"Z\")\n",
    "axs[2].set_xlabel(\"Sweep index\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "## Focused image\n",
    "\n",
    "After autofocus the point targets are well focused. The image is still in pseudo-polar coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_pga_db = 20 * torch.log10(torch.abs(img_pga)).detach()\n",
    "m = torch.max(img_pga_db)\n",
    "\n",
    "plt.figure()\n",
    "plt.title(\"Image after GPGA autofocus\")\n",
    "plt.imshow(img_pga_db.cpu().numpy().T, origin=\"lower\", vmin=m - 30, vmax=m,\n",
    "           extent=extent, aspect=\"auto\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Angle (sin radians)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "## Cartesian comparison\n",
    "\n",
    "Finally we resample both the uncorrected and the autofocused images from polar\n",
    "to Cartesian coordinates with `polar_to_cart` for a side-by-side comparison.\n",
    "\n",
    "The correct image should look like 5x5 grid of points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "origin = torch.mean(pos_error, axis=0)\n",
    "r0, r1 = grid_polar[\"r\"]\n",
    "y1 = 1.5 * (r1 - r0) / 2\n",
    "grid = {\"x\": grid_polar[\"r\"], \"y\": (-y1, y1),\n",
    "        \"nx\": 3 * grid_polar[\"nr\"], \"ny\": 3 * grid_polar[\"nr\"]}\n",
    "cart_extent = [grid[\"x\"][0], grid[\"x\"][1], grid[\"y\"][0], grid[\"y\"][1]]\n",
    "\n",
    "def to_cart(image):\n",
    "    c = torchbp.ops.polar_to_cart(torch.abs(image), origin, grid_polar, grid, fc)\n",
    "    return 20 * torch.log10(torch.abs(c) + 1e-10)\n",
    "\n",
    "img_cart = to_cart(img)\n",
    "img_pga_cart = to_cart(img_pga)\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=(12, 5))\n",
    "axs[0].set_title(\"Uncorrected\")\n",
    "axs[0].imshow(img_cart.cpu().numpy().T, origin=\"lower\", vmin=m - 30, vmax=m,\n",
    "              extent=cart_extent, aspect=\"auto\")\n",
    "axs[0].set_xlabel(\"X (m)\")\n",
    "axs[0].set_ylabel(\"Y (m)\")\n",
    "axs[0].grid(False)\n",
    "\n",
    "axs[1].set_title(\"After GPGA TDE autofocus\")\n",
    "axs[1].imshow(img_pga_cart.cpu().numpy().T, origin=\"lower\", vmin=m - 30, vmax=m,\n",
    "              extent=cart_extent, aspect=\"auto\")\n",
    "axs[1].set_xlabel(\"X (m)\")\n",
    "axs[1].set_ylabel(\"Y (m)\")\n",
    "axs[1].grid(False)"
   ]
  }
 ],
 "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
}