{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Polarimetric calibration (Ainsworth)\n",
    "\n",
    "This example synthesises fully-polarimetric SAR data with NUFFT projection,\n",
    "applies a known polarimetric distortion (channel imbalance + crosstalk), and\n",
    "recovers it with `torchbp.polarimetry.ainsworth` [1].\n",
    "\n",
    "The scene is distributed clutter (many independent scatterers\n",
    "obeying reflection symmetry) plus a single trihedral corner reflector. The corner reflector is\n",
    "used to resolve HH/VV channel imbalance (`k`), clutter correlation is used to estimate other parameters.\n",
    "\n",
    "[1] T. L. Ainsworth, L. Ferro-Famil and Jong-Sen Lee, \"Orientation angle\n",
    "preserving a posteriori polarimetric SAR calibration,\" in IEEE Transactions\n",
    "on Geoscience and Remote Sensing, vol. 44, no. 4, pp. 994-1003, April 2006."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from numpy import hamming\n",
    "import torchbp\n",
    "import torchbp.polarimetry as pol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(\"Device:\", device)\n",
    "torch.manual_seed(0);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Radar and imaging geometry\n",
    "\n",
    "The platform flies along `y` at constant altitude looking in `x`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "fc = 6e9            # RF center frequency (Hz)\n",
    "bw = 200e6          # Sweep bandwidth (Hz)\n",
    "tsweep = 100e-6     # Sweep length (s)\n",
    "fs = 20e6           # Sampling frequency (Hz)\n",
    "nsamples = int(fs * tsweep)\n",
    "nsweeps = 256       # Azimuth measurements\n",
    "wl = 3e8 / fc\n",
    "altitude = 100.0\n",
    "element_spacing = 0.25 * wl\n",
    "\n",
    "# Projection (scene) grid and imaging (polar) grid\n",
    "grid_proj = {\"x\": (60.0, 260.0), \"y\": (-90.0, 90.0), \"nx\": 256, \"ny\": 256}\n",
    "r0, r1, theta_limit = 90.0, 260.0, 0.35\n",
    "res = 3e8 / (2 * bw)\n",
    "nr = int(1.3 * (r1 - r0) / res)\n",
    "ntheta = int(1 + nsweeps * 1.3 * (element_spacing / wl) * theta_limit / 0.25)\n",
    "grid_polar = {\"r\": (r0, r1), \"theta\": (-theta_limit, theta_limit), \"nr\": nr, \"ntheta\": ntheta}\n",
    "\n",
    "pos = torch.zeros(nsweeps, 3, dtype=torch.float32, device=device)\n",
    "pos[:, 1] = torch.linspace(-nsweeps / 2, nsweeps / 2, nsweeps) * element_spacing\n",
    "pos[:, 2] = altitude\n",
    "print(\"nr\", nr, \"ntheta\", ntheta, \"nsamples\", nsamples)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## True scene\n",
    "\n",
    "Ainsworth's only assumption on the scene is scattering reciprocity and\n",
    "reflection symmetry of the distributed clutter, the cross-pol return is\n",
    "uncorrelated with the co-pol returns, $\\langle S_{hh}S_{hv}^*\\rangle =\n",
    "\\langle S_{vv}S_{hv}^*\\rangle = 0$, and $S_{vh} = S_{hv}$. We build the four\n",
    "channels in `[HH, HV, VH, VV]` order as independent speckle, with the cross-pol\n",
    "channel drawn independently of the co-pol channels.\n",
    "\n",
    "A single bright trihedral corner reflector ($S_{hh}=S_{vv}$, $S_{hv}=S_{vh}=0$)\n",
    "is placed at the scene centre to later pin down the HH/VV imbalance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "nx, ny = grid_proj['nx'], grid_proj['ny']\n",
    "def cn(scale):\n",
    "    return scale * (torch.randn(nx, ny, device=device)\n",
    "                    + 1j * torch.randn(nx, ny, device=device)) / np.sqrt(2)\n",
    "\n",
    "Shh, Svv = cn(1.0), cn(1.0)   # co-pol (independent)\n",
    "Shv = cn(0.4)                 # cross-pol, independent of co-pol\n",
    "Svh = Shv.clone()             # reciprocity\n",
    "\n",
    "cx, cy = nx // 2, ny // 2     # trihedral corner reflector\n",
    "Shh[cx, cy] = Svv[cx, cy] = 200.0\n",
    "Shv[cx, cy] = Svh[cx, cy] = 0.0\n",
    "\n",
    "S = torch.stack([Shh, Shv, Svh, Svv])   # [4, nx, ny], order HH,HV,VH,VV\n",
    "order = ['HH', 'HV', 'VH', 'VV']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Polarimetric distortion\n",
    "\n",
    "The system distortion mixes the four channels. We use the same parameterisation\n",
    "that `ainsworth` inverts: channel imbalances `alpha` (cross-pol), `k`\n",
    "(HH/VV), and crosstalk `u, v, w, z`.\n",
    "\n",
    "Ainsworth is orientation-angle preserving. It recovers only the crosstalk\n",
    "differences $u-z$ and $v-w$, because the common-mode part $(u+z)/2$,\n",
    "$(v+w)/2$ is mathematically indistinguishable from a real scene orientation\n",
    "angle. We therefore choose `z = -u`, `w = -v` (zero common-mode) so the\n",
    "distortion is fully recoverable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 1.15 * np.exp(1j * 0.20)   # cross-pol channel imbalance\n",
    "k = 0.85 * np.exp(-1j * 0.35)      # co-pol HH/VV imbalance\n",
    "u = 0.08 + 0.03j\n",
    "v = -0.05 + 0.06j\n",
    "w = -v                             # zero common-mode crosstalk\n",
    "z = -u\n",
    "M_true = torchbp.polarimetry.distortion_matrix(alpha, k, u, v, w, z, pol_order=order, device=device)\n",
    "\n",
    "# Observed (distorted) scene, per channel\n",
    "O = (M_true @ S.reshape(4, -1)).reshape(4, nx, ny)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Forward project and focus each channel\n",
    "\n",
    "Each channel is forward projected to raw FMCW data with\n",
    "`projection_cart_2d_nufft`, range compressed, and backprojected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "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",
    "win = torch.tensor(hamming(nsamples)[None, :], dtype=torch.float32, device=device)\n",
    "fmod_f = torch.exp(1j * data_fmod * torch.arange(n, device=device))[None, :]\n",
    "\n",
    "def make_image(scene):\n",
    "    data = torchbp.ops.projection_cart_2d_nufft(\n",
    "        scene, pos, grid_proj, fc, fs, bw / tsweep, nsamples,\n",
    "        use_rvp=False, normalization='gamma')[0]\n",
    "    data = torch.fft.ifft(data * win, dim=-1, n=n) * fmod_f\n",
    "    return torchbp.ops.backprojection_polar_2d(\n",
    "        data, grid_polar, fc, r_res, pos, dealias=False, data_fmod=data_fmod)[0]\n",
    "\n",
    "sar = torch.stack([make_image(O[c]) for c in range(4)])   # [4, nr, ntheta]\n",
    "print('Image shape:', tuple(sar.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Visualize the measured SAR images\n",
    "\n",
    "Corner reflector should not reflect cross-polarization, but it is slightly visible also in HV and VH images due to cross-talk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "extent = [*grid_polar['r'], *grid_polar['theta']]\n",
    "fig, axes = plt.subplots(1, 4, figsize=(16, 3.2))\n",
    "m = None\n",
    "for ax, c, name in zip(axes, range(4), order):\n",
    "    db = 20 * torch.log10(sar[c].abs() + 1e-9).cpu().numpy()\n",
    "    if m is None:\n",
    "        m = db.max()\n",
    "    ax.imshow(db.T, origin='lower', vmin=m - 60, vmax=m, extent=extent, aspect='auto')\n",
    "    ax.set_title(name); ax.set_xlabel('Range (m)')\n",
    "axes[0].set_ylabel('Angle (sin radians)')\n",
    "plt.tight_layout();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "## Solve k from the corner reflector\n",
    "\n",
    "The brightest pixel is the corner reflector. Its measured HH/VV ratio fixes the\n",
    "otherwise-ambiguous co-pol imbalance `k`. We also mask a small window around it\n",
    "with the `weight` argument so the corner does not bias the distributed-clutter\n",
    "crosstalk statistics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "mag = sar[0].abs()\n",
    "pk = torch.argmax(mag)\n",
    "pr, pt = int(pk // mag.shape[1]), int(pk % mag.shape[1])\n",
    "print(f'corner at ({pr},{pt}), peak/median = {(mag[pr,pt]/mag.median()).item():.1f}')\n",
    "corner_hh_vv = (sar[0, pr, pt] / sar[3, pr, pt]).item()\n",
    "print(f'measured corner HH/VV = {abs(corner_hh_vv):.3f} <{np.angle(corner_hh_vv):+.3f} rad')\n",
    "\n",
    "weight = torch.ones(nr, ntheta, device=device)\n",
    "ww = 6\n",
    "weight[max(0, pr-ww):pr+ww+1, max(0, pt-ww):pt+ww+1] = 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "## Calibration\n",
    "\n",
    "`ainsworth` returns the normalised correction matrix `Minv`. A perfect\n",
    "calibration makes `Minv @ M_true` a scalar multiple of the identity.\n",
    "\n",
    "Note that there is a possible 180 degree phase ambiguity with corner reflector `k` solving due to complex square root. A guess for `k` can also be specified and then the closer solution is used. `k` can also be specified directly if known and then `corner_hh_vv` should be `None` to disable corner reflector `k` solving."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "Minv = pol.ainsworth(sar, weight=weight, pol_order=order, corner_hh_vv=corner_hh_vv)\n",
    "\n",
    "R = Minv @ M_true\n",
    "R = R / R[0, 0]\n",
    "eye = torch.eye(4, dtype=R.dtype, device=device)\n",
    "offdiag = (R * (1 - eye)).abs().max().item()\n",
    "diagerr = (torch.diagonal(R) - 1).abs().max().item()\n",
    "print('Minv @ M_true:')\n",
    "for row in R.cpu().numpy():\n",
    "    print('  ' + '  '.join(f'{x.real:+.3f}{x.imag:+.3f}j' for x in row))\n",
    "print(f'\\nmax off-diagonal (residual crosstalk) = {offdiag:.4f}')\n",
    "print(f'max diagonal error (residual imbalance) = {diagerr:.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "Product is almost identity. The small residual error left on diagnal is from corner reflector measurement error."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "## Applying calibration\n",
    "\n",
    "`torchbp.polarimetry.apply_cal` can be used to apply calibration to image. After calibration the corner reflector isn't anymore visible in HV and VH images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "sar_cal = torchbp.polarimetry.apply_cal(sar, Minv)\n",
    "\n",
    "extent = [*grid_polar['r'], *grid_polar['theta']]\n",
    "fig, axes = plt.subplots(1, 4, figsize=(16, 3.2))\n",
    "for ax, c, name in zip(axes, range(4), order):\n",
    "    db = 20 * torch.log10(sar_cal[c].abs() + 1e-9).cpu().numpy()\n",
    "    ax.imshow(db.T, origin='lower', vmin=m - 60, vmax=m, extent=extent, aspect='auto')\n",
    "    ax.set_title(name); ax.set_xlabel('Range (m)')\n",
    "axes[0].set_ylabel('Angle (sin radians)')\n",
    "plt.tight_layout();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "## Orientation-angle and common-mode crosstalk\n",
    "\n",
    "Ainsworth's model allows the true scene to carry a non-zero\n",
    "$\\langle S_{co}S_{cross}^*\\rangle$, the scene orientation angle,\n",
    "requiring only that it be equal for HV and VH (reciprocity). A common-mode\n",
    "crosstalk $(u+z)/2$, $(v+w)/2$ produces exactly the same observed covariance as\n",
    "such an orientation angle, so the two cannot be separated from distributed\n",
    "targets alone. Ainsworth deliberately leaves the common-mode in place. The check below shows that with a\n",
    "non-zero common-mode crosstalk the residual equals common-mode.\n",
    "\n",
    "Common-mode crosstalk should also be calibrated for the best calibration quality, but Ainsworth can't correct it. It can be done with for example corner reflector measurements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Same distortion but with a non-zero common-mode crosstalk component.\n",
    "u2, v2 = 0.08 + 0.03j, -0.05 + 0.06j\n",
    "w2, z2 = 0.04 - 0.07j, 0.06 + 0.02j        # u2+z2, v2+w2 != 0\n",
    "M2 = torch.tensor([\n",
    "    [k*alpha,      v2/alpha,    w2*alpha,     v2*w2/(k*alpha)],\n",
    "    [z2*k*alpha,   1/alpha,     w2*z2*alpha,  w2/(k*alpha)],\n",
    "    [u2*k*alpha,   u2*v2/alpha, alpha,        v2/(k*alpha)],\n",
    "    [u2*z2*k*alpha,u2/alpha,    z2*alpha,     1/(k*alpha)],\n",
    "], dtype=torch.complex64, device=device)\n",
    "O2 = (M2 @ S.reshape(4, -1)).reshape(4, nx, ny)\n",
    "sar2 = torch.stack([make_image(O2[c]) for c in range(4)])\n",
    "chv2 = (sar2[0, pr, pt] / sar2[3, pr, pt]).item()\n",
    "Minv2 = pol.ainsworth(sar2, weight=weight, pol_order=order, corner_hh_vv=chv2)\n",
    "R2 = Minv2 @ M2; R2 = R2 / R2[0, 0]\n",
    "resid = (R2 * (1 - eye)).abs().max().item()\n",
    "print(f'common-mode (u+z)/2 = {abs((u2+z2)/2):.3f}, (v+w)/2 = {abs((v2+w2)/2):.3f}')\n",
    "print(f'residual off-diagonal = {resid:.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "## Calibration with rotation\n",
    "\n",
    "Apply rotation to antenna, calibrate, verify that polarimetric orientation angle is still correct after calibration and cross-talk is correctly calibrated. After rotation co- and cross-polarization channels correlate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "def copol_xpol_corr(img, weight):\n",
    "      \"\"\"Normalized co-pol/cross-pol correlation\n",
    "      ~0 for a reflection-symmetric scene\"\"\"\n",
    "      d = img[0] - img[3]                 # HH - VV\n",
    "      x = 0.5 * (img[1] + img[2])         # (HV + VH)/2\n",
    "      num = (weight * d * x.conj()).mean()\n",
    "      den = ((weight * d.abs()**2).mean() * (weight * x.abs()**2).mean()).sqrt()\n",
    "      return (num / den).item()\n",
    "\n",
    "# Give the clutter a known orientation angle, then distort, image and calibrate.\n",
    "theta_scene = np.radians(20.0)\n",
    "S_rot = pol.pol_antenna_rotation(S, theta_scene, pol_order=order)\n",
    "O_rot = (M_true @ S_rot.reshape(4, -1)).reshape(4, nx, ny)\n",
    "sar_rot = torch.stack([make_image(O_rot[c]) for c in range(4)])\n",
    "\n",
    "chv_rot = (sar_rot[0, pr, pt] / sar_rot[3, pr, pt]).item()\n",
    "Minv_rot = pol.ainsworth(sar_rot, weight=weight, pol_order=order, corner_hh_vv=chv_rot)\n",
    "cal_rot = pol.apply_cal(sar_rot, Minv_rot)\n",
    "\n",
    "R_rot = Minv_rot @ M_true\n",
    "R_rot = R_rot / R_rot[0, 0]\n",
    "resid = (R_rot * (1 - eye)).abs().max().item()\n",
    "\n",
    "deg = lambda im: np.degrees(torchbp.polarimetry.orientation_angle(im, weight, order).item())\n",
    "print(f\"scene orientation applied = {np.degrees(theta_scene):+.1f} deg\\n\")\n",
    "print(f\"{'':30}{'|corr(HH-VV, HV)|':>17}{'orientation':>14}\")\n",
    "for name, im in [(\"un-oriented, calibrated\", sar_cal), (\"+20 deg, calibrated\", cal_rot), (\"+20 deg, uncalibrated\", sar_rot)]:\n",
    "    print(f\"  {name:28}{abs(copol_xpol_corr(im, weight)):>15.3f}{deg(im):>11.1f} deg\")\n",
    "print(f\"\\nresidual system crosstalk = {resid:.4f}\")"
   ]
  }
 ],
 "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
}