{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Backprojection image formation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchbp\n",
    "import matplotlib.pyplot as plt\n",
    "from numpy import hamming"
   ]
  },
  {
   "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": [
    "Constant definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "nr = 128 # Range points\n",
    "ntheta = 128 # Azimuth points\n",
    "nsweeps = 128 # Number of measurements\n",
    "fc = 6e9 # RF center frequency\n",
    "bw = 100e6 # RF bandwidth. Negative for falling sweep.\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": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "Define target and radar positions. There is one point target at 100 m distance and zero azimuth angle.\n",
    "For polar image formation radar motion should be in direction of Y-axis.\n",
    "If this is not the case positions should be rotated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_pos = torch.tensor([[100, 0, 0]], dtype=torch.float32, device=device)\n",
    "target_rcs = torch.tensor([[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\n",
    "pos[:,2] = 50 # Platform height"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "Generate synthetic radar data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Oversampling input data decreases interpolation errors\n",
    "oversample = 2\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",
    "print(\"data_fmod\", data_fmod)\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",
    "w = torch.tensor(hamming(data.shape[-1])[None,:], dtype=torch.float32, device=device)\n",
    "# With rising sweep the IF frequencies are negative.\n",
    "if bw > 0:\n",
    "    data = torch.fft.ifft(data * w, dim=-1, n=nsamples * oversample)\n",
    "else:\n",
    "    data = torch.fft.ifft(data.conj() * w, dim=-1, n=nsamples * oversample).conj()\n",
    "    data_fmod = -data_fmod\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": "10",
   "metadata": {},
   "source": [
    "Plot the range spectrum of the raw data. With correct `data_fmod`, the spectrum should be centered around 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "freqs = torch.fft.fftshift(torch.fft.fftfreq(data.shape[-1]))\n",
    "plt.figure()\n",
    "plt.plot(freqs.cpu().numpy(), 20*torch.log10(torch.abs(torch.fft.fftshift(torch.fft.fft(data[0])))).cpu().numpy())\n",
    "plt.xlabel(\"Frequency (cycles/sample)\")\n",
    "plt.ylabel(\"Magnitude (dB)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "Image formation.\n",
    "Hamming window was applied in range direction so low sidelobes in range are expected.\n",
    "Azimuth direction has no windowing function and high sidelobes (Highest -13 dB) are expected.\n",
    "Azimuth sidelobes could be decreased by windowing the input data also in the other dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "r_res = 3e8 / (2 * abs(bw) * oversample) # Range bin size in input data\n",
    "\n",
    "# Calculate modulation frequency to center the spectrum of image around DC for decreased interpolation error.\n",
    "dr = (grid_polar[\"r\"][1] - grid_polar[\"r\"][0]) / grid_polar[\"nr\"]\n",
    "im_margin = oversample * r_res / dr - 1\n",
    "alias_fmod = -torch.pi * (1 - im_margin / (1 + im_margin))\n",
    "if bw < 0:\n",
    "    alias_fmod = -alias_fmod\n",
    "\n",
    "img = torchbp.ops.backprojection_polar_2d(data, grid_polar, fc, r_res, pos, dealias=True, data_fmod=data_fmod, alias_fmod=alias_fmod)\n",
    "img = img.squeeze() # Removes singular batch dimension\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-30, vmax=m, extent=extent, aspect=\"auto\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Angle (sin radians)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "Plot the SAR image spectrum. It should be centered around 0,0 with correct alias_fmod choice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "fimg = torch.fft.fftshift(torch.fft.fft2(img), (0, 1))\n",
    "fimg_db = (20*torch.log10(torch.abs(fimg)))\n",
    "vmax = torch.max(fimg_db)\n",
    "plt.imshow(fimg_db.cpu().numpy(), origin=\"lower\", aspect=\"auto\", vmin=vmax-40, vmax=vmax, extent=[-0.5, 0.5, -0.5, 0.5])\n",
    "plt.xlabel(\"Cross-range frequency (cycles/sample)\")\n",
    "plt.ylabel(\"Range frequency (cycles/sample)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "Image entropy. Can be used as a loss function for optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "entropy = torchbp.util.entropy(img)\n",
    "print(\"Entropy:\", entropy.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "Convert image to cartesian coordinates:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Origin of the polar coordinates\n",
    "origin = torch.mean(pos, axis=0)\n",
    "# Cartesian grid definition\n",
    "grid_cart = {\"x\": (90, 110), \"y\": (-10, 10), \"nx\": 128, \"ny\": 128}\n",
    "\n",
    "img_cart = torchbp.ops.polar_to_cart(img, origin, grid_polar, grid_cart, fc, rotation=0, alias_fmod=alias_fmod, method=(\"lanczos\", 10))\n",
    "\n",
    "img_db = 20*torch.log10(torch.abs(img_cart)).detach()\n",
    "\n",
    "m = torch.max(img_db)\n",
    "\n",
    "extent = [*grid_cart[\"x\"], *grid_cart[\"y\"]]\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(img_db.cpu().numpy().T, origin=\"lower\", vmin=m-30, vmax=m, extent=extent, aspect=\"equal\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Cross-range (m)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "Backprojection directly onto Cartesian grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_cart2 = torchbp.ops.backprojection_cart_2d(data, grid_cart, fc, r_res, pos, data_fmod=data_fmod)\n",
    "img_cart2 = img_cart2.squeeze() # Removes singular batch dimension\n",
    "\n",
    "img_db = 20*torch.log10(torch.abs(img_cart2)).detach()\n",
    "\n",
    "m = torch.max(img_db)\n",
    "\n",
    "extent = [*grid_cart[\"x\"], *grid_cart[\"y\"]]\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(img_db.cpu().numpy().T, origin=\"lower\", vmin=m-30, vmax=m, extent=extent, aspect=\"equal\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Cross-range (m)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "Difference between the results should be very small"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Phase difference\")\n",
    "plt.imshow(torch.angle(img_cart * torch.conj(img_cart2)).cpu().numpy().T, origin=\"lower\", extent=extent, aspect=\"equal\")\n",
    "plt.xlabel(\"Range (m)\")\n",
    "plt.ylabel(\"Cross-range (m)\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.linalg.norm(img_cart - img_cart2) / torch.linalg.norm(img_cart)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "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
}