gwcs-profiling.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4a8904aa-8426-41df-8e68-6818abe25c42",
   "metadata": {},
   "source": [
    "# Profiling GWCS vs FITS SIP for Roman WFI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2355d6fe-e050-44f6-9566-b0c271f9db18",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.auto import tqdm\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.lines import Line2D\n",
    "\n",
    "from astropy.wcs import WCS\n",
    "from gwcs import WCS as GWCS\n",
    "import roman_datamodels.datamodels as rdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "924312e4-81ff-45d4-90c0-09845b8d292d",
   "metadata": {},
   "source": [
    "Example WFI L2 images per SCA from Build 16:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "254aab2b-e92b-493d-88e3-fc47aa14b708",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://stsci.box.com/shared/static/frh9rvdy18ig7azjih3cgduojocjo251.asdf\n",
    "path_wfi01 = 'roman_data/r0000101001001001001_0001_wfi01_cal.asdf'\n",
    "\n",
    "# https://stsci.box.com/shared/static/z44o6xg5sver1qheku200q1v7fjl8lzh.asdf\n",
    "path_wfi14 = 'roman_data/r0000301001001001001_0001_wfi14_cal.asdf'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5915bb51-deba-4cea-88c9-a57974b64e49",
   "metadata": {},
   "source": [
    "Load the example data, get the GWCS coordinate representation, approximate with FITS SIP:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e507a7f5-20d1-4bfb-9639-070bea04324a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dm = rdm.open(path_wfi01)\n",
    "gwcs = dm.meta['wcs']\n",
    "sip_header = gwcs.to_fits_sip()\n",
    "fits_sip = WCS(sip_header)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bbfff4e-bd44-4ead-a91e-f76c1432c685",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2, 2, figsize=(12, 10), dpi=250, sharey='row')\n",
    "\n",
    "skips = [1, 3, 10, 30, 100, 300]\n",
    "\n",
    "fits_color = 'dodgerblue'\n",
    "gwcs_color = 'tomato'\n",
    "\n",
    "for skip in tqdm(skips):\n",
    "    \n",
    "    x, y = np.mgrid[0:sip_header['NAXIS1']:skip, 0:sip_header['NAXIS2']:skip]\n",
    "    \n",
    "    one_pixel_apart = gwcs.pixel_to_world([0, 0], [0, 1])\n",
    "\n",
    "    # estimate pixel scale\n",
    "    pixel_scale = one_pixel_apart[0].separation(one_pixel_apart[1])\n",
    "\n",
    "    # compute p2w\n",
    "    gwcs_world = gwcs.pixel_to_world(x, y)\n",
    "    fits_sip_world = fits_sip.pixel_to_world(x, y)\n",
    "\n",
    "    # measure difference in pixel units\n",
    "    fits_sip_p2w_2d = (gwcs_world.separation(fits_sip_world) / pixel_scale)\n",
    "    fits_sip_p2w_diff = fits_sip_p2w_2d.ravel()\n",
    "\n",
    "    # compute w2p using FITS SIP approx\n",
    "    fits_sip_w2p = fits_sip.world_to_pixel(gwcs_world)\n",
    "    fits_sip_w2p_diff = np.concatenate([\n",
    "        np.ravel(fits_sip_w2p[0] - x),\n",
    "        np.ravel(fits_sip_w2p[1] - y),\n",
    "    ])\n",
    "\n",
    "    # measure runtimes:\n",
    "    fits_sip_timeit_p2w = %timeit -o -r 3 -n 5 fits_sip.pixel_to_world(x, y)\n",
    "    fits_sip_timeit_w2p = %timeit -o -r 3 -n 5 fits_sip.world_to_pixel(fits_sip_world)\n",
    "    gwcs_timeit_p2w = %timeit -o -r 3 -n 5 gwcs.pixel_to_world(x, y)\n",
    "    gwcs_timeit_w2p = %timeit -o -r 3 -n 5 gwcs.world_to_pixel(gwcs_world)\n",
    "\n",
    "    # plot results:\n",
    "    ax[0, 0].errorbar(skip, gwcs_timeit_p2w.average, gwcs_timeit_p2w.stdev, fmt='o', color=gwcs_color)\n",
    "    ax[0, 0].errorbar(skip, fits_sip_timeit_p2w.average, fits_sip_timeit_p2w.stdev, fmt='o', color=fits_color)\n",
    "    \n",
    "    ax[0, 1].errorbar(skip, gwcs_timeit_w2p.average, gwcs_timeit_w2p.stdev, fmt='o', color=gwcs_color)\n",
    "    ax[0, 1].errorbar(skip, fits_sip_timeit_w2p.average, fits_sip_timeit_w2p.stdev, fmt='o', color=fits_color)\n",
    "\n",
    "    if skip == skips[0]:\n",
    "        for i, diff in enumerate([fits_sip_p2w_diff, fits_sip_w2p_diff]):\n",
    "            ax[1, i].bar(['mean', 'std', 'max'], [diff.mean(), diff.std(), diff.max()])\n",
    "\n",
    "ax[0, 0].set(\n",
    "    title='pixel to world',\n",
    "    xlabel='Downsampling factor',\n",
    "    ylabel='Runtime [s]',\n",
    "    yscale='log',\n",
    "    xscale='log',\n",
    ")\n",
    "\n",
    "ax[0, 1].set(\n",
    "    title='world to pixel',\n",
    "    xlabel='Downsampling factor',\n",
    "    ylabel='Runtime [s]',\n",
    "    yscale='log',\n",
    "    xscale='log',\n",
    ")\n",
    "\n",
    "\n",
    "ax[1, 0].set(\n",
    "    ylabel='Difference [pix]'\n",
    ")\n",
    "\n",
    "handles = [\n",
    "    Line2D([], [], marker='o', lw=0, color=color) \n",
    "    for color in [gwcs_color, fits_color]\n",
    "]\n",
    "ax[0, 0].legend(handles, ['GWCS', 'FITS SIP'])\n",
    "fig.savefig('plots/gwcs_fits_sip.pdf', bbox_inches='tight')\n",
    "fig.savefig('plots/gwcs_fits_sip.png', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e0777d8-7717-477a-b6a8-7cc61f48554c",
   "metadata": {},
   "outputs": [],
   "source": [
    "cax = plt.imshow(fits_sip_p2w_2d)\n",
    "plt.colorbar(cax, label='Difference [pix]')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11f5ba07-da5f-4c9e-be3d-7283ed41ce42",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}