{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Antenna pattern normalization\n",
    "\n",
    "In SAR image the pixels are brighter where the antenna is pointed and at closer distances. This is caused by the two-way antenna gain and the free-space spreading loss, not by the scene itself. To compare reflectivities across an image the effect of antenna and distance has to be removed.\n",
    "\n",
    "`torchbp` models the illumination directly with `backprojection_polar_2d_tx_power`. Given the antenna gain pattern, the trajectory and the antenna attitude, it forms a pseudo-polar image of the square-root power that a constant reflectivity ground would return. Dividing the SAR image by this map removes the antenna footprint and the range falloff, leaving an image whose brightness reflects the scene.\n",
    "\n",
    "This example:\n",
    "  1. Creates synthetic data\n",
    "  2. Simulates a pass over a uniform field of random scatterers with `projection_cart_2d_nufft`.\n",
    "  3. Forms the SAR image and shows the antenna footprint imprinted on it.\n",
    "  4. Models the illumination with `backprojection_polar_2d_tx_power` and normalizes the image with it."
   ]
  },
  {
   "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",
    "import time"
   ]
  },
  {
   "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",
    "A small FMCW radar is flown along the $y$ axis at 45 m altitude, looking in the $+x$ direction. The low altitude relative to the imaged range means the elevation angle changes a lot across the swath, so the elevation antenna pattern shows up strongly as a range-dependent brightness."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "fc = 6e9            # RF center frequency (Hz)\n",
    "bw = 200e6          # Sweep bandwidth (Hz)\n",
    "tsweep = 100e-6     # Sweep length (s)\n",
    "fs = 5e6           # Sampling frequency (Hz)\n",
    "nsamples = int(fs * tsweep)   # Time-domain samples per sweep\n",
    "nsweeps = 384       # Number of azimuth measurements\n",
    "wl = 3e8 / fc       # Wavelength\n",
    "\n",
    "altitude = 45.0     # Platform altitude (m)\n",
    "element_spacing = 0.25 * wl   # Azimuth sampling along the track"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "Polar imaging grid. `theta` is the sine of the azimuth angle. The number of range and azimuth bins is chosen from the range resolution and the synthetic aperture so the grid is slightly oversampled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "r0, r1 = 60.0, 250.0      # Ground-range extent (m)\n",
    "theta_limit = 0.75         # Azimuth half-extent (sine of angle)\n",
    "\n",
    "res = 3e8 / (2 * bw)      # Slant-range resolution\n",
    "res_interp = 1.3          # Grid oversampling factor\n",
    "nr = int(res_interp * (r1 - r0) / res)\n",
    "ntheta = int(1 + nsweeps * res_interp * (element_spacing / wl) * theta_limit / 0.25)\n",
    "\n",
    "grid_polar = {\"r\": (r0, r1), \"theta\": (-theta_limit, theta_limit),\n",
    "              \"nr\": nr, \"ntheta\": ntheta}\n",
    "print(\"nr:\", nr, \"ntheta:\", ntheta, \"nsamples:\", nsamples)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## Antenna pattern\n",
    "\n",
    "The operators take the square root of the two-way antenna gain `g`. If TX and RX antennas are identical, this is just the antenna pattern, but if they are not then this should be $\\sqrt{g_{tx} g_{rx}}$. It's sampled on a regular grid of elevation x azimuth angles in radians, with the beam center at angle $(0, 0)$. The corresponding angular extent is given as `g_extent = [el0, az0, el1, az1]`.\n",
    "\n",
    "Here we use a separable Gaussian beam: a narrow pencil beam in elevation and a wider beam in azimuth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "naz, nel = 64, 64\n",
    "az = np.linspace(-np.pi / 3, np.pi / 3, naz)\n",
    "el = np.linspace(-np.pi / 3, np.pi / 3, nel)\n",
    "az_bw = np.deg2rad(30)    # Azimuth beamwidth\n",
    "el_bw = np.deg2rad(14)    # Elevation beamwidth\n",
    "\n",
    "gain = np.exp(-(el[:, None] / el_bw) ** 2) * np.exp(-(az[None, :] / az_bw) ** 2)  # [el, az]\n",
    "g = torch.tensor(gain, dtype=torch.float32, device=device)\n",
    "g_extent = [el[0], az[0], el[-1], az[-1]]   # [el0, az0, el1, az1]\n",
    "\n",
    "plt.figure(figsize=(5, 4))\n",
    "plt.imshow(20 * np.log10(gain), origin=\"lower\", aspect=\"auto\",\n",
    "           extent=[np.rad2deg(az[0]), np.rad2deg(az[-1]),\n",
    "                   np.rad2deg(el[0]), np.rad2deg(el[-1])], vmin=-40, vmax=0)\n",
    "plt.colorbar(label=\"Gain (dB)\")\n",
    "plt.xlabel(\"Azimuth (deg)\")\n",
    "plt.ylabel(\"Elevation (deg)\")\n",
    "plt.title(\"Antenna pattern\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "## Scene\n",
    "\n",
    "To see the antenna footprint we need uniform scene: a dense field of random complex scatterers on a Cartesian grid on the ground ($z = 0$). Each cell has unit expected power, so any brightness variation in the SAR image comes from the illumination, not from the scene."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "grid_proj = {\"x\": (40.0, 270.0), \"y\": (-270.0, 270.0), \"nx\": 512, \"ny\": 512}\n",
    "img = torch.randn([grid_proj[\"nx\"], grid_proj[\"ny\"]], dtype=torch.complex64, device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "## Platform trajectory and attitude\n",
    "\n",
    "The platform flies straight along $y$ at $x = 0$. The antenna is rolled so the elevation beam center points at the middle of the swath; without this roll the narrow elevation beam would mostly miss the ground. The attitude `att` holds the [roll, pitch, yaw] Euler angles per sweep."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_coords = torch.linspace(-nsweeps / 2, nsweeps / 2, nsweeps) * element_spacing\n",
    "pos = torch.zeros([nsweeps, 3], dtype=torch.float32, device=device)\n",
    "pos[:, 1] = y_coords\n",
    "pos[:, 2] = altitude\n",
    "\n",
    "r_mid = 0.5 * (r0 + r1)\n",
    "roll = -np.arctan2(altitude, r_mid)   # tilt elevation beam down to mid-swath\n",
    "att = torch.zeros_like(pos)\n",
    "att[:, 0] = roll\n",
    "print(\"Antenna roll:\", np.rad2deg(roll), \"deg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "## Forward projection and image formation\n",
    "\n",
    "We simulate the raw radar data with `projection_cart_2d_nufft`, passing the antenna pattern (`g`, `g_extent`) and attitude (`att`) so that each scatterer is weighted by the two-way gain in the direction it is seen. The forward model is the point-target radar equation: amplitude $\\propto g / R^2$. After projection the raw data is range-compressed (windowed, oversampled IFFT) and backprojected onto the polar grid. The `data_fmod` term centers the range spectrum around DC to reduce interpolation error.\n",
    "\n",
    "`\"sigma\"` normalization means that input scene RCS is normalized to $1~m^2$ area.\n",
    "\n",
    "Alternatives are:\n",
    " * `\"beta\"` (normalized to slant-range area)\n",
    " * `\"gamma\"` (normalized to area projected towards radar)\n",
    " * `\"point\"` (point target with fixed RCS)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "oversample = 2\n",
    "data_fmod = -torch.pi * (1 - (oversample - 1) / oversample)\n",
    "n = nsamples * oversample\n",
    "r_res = 3e8 / (2 * bw * oversample)\n",
    "\n",
    "t = time.time()\n",
    "# Forward project the scene to raw FMCW data (stop-and-go NUFFT) with the antenna pattern\n",
    "data = torchbp.ops.projection_cart_2d_nufft(\n",
    "    img, pos, grid_proj, fc, fs, bw / tsweep, nsamples,\n",
    "    att=att, g=g, g_extent=g_extent, use_rvp=False, normalization=\"sigma\")[0]\n",
    "# Range compression: window + oversampled IFFT\n",
    "w = torch.tensor(hamming(data.shape[-1])[None, :], dtype=torch.float32, device=device)\n",
    "data = torch.fft.ifft(data * w, dim=-1, n=n)\n",
    "data = data * torch.exp(1j * data_fmod * torch.arange(data.shape[-1], device=device))[None, :]\n",
    "# Backproject to the polar grid\n",
    "sar = torchbp.ops.backprojection_polar_2d(\n",
    "    data, grid_polar, fc, r_res, pos, dealias=False, data_fmod=data_fmod)[0]\n",
    "print(f\"Image formed in {time.time() - t:.2f} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "The raw image is dominated by speckle, so to see the illumination we smooth the power with a small box filter. The smoothed envelope clearly shows the antenna footprint: a bright blob centered mid-swath where the rolled pencil beam illuminates the ground, fading toward the range edges (elevation pattern + range loss) and toward the azimuth edges (azimuth pattern)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "def envelope(x):\n",
    "    \"\"\"Smoothed magnitude (power box filter) to reveal the illumination under speckle.\"\"\"\n",
    "    p = (torch.abs(x) ** 2)[None, None]\n",
    "    p = torch.nn.functional.avg_pool2d(p, 9, stride=1, padding=4)[0, 0]\n",
    "    return torch.sqrt(p)\n",
    "\n",
    "extent = [*grid_polar[\"r\"], *grid_polar[\"theta\"]]\n",
    "env = envelope(sar)\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
    "for ax, x, title in zip(axes, [sar, env], [\"Raw image (speckle)\", \"Raw envelope (smoothed)\"]):\n",
    "    db = 20 * torch.log10(torch.abs(x) + 1e-12).cpu().numpy()\n",
    "    m = db.max()\n",
    "    im = ax.imshow(db.T, origin=\"lower\", vmin=m - 35, vmax=m, extent=extent, aspect=\"auto\")\n",
    "    ax.set_title(title)\n",
    "    ax.set_xlabel(\"Range (m)\")\n",
    "axes[0].set_ylabel(\"Angle (sin radians)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "## Modeling the illumination\n",
    "\n",
    "`backprojection_polar_2d_tx_power` forms this from the geometry. For every polar pixel it sums the two-way gain $g^2$ over all sweeps, weighted by the range loss, and returns the square-root power that a constant-reflectivity ground would return. Its inputs mirror the image formation:\n",
    "\n",
    "  * `wa` per-sweep amplitude weighting (transmit power and/or an azimuth window). Here it is flat.\n",
    "  * `g`, `g_extent`, `att`, `pos`, `grid_polar`, `r_res` the same antenna pattern, attitude and geometry used to form the image.\n",
    "  * `normalization` the radiometric convention. `\"beta\"` gives the raw illumination of the imaging plane. `\"sigma\"` and `\"gamma\"` additionally divide by the incidence-angle terms to estimate $\\sigma_0$ / $\\gamma_0$ backscatter coefficients. `\"point\"` normalizes for point targets.\n",
    "  * `azimuth_resolution` (default `True`) additionally folds in the azimuth-resolution normalization, so dividing by `tx_power` also flattens the residual azimuth brightness. Pass `False` to get the pure antenna/range illumination.\n",
    "\n",
    "The modeled footprint matches the measured envelope above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "wa = torch.ones(nsweeps, device=device)   # flat per-sweep weighting\n",
    "\n",
    "tx_power = torchbp.ops.backprojection_polar_2d_tx_power(\n",
    "    wa, g, g_extent, grid_polar, r_res, pos, att, normalization=\"sigma\")[0]\n",
    "\n",
    "plt.figure(figsize=(6, 4))\n",
    "db = 20 * torch.log10(tx_power + 1e-12).cpu().numpy()\n",
    "m = db.max()\n",
    "plt.imshow(db.T, origin=\"lower\", vmin=m - 35, vmax=m, extent=extent, aspect=\"auto\")\n",
    "plt.colorbar(label=\"Relative power (dB)\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Angle (sin radians)\")\n",
    "plt.title(\"Modeled illumination (tx_power)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "## Normalization\n",
    "\n",
    "Dividing the SAR image by the modeled illumination removes the antenna footprint and the range falloff. The residual still has speckle noise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "sar_norm = sar / tx_power\n",
    "env_norm = envelope(sar_norm)\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
    "for ax, x, title in zip(axes, [env, env_norm], [\"Before\", \"After normalization\"]):\n",
    "    db = 20 * torch.log10(torch.abs(x) + 1e-12).cpu().numpy()\n",
    "    m = db.max()\n",
    "    im = ax.imshow(db.T, origin=\"lower\", vmin=m - 20, vmax=m, extent=extent, aspect=\"auto\")\n",
    "    ax.set_title(title)\n",
    "    ax.set_xlabel(\"Range (m)\")\n",
    "    plt.colorbar(im, ax=ax, label=\"Relative power (dB)\")\n",
    "axes[0].set_ylabel(\"Angle (sin radians)\")\n",
    "plt.tight_layout();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "Range and azimuth profiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "r_axis = np.linspace(r0, r1, nr)\n",
    "th_axis = np.linspace(-theta_limit, theta_limit, ntheta)\n",
    "\n",
    "def db(x):\n",
    "    return 20 * torch.log10(x + 1e-12).cpu().numpy()\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
    "axes[0].plot(r_axis, db(env.mean(1)) - db(env.mean(1)).max(), label=\"raw\")\n",
    "axes[0].plot(r_axis, db(env_norm.mean(1)) - db(env_norm.mean(1)).max(), label=\"normalized\")\n",
    "axes[0].set_xlabel(\"Range (m)\")\n",
    "axes[0].set_ylabel(\"Relative power (dB)\")\n",
    "axes[0].set_title(\"Range profile\")\n",
    "axes[0].legend()\n",
    "\n",
    "axes[1].plot(th_axis, db(env.mean(0)) - db(env.mean(0)).max(), label=\"raw\")\n",
    "axes[1].plot(th_axis, db(env_norm.mean(0)) - db(env_norm.mean(0)).max(), label=\"normalized\")\n",
    "axes[1].set_xlabel(\"Angle (sin radians)\")\n",
    "axes[1].set_ylabel(\"Relative power (dB)\")\n",
    "axes[1].set_title(\"Azimuth profile\")\n",
    "axes[1].legend()\n",
    "plt.tight_layout();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "## SNR-aware normalization\n",
    "\n",
    "With real data the illumination is poor at the swath edges and antenna nulls. There `tx_power` is tiny, and dividing the SAR image by it divides receiver noise by an almost-zero number, so the normalized edges explode with amplified noise.\n",
    "\n",
    "`torchbp.util.wiener_normalize` replaces the division by the Wiener (MMSE) estimate\n",
    "\n",
    "$$\\hat{s} = \\frac{\\mathrm{sar}\\cdot\\mathrm{tx\\_power}}{\\mathrm{tx\\_power}^2 + \\varepsilon^2} = \\frac{\\mathrm{sar}}{\\mathrm{tx\\_power}}\\cdot\\frac{\\mathrm{SNR}}{1+\\mathrm{SNR}}, \\qquad \\mathrm{SNR} = \\left(\\frac{\\mathrm{tx\\_power}}{\\varepsilon}\\right)^2.$$\n",
    "\n",
    "Where the illumination is strong this is the full normalization; where it is weak the gain rolls off as $\\mathrm{tx\\_power}^2$, so the output goes to zero instead of blowing up.\n",
    "\n",
    "The regularization $\\varepsilon = \\sigma_n/\\sigma_s$ is the noise-to-signal amplitude ratio in `tx_power` units, the illumination level where the SNR is one. It isn't simply the noise level $\\sigma_n$: with the model $\\mathrm{sar} = s\\cdot\\mathrm{tx\\_power} + n$ the noise amplitude $\\sigma_n$ must be divided by the reflectivity scale $\\sigma_s$, which also absorbs the leftover radiometric calibration constant of `tx_power`. When `eps` is not given it is estimated from the data using $E|\\mathrm{sar}|^2 = \\sigma_s^2\\,\\mathrm{tx\\_power}^2 + \\sigma_n^2$: $\\sigma_s^2$ from the brightest pixels and $\\sigma_n^2$ from the dimmest. For real data prefer an explicit value from a shadow region and a calibration target.\n",
    "\n",
    "To show the effect we add receiver noise to the SAR image (constant noise, but the signal fades toward the edges) and compare plain division with the Wiener estimate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Emulate receiver noise constant complex noise added to the SAR image.\n",
    "torch.manual_seed(1)\n",
    "noise_std = 0.5 * sar.abs().mean() # sigma_n: noise amplitude\n",
    "sar_noisy = sar + noise_std * torch.randn_like(sar)\n",
    "\n",
    "# The correct regularization is eps = sigma_n / sigma_s, the noise-to-signal ratio\n",
    "# in tx_power units. sigma_s is the reflectivity scale relating\n",
    "# the image to the illumination (sar = s * tx_power + n), it also absorbs the\n",
    "# leftover calibration constant of tx_power. Measure it where illumination is strong.\n",
    "bright = tx_power > tx_power.quantile(0.8)\n",
    "sigma_s = ((sar / tx_power).abs()[bright] ** 2).mean().sqrt()\n",
    "eps = noise_std / sigma_s\n",
    "print(f\"sigma_n = {float(noise_std):.2e}, sigma_s = {float(sigma_s):.3f}, eps = {float(eps):.2e}\")\n",
    "\n",
    "naive = sar_noisy / tx_power\n",
    "wiener = torchbp.util.wiener_normalize(sar_noisy, tx_power, eps=eps)\n",
    "# Passing eps=None makes wiener_normalize estimate the same value from the data.\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
    "for ax, x, title in zip(axes, [naive, wiener], [\"Plain sar / tx_power\", \"Wiener normalize\"]):\n",
    "    d = 20 * torch.log10(envelope(x) + 1e-12).cpu().numpy()\n",
    "    m = np.nanmedian(d)\n",
    "    im = ax.imshow(d.T, origin=\"lower\", vmin=m - 15, vmax=m + 25, extent=extent, aspect=\"auto\")\n",
    "    ax.set_title(title)\n",
    "    ax.set_xlabel(\"Range (m)\")\n",
    "    plt.colorbar(im, ax=ax, label=\"Relative power (dB)\")\n",
    "axes[0].set_ylabel(\"Angle (sin radians)\")\n",
    "plt.tight_layout();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26",
   "metadata": {},
   "source": [
    "The calibration is still missing a constant factor. If calibrated RCS values are needed this constant factor needs to be also solved. It can be found by imaging a known RCS target, usually a corner reflector.\n",
    "\n",
    "For slant-range images (BP origin at sensor altitude, `pos` z = 0) use `backprojection_polar_2d_tx_power_slant`, which takes the sensor `altitude` and maps each polar pixel to its ground position before computing angles."
   ]
  }
 ],
 "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
}