{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5c24fc4b",
   "metadata": {},
   "source": [
    "# Notebook to mask persistence from saturation in MIRI imaging\n",
    "\n",
    "Authors: S. Alberts<br>\n",
    "Last Updated: December 5, 2025<br>\n",
    "Tested using JWST Pipeline v1.19\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "365c2241",
   "metadata": {},
   "source": [
    "## Purpose \n",
    "This notebook identifies saturated pixels in dithered MIRI imaging uncal files and sets the affected pixels' data quality (DQ) flag to DO_NOT_USE in subsequent exposures to mask persistence.  <span style=\"color:orange;\">This alogorithm only addresses persistence caused by saturating sources observed during the input MIRI exposures.  Persistence artifacts can also be caused by saturation from observations and/or slewing prior to the input MIRI observations.  If these pixels are identified by eye, a user-supplied pixel mask or list of effected pixels can be input.</span>\n",
    "\n",
    "## Persistence artifacts following saturation\n",
    "Persistence artifacts can occur when bright (saturating or non-saturating) sources fall on the MIRI detectors, leaving a residual charge in affected pixels.  This residual charge, called persistence, decays away on a timescale dependent on the brightness of the source, how long a pixel was exposed to that source, and other factors.  Because of this complexity, there is currently no correction for persistence; mitigation techniques include dithering during the observations and masking the affected pixels in post-processing.  In the case of persistence from saturation occuring during an observation, we can use the saturation flagging in the pipeline to identify and mask pixels that are likely effected by persistence, as shown in this notebook. \n",
    "\n",
    "Persistence can appear both as the affected pixels being over-luminous (positive persistence) or under-luminous (negative persistence).  Positive persistence decays more quickly, on the order of minutes.  The decay timescale for negative persistence is not well known, but it has been observed to last for tens of hours or more and persist across integrations, dithers and filter changes. Because of this, a given affected pixel may need to be masked across subsequent integrations and/or exposures over a long timescale. For more information on persistance, including different sources of persistance, visit the [JDocs MIRI Imaging Known Issues page](https://jwst-docs.stsci.edu/known-issues-with-jwst-data/miri-known-issues/miri-imaging-known-issues#MIRIImagingKnownIssues-Persistence) and see [Dicken et al. 2024](https://ui.adsabs.harvard.edu/abs/2024A%26A...689A...5D/abstract).\n",
    "\n",
    "## Method\n",
    "The input is a set of MIRI imaging uncal files, which can include multiple imaging filters. The uncal files are ordered by time at mid-exposure. The [DQInitStep](https://jwst-pipeline.readthedocs.io/en/latest/jwst/dq_init/index.html) and [SaturationStep](https://jwst-pipeline.readthedocs.io/en/stable/jwst/saturation/index.html) are then run on each uncal file to set up the [data quality (DQ) array](https://jwst-docs.stsci.edu/accessing-jwst-data/jwst-science-data-overview#JWSTScienceDataOverview-Dataqualityarrays(DQ)) and identify saturated or partially saturated pixels. Saturation flags are recorded group-by-group in the GROUPDQ array. This code uses the GROUPDQ array to then identify pixels to mask based on a required minimum number of unsaturated groups, set by the user.  For example, the user may specify that any pixel that has less than or equal to 3 unsaturated groups will be masked. Pixels that meet this threshold are recorded and masked in  ALL SUBSEQUENT exposures in the input uncal set by setting the PIXELDQ flag to DO_NOT_USE.  All input uncal exposures are checked for saturation in order by time at mid-exposure and the masking is cumulative.  The masked uncal files are written out and can then be processed through [calwebb detector1](https://jwst-pipeline.readthedocs.io/en/latest/jwst/pipeline/calwebb_detector1.html) as per normal pipeline procedure.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd2a9046",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Check whether the local CRDS cache directory has been set.\n",
    "# If not, set it to the user home directory\n",
    "if (os.getenv('CRDS_PATH') is None):\n",
    "    os.environ['CRDS_PATH'] = os.path.join(os.path.expanduser('~'), 'crds')\n",
    "# Check whether the CRDS server URL has been set.  If not, set it.\n",
    "if (os.getenv('CRDS_SERVER_URL') is None):\n",
    "    os.environ['CRDS_SERVER_URL'] = 'https://jwst-crds.stsci.edu'\n",
    "\n",
    "# Echo CRDS path and context in use\n",
    "print('CRDS local filepath:', os.environ['CRDS_PATH'])\n",
    "print('CRDS file server:', os.environ['CRDS_SERVER_URL'])\n",
    "\n",
    "\n",
    "import glob\n",
    "import numpy as np\n",
    "\n",
    "from astropy.table import Table\n",
    "\n",
    "from jwst import datamodels\n",
    "from jwst.dq_init import DQInitStep\n",
    "from jwst.saturation import SaturationStep\n",
    "\n",
    "print(\"JWST Calibration Pipeline Version = {}\".format(jwst.__version__))\n",
    "print(\"Using CRDS Context = {}\".format(crds.get_context_name('jwst')))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef438f8d",
   "metadata": {},
   "source": [
    "## <u>User-inputs</u>\n",
    "\n",
    "### ---Set directories with uncal files to flag---\n",
    "\n",
    "All uncal files in a given directory will be grabbed.  Multiple directories can be specified.\n",
    "\n",
    "### ---Set threshold for saturation---\n",
    "\n",
    "Pixels with less than or equal to ```n_unsat_threshold``` **unsaturated** groups will be flagged for masking (default ```n_unsat_threshold=3```).\n",
    "\n",
    "### ---Set mask for pixel identified by user---\n",
    "\n",
    "A custom mask can be used to mask user-identified pixels in all exposures.  This mask can either be supplied via 1) an input list of pixel coordinates (1-indexed) or 2) a boolean array supplied as a fits file.  \n",
    "\n",
    "To use a custom mask, set ```use_custom_mask = True``` and supply a ```user_input``` as a <span style=\"color:blue;\">**.txt**</span> or <span style=\"color:blue;\">**.fits**</span> file.  A pixel coordinate text file needs to have integer value columns <span style=\"color:blue;\">x</span> and <span style=\"color:blue;\">y</span>."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d18ca13",
   "metadata": {},
   "outputs": [],
   "source": [
    "# directories containing the uncal files to mask\n",
    "# uncal_dirs = ['/path/to/uncal_files1/', '/path/to/uncal_files2/']\n",
    "uncal_dirs = ['']\n",
    "\n",
    "# set the threshold for hard saturation\n",
    "n_unsat_threshold = 3\n",
    "\n",
    "# Define a custom mask, if desired\n",
    "# use_custom_mask = False and/or user_input = '' will skip using a custom mask\n",
    "use_custom_mask = False\n",
    "user_input = '' "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4809a958",
   "metadata": {},
   "source": [
    "## <u>Define Useful Functions</u>\n",
    "\n",
    "- **get_obs_time(file, verbose=False):**  \n",
    "    Returns the mid-exposure time (MJD) from a given uncal file.\n",
    "\n",
    "- **order_by_time(files):**  \n",
    "    Sorts a list of files by their mid-exposure time.\n",
    "\n",
    "- **find_saturated_pixels(dq, n_unsat_threshold=3, dq_value=2):**  \n",
    "    Identifies pixels with fewer than a threshold number of unsaturated groups, indicating saturation.\n",
    "\n",
    "- **check_user_input_type(user_input, uncal_file=None):**  \n",
    "    Determines if the user-supplied mask is a FITS file or a pixel list, validates its shape, and returns a boolean mask.\n",
    "\n",
    "- **create_user_mask_from_pix_list(x, y, uncal_file):**  \n",
    "    Creates a boolean mask from lists of x and y pixel coordinates (1-indexed)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2efde78d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_obs_time(file, verbose=False):\n",
    "    dm = datamodels.RampModel(file)\n",
    "    obs_time = dm.meta.exposure.mid_time_mjd\n",
    "    dm.close()\n",
    "    if verbose:\n",
    "        print(obs_time)\n",
    "    return obs_time\n",
    "\n",
    "def order_by_time(files):\n",
    "    return sorted(files, key=get_obs_time)\n",
    "\n",
    "def find_saturated_pixels(dq, n_unsat_threshold=3, dq_value=2):\n",
    "    # dq flag = 2 means saturation\n",
    "    # Ensure the array is 4D\n",
    "    assert dq.ndim == 4, \"Input array must be 4-dimensional\"\n",
    "\n",
    "    # If multi-int, just make a mask of the first integration\n",
    "    dq = dq[0,:,:,:]\n",
    "\n",
    "    num_unsat = (dq != dq_value).sum(axis=0)\n",
    "    mask = num_unsat <= n_unsat_threshold\n",
    "\n",
    "    return mask\n",
    "\n",
    "def check_user_input_type(user_input, uncal_file=None):\n",
    "    import re\n",
    "\n",
    "    # check if user input is a fits file\n",
    "    if user_input.lower().endswith('.fits'):\n",
    "        user_mask_dm = datamodels.open(user_input)\n",
    "\n",
    "        if uncal_file is None:\n",
    "            raise ValueError(f'Uncal file must be provided to check user mask dimensions.')\n",
    "            return None\n",
    "        \n",
    "        else:\n",
    "            dm = datamodels.RampModel(uncal_file)\n",
    "            nx, ny = dm.data.shape[3], dm.data.shape[2]\n",
    "            dm.close()\n",
    "        \n",
    "            if user_mask_dm.data.shape[0] != nx or user_mask_dm.data.shape[1] != ny:\n",
    "                raise ValueError(f'User mask dimensions {user_mask_dm.data.shape} do not match data dimensions {ny, nx}.')\n",
    "            else:\n",
    "                return user_mask_dm.data\n",
    "\n",
    "    # check if user input is a list of pixel coordinates\n",
    "    elif user_input.lower().endswith('.txt'):\n",
    "        # Check if file exists and has two columns\n",
    "        tab = Table.read(user_input, format='ascii')\n",
    "        if 'x' in tab.colnames and 'y' in tab.colnames:\n",
    "            x = tab['x'].data\n",
    "            y = tab['y'].data\n",
    "\n",
    "            user_mask = create_user_mask_from_pix_list(x, y, uncal_file)\n",
    "            return user_mask\n",
    "\n",
    "        else:\n",
    "            raise ValueError('Input text file must contain \"x\" and \"y\" columns.')\n",
    "            return None\n",
    "                \n",
    "    else:\n",
    "        print('User input type not recognized. Please provide a fits file or a list of pixel coordinates.')\n",
    "        return None\n",
    "\n",
    "def create_user_mask_from_pix_list(x, y, uncal_file):\n",
    "    dm = datamodels.RampModel(uncal_file)\n",
    "    nx, ny = dm.data.shape[3], dm.data.shape[2]\n",
    "    dm.close()\n",
    "\n",
    "    mask = np.zeros((ny, nx), dtype=bool)\n",
    "    \n",
    "    # Ensure x and y are numpy arrays of integers\n",
    "    # shift to zero indexed\n",
    "    x = np.asarray(x, dtype=int) - 1\n",
    "    y = np.asarray(y, dtype=int) - 1\n",
    "\n",
    "    valid = (x >= 0) & (x < nx) & (y >= 0) & (y < ny)\n",
    "    mask[y[valid], x[valid]] = True\n",
    "\n",
    "    return mask\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8b1a5d1",
   "metadata": {},
   "source": [
    "## Set up masking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc8ecce2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# gather all uncal files\n",
    "uncal_files = []\n",
    "for dirs in uncal_dirs: uncal_files += glob.glob(f'{dirs}/*mirimage_uncal.fits')\n",
    "\n",
    "# sort by time\n",
    "uncal_files = order_by_time(uncal_files)\n",
    "print(f'Found and sorted {len(uncal_files)} uncal files.')\n",
    "\n",
    "\n",
    "# check if user input is provided for custom mask\n",
    "if not use_custom_mask or user_input == '' or user_input.strip() == '':\n",
    "        user_mask = None\n",
    "else:\n",
    "    user_mask = check_user_input_type(user_input, uncal_file=uncal_files[0])\n",
    "\n",
    "# set up logging to dump pipeline output\n",
    "logcfg = 'stpipe-log.cfg'\n",
    "if not os.path.exists(logcfg): \n",
    "        with open(logcfg, 'w') as f:\n",
    "            f.write('[*]\\nlevel = INFO\\nhandler = append:stpipe.log')\n",
    "            f.close()\n",
    "if not os.path.exists('stpipe.log'): \n",
    "        with open('stpipe.log', 'w') as f:\n",
    "            f.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33110929",
   "metadata": {},
   "source": [
    "## Perform masking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16f7aed1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# go through uncal files in order\n",
    "\n",
    "masks = {}\n",
    "num_masked = []\n",
    "\n",
    "for i, file in enumerate(uncal_files):\n",
    "    # loop through uncal files in order of time of mid-exposure\n",
    "    print(f'Processing file {i+1}/{len(uncal_files)}: {file}')\n",
    "    \n",
    "    dm = datamodels.RampModel(file)\n",
    "    ngroups = dm.meta.exposure.ngroups\n",
    "\n",
    "    if i==0: \n",
    "        # should be 1024, 1032\n",
    "        ny, nx = dm.data.shape[2], dm.data.shape[3]\n",
    "\n",
    "        # check if user mask is provided \n",
    "        if user_mask is not None:\n",
    "            print(f'Using user-supplied custom mask.')\n",
    "\n",
    "            masks['user_mask'] = user_mask\n",
    "            num_masked.append(np.count_nonzero(masks['user_mask']))\n",
    "            print(f'User mask has {num_masked[-1]} pixels masked.')\n",
    "        else:\n",
    "            masks['user_mask'] = None\n",
    "    \n",
    "    else:\n",
    "        # check if the dimensions of the current file match the first file\n",
    "        assert ny == dm.data.shape[2] and nx == dm.data.shape[3], \\\n",
    "            f'Warning: {file} has different dimensions than previous files. Skipping.'\n",
    "    \n",
    "    # Initialize the data quality array\n",
    "    dm = DQInitStep.call(dm, logcfg=logcfg)\n",
    "\n",
    "    # run the saturation step\n",
    "    dm = SaturationStep.call(dm, logcfg=logcfg)\n",
    "\n",
    "    # find saturated pixels\n",
    "    mask = find_saturated_pixels(dm.groupdq, n_unsat_threshold=n_unsat_threshold)\n",
    "    \n",
    "    # check if any saturated pixels were found, add mask to masks dict\n",
    "    masks[file] = {}\n",
    "    if np.count_nonzero(mask) == 0:\n",
    "        print('No saturated pixels found.')\n",
    "        masks[file]['mask'] = []\n",
    "    else:\n",
    "        masks[file]['mask'] = mask\n",
    "        num_masked.append(np.count_nonzero(mask))\n",
    "        print(f'Found {np.count_nonzero(mask)} saturated pixels ({(100 * num_masked[-1] / (nx * ny)):.2f}% of total pixels)\\n\\n.')\n",
    "\n",
    "    # apply user mask if provided\n",
    "    if masks['user_mask'] is not None:\n",
    "        dm.pixeldq[masks['user_mask']] = 1\n",
    "    \n",
    "    # for all exposure after the first, \n",
    "    # mask pixels that were saturated in previous exposures\n",
    "    if i!=0:\n",
    "        combined_mask = np.logical_or.reduce([masks[f]['mask'] for f in uncal_files[0:i] if isinstance(masks[f]['mask'], np.ndarray)])\n",
    "        dm.pixeldq[combined_mask] = 1\n",
    "  \n",
    "\n",
    "    dm.to_fits(file.replace('_uncal.fits', '_masked_persistence_uncal.fits'), overwrite=True)\n",
    "    dm.close()\n",
    "\n",
    "print(f'Cumulative percentage of pixels masked: {(100 * np.count_nonzero(combined_mask) / (nx * ny)):.2f}%')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f1a7c5f",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jwst-1.19_dev",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
