Merge pull request #6 from /issues/1

Issues/1: USe Roman PSF realized at locations

notebook/RomanDESCForwardModelLightcurves.ipynb

@@ -7,7 +7,7 @@
     "# RomanDESC SN Simulation modeling with AstroPhot\n",
     "\n",
     "Author: Michael Wood-Vasey <[email protected]>  \n",
-    "Last Verified to run: 2024-03-08\n",
+    "Last Verified to run: 2024-06-06\n",
     "\n",
     "Use the [AstroPhot](https://autostronomy.github.io/AstroPhot/) package to model lightcurve of SN in Roman+Rubin DESC simulations\n",
     "\n",
@@ -17,7 +17,7 @@
     "torch  \n",
     "\n",
     "Major TODO:\n",
-    "  * [ ] Use full Roman PSF including variation across focal plane.\n",
+    "  * [x] Use full Roman PSF including variation across focal plane.\n",
     "  * [ ] Start utility support Python file as developing package\n",
     "  * [ ] Write tests for package.  Decide on test data.\n",
     "  * [ ] Write logic into functions that can be more readily called from Python script\n",
@@ -34,12 +34,14 @@
     "This Notebook was developed and tested within a conda environment.  You can create this environment with:\n",
     "\n",
     "```\n",
-    "conda create --name astrophot -c conda-forge python astropy cudatoolkit h5py  ipykernel jupyter matplotlib numpy pandas pyyaml pyarrow scipy requests tqdm \n",
+    "conda create --name astrophot -c conda-forge python astropy cudatoolkit h5py ipykernel jupyter matplotlib numpy pandas pyyaml pyarrow scipy requests tqdm webbpsf\n",
     "conda activate astrophot\n",
     "pip install astrophot pyro-ppl torch\n",
     "ipython kernel install --user --name=astrophot\n",
     "```\n",
     "\n",
+    "And then download the `webbpsf` extra data files.\n",
+    "\n",
     "This requires astrophot >= v0.15.2"
    ]
   },
@@ -74,20 +76,15 @@
     "\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
-    "from scipy.stats import iqr\n",
-    "\n",
-    "import pandas as pd\n",
-    "import torch\n",
     "\n",
     "import astropy.units as u\n",
     "from astropy.coordinates import SkyCoord\n",
     "from astropy.io import fits\n",
     "from astropy.table import Table\n",
-    "from astropy.time import Time\n",
     "from astropy.wcs import WCS\n",
     "\n",
     "import astrophot as ap\n",
