From cd9c19e3f049650310cba15d21e771d2a32f565e Mon Sep 17 00:00:00 2001 From: Howard Bushouse Date: Tue, 5 Dec 2023 15:52:55 -0500 Subject: [PATCH] JP-3235: Implement NSClean step for NIRSpec 1/f noise (#8000) --- CHANGES.rst | 16 +- CODEOWNERS.txt | 1 + docs/jwst/nsclean/arguments.rst | 24 ++ docs/jwst/nsclean/index.rst | 13 + docs/jwst/nsclean/main.rst | 156 +++++++ docs/jwst/nsclean/reference_files.rst | 3 + docs/jwst/package_index.rst | 1 + docs/jwst/pipeline/calwebb_spec2.rst | 12 +- jwst/lib/suffix.py | 262 ++++++------ jwst/nsclean/__init__.py | 3 + jwst/nsclean/lib.py | 535 +++++++++++++++++++++++++ jwst/nsclean/nsclean.py | 386 ++++++++++++++++++ jwst/nsclean/nsclean_step.py | 67 ++++ jwst/pipeline/calwebb_spec2.py | 45 ++- jwst/pipeline/nsclean.cfg | 3 + jwst/regtest/test_nirspec_fs_spec2.py | 4 +- jwst/regtest/test_nirspec_ifu_spec2.py | 4 +- jwst/regtest/test_nirspec_mos_spec2.py | 6 +- jwst/step.py | 5 +- jwst/stpipe/integration.py | 1 + 20 files changed, 1385 insertions(+), 162 deletions(-) create mode 100644 docs/jwst/nsclean/arguments.rst create mode 100644 docs/jwst/nsclean/index.rst create mode 100644 docs/jwst/nsclean/main.rst create mode 100644 docs/jwst/nsclean/reference_files.rst create mode 100644 jwst/nsclean/__init__.py create mode 100644 jwst/nsclean/lib.py create mode 100644 jwst/nsclean/nsclean.py create mode 100644 jwst/nsclean/nsclean_step.py create mode 100644 jwst/pipeline/nsclean.cfg diff --git a/CHANGES.rst b/CHANGES.rst index 65e47b07d4..52d05b5115 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -61,12 +61,26 @@ imprint - Updated the logging to report which imprint image is being subtracted from the science image. [#8041] +nsclean +------- + +- Implemented this new step, which is used to remove 1/f noise from NIRSpec + images. [#8000] + photom --------- +------ - Added time-dependent correction for MIRI Imager data. [#8096, #8102, spacetelescope/stdatamodels#235] +pipeline +-------- + +- Updated the ``calwebb_spec2`` pipeline to add in calling the ``nsclean`` step + for NIRSpec exposures. Also rearranged the order of the steps, so that + ``msa_flagging`` immediately follows ``assign_wcs``, so that both steps have + been applied before calling ``nsclean``. [#8000] + pixel_replace ------------- diff --git a/CODEOWNERS.txt b/CODEOWNERS.txt index aa6e4195ff..88aab08b8b 100644 --- a/CODEOWNERS.txt +++ b/CODEOWNERS.txt @@ -24,6 +24,7 @@ jwst/ipc @philhodge @stscirij jwst/jump @dmggh @hbushouse jwst/lastframe @jemorrison @dmggh jwst/linearity @hbushouse @stscirij +jwst/nsclean @hbushouse @jemorrison jwst/outlier_detection @stsci-hack @jdavies-st jwst/persistence @philhodge @stscirij jwst/photom @hbushouse @dmggh diff --git a/docs/jwst/nsclean/arguments.rst b/docs/jwst/nsclean/arguments.rst new file mode 100644 index 0000000000..c91bfaba3e --- /dev/null +++ b/docs/jwst/nsclean/arguments.rst @@ -0,0 +1,24 @@ +.. _nsclean_arguments: + +Step Arguments +============== + +The ``nsclean`` step has the following optional arguments to control +the behavior of the processing. + +``--mask_spectral_regions`` (boolean, default=True) + Mask regions in IFU and MOS images that are within the bounding boxes + for each slice or slitlet defined in the WCS object of the image. + +``--n_sigma`` (float, default=5.0) + The sigma-clipping threshold to use when searching for outliers + and illuminated pixels to be excluded from use in the fitting + process. + +``--save_mask`` (boolean, default=False) + A flag to indicate whether the mask constructed by the step + should be saved to a file. + +``--user_mask`` (string, default=None) + Path to a user-supplied mask file. If supplied, the mask is used + directly and the process of creating a mask in the step is skipped. diff --git a/docs/jwst/nsclean/index.rst b/docs/jwst/nsclean/index.rst new file mode 100644 index 0000000000..e4a2b784bb --- /dev/null +++ b/docs/jwst/nsclean/index.rst @@ -0,0 +1,13 @@ +.. _nsclean_step: + +====================== +NSClean 1/f Correction +====================== +.. toctree:: + :maxdepth: 2 + + main.rst + reference_files.rst + arguments.rst + +.. automodapi:: jwst.nsclean diff --git a/docs/jwst/nsclean/main.rst b/docs/jwst/nsclean/main.rst new file mode 100644 index 0000000000..7a377b8fb8 --- /dev/null +++ b/docs/jwst/nsclean/main.rst @@ -0,0 +1,156 @@ +Description +=========== + +:Class: `jwst.nsclean.NSCleanStep` +:Alias: nsclean + +Overview +======== +The ``nsclean`` step applies an algorithm for removing correlated read +noise from NIRSpec images. The noise often appears as faint vertical +banding and so-called "picture frame noise." The algorithm uses dark +(unilluminated) areas of an image to fit a background model in Fourier +space. When the fit is subtracted, it removes nearly all correlated noise. +Compared to simpler strategies, like subtracting a rolling median, this +algorithm is more thorough and uniform. It is also computationally +undemanding, typically requiring only a few seconds to clean a full-frame +image. + +The correction can be applied to any type of NIRSpec exposure, including +IFU, MOS, fixed slit, and Bright Object Time Series (BOTS), in both full-frame +and subarray readouts. Time series (3D) data are corrected one integration +at a time. + +.. note:: + + The step is currently not capable of processing images taken using the + "ALLSLITS" subarray. Other subarray types are allowed. + +Details on the source of the correlated noise and the algorithm used +in the ``nsclean`` step to fit and remove it can be found in +`Rauscher 2023 `_. + +Upon completion of the step, the step status keyword "S_NSCLEN" gets set +to "COMPLETE" in the output science data. + +Assumptions +=========== +As described below, the creation of a pixel mask depends on the presence +of a World Coordinate System (WCS) object for the image, which is +constructed by the :ref:`assign_wcs ` step. +In addition, creating a mask for IFU and MOS images depends on +the presence of DQ flags assigned by the +:ref:`msaflagopen ` step. +It is therefore required that those steps be run before attempting to +apply ``nsclean``. + +Creation of an image mask +========================= +One of the key components of the correction is knowing which pixels are +unilluminated and hence can be used in fitting the background noise. +The step builds a mask on the fly for each image, which is used to mark +useable and unuseable pixels. The mask is a 2D boolean array, having the same +size as the image, with pixels set to True interpreted as being OK to use. +The process of building the mask varies somewhat depending on the +observing mode of the image being processed. Some features are common +to all modes, while others are mode-specific. The following sections +describe each type of masking that can be applied and at the end there +is a summary of the types applied to each image mode. + +The user-settable step parameter `save_mask` can be used to save the +mask to a file, if desired (see :ref:`nsclean step arguments `). + +Note that a user may supply their own mask image as input to the step, +in which case the process of creating a mask is skipped. The step parameter +`user_mask` is used to specify an input mask. + +IFU Slices +---------- +For IFU images the majority of the mask is based on knowing which +pixels are contained within the footprints of the IFU slices. To do +this, the image's World Coordinate System (WCS) object is queried in +order to determine which pixels are contained within each of the 30 +slices. Pixels within each slice are set to False (do not use) in the +mask. + +MOS/FS Slits +------------ +The footprints of each open MOS slitlet or fixed slit are flagged in +a similar way as IFU slices. For MOS and FS images, the WCS object is +queried to determine which pixels are contained within each open +slit/slitlet and they are set to False in the mask. + +MSA Failed Open Shutters +------------------------ +Pixels affected by stuck open MSA shutters are masked, because they +may contain signal. This is accomplished by setting all pixels flagged by the +:ref:`msaflagopen ` step with DQ value "MSA_FAILED_OPEN" +to False in the mask. + +NaN Pixels +---------- +Any pixel in the input image that has a value of NaN is temporarily reset +to zero for input to the fitting routine and flagged as False in the mask. +Upon completion of the noise subtraction, this population of pixels is +set back to NaN again in the output (corrected) image. + +Fixed-Slit Region Pixels +------------------------ +Full-frame MOS and IFU images may contain signal from the always open +fixed slits, which appear in fixed region in the middle of each image. +The entire region containing the fixed slits is masked out when +processing MOS and IFU images. The masked region is currently hardwired +in the step to image indexes [1:2048, 923:1116], where the indexes are +in x, y order and in 1-indexed values. + +Left/Right Reference Pixel Columns +---------------------------------- +Full-frame images contain 4 columns of reference pixels on the left and +right edges of the image. These are not to be used in the fitting +algorithm and hence are set to False in the mask. + +Outliers +-------- +Pixels in the unilluminated regions of the region can contain anomalous +signal, due to uncaught Cosmic Rays, hot pixels, etc. A sigma-clipping +routine is employed to find such outliers within the input image and set +them to False in the mask. All pixels with values greater than +:math:`median+n_sigma*sigma` are set to False in the mask. +Here `median` and `sigma` are computed +from the image using the astropy.stats `sigma_clipped_stats` routine, +using the image mask to exclude pixels that have already been flagged +and a clipping level of 5 sigma. `n_sigma` is a user-settable step +parameter, with a default value of 5.0 +(see :ref:`nsclean step arguments `). + +Mode-Specific Masking Steps +--------------------------- +The following table indicates which flavors of masking are applied to +images from each type of observing mode. + +.. |c| unicode:: U+2713 .. checkmark + ++--------------------------+-----+-----+-----+ +| | | Mode| | ++--------------------------+-----+-----+-----+ +| Masking | IFU | MOS | FS | ++==========================+=====+=====+=====+ +| IFU Slices\ :sup:`1` | |c| | | | ++--------------------------+-----+-----+-----+ +| Slits/Slitlets\ :sup:`1` | | |c| | |c| | ++--------------------------+-----+-----+-----+ +| MSA_FAILED_OPEN | |c| | |c| | |c| | ++--------------------------+-----+-----+-----+ +| NaN Pixels | |c| | |c| | |c| | ++--------------------------+-----+-----+-----+ +| FS Region | |c| | |c| | | ++--------------------------+-----+-----+-----+ +| Reference Pix | |c| | |c| | |c| | ++--------------------------+-----+-----+-----+ +| Outliers | |c| | |c| | |c| | ++--------------------------+-----+-----+-----+ + +:sup:`1`\ The application of these steps can be turned on and off via +the step parameter `mask_spectral_regions`. This parameter controls +whether the "IFU Slices" and "Slits/Slitlets" portions of the masking +are applied. diff --git a/docs/jwst/nsclean/reference_files.rst b/docs/jwst/nsclean/reference_files.rst new file mode 100644 index 0000000000..0a02d1bfd5 --- /dev/null +++ b/docs/jwst/nsclean/reference_files.rst @@ -0,0 +1,3 @@ +Reference Files +=============== +The ``nsclean`` step does not use any reference files. diff --git a/docs/jwst/package_index.rst b/docs/jwst/package_index.rst index 19180c5552..1a0ce623ef 100644 --- a/docs/jwst/package_index.rst +++ b/docs/jwst/package_index.rst @@ -41,6 +41,7 @@ Package Index model_blender/index.rst mrs_imatch/index.rst msaflagopen/index.rst + nsclean/index.rst outlier_detection/index.rst pathloss/index.rst persistence/index.rst diff --git a/docs/jwst/pipeline/calwebb_spec2.rst b/docs/jwst/pipeline/calwebb_spec2.rst index 094af07ffe..daddc6e2ff 100644 --- a/docs/jwst/pipeline/calwebb_spec2.rst +++ b/docs/jwst/pipeline/calwebb_spec2.rst @@ -47,12 +47,14 @@ TSO exposures. The instrument mode abbreviations used in the table are as follow +==========================================================+=====+=====+=====+=====+=====+=====+=================+======+========+=====+ | :ref:`assign_wcs ` | |c| | |c| | |c| | |c| | |c| | |c| | |c| | |c| | |c| | |c| | +----------------------------------------------------------+-----+-----+-----+-----+-----+-----+-----------------+------+--------+-----+ +| :ref:`msaflagopen ` | | |c| | |c| | | | | | | | | ++----------------------------------------------------------+-----+-----+-----+-----+-----+-----+-----------------+------+--------+-----+ +| :ref:`nsclean ` | |c| | |c| | |c| | | | | | | | | ++----------------------------------------------------------+-----+-----+-----+-----+-----+-----+-----------------+------+--------+-----+ | :ref:`imprint ` | | |c| | |c| | | | | | | | | +----------------------------------------------------------+-----+-----+-----+-----+-----+-----+-----------------+------+--------+-----+ | :ref:`background ` | |c| | |c| | |c| | |c| | | |c| | |c| | |c| | |c| | | +----------------------------------------------------------+-----+-----+-----+-----+-----+-----+-----------------+------+--------+-----+ -| :ref:`msaflagopen ` | | |c| | |c| | | | | | | | | -+----------------------------------------------------------+-----+-----+-----+-----+-----+-----+-----------------+------+--------+-----+ | :ref:`extract_2d `\ :sup:`1` | |c| | |c| | | | | | | |c| | |c| | |c| | +----------------------------------------------------------+-----+-----+-----+-----+-----+-----+-----------------+------+--------+-----+ | :ref:`srctype `\ :sup:`1` | |c| | |c| | |c| | |c| | |c| | |c| | |c| | |c| | |c| | |c| | @@ -181,12 +183,14 @@ abbreviations used in the table are as follows: +=======================================+============+==============+=================+==============+ | :ref:`assign_wcs ` | ALL | ALL | ALL | ALL | +---------------------------------------+------------+--------------+-----------------+--------------+ +| :ref:`msaflagopen ` | MOS, IFU | MOS, IFU | MOS, IFU | MOS | ++---------------------------------------+------------+--------------+-----------------+--------------+ +| :ref:`nsclean ` | NONE | NONE | NONE | NONE | ++---------------------------------------+------------+--------------+-----------------+--------------+ | :ref:`imprint ` | NONE | IFU | NONE | NONE | +---------------------------------------+------------+--------------+-----------------+--------------+ | :ref:`background ` | NONE | NONE | NONE | NONE | +---------------------------------------+------------+--------------+-----------------+--------------+ -| :ref:`msaflagopen ` | MOS, IFU | MOS, IFU | MOS, IFU | MOS | -+---------------------------------------+------------+--------------+-----------------+--------------+ | :ref:`extract_2d ` | MOS, FS | MOS, FS | MOS, FS | MOS | +---------------------------------------+------------+--------------+-----------------+--------------+ | :ref:`srctype ` | NONE | NONE | NONE | NONE | diff --git a/jwst/lib/suffix.py b/jwst/lib/suffix.py index 8c93443223..778c5db350 100644 --- a/jwst/lib/suffix.py +++ b/jwst/lib/suffix.py @@ -58,139 +58,135 @@ # Calculated suffixes. # This is produced by the `find_suffixes` function below _calculated_suffixes = { - 'masterbackgroundmosstep', - 'ami3pipeline', - 'whitelightstep', - 'ami_average', - 'spec3pipeline', - 'wfscombine', - 'fringe', - 'resamplestep', - 'resample_spec', - 'saturationstep', - 'firstframestep', - 'cat', - 'systemcall', - 'alignrefsstep', - 'functionwrapper', - 'darkcurrentstep', - 'imprintstep', - 'source_catalog', - 'straylight', - 'amiaveragestep', - 'ami_normalize', - 'jumpstep', - 'resample', - 'sourcetypestep', - 'spec2pipeline', - 'tweakreg', - 'msaflagopenstep', - 'outlierdetectionstep', - 'saturation', - 'pathloss', - 'groupscalestep', - 'ramp_fit', - 'lastframe', - 'darkpipeline', - 'image2pipeline', - 'outlierdetectionstackstep', - 'tso3pipeline', - 'straylightstep', - 'sourcecatalogstep', - 'dark_current', - 'mrs_imatch', - 'assignwcsstep', - 'skymatch', - 'extract_2d', - 'cubebuildstep', - 'residualfringestep', - 'residual_fringe', - 'spec2nrslamp', - 'ipc', - 'refpix', - 'image3pipeline', - 'superbiasstep', - 'hlspstep', - 'reset', - 's2d', - 'ami_analyze', - 'flatfieldstep', - 'tsophotometrystep', - 'combine_1d', - 'step', - 'cubeskymatchstep', - 'i2d', - 'group_scale', - 'rscdstep', - 'stackrefsstep', - 'flat_field', - 'guidercdsstep', - 'mrsimatchstep', - 'align_refs', - 'dqinitstep', - 'outlierdetectionscaledstep', - 'superbias', - 'assign_wcs', - 'guider_cds', - 'firstframe', - 'masterbackgroundstep', - 'master_background', - 'skymatchstep', - 'white_light', - 'persistencestep', - 'amianalyzestep', - 'backgroundstep', - 'photomstep', - 'background', - 'photom', - 'extract_1d', - 'cube_build', - 'wfscombinestep', - 'lastframestep', - 'aminormalizestep', - 'linearity', - 'rscd', - 'rampfitstep', - 'pipeline', - 'engdblog', - 'resamplespecstep', - 'persistence', - 'klip', - 'dq_init', - 'barshadowstep', - 'klipstep', - 'linearitystep', - 'hlsp', - 'pathlossstep', - 'refpixstep', - 'gainscalestep', - 'extract2dstep', - 'detector1pipeline', - 'fringestep', - 'dark', - 'whtlt', - 'guiderpipeline', - 'stackrefs', - 'imprint', - 'coron3pipeline', - 'resetstep', - 'combine1dstep', - 'outlier_detection_scaled', - 'outlier_detection_stack', - 'srctype', - 'outlier_detection', - 'engdblogstep', - 'gain_scale', - 'ipcstep', - 'jump', - 'extract1dstep', - 'tweakregstep', - 'assignmtwcsstep', - 'assign_mtwcs', - 'wavecorrstep', - 'wfsscontamstep', - 'chargemigrationstep', - 'pixelreplacestep', + 'lastframe', + 'saturation', + 'pixelreplacestep', + 'reset', + 'wavecorrstep', + 'lastframestep', + 'residual_fringe', + 'cat', + 'extract2dstep', + 'guider_cds', + 'gain_scale', + 'background', + 'outlier_detection_stack', + 'resetstep', + 'imprintstep', + 'extract1dstep', + 'pathlossstep', + 'spec2pipeline', + 'darkpipeline', + 'mrsimatchstep', + 'jumpstep', + 'masterbackgroundmosstep', + 's3d', + 'mrs_imatch', + 'saturationstep', + 'jump', + 'ipc', + 'wfscmb', + 'rampfitstep', + 'sourcecatalogstep', + 'outlier_detection', + 'assignwcsstep', + 'persistencestep', + 'chargemigrationstep', + 'assign_wcs', + 'whtlt', + 'assignmtwcsstep', + 'engdblog', + 'outlierdetectionstackstep', + 'align_refs', + 'master_background', + 'tso3pipeline', + 'backgroundstep', + 'tweakregstep', + 'i2d', + 'whitelightstep', + 'photom', + 'guidercdsstep', + 'imprint', + 'resample_spec', + 'ami_average', + 'white_light', + 'resample', + 'sourcetypestep', + 'flat_field', + 'extract_1d', + 'skymatchstep', + 'residualfringestep', + 'superbias', + 'fringestep', + 'photomstep', + 'klipstep', + 'tweakreg', + 'fringe', + 'stackrefsstep', + 'stackrefs', + 'extract_2d', + 'source_catalog', + 'resamplestep', + 'refpix', + 'amiaveragestep', + 'superbiasstep', + 'dq_init', + 'wfsscontamstep', + 'straylightstep', + 'rscdstep', + 'barshadowstep', + 'msaflagopenstep', + 'refpixstep', + 'ami_normalize', + 'wfscombinestep', + 'aminormalizestep', + 'rscd', + 'firstframestep', + 'ami_analyze', + 'resamplespecstep', + 'group_scale', + 'linearitystep', + 'combine1dstep', + 'tsophotometrystep', + 'spec3pipeline', + 'nscleanstep', + 'outlierdetectionstep', + 'groupscalestep', + 'hlspstep', + 'masterbackgroundstep', + 'alignrefsstep', + 'flatfieldstep', + 'skymatch', + 'cube_build', + 'ramp_fit', + 'firstframe', + 'srctype', + 'outlierdetectionscaledstep', + 'outlier_detection_scaled', + 'straylight', + 'spec2nrslamp', + 'pathloss', + 'guiderpipeline', + 'linearity', + 'assign_mtwcs', + 'hlsp', + 'wfscombine', + 'detector1pipeline', + 'ami3pipeline', + 'gainscalestep', + 'coron3pipeline', + 'image3pipeline', + 'combine_1d', + 'amianalyzestep', + 'ipcstep', + 'dark_current', + 'dqinitstep', + 'cubeskymatchstep', + 'cubebuildstep', + 'darkcurrentstep', + 'persistence', + 'image2pipeline', + 'klip' } diff --git a/jwst/nsclean/__init__.py b/jwst/nsclean/__init__.py new file mode 100644 index 0000000000..23103562ce --- /dev/null +++ b/jwst/nsclean/__init__.py @@ -0,0 +1,3 @@ +from .nsclean_step import NSCleanStep + +__all__ = ['NSCleanStep'] diff --git a/jwst/nsclean/lib.py b/jwst/nsclean/lib.py new file mode 100644 index 0000000000..c53c09ef2b --- /dev/null +++ b/jwst/nsclean/lib.py @@ -0,0 +1,535 @@ +import numpy as np + + +class NSClean: + """ + NSClean is the base class for removing residual correlated + read noise from JWST NIRSpec images. It is intended for use + on Level 2a pipeline products, i.e. IRS2 corrected slope + images. All processing is done in detector coordinates, with + arrays transposed and flipped so that the IRS2 "zipper" + appears along the bottom. For most users, NSClean's `clean` + method automatically transposes and flips the data as + necessary. + """ + + def __init__(self, detector, mask, fc=1/(2*2048/16), kill_width=1/(2*2048/16)/4, + buffer_sigma=1.5, sigrej=3.0, weights_kernel_sigma=32): + """ + JWST NIRSpec background modeling and subtraction -AKA "clean" (NSClean) + + Parameters + ---------- + detector : str + A string selected from {'NRS1','NRS2'} + + mask : bool array + The background model is fitted to pixels set to True. + Pixels set to False are ignored. + + fc : float + Critical frequency. This is the 1/2 power point for NSClean's + low pass filter. The units are 1/pixels. + + kill_width : float + "Kill width". The low pass filter's cutoff is 1/2 of a cosine. + The filter function goes from 1x to 0x gain over this bandwidth. + The units are 1/pixels. + + buffer_sigma : float + Standard deviation of the buffer's Gaussian smooth. This is an + optional 1-dimensional smooth in the spectral dispersion direction + only. If selected, buffing is done after modeling the background + but before subtracting it. + + sigrej : float + Standard deviation threshold for flagging statistical outliers, + to exclude them from use in fitting the background model. + + weights_kernel_sigma : int + Used to assign weights for MASK mode fitting. + """ + self.detector = detector + self.mask = mask + self.fc = fc + self.kill_width = kill_width + self.buffer_sigma = buffer_sigma + self.sigrej = sigrej + self.weights_kernel_sigma = weights_kernel_sigma + + # Transpose and flip mask to detector coordinates with the IRS2 + # zipper running along the bottom as displayed in ds9. + if self.detector == 'NRS1': + # NRS1 requires transpose only + self.mask = self.mask.transpose() + else: + # NRS2 requires transpose and flip + self.mask = self.mask.transpose()[::-1] + + self.ny = self.mask.shape[0] + self.nx = self.mask.shape[1] + + # FFT frequencies + self.rfftfreq = np.fft.rfftfreq(self.ny) + + # Extract other necessary information from model_parameters and build the + # low pass filter. This alse sets the number of Fourier vectors that + # we need to project out. + self.apodizer = np.array(make_lowpass_filter(self.fc, self.kill_width, self.nx)) + self.nvec = np.sum(self.apodizer > 0) # Project out this many frequencies + + # Only MASK mode uses a weighted fit. Compute the weights here. The aim is + # to weight by the reciprocal of the local background sample density along + # each line. Roughly approximate the local density, P, using the reciprocal of + # convolution by a Gaussian kernel for now. For now, hard code the kernel. We + # will optimize this later. + W = np.zeros((self.ny, self.nx), dtype=np.float32) # Build the kernel here + _x = np.arange(self.nx) + _mu = self.nx//2 + 1 + _sigma = self.weights_kernel_sigma + W[self.ny//2+1] = np.exp(-(_x - _mu)**2 / _sigma**2/2) / _sigma / np.sqrt(2*np.pi) + FW = np.fft.rfft2(np.fft.ifftshift(W)) + with np.errstate(divide='ignore'): + self.P = 1 / np.fft.irfft2(np.fft.rfft2(np.array(self.mask, dtype=np.float32)) * FW, (self.ny,self.nx)) + self.P = np.where(self.mask==True, self.P, 0.) # Illuminated areas carry no weight + + # Build a 1-dimensional Gaussian kernel for "buffing". Buffing is in the + # dispersion direction only. In detector coordinates, this is axis zero. Even though + # the kernel is 1-dimensional, we must still use a 2-dimensional array to + # represent it. I tried broadcasting a vector, but that made a kernel 2048 + # columns wide (in detector space). + _y = np.arange(self.ny) + _mu = self.nx//2 + 1 + _sigma = self.buffer_sigma # Standard deviation of kernel + _gkern = np.exp(-((_y-_mu) / _sigma)**2 / 2) / _sigma / np.sqrt(2*np.pi) # Centered kernel as a vector + gkern = np.zeros((self.ny, self.nx), dtype=np.float32) # 2D kernel template + gkern[:, _mu] = _gkern # Copy in the kernel. Normalization is already correct. + gkern = np.fft.ifftshift(gkern) # Shift for Numpy + self.fgkern = np.array(np.fft.rfft2(gkern), dtype=np.complex64) # FFT for fast convolution + + + def fit(self, data): + """ + Fit a background model to the supplied frame of data. + + Parameters + ---------- + data : float array + A NIRSpec image. The image must be in the detector-space + orientation with the IRS2 zipper running along the bottom as + displayed in SAOImage DS9. + + Returns + ------- + Bkg : float array + The fitted background model. + + Notes + ----- + Fitting is done line by line because the matrices get very big if one + tries to project out Fourier vectors from the entire 2K x 2K image area. + """ + model = np.zeros((self.ny, self.nx), dtype=np.float32) # Build the model here + for y in np.arange(self.ny)[4:-4]: + + # Get data and weights for this line + d = data[y][self.mask[y]] # unmasked (useable) data + p = np.diag(self.P[y][self.mask[y]]) # Weights + + # If none of the pixels in this line is useable (all masked out), + # skip and move on to the next line. + if len(d) == 0: + continue + + # Fill statistical outliers with line median. We know that the rolling + # median cleaning technique worked reasonably well, so this is a fast + # justifiable approximation. + _mu = np.median(d) # Robust estimate of mean + _sigma = 1.4826 * np.median(np.abs(d - _mu)) # Robust estimate of standard deviation + + # Fill outliers + d = np.where(np.logical_and(_mu - self.sigrej * _sigma <= d, + d <= _mu + self.sigrej * _sigma), d, _mu) + + # Build the Fourier basis matrix for this line + m = np.arange(self.nx)[self.mask[y]].reshape((-1, 1)) # Must be a column vector to broadcast + k = np.arange(self.nvec).reshape((1,-1)) # Must be a row vector to broadcast. We can optimize + # the code later by putting this into object instantiation + # since it is the same for every line. For now, leave it + # here since the cost is negligible and it may aid + # comprehension. + + # Build the basis matrix + B = np.array(np.exp(2 * np.pi * 1J * m * k / self.nx) / m.shape[0], dtype=np.complex64) + + # Compute the Moore-Penrose inverse of A = P*B. + # $A^+ = (A^H A)^{-1} A^H$ + A = np.matmul(p, B) + AH = np.conjugate(A.transpose()) # Hermitian transpose of A + pinv_PB = np.matmul(np.linalg.inv(np.matmul(AH, A)), AH) + + # Solve for the Fourier transform of this line's background samples. + # The way that we have done it, this multiplies the input data by the + # number of samples used for the fit. + rfft = np.zeros(self.nx//2 + 1, dtype=np.complex64) + rfft[:k.shape[1]] = np.matmul(np.matmul(pinv_PB, p), d) + + # Numpy requires that the forward transform multiply + # the data by n. Correct normalization. + rfft *= self.nx / m.shape[0] + + # Apodize if necessary + if self.kill_width > 0: + rfft[:self.nvec] *= self.apodizer[:self.nvec] + + # Invert the FFT to build the background model for this line + model[y] = np.fft.irfft(rfft, self.nx) + + # Done! + return model + + + def clean(self, data, buff=True): + """ + "Clean" NIRspec images by fitting and subtracting the + instrumental background. This is intended to improve the + residual correlated noise (vertical banding) that is + sometimes seen. Because the banding is not seen by the + reference pixels, normal IRS^2 processing does not + remove it. + + This is an ad-hoc correction. We model the background by + fitting it in Fourier space. There is an option + to "improve" the result by "buffing" in the spectral + dispersion direction. "Buffing" is light smoothing using + a 1-dimensional Gaussian. + + Parameters + ---------- + data : array_like + The input data. This should be the normal end result of Stage 1 processing. + + buff : bool + "Buff" the fitted spectrum by applying a slight Gaussian blur + in the spectral dispersion direction. + + Returns + ------- + data : array_like + The data, but with less striping and the background subtracted. + """ + + # Transform the data to detector space with the IRS2 zipper running along the bottom. + if self.detector == 'NRS2': + # Transpose and flip for NRS2 + data = data.transpose()[::-1] + else: + # Transpose (no flip) for NRS1 + data = data.transpose() + + # Fit the background model + Bkg = self.fit(data) # Background model + + # Buff, if requested + if buff: + Bkg = np.fft.irfft2(np.fft.rfft2(Bkg) * self.fgkern, s=Bkg.shape) + + # Subtract the background model from the data + data -= Bkg + + # Transform back to DMS space + if self.detector=='NRS2': + data = data[::-1].transpose() + else: + data = data.transpose() + + # Done + return data + + +def make_lowpass_filter(f_half_power, w_cutoff, n, d=1.0): + """ + Make a lowpass Fourier filter + + Parameters + ---------- + f_half_power : float + Half power frequency + + w_cutoff : float + Width of cosine cutoff. The response transitions + from 1x to 0x over this range of frequencies + + n : int + Number of samples in the timeseries + + d : float + Sample spacing (inverse of the sampling rate). Defaults to 1. + + Returns + ------- + filt : array + Filter array + """ + + # Make frequencies vector + freq = np.fft.rfftfreq(n, d=d) + + # Build a cosine wave that is approriately shifted + cos = (1 + np.cos(np.pi * (freq - f_half_power) / w_cutoff + np.pi / 2)) / 2 + + # Construct low-pass filter with cosine rolloff + filt = np.where(freq <= f_half_power - w_cutoff / 2, 1, cos) + filt = np.where(freq <= f_half_power + w_cutoff / 2, filt, 0) + + # Done + return filt + + +def med_abs_deviation(d, median=True): + """ + Median absolute deviation + + Computes the median and the median absolute deviation (MAD). For normally + distributed data, multiply the MAD by 1.4826 to approximate standard deviation. + If median=True (the default), this returns a tuple containing (median, MAD). + Otherwise, only the MAD is returned. + + Parameters + ---------- + d : ndarray + The input data + + median : bool + Return both the median and MAD + + Returns + ------- + m : float + median + + mad : float + median absolute deviation + """ + d = d[np.isfinite(d)] # Exclude NaNs + m = np.median(d) + mad = np.median(np.abs(d - m)) + if median is True: + return(m, mad) + else: + return(mad) + + +class NSCleanSubarray: + """ + NSCleanSubarray is the base class for removing residual correlated + read noise from generic JWST near-IR Subarray images. It is + intended for use on Level 2a pipeline products, i.e. slope images. + """ + + # Class variables. These are the same for all instances. + nloh = np.int32(12) # New line overhead in pixels + tpix = np.float32(10.e-6) # Pixel dwell time in seconds + sigrej = np.float32(4.0) # Standard deviation threshold for flagging + # statistical outliers. + + def __init__(self, data, mask, fc=(1061, 1211, 49943, 49957), + exclude_outliers=True, weights_kernel_sigma=None): + """ + Background modeling and subtraction for generic JWST near-IR subarrays. + + Parameters + ---------- + data : float array + The 2D input image data array to be operated on. + + mask : bool array + The background model is fitted to pixels set to True. + Pixels set to False are ignored. + + fc : tuple + Apodizing filter definition. These parameters are tunable. They + happen to work well for NIRSpec BOTS exposures. + 1) Unity gain for f < fc[0] + 2) Cosine roll-off from fc[0] to fc[1] + 3) Zero gain from fc[1] to fc[2] + 4) Cosine roll-on from fc[2] to fc[3] + + exclude_outliers : bool + Exclude statistical outliers and their nearest neighbors + from the background pixels mask. + + weights_kernel_sigma : float + Standard deviation of the 1-dimensional Gaussian kernel + that is used to approximate background sample density. This + is ad-hoc. See the NSClean journal article for more information. The + default for subarrays results in nearly equal weighting of all background + samples. + + Notes: + 1) NSCleanSubarray works in detector coordinates. Both the data and mask + need to be transposed and flipped so that slow-scan runs from bottom + to top as displayed in SAOImage DS9. The fast scan direction is + required to run from left to right. + """ + # Definitions + self.data = np.array(data, dtype=np.float32) + self.mask = np.array(mask, dtype=np.bool_) + self.ny = np.int32(data.shape[0]) # Number of pixels in slow scan direction + self.nx = np.int32(data.shape[1]) # Number of pixels in fast scan direction + self.fc = np.array(fc, dtype=np.float32) + self.n = np.int32(self.ny * (self.nx + self.nloh)) # Number of ticks in clocking pattern + self.rfftfreq = np.array(np.fft.rfftfreq(self.n, self.tpix), + dtype=np.float32) # Fourier frequencies in clocking pattern + + # We will weight by the inverse of the local sample density in time. We compute the local + # sample density by convolution using a Gaussian. Define the standard deviation of the + # Gaussian here. This is ad-hoc. + if weights_kernel_sigma is None: + self.weights_kernel_sigma = 1 / ((fc[0]+fc[1]) / 2) / self.tpix / 2 / 4 + else: + self.weights_kernel_sigma = weights_kernel_sigma + + # The mask potentially contains NaNs. Exclude them. + self.mask[np.isnan(self.data)] = False + + # The mask potentially contains statistical outliers. + # Optionally exclude them. + if exclude_outliers is True: + m, s = med_abs_deviation(self.data[self.mask]) # Compute median and median absolute deviation + s *= 1.4826 # Convert MAD to std + vmin = m - self.sigrej*s # Minimum value to keep + vmax = m + self.sigrej*s # Maximum value to keep + + # Temporarily change NaNs to inf, so that they don't cause problems. + self.data[np.isnan(self.data)] = np.inf + + # Flag statistical outliers + bdpx = np.array(np.where(np.logical_or(self.data < vmin, self.data > vmax), + 1, 0), dtype=np.float32) + self.data[np.isinf(self.data)] = np.nan # Restore NaNs + + bdpx[np.logical_not(self.mask)] = 0 # We don't need to worry about non-background pixels + # Also flag 4 nearest neighbors + bdpx = bdpx +\ + np.roll(bdpx, (+1,0), axis=(0,1)) +\ + np.roll(bdpx, (-1,0), axis=(0,1)) +\ + np.roll(bdpx, (0,+1), axis=(0,1)) +\ + np.roll(bdpx, (0,-1), axis=(0,1)) + # bdpx now contains the pixels to exclude from the background pixels + # mask. Exclude them. + self.mask[bdpx != 0] = False + self.data -= m # STUB - Median subtract + + # Build the apodizing filter. This has unity gain at low frequency to capture 1/f. It + # also has unity gain at Nyquist to capture alternating column noise. + # At mid frequencies, the gain is zero. Unity gain at low frequencies. + self.apodizer = np.zeros(len(self.rfftfreq), dtype=np.float32) + self.apodizer[self.rfftfreq < self.fc[0]] = 1.0 + # Cosine roll-off + here = np.logical_and(self.fc[0] <= self.rfftfreq, self.rfftfreq < self.fc[1]) + self.apodizer[here] = 0.5*(np.cos((np.pi/(self.fc[1]-self.fc[0])) * + (self.rfftfreq[here] - self.fc[0]))+1) + # Cosine roll-on + here = np.logical_and(self.fc[2] <= self.rfftfreq, self.rfftfreq < self.fc[3]) + self.apodizer[here] = 0.5*(np.cos((np.pi/(self.fc[3]-self.fc[2])) * + (self.rfftfreq[here] - self.fc[2]) + np.pi)+1) + # Unity gain between f[3] and end + self.apodizer[self.rfftfreq >= self.fc[3]] = 1.0 + + + def fit(self, return_fit=False, weight_fit=False): + """ + Fit a background model to the data. + + Parameters + ---------- + return_fit : bool + Return the Fourier transform. + + weight_fit : bool + Use weighted least squares as described in the NSClean paper. + Turn off by default. For subarrays it is TBD if this is necessary. + + Returns + ------- + rfft : numpy array + The computed Fourier transform. + """ + + # To build the incomplete Fourier matrix, we require the index of each + # clock tick of each valid pixel in the background samples. For consistency with + # numpy's notation, we call this 'm' and require it to be a column vector. + _x = np.arange(self.nx).reshape((1,-1)) + _y = np.arange(self.ny).reshape((-1,1)) + m = (_y*(self.nx+self.nloh) + _x)[self.mask].reshape((-1,1)) + + # Define which Fourier vectors to fit. For consistency with numpy, call this k. + k = np.arange(len(self.rfftfreq))[self.apodizer>0.].reshape((1,-1)) + + # Build the incomplete Fourier matrix + B = np.array(np.exp(2*np.pi*1J*m*k/self.n)/m.shape[0], dtype=np.complex64) + + # Weighted NSClean fitting + if weight_fit: + + # Build the weight matrix. Weight by the reciprocal of the local background + # sample density in time. Roughly approximate the local density, P, using + # the reciprocal of convolution by a Gaussian kernel. + _x = np.arange(self.n, dtype=np.float32) # x-values for building a 1-D Gaussian + _mu = np.float32(self.n//2+1) # Center point of Gaussian + _sigma = np.float32(self.weights_kernel_sigma) # Standard deviation of Gaussian + W = np.exp(-(_x-_mu)**2/_sigma**2/2)/_sigma/np.sqrt(2*np.pi) # Build centered Gaussian + FW = np.fft.rfft(np.fft.ifftshift(W)) # Forward FFT + _M = np.hstack((self.mask, np.zeros((self.ny,self.nloh), dtype=np.bool_))).flatten() # Add new line overhead to mask + with np.errstate(divide='ignore'): + P = 1/np.fft.irfft(np.fft.rfft(np.array(_M, dtype=np.float32)) * FW, self.n) # Compute weights + P = P[_M==True] # Keep only background samples + + # NSClean's weighting requires the Moore-Penrose invers of A = P*B. + # $A^+ = (A^H A)^{-1} A^H$ + A = P.reshape((-1,1)) * B # P is diagonal. Hadamard product is most RAM efficient + AH = np.conjugate(A.transpose()) # Hermitian transpose of A + pinv_PB = np.matmul(np.linalg.inv(np.matmul(AH, A)), AH) + + else: + # Unweighted fit + pinvB = np.linalg.pinv(B) + + # Solve for the (approximate) Fourier transform of the background samples. + rfft = np.zeros(len(self.rfftfreq), dtype=np.complex64) + if weight_fit is True: + rfft[self.apodizer>0.] = np.matmul(pinv_PB * P.reshape((1,-1)), self.data[self.mask]) + else: + rfft[self.apodizer>0.] = np.matmul(pinvB, self.data[self.mask]) + + # Numpy requires that the forward transform multiply + # the data by n. Correct normalization. + rfft *= self.n / m.shape[0] + + # Invert the apodized Fourier transform to build the background model for this integration + self.model = np.fft.irfft(rfft*self.apodizer, self.n).reshape((self.ny,-1))[:,:self.nx] + + # Done + if return_fit: + return(rfft) + + + def clean(self, weight_fit=True): + """ + Clean the data + + Parameters + ---------- + weight_fit : bool + Use weighted least squares as described in the NSClean paper. + Otherwise, it is a simple unweighted fit. + + Returns + ------- + data : 2D float array + The cleaned data array. + """ + self.fit(weight_fit=weight_fit) # Fit the background model + self.data -= self.model # Overwrite data with cleaned data + return(self.data) diff --git a/jwst/nsclean/nsclean.py b/jwst/nsclean/nsclean.py new file mode 100644 index 0000000000..b867787217 --- /dev/null +++ b/jwst/nsclean/nsclean.py @@ -0,0 +1,386 @@ +import logging +import numpy as np + +import gwcs +from astropy.stats import sigma_clipped_stats + +from jwst import datamodels +from jwst.assign_wcs import nirspec +from stdatamodels.jwst.datamodels import dqflags + +from jwst.nsclean.lib import NSClean, NSCleanSubarray + +log = logging.getLogger(__name__) +log.setLevel(logging.DEBUG) + + +def mask_ifu_slices(input_model, mask): + + """Find pixels located within IFU slices and flag them in the + mask, so that they do not get used. + + Parameters + ---------- + input_model : data model object + science data + + mask : 2D bool array + input mask that will be updated + + Returns + ------- + mask : 2D bool array + output mask with additional flags for science pixels + """ + + log.info("Finding slice pixels for an IFU image") + + # Initialize global DQ map to all zero (OK to use) + dqmap = np.zeros_like(input_model.dq) + + # Get the wcs objects for all IFU slices + list_of_wcs = nirspec.nrs_ifu_wcs(input_model) + + # Loop over the IFU slices, finding the valid region for each + for (k, ifu_wcs) in enumerate(list_of_wcs): + + # Construct array indexes for pixels in this slice + x, y = gwcs.wcstools.grid_from_bounding_box(ifu_wcs.bounding_box, + step=(1, 1), + center=True) + # Get the world coords for all pixels in this slice; + # all we actually need are wavelengths + coords = ifu_wcs(x, y) + dq = dqmap[y.astype(int), x.astype(int)] + wl = coords[2] + # set non-NaN wavelength locations as do not use (one) + valid = ~np.isnan(wl) + dq = dq[valid] + x = x[valid] + y = y[valid] + dq[:] = 1 + + # Copy DQ for this slice into global DQ map + dqmap[y.astype(int), x.astype(int)] = dq + + # Now set all non-zero locations in the mask to False (do not use) + mask[dqmap==1] = False + + return mask + + +def mask_slits(input_model, mask): + + """Find pixels located within MOS or fixed slit footprints + and flag them in the mask, so that they do not get used. + + Parameters + ---------- + input_model : data model object + science data + + mask : 2D bool array + input mask that will be updated + + Returns + ------- + mask : 2D bool array + output mask with additional flags for slit pixels + """ + + from jwst.extract_2d.nirspec import offset_wcs + + log.info("Finding slit/slitlet pixels") + + # Get the slit-to-msa frame transform from the WCS object + slit2msa = input_model.meta.wcs.get_transform('slit_frame', 'msa_frame') + + # Loop over the slits, marking all the pixels within each bounding + # box as False (do not use) in the mask. + # Note that for 3D masks (TSO mode), all planes will be set to the same value. + for slit in slit2msa.slits: + slit_wcs = nirspec.nrs_wcs_set_input(input_model, slit.name) + xlo, xhi, ylo, yhi = offset_wcs(slit_wcs) + mask[..., ylo:yhi, xlo:xhi] = False + + return mask + + +def create_mask(input_model, mask_spectral_regions, n_sigma): + """Create the pixel mask needed for setting which pixels to use + for measuring 1/f noise. + + Parameters + ---------- + input_model : data model object + science data + + mask_spectral_regions : bool + mask slit/slice regions defined in WCS + + n_sigma : float + sigma threshold for masking outliers + + Returns + ------- + mask : 2D or 3D bool array + image mask + + nan_pix : array + indexes of image locations with NaN values + """ + exptype = input_model.meta.exposure.type.lower() + + # Initialize mask to all True. Subsequent operations will mask + # out pixels that contain signal. + # Note: mask will be 3D for BOTS mode data + mask = np.full(np.shape(input_model.dq), True) + + # If IFU, mask all pixels contained in the IFU slices + if exptype == 'nrs_ifu' and mask_spectral_regions: + mask = mask_ifu_slices(input_model, mask) + + # If MOS or FS, mask all pixels affected by open slitlets + if exptype in ['nrs_fixedslit', 'nrs_brightobj', 'nrs_msaspec'] and mask_spectral_regions: + mask = mask_slits(input_model, mask) + + # If IFU or MOS, mask pixels affected by failed-open shutters + if exptype in ['nrs_ifu', 'nrs_msaspec']: + open_pix = input_model.dq & dqflags.pixel['MSA_FAILED_OPEN'] + mask[open_pix > 0] = False + + # Temporarily reset NaN pixels and mask them. + # Save the list of NaN pixel coords, so that they can be reset at the end. + nan_pix = np.isnan(input_model.data) + input_model.data[nan_pix] = 0 + mask[nan_pix] = False + + # If IFU or MOS, mask the fixed-slit area of the image; uses hardwired indexes + if exptype in ['nrs_ifu', 'nrs_msaspec']: + mask[922:1116, :] = False + + # Use left/right reference pixel columns (first and last 4). Can only be + # applied to data that uses all 2048 columns of the detector. + if mask.shape[-1] == 2048: + mask[..., :, :5] = True + mask[..., :, -4:] = True + + # Mask outliers using sigma clipping stats. + # For BOTS mode, which uses 3D data, loop over each integration separately. + if len(input_model.data.shape) == 3: + for i in range(input_model.data.shape[0]): + _, median, sigma = sigma_clipped_stats(input_model.data[i], mask=~mask[i], mask_value=0, sigma=5.0) + outliers = input_model.data[i] > (median + n_sigma * sigma) + mask[i][outliers] = False + else: + _, median, sigma = sigma_clipped_stats(input_model.data, mask=~mask, mask_value=0, sigma=5.0) + outliers = input_model.data > (median + n_sigma * sigma) + mask[outliers] = False + + + # Return the mask and the record of which pixels were NaN in the input; + # it'll be needed later + return mask, nan_pix + + +def clean_full_frame(detector, image, mask): + """Clean a full-frame (2048x2048) image. + + Parameters + ---------- + detector : str + The name of the detector from which the data originate. + + image : 2D float array + The image to be cleaned. + + mask : 2D bool array + The mask that indicates which pixels are to be used in fitting. + + Returns + ------- + cleaned_image : 2D float array + The cleaned image. + """ + + # Instantiate the cleaner + cleaner = NSClean(detector, mask) + + # Clean the image + try: + cleaned_image = cleaner.clean(image, buff=True) + except np.linalg.LinAlgError: + log.warning("Error cleaning image; step will be skipped") + return None + + return cleaned_image + + +def clean_subarray(detector, image, mask): + """Clean a subarray image. + + Parameters + ---------- + detector : str + The name of the detector from which the data originate. + + image : 2D float array + The image to be cleaned. + + mask : 2D bool array + The mask that indicates which pixels are to be used in fitting. + + Returns + ------- + cleaned_image : 2D float array + The cleaned image. + """ + + # Flip the image to detector coords. NRS1 requires a transpose + # of the axes, while NRS2 requires a transpose and flip. + if detector == "NRS1": + image = image.transpose() + mask = mask.transpose() + else: + image = image.transpose()[::-1] + mask = mask.transpose()[::-1] + + # Instantiate the cleaner + cleaner = NSCleanSubarray(image, mask) + + # Clean the image + try: + cleaned_image = cleaner.clean() + except np.linalg.LinAlgError: + log.warning("Error cleaning image; step will be skipped") + return None + + # Restore the cleaned image to the science frame + if detector == "NRS1": + cleaned_image = cleaned_image.transpose() + else: + cleaned_image = cleaned_image[::-1].transpose() + + return cleaned_image + + +def do_correction(input_model, mask_spectral_regions, n_sigma, save_mask, user_mask): + + """Apply the NSClean 1/f noise correction + + Parameters + ---------- + input_model : data model object + science data to be corrected + + mask_spectral_regions : bool + Mask slit/slice regions defined in WCS + + n_sigma : float + n-sigma rejection level for finding outliers + + save_mask : bool + switch to indicate whether the mask should be saved + + user_mask : str or None + Path to user-supplied mask image + + Returns + ------- + output_model : `~jwst.datamodel.JwstDataModel` + corrected data + + mask_model : `~jwst.datamodel.JwstDataModel` + pixel mask to be saved or None + """ + + detector = input_model.meta.instrument.detector.upper() + exp_type = input_model.meta.exposure.type + log.info(f'Input exposure type is {exp_type}, detector={detector}') + + # Check for a valid input that we can work on + if input_model.meta.subarray.name.upper() == "ALLSLITS": + log.warning("Step cannot be applied to ALLSLITS subarray images") + log.warning("Step will be skipped") + input_model.meta.cal_step.nsclean = 'SKIPPED' + return input_model, None + + output_model = input_model.copy() + + # Check for a user-supplied mask image. If so, use it. + if user_mask is not None: + mask_model = datamodels.open(user_mask) + Mask = (mask_model.data.copy()).astype(np.bool_) + + # Reset and save list of NaN pixels in the input image + nan_pix = np.isnan(input_model.data) + input_model.data[nan_pix] = 0 + Mask[nan_pix] = False + + else: + # Create the pixel mask that'll be used to indicate which pixels + # to include in the 1/f noise measurements. Basically, we're setting + # all illuminated pixels to False, so that they do not get used, and + # setting all unilluminated pixels to True (so they DO get used). + # For BOTS mode the mask will be 3D, to accommodate changes in masked + # pixels per integration. + log.info("Creating mask") + Mask, nan_pix = create_mask(input_model, mask_spectral_regions, n_sigma) + + # Store the mask image in a model, if requested + if save_mask: + if len(Mask.shape) == 3: + mask_model = datamodels.CubeModel(data=Mask) + else: + mask_model = datamodels.ImageModel(data=Mask) + else: + mask_model = None + + log.info(f"Cleaning image {input_model.meta.filename}") + + # Setup for handling 2D or 3D inputs + if len(input_model.data.shape) == 3: + nints = input_model.data.shape[0] + # Check for 3D mask + if len(Mask.shape) == 2: + log.warning("Data are 3D, but mask is 2D. Step will be skipped.") + output_model.meta.cal_step.nsclean = 'SKIPPED' + return output_model, None + else: + nints = 1 + + # Loop over integrations (even if there's only 1) + for i in range(nints): + log.debug(f" working on integration {i+1}") + if len(input_model.data.shape) == 3: + image = np.float32(input_model.data[i]) + mask = Mask[i] + else: + image = np.float32(input_model.data) + mask = Mask + + if input_model.data.shape[-2:] == (2048, 2048): + # Clean a full-frame image + cleaned_image = clean_full_frame(detector, image, mask) + else: + # Clean a subarray image + cleaned_image = clean_subarray(detector, image, mask) + + # Check for failure + if cleaned_image is None: + output_model.meta.cal_step.nsclean = 'SKIPPED' + break + else: + # Store the cleaned image in the output model + if len(output_model.data.shape) == 3: + output_model.data[i] = cleaned_image + else: + output_model.data = cleaned_image + + # Restore NaN's from original image + output_model.data[nan_pix] = np.nan + + # Set completion status + output_model.meta.cal_step.nsclean = 'COMPLETE' + + return output_model, mask_model diff --git a/jwst/nsclean/nsclean_step.py b/jwst/nsclean/nsclean_step.py new file mode 100644 index 0000000000..79a61128e1 --- /dev/null +++ b/jwst/nsclean/nsclean_step.py @@ -0,0 +1,67 @@ +from stdatamodels.jwst import datamodels + +from ..stpipe import Step +from . import nsclean + +__all__ = ["NSCleanStep"] + + +class NSCleanStep(Step): + """ + NSCleanStep: This step performs 1/f noise correction ("cleaning") + of NIRSpec images, using the "NSClean" method. + """ + + class_alias = "nsclean" + + spec = """ + mask_spectral_regions = boolean(default=True) # Mask WCS-defined regions + n_sigma = float(default=5.0) # Clipping level for outliers + save_mask = boolean(default=False) # Save the created mask + user_mask = string(default=None) # Path to user-supplied mask + skip = boolean(default=True) # By default, skip the step + """ + + def process(self, input): + """ + Fit and subtract 1/f background noise from a NIRSpec image + + Parameters + ---------- + input : `~jwst.datamodels.ImageModel`, `~jwst.datamodels.IFUImageModel` + Input datamodel to be corrected + + n_sigma : float, optional + Sigma clipping threshold to be used in detecting outliers in the image + + save_mask : bool, optional + Save the computed mask image + + user_mask : None, str, or `~jwst.datamodels.ImageModel` + Optional user-supplied mask image; path to file or opened datamodel + + mask_spectral_regions : bool, optional + Mask regions of the image defined by WCS bounding boxes for slits/slices + + Returns + ------- + output_model : `~jwst.datamodels.ImageModel`, `~jwst.datamodels.IFUImageModel` + The 1/f corrected datamodel + """ + + # Open the input data model + with datamodels.open(input) as input_model: + + # Do the NSClean correction + result = nsclean.do_correction(input_model, self.mask_spectral_regions, self.n_sigma, + self.save_mask, self.user_mask) + output_model, mask_model = result + + # Save the mask, if requested + if self.save_mask and mask_model is not None: + mask_path = self.make_output_path(basepath=input_model.meta.filename, suffix='mask') + self.log.info(f"Saving mask file {mask_path}") + mask_model.save(mask_path) + mask_model.close() + + return output_model diff --git a/jwst/pipeline/calwebb_spec2.py b/jwst/pipeline/calwebb_spec2.py index cedad6d99d..4ecd6ff5fc 100644 --- a/jwst/pipeline/calwebb_spec2.py +++ b/jwst/pipeline/calwebb_spec2.py @@ -23,6 +23,7 @@ from ..imprint import imprint_step from ..master_background import master_background_mos_step from ..msaflagopen import msaflagopen_step +from ..nsclean import nsclean_step from ..pathloss import pathloss_step from ..photom import photom_step from ..pixel_replace import pixel_replace_step @@ -47,8 +48,8 @@ class Spec2Pipeline(Pipeline): Accepts a single exposure or an association as input. Included steps are: - assign_wcs, background subtraction, NIRSpec MSA imprint subtraction, - NIRSpec MSA bad shutter flagging, 2-D subwindow extraction, flat field, + assign_wcs, NIRSpec MSA bad shutter flagging, nsclean, background subtraction, + NIRSpec MSA imprint subtraction, 2-D subwindow extraction, flat field, source type decision, straylight, fringe, residual_fringe, pathloss, barshadow, photom, resample_spec, cube_build, and extract_1d. """ @@ -63,10 +64,11 @@ class Spec2Pipeline(Pipeline): # Define aliases to steps step_defs = { - 'bkg_subtract': background_step.BackgroundStep, 'assign_wcs': assign_wcs_step.AssignWcsStep, - 'imprint_subtract': imprint_step.ImprintStep, 'msa_flagging': msaflagopen_step.MSAFlagOpenStep, + 'nsclean': nsclean_step.NSCleanStep, + 'bkg_subtract': background_step.BackgroundStep, + 'imprint_subtract': imprint_step.ImprintStep, 'extract_2d': extract_2d_step.Extract2dStep, 'master_background_mos': master_background_mos_step.MasterBackgroundMosStep, 'wavecorr': wavecorr_step.WavecorrStep, @@ -194,9 +196,7 @@ def process_exposure_product( suffix = 'cal' self.extract_1d.suffix = 'x1d' - # Apply WCS info - # check the datamodel to see if it's - # a grism image, if so get the catalog + # Check the datamodel to see if it's a grism image, if so get the catalog # name from the asn and record it to the meta if exp_type in WFSS_TYPES: try: @@ -216,8 +216,7 @@ def process_exposure_product( self._step_verification(exp_type, science, members_by_type, multi_int) # Start processing the individual steps. - # `assign_wcs` is the critical step. Without it, processing - # cannot proceed. + # `assign_wcs` is the critical step. Without it, processing cannot proceed. assign_wcs_exception = None try: calibrated = self.assign_wcs(science) @@ -240,7 +239,13 @@ def process_exposure_product( else: raise RuntimeError('Cannot determine WCS.') - # Steps whose order is the same for all types of input. + # Steps whose order is the same for all types of input: + + # apply msa_flagging (flag stuck open shutters for NIRSpec IFU and MOS) + calibrated = self.msa_flagging(calibrated) + + # apply the "nsclean" 1/f correction to NIRSpec images + calibrated = self.nsclean(calibrated) # Leakcal subtraction (imprint) occurs before background subtraction on a per-exposure basis. # If there is only one `imprint` member, this imprint exposure is subtracted from all the @@ -256,8 +261,6 @@ def process_exposure_product( calibrated = self.bkg_subtract(calibrated, members_by_type['background']) - calibrated = self.msa_flagging(calibrated) - # The order of the next few steps is tricky, depending on mode: # WFSS/Grism data need flat_field before extract_2d, but other modes # need extract_2d first. Furthermore, NIRSpec MOS and FS need @@ -351,6 +354,18 @@ def _step_verification(self, exp_type, input, members_by_type, multi_int): logic can be removed. """ + # Check for NIRSpec MSA bad shutter flagging. + if not self.msa_flagging.skip and exp_type not in ['NRS_MSASPEC', 'NRS_IFU', 'NRS_LAMP', + 'NRS_AUTOFLAT', 'NRS_AUTOWAVE']: + self.log.debug('Science data does not allow MSA flagging. Skipping "msa_flagging".') + self.msa_flagging.skip = True + + # Check for NIRSpec "nsclean" correction. Only attempt to apply to + # IFU, MOS, and FIXEDSLIT modes, for now. + if not self.nsclean.skip and exp_type not in ['NRS_MSASPEC', 'NRS_IFU', 'NRS_FIXEDSLIT']: + self.log.debug('Science data does not allow NSClean correction. Skipping "nsclean".') + self.nsclean.skip = True + # Check for image-to-image background subtraction can be done. if not self.bkg_subtract.skip: if exp_type in WFSS_TYPES or len(members_by_type['background']) > 0: @@ -383,12 +398,6 @@ def _step_verification(self, exp_type, input, members_by_type, multi_int): self.log.debug('Science data does not allow imprint processing. Skipping "imprint_subtraction".') self.imprint_subtract.skip = True - # Check for NIRSpec MSA bad shutter flagging. - if not self.msa_flagging.skip and exp_type not in ['NRS_MSASPEC', 'NRS_IFU', 'NRS_LAMP', - 'NRS_AUTOFLAT', 'NRS_AUTOWAVE']: - self.log.debug('Science data does not allow MSA flagging. Skipping "msa_flagging".') - self.msa_flagging.skip = True - # Check for straylight correction for MIRI MRS. if not self.straylight.skip and exp_type != 'MIR_MRS': self.log.debug('Science data does not allow stray light correction. Skipping "straylight".') diff --git a/jwst/pipeline/nsclean.cfg b/jwst/pipeline/nsclean.cfg new file mode 100644 index 0000000000..0fdf3542f4 --- /dev/null +++ b/jwst/pipeline/nsclean.cfg @@ -0,0 +1,3 @@ +name = "nsclean" +class = "jwst.nsclean.NSCleanStep" + diff --git a/jwst/regtest/test_nirspec_fs_spec2.py b/jwst/regtest/test_nirspec_fs_spec2.py index 07cfeee30c..3b3f14ebf1 100644 --- a/jwst/regtest/test_nirspec_fs_spec2.py +++ b/jwst/regtest/test_nirspec_fs_spec2.py @@ -52,6 +52,8 @@ def run_pipeline(jail, rtdata_module, request): # Run the calwebb_spec2 pipeline; save results from intermediate steps args = ["calwebb_spec2", rtdata.input, "--steps.assign_wcs.save_results=true", + "--steps.nsclean.skip=False", + "--steps.nsclean.save_results=true", "--steps.extract_2d.save_results=true", "--steps.wavecorr.save_results=true", "--steps.srctype.save_results=true", @@ -64,7 +66,7 @@ def run_pipeline(jail, rtdata_module, request): @pytest.mark.bigdata @pytest.mark.parametrize("suffix", [ - "assign_wcs", "extract_2d", "wavecorr", "flat_field", "pathloss", "srctype", + "assign_wcs", "nsclean", "extract_2d", "wavecorr", "flat_field", "pathloss", "srctype", "cal", "s2d", "x1d"]) def test_nirspec_fs_spec2(run_pipeline, fitsdiff_default_kwargs, suffix): """Regression test of the calwebb_spec2 pipeline on a diff --git a/jwst/regtest/test_nirspec_ifu_spec2.py b/jwst/regtest/test_nirspec_ifu_spec2.py index 0bc447de58..5eeb8516a3 100644 --- a/jwst/regtest/test_nirspec_ifu_spec2.py +++ b/jwst/regtest/test_nirspec_ifu_spec2.py @@ -24,6 +24,8 @@ def run_spec2(jail, rtdata_module): 'args': [ '--steps.assign_wcs.save_results=true', '--steps.msa_flagging.save_results=true', + '--steps.nsclean.skip=False', + '--steps.nsclean.save_results=true', '--steps.srctype.save_results=true', '--steps.flat_field.save_results=true', '--steps.pathloss.save_results=true', @@ -39,7 +41,7 @@ def run_spec2(jail, rtdata_module): @pytest.mark.parametrize( 'suffix', ['assign_wcs', 'cal', 'flat_field', 'msa_flagging', - 'pathloss', 's3d', 'srctype', 'x1d'] + 'nsclean', 'pathloss', 's3d', 'srctype', 'x1d'] ) def test_spec2(run_spec2, fitsdiff_default_kwargs, suffix): """Regression test matching output files""" diff --git a/jwst/regtest/test_nirspec_mos_spec2.py b/jwst/regtest/test_nirspec_mos_spec2.py index 0b9966501a..8bc6f909d6 100644 --- a/jwst/regtest/test_nirspec_mos_spec2.py +++ b/jwst/regtest/test_nirspec_mos_spec2.py @@ -20,8 +20,10 @@ def run_pipeline(rtdata_module): # Run the calwebb_spec2 pipeline; save results from intermediate steps args = ["calwebb_spec2", rtdata.input, "--steps.assign_wcs.save_results=true", - "--steps.bkg_subtract.save_results=true", "--steps.msa_flagging.save_results=true", + "--steps.nsclean.skip=False", + "--steps.nsclean.save_results=true", + "--steps.bkg_subtract.save_results=true", "--steps.extract_2d.save_results=true", "--steps.srctype.save_results=true", "--steps.wavecorr.save_results=true", @@ -35,7 +37,7 @@ def run_pipeline(rtdata_module): @pytest.mark.bigdata @pytest.mark.parametrize("suffix", [ - "assign_wcs", "bsub", "msa_flagging", "extract_2d", "wavecorr", "flat_field", + "assign_wcs", "msa_flagging", "nsclean", "bsub", "extract_2d", "wavecorr", "flat_field", "srctype", "pathloss", "barshadow", "cal", "s2d", "x1d"]) def test_nirspec_mos_spec2(run_pipeline, fitsdiff_default_kwargs, suffix): """Regression test of the calwebb_spec2 pipeline on a diff --git a/jwst/step.py b/jwst/step.py index df270647d8..7496e8f763 100644 --- a/jwst/step.py +++ b/jwst/step.py @@ -5,6 +5,7 @@ from .assign_wcs.assign_wcs_step import AssignWcsStep from .background.background_step import BackgroundStep from .barshadow.barshadow_step import BarShadowStep +from .charge_migration.charge_migration_step import ChargeMigrationStep from .combine_1d.combine_1d_step import Combine1dStep from .coron.stack_refs_step import StackRefsStep from .coron.align_refs_step import AlignRefsStep @@ -31,6 +32,7 @@ from .master_background.master_background_mos_step import MasterBackgroundMosStep from .mrs_imatch.mrs_imatch_step import MRSIMatchStep from .msaflagopen.msaflagopen_step import MSAFlagOpenStep +from .nsclean.nsclean_step import NSCleanStep from .outlier_detection.outlier_detection_step import OutlierDetectionStep from .outlier_detection.outlier_detection_scaled_step import OutlierDetectionScaledStep from .outlier_detection.outlier_detection_stack_step import OutlierDetectionStackStep @@ -53,8 +55,6 @@ from .superbias.superbias_step import SuperBiasStep from .tso_photometry.tso_photometry_step import TSOPhotometryStep from .tweakreg.tweakreg_step import TweakRegStep -from .charge_migration.charge_migration_step import ChargeMigrationStep - from .wavecorr.wavecorr_step import WavecorrStep from .wfs_combine.wfs_combine_step import WfsCombineStep from .wfss_contam.wfss_contam_step import WfssContamStep @@ -95,6 +95,7 @@ "MasterBackgroundMosStep", "MRSIMatchStep", "MSAFlagOpenStep", + "NSCleanStep", "OutlierDetectionStep", "OutlierDetectionScaledStep", "OutlierDetectionStackStep", diff --git a/jwst/stpipe/integration.py b/jwst/stpipe/integration.py index 84daef34c2..ddfd823838 100644 --- a/jwst/stpipe/integration.py +++ b/jwst/stpipe/integration.py @@ -64,6 +64,7 @@ def get_steps(): ("jwst.step.MasterBackgroundMosStep", 'master_background_mos', False), ("jwst.step.MRSIMatchStep", 'mrs_imatch', False), ("jwst.step.MSAFlagOpenStep", 'msa_flagging', False), + ("jwst.step.NSCleanStep", 'nsclean', False), ("jwst.step.OutlierDetectionStep", 'outlier_detection', False), ("jwst.step.OutlierDetectionScaledStep", 'outlier_detection_scaled', False), ("jwst.step.OutlierDetectionStackStep", 'outlier_detection_stack', False),