Demo for @PaulHuwe + @tddesjardins

imviz-roman-all-SCAs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e024fdd8-a348-459f-bfaf-6a997baac6ef",
   "metadata": {},
   "source": [
    "# Quickly load a full Roman exposure into Imviz\n",
    "\n",
    "by dodging 18 WCS links."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95e556f5-f3ec-47c4-8365-79bd1b179b04",
   "metadata": {
    "tags": []
   },
   "source": [
    "Use the matching version of rdm/rad:\n",
    "```bash\n",
    "pip install roman_datamodels==0.21.0 rad==0.21.0\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a105dad-7107-440e-8d07-9b6eddc450a7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import asdf\n",
    "from roman_datamodels import __version__ as rdm_version_pkg\n",
    "from packaging.version import Version\n",
    "\n",
    "\n",
    "path = '/grp/roman/SCIENCE_PLATFORM_DATA/ROMANISIM/DENSE_REGION/R0.5_DP0.5_PA0/r0000101001001001001_01101_0001_WFI01_cal.asdf'\n",
    "file = asdf.open(path)\n",
    "\n",
    "[rdm_version_file] = [\n",
    "    ext['software']['version'] \n",
    "    for ext in file.tree['history']['extensions'] \n",
    "    if ext['software']['name'] == 'roman_datamodels'\n",
    "]\n",
    "\n",
    "print(f\"{rdm_version_file=}, {rdm_version_pkg=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e32a644b-67a8-4c35-8120-ef6d60709bf2",
   "metadata": {},
   "source": [
    "Retrieving large files from Central Storage is a bit slow, so let's avoid repeated downloads. For this demo, download local copies of the files at: \n",
    "```\n",
    "/grp/roman/SCIENCE_PLATFORM_DATA/ROMANISIM/DENSE_REGION/R0.5_DP0.5_PA0/r0000101001001001001*_cal.asdf\n",
    "```\n",
    "\n",
    "For a minimal example with two SCAs, run:\n",
    "```bash\n",
    "mdkir ~/Downloads/R0.5_DP0.5_PA0 \n",
    "cd ~/Downloads/R0.5_DP0.5_PA0/\n",
    "cp /grp/roman/SCIENCE_PLATFORM_DATA/ROMANISIM/DENSE_REGION/R0.5_DP0.5_PA0/r0000101001001001001_01101_0001_WFI01_cal.asdf .\n",
    "cp /grp/roman/SCIENCE_PLATFORM_DATA/ROMANISIM/DENSE_REGION/R0.5_DP0.5_PA0/r0000101001001001001_01101_0001_WFI02_cal.asdf .\n",
    "```\n",
    "\n",
    "In the following cell, we assume those files are at: `~/Downloads/R0.5_DP0.5_PA0/*`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da8dc13e-2863-416b-b6da-6468b8b89cf8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from copy import deepcopy\n",
    "from glob import glob\n",
    "import numpy as np\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from asdf.exceptions import AsdfWarning\n",
    "from erfa import ErfaWarning\n",
    "import warnings\n",
    "from astropy.wcs import WCS\n",
    "from astropy.nddata import NDData\n",
    "import roman_datamodels.datamodels as rdd\n",
    "from jdaviz import Imviz\n",
    "from astropy.modeling import models\n",
    "\n",
    "def convert_to_fits_wcs(gwcs):\n",
    "    \"\"\"\n",
    "    Get a FITS WCS approximation to the gwcs:\n",
    "    \"\"\"\n",
    "    fits_wcs_header = gwcs.to_fits(((0, 4088), (0, 4088)))[0]\n",
    "    fits_wcs = WCS(header=fits_wcs_header)\n",
    "    return fits_wcs\n",
    "\n",
    "# paths = sorted(glob(\"/grp/roman/SCIENCE_PLATFORM_DATA/ROMANISIM/DENSE_REGION/R0.5_DP0.5_PA0/r0000101001001001001*_cal.asdf\"))\n",
    "paths = sorted(glob(os.path.expanduser(\"~/Downloads/R0.5_DP0.5_PA0/*\")))\n",
    "\n",
    "# all pixel coords will be estimated in reference to WFI01:\n",
    "dm = rdd.open(paths[0])\n",
    "ref_wcs = convert_to_fits_wcs(dm.meta.wcs)\n",
    "\n",
    "# we'll save a copy of the WCS from SCA01 for later:\n",
    "new_wcs = deepcopy(dm.meta.wcs)\n",
    "\n",
    "image_shape = (4088, 4088)\n",
    "init_offset = (-4088//2, -5000)\n",
    "\n",
    "# set up the memmap array:\n",
    "outarray = np.memmap(\n",
    "    # this path can be whatever you like:\n",
    "    filename='output.mmap', \n",
    "    # this will overwrite and update the file:\n",
    "    mode='w+', \n",
    "    # this is about big enough to fit all SCAs:\n",
    "    shape=(30_000, 25_000), \n",
    "    dtype=np.float32\n",
    ")\n",
    "# This will help center the placement of the SCAs in the memmap array:\n",
    "center_reference = np.array(outarray.shape) // 2 - np.array(image_shape) // 2\n",
    "\n",
    "# you can optionally add DC offsets to the flux in each SCA. This is helpful\n",
    "# when trying to confirm the orientation of the viewer.\n",
    "add_offset_per_sca = True\n",
    "\n",
    "sky_centers = []\n",
    "\n",
    "for i, path in tqdm(enumerate(paths), total=len(paths)):\n",
    "\n",
    "    data_label = path.split('_')[-2]\n",
    "\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.simplefilter('ignore')\n",
    "        dm = rdd.open(path)\n",
    "    \n",
    "    dm_wcs = convert_to_fits_wcs(dm.meta.wcs)\n",
    "    \n",
    "    xmin, xmax = 0, 4088\n",
    "    ymin, ymax = 0, 4088\n",
    "    sky_corners = dm_wcs.pixel_to_world([xmin, xmin, xmax, xmax], [ymin, ymax, ymin, ymax])\n",
    "    sky_centers.append(dm_wcs.pixel_to_world(np.mean([xmin, xmax]), np.mean([ymin, ymax])))\n",
    "\n",
    "    pixel_corners = ref_wcs.world_to_pixel(sky_corners)\n",
    "    \n",
    "    xmin_prime = np.mean(pixel_corners[0][:2])\n",
    "    ymin_prime = np.mean(pixel_corners[1][1::2])\n",
    "    \n",
    "    xmin_pprime = int(np.round(xmin_prime)) + center_reference[0] + init_offset[0]\n",
    "    xmax_pprime = xmin_pprime + 4088\n",
    "    ymin_pprime = int(np.round(ymin_prime)) + center_reference[1] + init_offset[1]\n",
    "    ymax_pprime = ymin_pprime + 4088\n",
    "\n",
    "    if i == 0:\n",
    "        reference_transformation = [\n",
    "            xmin_pprime, ymin_pprime\n",
    "        ]\n",
    "    \n",
    "    outarray[\n",
    "        xmin_pprime:xmax_pprime, \n",
    "        ymin_pprime:ymax_pprime\n",
    "    ] = dm.data.value + (i / 10 if add_offset_per_sca else 0.0)\n",
    "    dm.close()\n",
    "    \n",
    "outarray.flush()\n",
    "\n",
    "dx, dy = reference_transformation\n",
    "new_wcs.pipeline[0].transform = (models.Shift(-dx) & models.Shift(-dy)) | new_wcs.pipeline[0].transform"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd9f11cf-dbae-4460-b76b-e9846ee1f7b6",
   "metadata": {},
   "source": [
    "Load the memmap array into Imviz, using the \"shifted\" reference WCS (shifted to account for its new pixel position in the memmap array):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2a547fb-45bf-49d5-92cb-173cbb665325",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from astropy.coordinates import SkyCoord\n",
    "\n",
    "imviz = Imviz()\n",
    "\n",
    "ndd = NDData(outarray.T, wcs=convert_to_fits_wcs(new_wcs))\n",
    "imviz.load_data(ndd)\n",
    "\n",
    "# set a good colormap stretch\n",
    "po = imviz.plugins['Plot Options']\n",
    "po.stretch_vmin = 1\n",
    "po.stretch_vmax = 3.5\n",
    "po.image_colormap = 'viridis'\n",
    "# add markers based on the sky coordinates of the \n",
    "# centers of each SCA. These should appear near the centers\n",
    "# of the SCAs if the WCS has been transformed correctly.\n",
    "# Note: these markers won't appear in the Markers plugin table.\n",
    "pixel_centers = ndd.wcs.world_to_pixel(SkyCoord(sky_centers))\n",
    "\n",
    "markers = imviz.plugins['Markers']._obj\n",
    "markers.plugin_opened = True\n",
    "viewer = imviz.app.get_viewer('imviz-0')\n",
    "viewer.figure.marks[-1].x = pixel_centers[0]\n",
    "viewer.figure.marks[-1].y = pixel_centers[1]\n",
    "\n",
    "imviz.show(height=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a178df8e-b176-42df-aa6b-07d397a4fd62",
   "metadata": {},
   "source": [
    "In the Imviz rendering above, if `add_offset_per_sca=True`, the dimmest SCA is WFI01, and the brightest is WFI18. You can compare with the following graphic from the [WebbPSF docs](https://webbpsf.readthedocs.io/en/latest/roman.html#field-dependence-in-the-wfi-model) to see SCA names:\n",
    "\n",
    "<img src=\"https://webbpsf.readthedocs.io/en/latest/_images/field_layout.png\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2a4cbd1-2f0f-45bb-b061-07dba72465f1",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}