-    "from astrophot.image import PSF_Image, Window"
+    "import webbpsf\n"
    ]
   },
   {
@@ -103,8 +100,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "data_dir = \"data/RomanDESC\"\n",
-    "\n",
     "image_file_format = (\n",
     "    \"images/{band}/{visit}/Roman_TDS_simple_model_{band}_{visit}_{sca}.fits.gz\"\n",
     ")\n",
@@ -441,22 +436,10 @@
    "source": [
     "## Roman PSF\n",
     "\n",
-    "The Roman Wide-Field Imager (WFI) PSF is available in detail through the `webbpsf` package.  We here just used a simplified pre-stored model of the PSF.  This model was generated with the following code:\n",
-    "\n",
-    "```\n",
-    "import webbpsf\n",
-    "\n",
-    "wfi = webbpsf.roman.WFI()\n",
-    "wfi.filter = \"F158\"\n",
-    "default_psf = wfi.calc_psf(outfile=\"roman_psf_nominal_F158.fits\")\n",
-    "```\n",
-    "\n",
-    "For more information about Roman PSFs for the Wide-Field Imager, see\n",
+    "The Roman Wide-Field Imager (WFI) PSF is available in detail through the `webbpsf` package for band, sca, x, y and SED.  For more information about Roman PSFs see:\n",
     "\n",
     "https://roman-docs.stsci.edu/simulation-tools-handbook-home/webbpsf-for-roman/webbpsf-tutorials\n",
-    "https://github.com/spacetelescope/webbpsf/blob/develop/notebooks/WebbPSF-Roman_Tutorial.ipynb\n",
-    "\n",
-    "For actual usage, you should calculate the PSF for the specific observation, including detector (SCA).  Most notably, there is spatial variation of the PSF, and different roll angles will put the diffraction spikes in different places.  But we're just using simple approximations for now."
+    "https://github.com/spacetelescope/webbpsf/blob/develop/notebooks/WebbPSF-Roman_Tutorial.ipynb\n"
    ]
   },
   {
@@ -465,10 +448,36 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "if DATASET == \"RomanDESC\":\n",
-    "    with fits.open(\"roman_psf_nominal_F158.fits\") as hdu:\n",
-    "        pixstep = hdu[\"OVERSAMP\"].header[\"PIXELSCL\"]\n",
-    "        psf = hdu[\"OVERSAMP\"]"
+    "def get_roman_psf(band, sca, x, y):\n",
+    "    \"\"\"\n",
+    "    Return the Roman WFI PSF for the given band, sca at the detector position x, y\n",
+    "    \n",
+    "    Does not pass an SED, so default of 5700K will be used by webbpsf.roman.\n",
+    "    \"\"\"\n",
+    "    # translate from colloquial R, Z, Y, H to standard \"F*\" filter names\n",
+    "    standard_band_names = {\n",
+    "        \"F062\": \"F062\",\n",
+    "        \"F087\": \"F087\",\n",
+    "        \"F106\": \"F106\",\n",
+    "        \"F129\": \"F129\",\n",
+    "        \"F146\": \"F146\",\n",
+    "        \"F158\": \"F158\",\n",
+    "        \"F184\": \"F184\",\n",
+    "        \"R062\": \"F062\",\n",
+    "        \"Z087\": \"F087\",\n",
+    "        \"Y106\": \"F106\",\n",
+    "        \"J129\": \"F129\",\n",
+    "        \"H158\": \"F158\",\n",
+    "    }\n",
+    "    wfi = webbpsf.roman.WFI()\n",
+    "    wfi.filter = standard_band_names[band]\n",
+    "    wfi.detector = f\"SCA{sca:02d}\"\n",
+    "    wfi.detector_position = (x, y)\n",
+    "\n",
+    "    psf_hdu = wfi.calc_psf()\n",
+    "    psf = psf_hdu[\"DET_SAMP\"].data\n",
+    "\n",
+    "    return psf"
    ]
   },
   {
@@ -597,6 +606,7 @@
     "\n",
     "    mjd = header[\"MJD-OBS\"]\n",
     "    band = header[\"FILTER\"]\n",
+    "    sca = header[\"SCA_NUM\"]\n",
     "    \n",
     "    zp_band = {\"H158\": 32.603}\n",
     "    sigma_to_fwhm = 2.355\n",
@@ -617,22 +627,10 @@
     "            zeropoint = DEFAULT_ZP\n",
     "\n",
     "    wcs = WCS(hdu[hdu_idx[\"image\"]].header)\n",
+    "    x, y = wcs.world_to_pixel(coord)\n",
     "\n",
-    "    # If a PSF image is available, use it to calculate FWHM\n",
-    "    if \"psfex_info\" in hdu_idx.keys():\n",
-    "        pixstep, psf = read_psfex_image(\n",
-    "            hdu[hdu_idx[\"psfex_info\"]], hdu[hdu_idx[\"psfex_data\"]], resample=True,\n",
-    "        )\n",
-    "        pixel_scale = 3600 * wcs.pixel_scale_matrix\n",
-    "\n",
-    "        fwhm = calc_fwhm_from_psf_image(psf)\n",
-    "    else:\n",
-    "        # we construct a basic gaussian psf for each image\n",
-    "        # by giving the simga (arcsec), image width (pixels), and pixelscale (arcsec/pixel)\n",
-    "        psf = ap.utils.initialize.gaussian_psf(\n",
-    "            fwhm / sigma_to_fwhm, psf_size, pixel_scale\n",
-    "        )\n",
-    "\n",
+    "    psf = get_roman_psf(band, sca, x, y)\n",
+    "    \n",
     "    target_kwargs = {\n",
     "        \"data\": np.array(img, dtype=np.float64),\n",
     "        \"variance\": var,\n",
@@ -710,7 +708,7 @@
     "    return window\n",
     "\n",
     "def make_windows_for_targets(targets, ra, dec, npix=npix):\n",
-    "    windows = [make_window_for_target(t, ra, dec) for t in targets]\n",
+    "    windows = [make_window_for_target(t, ra, dec, npix) for t in targets]\n",
     "    return windows"
    ]
   },
@@ -985,7 +983,7 @@
     "        if model.name == model_sn[model_static_band[b]].name:\n",
     "            continue\n",
     "        for parameter in [\"center\"]:\n",
-    "            model[parameter].value = model_host[model_static_band[b]][parameter]"
+    "            model[parameter].value = model_static[model_static_band[b]][parameter]"
    ]
   },
   {
@@ -1154,7 +1152,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sn_flux_starts_at_parameter_idx = -len(targets.image_list)\n",
+    "sn_flux_starts_at_parameter_idx = -sum(live_sn)\n",
     "covar = result.covariance_matrix.detach().cpu().numpy()\n",
     "plt.imshow(\n",
     "    covar[sn_flux_starts_at_parameter_idx:, sn_flux_starts_at_parameter_idx:],\n",