{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Phase gradient autofocus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchbp\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.signal import get_window\n",
    "import numpy as np\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    device = \"cuda\"\n",
    "else:\n",
    "    device = \"cpu\"\n",
    "print(\"Device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "Generate synthetic radar data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "nr = 100 # Range points\n",
    "ntheta = 128 # Azimuth points\n",
    "nsweeps = 128 # Number of measurements\n",
    "fc = 6e9 # RF center frequency\n",
    "bw = 100e6 # RF bandwidth\n",
    "tsweep = 100e-6 # Sweep length\n",
    "fs = 1e6 # Sampling frequency\n",
    "nsamples = int(fs * tsweep) # Time domain samples per sweep\n",
    "\n",
    "# Imaging grid definition. Azimuth angle \"theta\" is sine of radians. 0.2 = 11.5 degrees.\n",
    "grid_polar = {\"r\": (90, 110), \"theta\": (-0.2, 0.2), \"nr\": nr, \"ntheta\": ntheta}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_pos = torch.tensor([[100, 0, 0], [105, 10, 0], [97, -5, 0], [102, -10, 0], [95, 5, 0]], dtype=torch.float32, device=device)\n",
    "target_rcs = torch.tensor([1,1,1,1,1], dtype=torch.float32, device=device)\n",
    "pos = torch.zeros([nsweeps, 3], dtype=torch.float32, device=device)\n",
    "pos[:,1] = torch.linspace(-nsweeps/2, nsweeps/2, nsweeps) * 0.25 * 3e8 / fc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Oversampling input data decreases interpolation errors\n",
    "oversample = 3\n",
    "\n",
    "# Modulation frequency in range direction to center the spectrum at DC\n",
    "# for more accurate interpolation.\n",
    "data_fmod = -torch.pi * (1 - (oversample-1) / oversample)\n",
    "\n",
    "data = torchbp.util.generate_fmcw_data(target_pos, target_rcs, pos, fc, bw, tsweep, fs)\n",
    "# Apply windowing function in range direction\n",
    "wr = torch.tensor(get_window((\"taylor\", 3, 30), data.shape[-1])[None,:], dtype=torch.float32, device=device)\n",
    "wa = torch.tensor(get_window((\"taylor\", 3, 30), data.shape[0])[:,None], dtype=torch.float32, device=device)\n",
    "data = torch.fft.ifft(data * wa * wr, dim=-1, n=nsamples * oversample)\n",
    "\n",
    "data_fmod_f = torch.exp(1j*data_fmod*torch.arange(data.shape[-1], device=device))[None,:]\n",
    "data = data * data_fmod_f\n",
    "\n",
    "data_db = 20*torch.log10(torch.abs(data)).detach()\n",
    "m = torch.max(data_db)\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(data_db.cpu().numpy(), origin=\"lower\", vmin=m-30, vmax=m, aspect=\"auto\")\n",
    "plt.xlabel(\"Range samples\")\n",
    "plt.ylabel(\"Azimuth samples\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "Focused image without motion error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "r_res = 3e8 / (2 * bw * oversample) # Range bin size in input data\n",
    "\n",
    "# dealias=True removes range spectrum aliasing\n",
    "img = torchbp.ops.backprojection_polar_2d(data, grid_polar, fc, r_res, pos, dealias=True, data_fmod=data_fmod)\n",
    "img = img.squeeze(0) # Removes singular batch dimension\n",
    "# Backprojection image has spectrum with DC at zero index.\n",
    "# Shifting the spectrum shifts the DC to center bin.\n",
    "# This makes the solved phase to have same order as the position vector\n",
    "# Without shifting of the image, fftshift needs to be applied to\n",
    "# the solved phase for it to be in the same order as the position vector.\n",
    "# This doesn't affect the absolute value of the image.\n",
    "img = torchbp.util.shift_spectrum(img)\n",
    "\n",
    "img_db = 20*torch.log10(torch.abs(img)).detach()\n",
    "\n",
    "m = torch.max(img_db)\n",
    "\n",
    "extent = [*grid_polar[\"r\"], *grid_polar[\"theta\"]]\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(img_db.cpu().numpy().T, origin=\"lower\", vmin=m-40, vmax=m, extent=extent, aspect=\"auto\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Angle (sin radians)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "Create corrupted image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "phase_error = torch.exp(1j*2*torch.pi*torch.linspace(-3, 3, ntheta, dtype=torch.float32, device=device)[None,:]**2)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(torch.angle(phase_error.squeeze()).cpu().numpy())\n",
    "plt.xlabel(\"Azimuth sample\")\n",
    "plt.ylabel(\"Phase error (radians)\")\n",
    "\n",
    "img_corrupted = torch.fft.ifft(torch.fft.fft(img, dim=-1) * phase_error, dim=-1)\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(20*torch.log10(torch.abs(img_corrupted)).cpu().numpy().T, origin=\"lower\", vmin=m-40, vmax=m, extent=extent, aspect=\"auto\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Angle (sin radians)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "Phase gradient autofocus with phase difference estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_pga, phi = torchbp.autofocus.pga(img_corrupted, remove_trend=False, estimator=\"pd\")\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(20*torch.log10(torch.abs(img_pga)).cpu().numpy().T, origin=\"lower\", vmin=m-40, vmax=m, extent=extent, aspect=\"auto\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Angle (sin radians)\");\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(torch.angle(torch.exp(1j*phi)).cpu().numpy())\n",
    "plt.xlabel(\"Azimuth samples\")\n",
    "plt.ylabel(\"Phase error (radians)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "Apply maximum likelihood phase gradient autofocus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_pga, phi = torchbp.autofocus.pga(img_corrupted, remove_trend=False, estimator=\"ml\")\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(20*torch.log10(torch.abs(img_pga)).cpu().numpy().T, origin=\"lower\", vmin=m-40, vmax=m, extent=extent, aspect=\"auto\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Angle (sin radians)\");\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(torch.angle(torch.exp(1j*phi)).cpu().numpy())\n",
    "plt.xlabel(\"Azimuth samples\")\n",
    "plt.ylabel(\"Phase error (radians)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "Multiplying the solved phase with FFT of the corrupted image gives the focused image and taking inverse FFT gives the focused image. This should be identical to the image returned by `pga_ml`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_focused = torch.fft.ifft(torch.fft.fft(img_corrupted, dim=-1) * torch.exp(-1j*phi), dim=-1)\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(20*torch.log10(torch.abs(img_focused)).cpu().numpy().T, origin=\"lower\", vmin=m-40, vmax=m, extent=extent, aspect=\"auto\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Angle (sin radians)\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}