{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Backprojection on DEM\n",
    "\n",
    "Backprojection normally assumes that every image pixel lies on the $z = 0$\n",
    "plane. A scatterer at height above the zero-height plane has a shorter slant range than a flat\n",
    "ground point at the same image coordinates, so on a flat-earth image it lands\n",
    "displaced towards the radar (layover). The displacement grows with target\n",
    "height and with the look angle, and since it varies across the scene with the\n",
    "terrain the image gets warped.\n",
    "\n",
    "`backprojection_polar_2d` and `ffbp` take an optional `dem` argument that\n",
    "replaces the flat-earth assumption with per-pixel terrain heights:\n",
    "\n",
    "- `dem` is a float32 tensor of shape `[dem_nr, dem_ntheta]` covering the\n",
    "  same r/theta extent as the image grid. It can be much coarser than the\n",
    "  image grid, heights are bilinearly interpolated at each pixel.\n",
    "- Values are pixel z coordinates in the same frame as the platform positions.\n",
    "- DEM sample `[i, j]` corresponds to image pixel `(i * nr / dem_nr,\n",
    "  j * ntheta / dem_ntheta)`. Sample `i` sits at\n",
    "  `r0 + i * (r1 - r0) / dem_nr`.\n",
    "- `torchbp.util.dem_to_polar` resamples a real Cartesian DEM (e.g. from a\n",
    "  GeoTIFF) to a polar grid. In this example a synthetic DEM is built directly.\n",
    "\n",
    "This notebook builds a synthetic terrain with a grid of point targets on it,\n",
    "and shows that with the DEM the targets focus at their correct image\n",
    "coordinates while the flat-earth image misplaces them."
   ]
  },
  {
   "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",
    "from torchbp.ops import backprojection_polar_2d, ffbp\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(\"Device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Scene: terrain and a grid of point targets\n",
    "\n",
    "The terrain is a smooth 15 m high hill in the middle of the scene.\n",
    "A 5x5 grid of equal point targets is placed on the terrain surface, so\n",
    "targets near the scene center are elevated and the ones at the edges are close\n",
    "to the ground plane. The platform flies a straight track along y at 50 m\n",
    "height, so the look angle is fairly steep and the layover displacement is\n",
    "several meters at the hilltop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "fc = 6e9          # RF center frequency (Hz)\n",
    "bw = 200e6        # RF bandwidth (Hz)\n",
    "tsweep = 100e-6   # Sweep length (s)\n",
    "fs = 2e6          # Sampling frequency (Hz)\n",
    "nsweeps = 512     # Number of pulses in the aperture\n",
    "oversample = 2    # Range FFT oversampling\n",
    "nsamples = int(fs * tsweep)\n",
    "\n",
    "nr, ntheta = 320, 384\n",
    "grid = {\"r\": (70.0, 130.0), \"theta\": (-0.3, 0.3), \"nr\": nr, \"ntheta\": ntheta}\n",
    "r0, r1 = grid[\"r\"]\n",
    "t0, t1 = grid[\"theta\"]\n",
    "\n",
    "def terrain(x, y):\n",
    "    \"\"\"Terrain height (m) as a function of world x, y.\"\"\"\n",
    "    return 15.0 * torch.exp(-((x - 100.0)**2 + y**2) / (2 * 25.0**2))\n",
    "\n",
    "# 5x5 grid of point targets on the terrain surface.\n",
    "rt = torch.linspace(80.0, 120.0, 5)\n",
    "tt = torch.linspace(-0.24, 0.24, 5)\n",
    "R, T = torch.meshgrid(rt, tt, indexing=\"ij\")\n",
    "xt = (R * torch.sqrt(1 - T**2)).flatten()\n",
    "yt = (R * T).flatten()\n",
    "zt = terrain(xt, yt)\n",
    "target_pos = torch.stack([xt, yt, zt], dim=1).to(torch.float32).to(device)\n",
    "target_rcs = torch.ones((target_pos.shape[0], 1), dtype=torch.float32, device=device)\n",
    "\n",
    "# Straight track along y at 50 m height, lambda/4 element spacing.\n",
    "pos = torch.zeros((nsweeps, 3), dtype=torch.float32, device=device)\n",
    "pos[:, 1] = torch.linspace(-nsweeps / 2, nsweeps / 2, nsweeps, device=device) * 0.25 * 3e8 / fc\n",
    "pos[:, 2] = 50.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Synthetic DEM on the polar grid\n",
    "\n",
    "The DEM is sampled at 64x64 points, much coarser than the 320x384\n",
    "image grid, to demonstrate that the kernel's bilinear interpolation handles\n",
    "the resolution difference. Sample `[i, j]` is placed at the cell-start\n",
    "coordinates `r = r0 + i*(r1-r0)/dem_nr`, `theta = t0 + j*(t1-t0)/dem_ntheta`\n",
    "to match the kernel's pixel-to-DEM index map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "dem_nr, dem_ntheta = 64, 64\n",
    "r_dem = r0 + (r1 - r0) * torch.arange(dem_nr, device=device) / dem_nr\n",
    "t_dem = t0 + (t1 - t0) * torch.arange(dem_ntheta, device=device) / dem_ntheta\n",
    "x_dem = r_dem[:, None] * torch.sqrt(1 - t_dem[None, :]**2)\n",
    "y_dem = r_dem[:, None] * t_dem[None, :]\n",
    "dem = terrain(x_dem, y_dem).to(torch.float32)\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(dem.cpu().numpy().T, origin=\"lower\", extent=[r0, r1, t0, t1], aspect=\"auto\")\n",
    "plt.colorbar(label=\"Height (m)\")\n",
    "plt.plot((xt**2 + yt**2).sqrt().cpu(), (yt / (xt**2 + yt**2).sqrt()).cpu(), \"r+\", ms=10, label=\"Targets\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Angle (sin)\")\n",
    "plt.title(\"DEM with point target locations\")\n",
    "plt.legend(loc=\"best\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## Simulate and range-compress the data\n",
    "\n",
    "`generate_fmcw_data` uses\n",
    "the true 3D target positions, so the raw data contains the real\n",
    "terrain geometry. Range window, FFT range compression, and the `data_fmod`\n",
    "spectrum shift that reduces interpolation error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw = torchbp.util.generate_fmcw_data(target_pos, target_rcs, pos, fc, bw, tsweep, fs)\n",
    "w = torch.tensor(hamming(raw.shape[-1])[None, :], dtype=torch.float32, device=device)\n",
    "data = torch.fft.ifft(raw * w, dim=-1, n=nsamples * oversample)\n",
    "\n",
    "data_fmod = -math.pi * (1 - (oversample - 1) / oversample)\n",
    "data = data * torch.exp(1j * data_fmod * torch.arange(data.shape[-1], device=device))[None, :]\n",
    "\n",
    "r_res = 3e8 / (2 * abs(bw) * oversample)\n",
    "dr = (grid[\"r\"][1] - grid[\"r\"][0]) / grid[\"nr\"]\n",
    "im_margin = oversample * r_res / dr - 1\n",
    "alias_fmod = -math.pi * (1 - im_margin / (1 + im_margin))\n",
    "print(f\"r_res = {r_res:.3f} m, image dr = {dr:.3f} m\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "extent = [r0, r1, t0, t1]\n",
    "\n",
    "def show_polar(img, title, ax=None, dyn=35, ref_max=None):\n",
    "    img = img.detach().cpu()\n",
    "    db = 20 * torch.log10(torch.abs(img) + 1e-12)\n",
    "    m = db.max() if ref_max is None else ref_max\n",
    "    if ax is None:\n",
    "        _, ax = plt.subplots()\n",
    "    ax.imshow(db.numpy().T, origin=\"lower\", vmin=m - dyn, vmax=m,\n",
    "              extent=extent, aspect=\"auto\")\n",
    "    # True target positions in image coordinates.\n",
    "    r_t = (xt**2 + yt**2).sqrt().cpu()\n",
    "    ax.plot(r_t, (yt.cpu() / r_t), \"o\", ms=10, mfc=\"none\", mec=\"r\", label=\"True position\")\n",
    "    ax.set_title(title)\n",
    "    ax.set_xlabel(\"Range (m)\")\n",
    "    ax.set_ylabel(\"Angle (sin)\")\n",
    "    ax.legend(loc=\"upper right\")\n",
    "    return m"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Image with and without the DEM\n",
    "\n",
    "Both images are formed with `ffbp` from the same data, the only difference is\n",
    "the `dem` argument. The red circles mark the true target coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_flat = ffbp(data, grid, fc, r_res, pos, stages=4, dealias=True,\n",
    "                data_fmod=data_fmod, alias_fmod=alias_fmod).squeeze()\n",
    "img_dem = ffbp(data, grid, fc, r_res, pos, stages=4, dealias=True,\n",
    "               data_fmod=data_fmod, alias_fmod=alias_fmod, dem=dem).squeeze()\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=(12, 4.5))\n",
    "m = show_polar(img_flat, \"FFBP, flat earth\", ax=axs[0])\n",
    "show_polar(img_dem, \"FFBP with DEM\", ax=axs[1], ref_max=m)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "Without the DEM every elevated target is displaced towards the radar. The\n",
    "higher the terrain under it, the larger the shift, so the regular target grid\n",
    "appears warped around the hill. The near-ground targets at the scene edges\n",
    "sit close to their circles in both images. With the DEM all 25 targets focus\n",
    "on their true coordinates.\n",
    "\n",
    "For a straight track the flat-earth error is almost purely a position\n",
    "error and the point-spread functions stay compact, they are just at wrong places."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "## FFBP against direct backprojection\n",
    "\n",
    "The DEM-aware FFBP merge should match direct backprojection on the same DEM.\n",
    "With a DEM the dealias carrier is always referenced to the DEM surface\n",
    "instead of the z=0 plane: a flat-plane carrier would leave a\n",
    "terrain-dependent residual phase that aliases the image spectrum on any\n",
    "significant topography, defeating the purpose of dealiasing.\n",
    "`torchbp.util.bp_polar_range_dealias` and `bp_polar_range_alias` take the\n",
    "same `dem` argument to apply or remove this carrier after image formation.\n",
    "These functions could be required for example with interferometric SAR processing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_bp_dem = backprojection_polar_2d(data, grid, fc, r_res, pos, dealias=True,\n",
    "                                     data_fmod=data_fmod, alias_fmod=alias_fmod,\n",
    "                                     dem=dem).squeeze()\n",
    "\n",
    "def relerr(a, b):\n",
    "    return (torch.linalg.norm(a - b) / torch.linalg.norm(a)).item()\n",
    "\n",
    "print(f\"FFBP+DEM vs direct BP+DEM relative error: {relerr(img_bp_dem, img_dem):.4f}\")\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=(12, 4.5))\n",
    "show_polar(img_bp_dem, \"Direct BP with DEM\", ax=axs[0], ref_max=m)\n",
    "show_polar(img_bp_dem - img_dem, \"Difference (BP - FFBP)\", ax=axs[1], ref_max=m, dyn=60)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "## Notes\n",
    "\n",
    "- With real measurement data, resample the DEM onto the polar grid with\n",
    "  `torchbp.util.dem_to_polar` and make sure its heights are in the same z\n",
    "  frame as the platform positions.\n",
    "- The autofocus functions `torchbp.autofocus.gpga` and `gpga_tde` also accept `dem`."
   ]
  }
 ],
 "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
}