diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4c46a1f25f..9dceff2eff 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -20,8 +20,8 @@ on: type: string required: false default: '' - crds_server: - description: CRDS server + crds_server_url: + description: CRDS server URL type: string required: false default: https://jwst-crds.stsci.edu @@ -46,28 +46,17 @@ jobs: envs: | - linux: check-dependencies - linux: check-types - latest_crds_contexts: - uses: ./.github/workflows/contexts.yml - crds_context: - needs: [ latest_crds_contexts ] - runs-on: ubuntu-latest - steps: - - id: context - run: echo context=${{ github.event_name == 'workflow_dispatch' && (inputs.crds_context != '' && inputs.crds_context || needs.latest_crds_contexts.outputs.jwst) || needs.latest_crds_contexts.outputs.jwst }} >> $GITHUB_OUTPUT - outputs: - context: ${{ steps.context.outputs.context }} test: uses: OpenAstronomy/github-actions-workflows/.github/workflows/tox.yml@9f1f43251dde69da8613ea8e11144f05cdea41d5 # v1.15.0 - needs: [ crds_context ] with: setenv: | CRDS_PATH: /tmp/data/crds_cache - CRDS_SERVER_URL: ${{ github.event_name == 'workflow_dispatch' && inputs.crds_server || 'https://jwst-crds.stsci.edu' }} - CRDS_CONTEXT: ${{ needs.crds_context.outputs.context }} + CRDS_SERVER_URL: ${{ inputs.crds_server_url || 'https://jwst-crds.stsci.edu' }} + CRDS_CONTEXT: ${{ inputs.crds_context || 'jwst-edit' }} CRDS_CLIENT_RETRY_COUNT: 3 CRDS_CLIENT_RETRY_DELAY_SECONDS: 20 cache-path: /tmp/data/crds_cache - cache-key: crds-${{ needs.crds_context.outputs.context }} + cache-key: crds-${{ inputs.crds_context || 'jwst-edit' }} envs: | - linux: py310-oldestdeps-xdist-cov pytest-results-summary: true diff --git a/.github/workflows/ci_cron.yml b/.github/workflows/ci_cron.yml index c0db35ced5..1822318bbf 100644 --- a/.github/workflows/ci_cron.yml +++ b/.github/workflows/ci_cron.yml @@ -11,8 +11,8 @@ on: type: string required: false default: '' - crds_server: - description: CRDS server + crds_server_url: + description: CRDS server URL type: string required: false default: https://jwst-crds.stsci.edu @@ -22,25 +22,14 @@ concurrency: cancel-in-progress: true jobs: - latest_crds_contexts: - uses: ./.github/workflows/contexts.yml - crds_context: - needs: [ latest_crds_contexts ] - runs-on: ubuntu-latest - steps: - - id: context - run: echo context=${{ github.event_name == 'workflow_dispatch' && (inputs.crds_context != '' && inputs.crds_context || needs.latest_crds_contexts.outputs.jwst) || needs.latest_crds_contexts.outputs.jwst }} >> $GITHUB_OUTPUT - outputs: - context: ${{ steps.context.outputs.context }} test: if: (github.repository == 'spacetelescope/jwst' && (github.event_name == 'schedule' || github.event_name == 'push' || github.event_name == 'workflow_dispatch' || contains(github.event.pull_request.labels.*.name, 'run scheduled tests'))) uses: OpenAstronomy/github-actions-workflows/.github/workflows/tox.yml@9f1f43251dde69da8613ea8e11144f05cdea41d5 # v1.15.0 - needs: [ crds_context ] with: setenv: | CRDS_PATH: /tmp/crds_cache - CRDS_SERVER_URL: ${{ github.event_name == 'workflow_dispatch' && inputs.crds_server || 'https://jwst-crds.stsci.edu' }} - CRDS_CONTEXT: ${{ needs.crds_context.outputs.context }} + CRDS_SERVER_URL: ${{ inputs.crds_server_url || 'https://jwst-crds.stsci.edu' }} + CRDS_CONTEXT: ${{ inputs.crds_context || 'jwst-edit' }} CRDS_CLIENT_RETRY_COUNT: 3 CRDS_CLIENT_RETRY_DELAY_SECONDS: 20 cache-path: /tmp/crds_cache diff --git a/.github/workflows/contexts.yml b/.github/workflows/contexts.yml deleted file mode 100644 index b3047bdb63..0000000000 --- a/.github/workflows/contexts.yml +++ /dev/null @@ -1,26 +0,0 @@ -name: contexts - -on: - workflow_call: - outputs: - jwst: - value: ${{ jobs.contexts.outputs.jwst }} - workflow_dispatch: - -jobs: - contexts: - name: retrieve latest CRDS contexts - runs-on: ubuntu-latest - outputs: - jwst: ${{ steps.jwst_crds_context.outputs.pmap }} - steps: - - id: jwst_crds_context - env: - OBSERVATORY: jwst - CRDS_SERVER_URL: https://jwst-crds.stsci.edu - run: > - echo "pmap=$( - curl -s -X POST -d '{"jsonrpc": "1.0", "method": "get_default_context", "params": ["${{ env.OBSERVATORY }}", null], "id": 1}' ${{ env.CRDS_SERVER_URL }}/json/ --retry 8 --connect-timeout 10 | - python -c "import sys, json; print(json.load(sys.stdin)['result'])" - )" >> $GITHUB_OUTPUT - - run: if [[ ! -z "${{ steps.jwst_crds_context.outputs.pmap }}" ]]; then echo ${{ steps.jwst_crds_context.outputs.pmap }}; else exit 1; fi diff --git a/.github/workflows/tests_devdeps.yml b/.github/workflows/tests_devdeps.yml index dda0e41f77..5f79a9982d 100644 --- a/.github/workflows/tests_devdeps.yml +++ b/.github/workflows/tests_devdeps.yml @@ -20,8 +20,8 @@ on: type: string required: false default: '' - crds_server: - description: CRDS server + crds_server_url: + description: CRDS server URL type: string required: false default: https://jwst-crds.stsci.edu @@ -31,25 +31,14 @@ concurrency: cancel-in-progress: true jobs: - latest_crds_contexts: - uses: ./.github/workflows/contexts.yml - crds_context: - needs: [ latest_crds_contexts ] - runs-on: ubuntu-latest - steps: - - id: context - run: echo context=${{ github.event_name == 'workflow_dispatch' && (inputs.crds_context != '' && inputs.crds_context || needs.latest_crds_contexts.outputs.jwst) || needs.latest_crds_contexts.outputs.jwst }} >> $GITHUB_OUTPUT - outputs: - context: ${{ steps.context.outputs.context }} test: if: (github.repository == 'spacetelescope/jwst' && (github.event_name == 'schedule' || github.event_name == 'push' || github.event_name == 'workflow_dispatch' || contains(github.event.pull_request.labels.*.name, 'run devdeps tests'))) uses: OpenAstronomy/github-actions-workflows/.github/workflows/tox.yml@9f1f43251dde69da8613ea8e11144f05cdea41d5 # v1.15.0 - needs: [ crds_context ] with: setenv: | CRDS_PATH: /tmp/data/crds_cache - CRDS_SERVER_URL: ${{ github.event_name == 'workflow_dispatch' && inputs.crds_server || 'https://jwst-crds.stsci.edu' }} - CRDS_CONTEXT: ${{ needs.crds_context.outputs.context }} + CRDS_SERVER_URL: ${{ inputs.crds_server_url || 'https://jwst-crds.stsci.edu' }} + CRDS_CONTEXT: ${{ inputs.crds_context || 'jwst-edit' }} CRDS_CLIENT_RETRY_COUNT: 3 CRDS_CLIENT_RETRY_DELAY_SECONDS: 20 cache-path: /tmp/data/crds_cache diff --git a/changes/8961.extract_1d.rst b/changes/8961.extract_1d.rst new file mode 100644 index 0000000000..ce89d4c2b9 --- /dev/null +++ b/changes/8961.extract_1d.rst @@ -0,0 +1 @@ +Refactor the core extraction algorithm and aperture definition modules for slit and slitless extractions, for greater efficiency and maintainability. Extraction reference files in FITS format are no longer supported. Current behavior for extractions proceeding from extract1d reference files in JSON format is preserved, with minor improvements: DQ arrays are populated and error propagation is improved for some aperture types. diff --git a/docs/jwst/extract_1d/arguments.rst b/docs/jwst/extract_1d/arguments.rst index 7b338a86f4..7311608d9a 100644 --- a/docs/jwst/extract_1d/arguments.rst +++ b/docs/jwst/extract_1d/arguments.rst @@ -3,6 +3,33 @@ Step Arguments The ``extract_1d`` step has the following step-specific arguments. +General Step Arguments +---------------------- +The following arguments apply to all modes unless otherwise specified. + +``--subtract_background`` + Flag to specify whether the background should be subtracted. If None or True, + background subtraction will be performed if there are background regions + specified in the reference file. If False, no background subtraction will be + performed. Has no effect for NIRISS SOSS data. + +``--apply_apcorr`` + Switch to select whether or not to apply an APERTURE correction during the + Extract1dStep processing. Default is ``True``. Has no effect for NIRISS SOSS data. + +Step Arguments for Slit and Slitless Spectroscopic Data +------------------------------------------------------- + +``--use_source_posn`` + Specify whether the target and background extraction + region locations specified in the :ref:`EXTRACT1D ` reference + file should be shifted to account for the expected position of the source. If None (the default), + the step will decide whether to use the source position based + on the observing mode and the source type. By default, source position corrections + are attempted only for NIRSpec MOS and NIRSpec and MIRI LRS fixed-slit point sources. + Set to False to ignore position estimates for all modes; set to True to additionally attempt + source position correction for NIRSpec BOTS data or extended sources. + ``--smoothing_length`` If ``smoothing_length`` is greater than 1 (and is an odd integer), the image data used to perform background extraction will be smoothed in the @@ -20,7 +47,7 @@ The ``extract_1d`` step has the following step-specific arguments. ``--bkg_fit`` The type of fit to perform to the background data in each image column (or row, if the dispersion is vertical). There are four allowed values: - "poly", "mean", and "median", and None (the default value). If left as None, + "poly", "mean", "median", and None (the default value). If left as None, the step will search the reference file for a value - if none is found, ``bkg_fit`` will be set to "poly". If set to "poly", the background values for each pixel within all background regions in a given column (or @@ -45,59 +72,28 @@ The ``extract_1d`` step has the following step-specific arguments. 0 for that particular column (or row). If "bkg_fit" is not "poly", this parameter will be ignored. -``--bkg_sigma_clip`` - The background values will be sigma-clipped to remove outlier values from - the determination of the background. The default value is a 3.0 sigma clip. - ``--log_increment`` - Most log messages are suppressed while looping over integrations, i.e. when - the input is a CubeModel or a 3-D SlitModel. Messages will be logged while - processing the first integration, but since they would be the same for - every integration, most messages will only be written once. However, since - there can be hundreds or thousands of integrations, which can take a long - time to process, it would be useful to log a message every now and then to - let the user know that the step is still running. - - ``log_increment`` is an integer, with default value 50. If it is greater - than 0, an INFO message will be printed every ``log_increment`` - integrations, e.g. "... 150 integrations done". + For multi-integration extractions, if this parameter is set to a value greater + than zero, an INFO-level log message will be printed every `log_increment` integrations + to report on progress. Default value is 50. -``--subtract_background`` - This is a boolean flag to specify whether the background should be - subtracted. If None, the value in the :ref:`EXTRACT1D ` - reference file (if any) will be used. If not None, this parameter overrides - the value in the reference file. +``--save_profile`` + Flag to enable saving the spatial profile representing the extraction aperture. + If True, the profile is saved to disk with suffix "profile". -``--use_source_posn`` - This is a boolean flag to specify whether the target and background extraction - region locations specified in the :ref:`EXTRACT1D ` reference - file should be shifted - to account for the expected position of the source. If None (the default), - the step will make the decision of whether to use the source position based - on the observing mode and the source type. The source position will only be - used for point sources and for modes where the source could be located - off-center due to things like nodding or dithering. If turned on, the position - of the source is used in conjunction with the World Coordinate System (WCS) to - compute the x/y source location. For NIRSpec non-IFU modes, the source position - is determined from the ``source_xpos/source_ypos`` parameters. For MIRI LRS fixed slit, - the dither offset is applied to the sky pointing location to determine source position. - All other modes use ``targ_ra/targ_dec``. If this parameter is specified in the - :ref:`EXTRACT1D ` reference file, the reference file value will - override any automatic settings based on exposure and source type. As always, a value - given by the user as an argument to the step overrides all settings in the reference - file or within the step code. +``--save_scene_model`` + Flag to enable saving a model of the 2D flux as defined by the extraction aperture. + If True, the model is saved to disk with suffix "scene_model". + +Step Arguments for IFU Data +--------------------------- ``--center_xy`` A list of two integer values giving the desired x/y location for the center of the circular extraction aperture used for extracting spectra from 3-D - IFU cubes. Ignored for non-IFU modes and non-point sources. Must be given in - x,y order and in units of pixels along the x,y axes of the 3-D IFU cube, e.g. - ``--center_xy="27,28"``. If given, the values override any position derived - from the use of the ``use_source_posn`` argument. Default is None. - -``--apply_apcorr`` - Switch to select whether or not to apply an APERTURE correction during the - Extract1dStep processing. Default is ``True`` + IFU cubes. Must be given in x,y order and in units of pixels along the x,y + axes of the 3-D IFU cube, e.g. ``--center_xy="27,28"``. + Default is None. ``--ifu_autocen`` Switch to select whether or not to enable auto-centroiding of the extraction @@ -106,6 +102,10 @@ The ``extract_1d`` step has the following step-specific arguments. becomes extremely low) and using DAOStarFinder to locate the brightest source in the field. Default is ``False``. +``--bkg_sigma_clip`` + The background values will be sigma-clipped to remove outlier values from + the determination of the background. The default value is a 3.0 sigma clip. + ``--ifu_rfcorr`` Switch to select whether or not to run 1d residual fringe correction on the extracted 1d spectrum (MIRI MRS only). Default is ``False``. @@ -128,65 +128,66 @@ The ``extract_1d`` step has the following step-specific arguments. for covariance between adjacent spaxels in the IFU data cube. The default value is 1.0 (i.e., no correction) unless set by a user or a parameter reference file. This parameter only affects MIRI and NIRSpec IFU spectroscopy. - + +Step Arguments for NIRISS SOSS Data +----------------------------------- + ``--soss_atoca`` - This is a NIRISS-SOSS algorithm-specific parameter; if True, use the ATOCA - algorithm to treat order contamination. Default is ``True``. + Flag to enable using the ATOCA algorithm to treat order contamination. Default is ``True``. ``--soss_threshold`` - This is a NIRISS-SOSS algorithm-specific parameter; this sets the threshold - value for a pixel to be included when modelling the spectral trace. The default + Threshold value for a pixel to be included when modeling the spectral trace. The default value is 0.01. ``--soss_n_os`` - This is a NIRISS-SOSS algorithm-specific parameter; this is an integer that sets + An integer that sets the oversampling factor of the underlying wavelength grid used when modeling the trace. The default value is 2. -``--soss_estimate`` - This is a NIRISS-SOSS algorithm-specific parameter; filename or SpecModel of the - estimate of the target flux. The estimate must be a SpecModel with wavelength and - flux values. - ``--soss_wave_grid_in`` - This is a NIRISS-SOSS algorithm-specific parameter; filename or SossWaveGridModel + Filename or SossWaveGridModel containing the wavelength grid used by ATOCA to model each valid pixel of the detector. If not given, the grid is determined based on an estimate of the flux (soss_estimate), the relative tolerance (soss_rtol) required on each pixel model and the maximum grid size (soss_max_grid_size). ``--soss_wave_grid_out`` - This is a NIRISS-SOSS algorithm-specific parameter; filename to hold the wavelength + Filename to hold the wavelength grid calculated by ATOCA, stored in a SossWaveGridModel. +``--soss_estimate`` + Filename or SpecModel of the + estimate of the target flux. The estimate must be a SpecModel with wavelength and + flux values. + ``--soss_rtol`` - This is a NIRISS-SOSS algorithm-specific parameter; the relative tolerance needed on a + The relative tolerance needed on a pixel model. It is used to determine the sampling of the soss_wave_grid when not directly given. Default value is 1.e-4. ``--soss_max_grid_size`` - This is a NIRISS-SOSS algorithm-specific parameter; the maximum grid size allowed. It is + The maximum grid size allowed. It is used when soss_wave_grid is not provided to make sure the computation time or the memory used stays reasonable. Default value is 20000. ``--soss_tikfac`` - This is a NIRISS-SOSS algorithm-specific parameter; this is the regularization + This is the regularization factor used in the SOSS extraction. If not specified, ATOCA will calculate a best-fit value for the Tikhonov factor. ``--soss_width`` - This is a NIRISS-SOSS algorithm-specific parameter; this specifies the aperture + This specifies the aperture width used to extract the 1D spectrum from the decontaminated trace. The default value is 40.0 pixels. ``--soss_bad_pix`` - This is a NIRISS-SOSS algorithm-specific parameter; this parameter sets the method + This parameter sets the method used to handle bad pixels. There are currently two options: "model" will replace the bad pixel values with a modeled value, while "masking" will omit those pixels from the spectrum. The default value is "model". ``--soss_modelname`` - This is a NIRISS-SOSS algorithm-specific parameter; if set, this will provide + If set, this will provide the optional ATOCA model output of traces and pixel weights, with the filename set by this parameter. By default this is set to None and this output is not provided. diff --git a/docs/jwst/extract_1d/description.rst b/docs/jwst/extract_1d/description.rst index f82e39122f..bd31735236 100644 --- a/docs/jwst/extract_1d/description.rst +++ b/docs/jwst/extract_1d/description.rst @@ -7,9 +7,9 @@ Description Overview -------- The ``extract_1d`` step extracts a 1D signal from a 2D or 3D dataset and -writes spectral data to an "x1d" product. This works on all JWST spectroscopic -modes, including MIRI LRS (slit and slitless) and MRS, NIRCam WFSS and -TSGRISM, NIRISS WFSS and SOSS, and NIRSpec fixed-slit, IFU, and MOS. +writes spectral data to an "x1d" product (or "x1dints" for time series data). +This step works on all JWST spectroscopic modes, including MIRI LRS (slit and slitless) +and MRS, NIRCam WFSS and TSGRISM, NIRISS WFSS and SOSS, and NIRSpec fixed-slit, IFU, and MOS. An EXTRACT1D reference file is used for most modes to specify the location and size of the target and background extraction apertures. @@ -52,8 +52,8 @@ CubeModel, SlitModel, IFUCubeModel, ImageModel, MultiSlitModel, or a ModelContai For some JWST modes this is usually a resampled product, such as the "s2d" products for MIRI LRS fixed-slit, NIRSpec fixed-slit, and NIRSpec MOS, or the "s3d" products for MIRI MRS and NIRSpec IFU. For other modes that are not resampled (e.g. MIRI -LRS slitless, NIRISS SOSS, NIRSpec BrightObj, and NIRCam and NIRISS WFSS), this will -be a "cal" product. +LRS slitless, NIRISS SOSS, NIRSpec BOTS, and NIRCam and NIRISS WFSS), this will +be a "cal" or "calints" product. For modes that have multiple slit instances (NIRSpec fixed-slit and MOS, WFSS), the SCI extensions should have the keyword SLTNAME to specify which slit was extracted, though if there is only one slit (e.g. MIRI LRS and NIRISS SOSS), the slit name can @@ -62,9 +62,9 @@ be taken from the EXTRACT1D reference file instead. Normally the :ref:`photom ` step should be applied before running ``extract_1d``. If ``photom`` has not been run, a warning will be logged and the output of ``extract_1d`` will be in units of count rate. The ``photom`` step -converts data to units of either surface brightness (MegaJanskys per steradian) or, +converts data to units of either surface brightness (megajanskys per steradian) or, for point sources observed with NIRSpec and NIRISS SOSS, units of flux density -(MegaJanskys). +(megajanskys). Output ------ @@ -73,41 +73,44 @@ be an output table extension with the name EXTRACT1D. This extension will have columns WAVELENGTH, FLUX, FLUX_ERROR, FLUX_VAR_POISSON, FLUX_VAR_RNOISE, FLUX_VAR_FLAT, SURF_BRIGHT, SB_ERROR, SB_VAR_POISSON, SB_VAR_RNOISE, SB_VAR_FLAT, DQ, BACKGROUND, BKGD_ERROR, BKGD_VAR_POISSON, BKGD_VAR_RNOISE, -BKGD_VAR_FLAT and NPIXELS. -Some meta data will be written to the table header, mostly copied from the -input header. For slit-like modes the extraction region is -recorded in the meta data of the table header as EXTRXSTR (x start of extraction), -EXTRXSTP (x end of extraction), EXTRYSTR (y start of extraction), and -EXTRYSTP (y end of extraction). For MIRI and NIRSpec IFU data the center of -the extraction region is recorded in the meta data EXTR_X (x center of extraction region) +BKGD_VAR_FLAT and NPIXELS. Some metadata for the slit will be written to the header for +the table extension, mostly copied from the input SCI extension headers. + +For slit-like modes, the extraction region is +recorded in the metadata of the table header as EXTRXSTR (x start of extraction), +EXTRXSTP (x end of extraction), EXTRYSTR (y start of extraction), and +EXTRYSTP (y end of extraction). For MIRI and NIRSpec IFU data, the center of +the extraction region is recorded in the metadata EXTR_X (x center of extraction region) and EXTR_Y (y center of extraction region). The NIRISS SOSS algorithm is a specialized extraction -algorithm that does not use fixed limits, therefore no extraction limits are provided for this mode. - +algorithm that does not use fixed limits; therefore, no extraction limits are provided for this mode. +Note that the pipeline takes input start/stop values from the reference files to be +zero-indexed positions, but all extraction values are recorded in the headers as one-indexed +values, following FITS header conventions. The output WAVELENGTH data is copied from the wavelength array of the input 2D data, -if that attribute exists and was populated, otherwise it is calculated from the WCS. -FLUX is the flux density in Janskys; see keyword TUNIT2 if the data are -in a FITS BINTABLE. FLUX_ERROR is the error estimate for FLUX, and it has the +if that attribute exists and was populated. Otherwise, it is calculated from the WCS. +FLUX is the summed flux density in janskys (see keyword TUNIT2 in the FITS table header). +FLUX_ERROR is the error estimate for FLUX; it has the same units as FLUX. The error is calculated as the square root of the sum of the three variance arrays: Poisson, read noise (RNOISE), and flat field (FLAT). SURF_BRIGHT is the surface brightness in MJy / sr, except that for point sources observed with NIRSpec and NIRISS SOSS, SURF_BRIGHT will be set to -zero, because there's no way to express the extracted results from those modes +zero, because there is no way to express the extracted results from those modes as a surface brightness. SB_ERROR is the error estimate for SURF_BRIGHT, calculated in the same fashion as FLUX_ERROR but using the SB_VAR arrays. While it's expected that a user will make use of the FLUX column for point-source data and the SURF_BRIGHT column for an extended source, both columns are populated (except for NIRSpec and NIRISS SOSS point sources, as mentioned above). + The ``extract_1d`` step collapses the input data from 2-D to 1-D by summing one or more rows (or columns, depending on the dispersion direction). -A background may optionally be subtracted, but -there are also other options for background subtraction prior to ``extract_1d``. +A residual background may optionally be subtracted, in addition to any +background subtraction performed prior to ``extract_1d``. For the case of input data in units of MJy / sr, the SURF_BRIGHT -and BACKGROUND columns are -populated by dividing the sum by the number of pixels (see the NPIXELS column, -described below) that were added together. The FLUX column is populated -by multiplying the sum by the solid angle of a pixel, and also multiplying -by 10^6 to convert from MJy to Jy. +and BACKGROUND columns are populated by dividing the sum by the number of pixels +(see the NPIXELS column, described below) summed over during extraction. +The FLUX column is populated by multiplying the sum by the solid angle of a pixel, +and also multiplying by 10^6 to convert from MJy to Jy. For the case of input data in units of MJy (i.e. point sources, NIRSpec or NIRISS SOSS), the SURF_BRIGHT column is set to zero, the FLUX column is just multiplied by 10^6, and the BACKGROUND column is @@ -115,12 +118,18 @@ divided by NPIXELS and by the solid angle of a pixel to convert to surface brightness (MJy / sr). NPIXELS is the number of pixels that were added together for the source -extraction region. Note that this is not necessarily a constant, and -the value is not necessarily an integer (the data type is float). +extraction region. Note that this is not necessarily a constant, since some +pixels might be excluded for some wavelengths and included for others, and +the value is not necessarily an integer, since partial pixels may have been +included in the extraction aperture. + BACKGROUND is the measured background, scaled to the extraction width used for FLUX and SURF_BRIGHT. BACKGROUND will be zero if background subtraction is not requested. BKGD_ERROR is calculated as the square root of the sum of the -BKGD_VAR arrays. DQ is not populated with useful values yet. +BKGD_VAR arrays. + +The DQ array is set to DO_NOT_USE for pixels with NaN flux values and zero +otherwise. .. _extract-1d-for-slits: @@ -136,114 +145,133 @@ Source Extraction Region As described in the documentation for the :ref:`EXTRACT1D ` reference file, the characteristics of the source extraction region can be specified in one -of two different ways. -The simplest approach is to use the ``xstart``, ``xstop``, ``ystart``, -``ystop``, and ``extract_width`` parameters. Note that all of these values are -zero-indexed integers, the start and stop limits are inclusive, and the values -are in the frame of the image being operated on (which could be a cutout of a -larger original image). -If ``dispaxis=1``, the limits in the dispersion direction are ``xstart`` -and ``xstop`` and the limits in the cross-dispersion direction are ``ystart`` -and ``ystop``. If ``dispaxis=2``, the rolls are reversed. - -If ``extract_width`` is also given, that takes priority over ``ystart`` and -``ystop`` (for ``dispaxis=1``) for the extraction width, but ``ystart`` and -``ystop`` will still be used to define the centering of the extraction region -in the cross-dispersion direction. For point source data, -then the ``xstart`` and ``xstop`` values (dispaxis = 2) are shifted to account -for the expected location of the source. If dispaxis=1, then the ``ystart`` and ``ystop`` values -are modified. The offset amount is calculated internally. If it is not desired to apply this -offset, then set ``use_source_posn`` = False. If the ``use_source_posn`` parameter is None (default), -the values of ``xstart/xstop`` or ``ystart/ystop`` in the ``extract_1d`` reference file will be used -to determine the center position of the extraction aperture. If these values are not set in the -reference file, the ``use_source_posn`` will be set internally to True for point source data -according to the table given in :ref:`srctype `. -Any of the extraction location parameters will be modified internally by the step code if the -extraction region would extend outside the limits of the input image or outside -the domain specified by the WCS. - -A more flexible way to specify the source extraction region is via the ``src_coeff`` -parameter. ``src_coeff`` is specified as a list of lists of floating-point +of two different ways. + +The simplest approach is to use the `xstart`, `xstop`, `ystart`, +`ystop`, and `extract_width` parameters. Note that all of these values are +zero-indexed floating point values, the start and stop limits are inclusive, and +the values are in the frame of the image being operated on (which could be a cutout +of a larger original image). +If `dispaxis=1`, the limits in the dispersion direction are `xstart` +and `xstop` and the limits in the cross-dispersion direction are `ystart` +and `ystop`. If `dispaxis=2`, the roles are reversed. + +If `extract_width` is also given, the start and stop values are used to define +the center of the extraction region in the cross-dispersion direction, but the +width of the aperture is set by the `extract_width` value. + +For some instruments and modes, the cross-dispersion start and stop values may be shifted +to account for the expected location of the source. This option +is available for NIRSpec MOS, fixed-slit, and BOTS data, as well as MIRI LRS fixed-slit. +If `use_source_posn` is set to None via the reference file or input parameters, +it is turned on by default for all point sources in these modes, except NIRSpec BOTS. +To turn it on for NIRSpec BOTS or extended sources, set `use_source_posn` to True. +To turn it off for any mode, set `use_source_posn` to False. +If source position correction is enabled, the planned location for the source is +calculated internally, via header metadata recording the source position and the +spectral WCS transforms, then used to offset the extraction start and stop values +in the cross-dispersion direction. + +A more flexible way to specify the source extraction region is via the `src_coeff` +parameter. `src_coeff` is specified as a list of lists of floating-point polynomial coefficients that define the lower and upper limits of the source extraction region as a function of dispersion. This allows, for example, following a tilted or curved spectral trace or simply following the variation in cross-dispersion FWHM as a function of wavelength. -If both ``src_coeff`` and ``ystart``/``ystop`` values are given, ``src_coeff`` -takes precedence. The ``xstart`` and ``xstop`` values can still be used to +If both `src_coeff` and cross-dispersion start/stop values are given, `src_coeff` +takes precedence. The start/stop values can still be used to limit the range of the extraction in the dispersion direction. More details on the specification and use of polynomial coefficients is given below. +Note that if source position correction is enabled, the position offset is applied to +any supplied `src_coeff` values, as well as the cross-dispersion start/stop values. +To ensure the provided `src_coeff` values are used as-is, set `use_source_posn` +to False. + + Background Extraction Regions ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ One or more background extraction regions for a given aperture instance can -be specified using the ``bkg_coeff`` parameter in the EXTRACT1D reference file. -This is directly analogous to the use of ``src_coeff`` for specifying source +be specified using the `bkg_coeff` parameter in the EXTRACT1D reference file. +This is directly analogous to the use of `src_coeff` for specifying source extraction regions and functions in exactly the same way. More details on the use of polynomial coefficients is given in the next section. -Background subtraction will be done if and only if ``bkg_coeff`` is given in -the EXTRACT1D reference file. The background is determined independently for + +By default, background subtraction will be done if `bkg_coeff` is set in +the EXTRACT1D reference file. To turn it off without modifying the reference +file, set `subtract_background` to False in the input step parameters. + +The background values are determined independently for each column (or row, if dispersion is vertical), using pixel values from all background regions within each column (or row). +Parameters related to background fitting are `smoothing_length`, +`bkg_fit`, and `bkg_order`: -Parameters related to background subtraction are ``smoothing_length``, -``bkg_fit``, and ``bkg_order``: - -#. If ``smoothing_length`` is specified, the 2D image data used to perform +#. If `smoothing_length` is specified, the 2D image data used to perform background extraction will be smoothed along the dispersion direction using - a boxcar of width ``smoothing_length`` (in pixels). If not specified, no + a boxcar of width `smoothing_length` (in pixels). If not specified, no smoothing of the input 2D image data is performed. -#. ``bkg_fit`` specifies the type of background computation to be performed +#. `bkg_fit` specifies the type of fit to the background data, to be performed within each column (or row). The default value is None; if not set by the user, the step will search the reference file for a value. If no value - is found, ``bkg_fit`` will be set to "poly". The "poly" mode fits a - polynomial of order ``bkg_order`` to the background values within + is found, `bkg_fit` will be set to "poly". The "poly" mode fits a + polynomial of order `bkg_order` to the background values within the column (or row). Alternatively, values of "mean" or "median" can be specified in order to compute the simple mean or median of the background - values in each column (or row). Note that using "bkg_fit=mean" is - mathematically equivalent to "bkg_fit=poly" with "bkg_order=0". If ``bkg_fit`` - is provided both by a reference file and by the user, e.g. - ``steps.extract_1d.bkg_fit='poly'``, the user-supplied value will override - the reference file value. + values in each column (or row). Note that using `bkg_fit=mean` is + mathematically equivalent to `bkg_fit=poly` with `bkg_order=0`. -#. If ``bkg_fit=poly`` is specified, ``bkg_order`` is used to indicate the +#. If `bkg_fit=poly` is specified, `bkg_order` is used to indicate the polynomial order to be used. The default value is zero, i.e. a constant. During source extraction, the background fit is evaluated at each pixel within the -source extraction region for that column (row), and the fitted values will -be subtracted (pixel by pixel) from the source count rate. +source extraction region for that column/row, and the fitted values will +be subtracted (pixel by pixel) from the source count rate, prior to summing +over the aperture. + +If source position correction is enabled, the calculated position offset is applied to +any supplied `bkg_coeff` values, as well as the source aperture limit values. +To ensure the provided `bkg_coeff` values are used as-is, set `use_source_posn` +to False. Source and Background Coefficient Lists ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ The interpretation and use of polynomial coefficients to specify source and -background extraction regions via ``src_coeff`` and ``bkg_coeff`` is the same. -The coefficients are specified as a list of an even number of lists (an -even number because both the lower and upper limits of each extraction region -must be specified). The source extraction coefficients will normally be -a list of just two lists, the coefficients for the lower limit function +background extraction regions is the same for both source coefficients (`src_coeff`) +and background coefficients (`bkg_coeff`). + +Polynomials specified via `src_coeff` and `bkg_coeff` are functions of either wavelength +(in microns) or pixel number (pixels in the dispersion direction, with respect to +the input 2D slit image), which is specified by the parameter `independent_var`. +The default is "pixel"; the alternative is "wavelength". The dependent values of these +polynomial functions are always pixel numbers (zero-indexed) in the cross-dispersion +direction, with respect to the input 2D slit image. + +The coefficients for the polynomial functions are specified as a list of an +even number of lists (an even number because both the lower and upper limits of each +extraction region must be specified). The source extraction coefficients will normally +be a list of just two lists: the coefficients for the lower limit function and the coefficients for the upper limit function of one extraction region. The limits could just be constant values, -e.g. \[\[324.5\], \[335.5\]\]. Straight but tilted lines are linear functions: +e.g. `[[324.5], [335.5]]`. Straight but tilted lines are linear functions, e.g. +`[[324.5, 0.0137], [335.5, 0.0137]]`. -\[\[324.5, 0.0137\], \[335.5, 0.0137\]\] - -Multiple regions may be specified for either the source or background, or -both. It will be common to specify more than one background region. Here +Multiple regions may be specified for either the source or background, but it is +more common to specify more than one background region. Here is an example for specifying two background regions: -\[\[315.2, 0.0135\], \[320.7, 0.0135\], \[341.1, 0.0139\], \[346.8, 0.0139\]\] +`[[315.2, 0.0135], [320.7, 0.0135], [341.1, 0.0139], [346.8, 0.0139]]` This is interpreted as follows: -* \[315.2, 0.0135\]: lower limit for first background region -* \[320.7, 0.0135\]: upper limit for first background region -* \[341.1, 0.0139\]: lower limit for second background region -* \[346.8, 0.0139\]: upper limit for second background region +* `[315.2, 0.0135]`: lower limit for first background region +* `[320.7, 0.0135]`: upper limit for first background region +* `[341.1, 0.0139]`: lower limit for second background region +* `[346.8, 0.0139]`: upper limit for second background region -Note: If the dispersion direction is vertical, replace "lower" with "left" and -"upper" with "right" in the above description. -Notice especially that ``src_coeff`` and ``bkg_coeff`` contain floating-point +Note that `src_coeff` and `bkg_coeff` contain floating-point values. For interpreting fractions of a pixel, the convention used here is that the pixel number at the center of a pixel is a whole number. Thus, if a lower or upper limit is a whole number, that limit splits the pixel @@ -251,7 +279,14 @@ in two, so the weight for that pixel will be 0.5. To include all the pixels between 325 and 335 inclusive, for example, the lower and upper limits would be given as 324.5 and 335.5 respectively. -The order of a polynomial is specified implicitly to be one less than the +Please note that this is different from the convention used for the cross-dispersion +start/stop values, which are expected to be inclusive index values. For the example here, +for horizontal dispersion, `ystart = 325`, `ystop = 335` is equivalent +to `src_coeff = [[324.5],[335.5]]`. To include half a pixel more at the top +and bottom of the aperture, `ystart = 324.5`, `ystop = 335.5` is equivalent +to `src_coeff = [[324],[336]]`. + +The order of the polynomial is specified implicitly to be one less than the number of coefficients. The number of coefficients for a lower or upper extraction region limit must be at least one (i.e. zeroth-order polynomial). There is no predefined upper limit on the number of coefficients (and hence polynomial order). @@ -260,45 +295,41 @@ not need to have the same number of coefficients; each of the inner lists specif a separate polynomial. However, the independent variable (wavelength or pixel) does need to be the same for all polynomials for a given slit. -Polynomials specified via ``src_coeff`` and ``bkg_coeff`` are functions of either wavelength -(in microns) or pixel number (pixels in the dispersion direction, with respect to -the input 2D slit image), which is specified by the parameter ``independent_var``. -The default is "pixel". The values of these polynomial functions are pixel numbers in the -direction perpendicular to dispersion. .. _extract-1d-for-ifu: Extraction for 3D IFU Data -------------------------- In IFU cube data, 1D extraction is controlled by a different set of EXTRACT1D -reference file parameters. For point source data the extraction -aperture is centered at the RA/DEC target location indicated by the header. If the target location is undefined in the header, then the extraction +reference file parameters. For point source data, the extraction +aperture is centered at the RA/Dec target location indicated by the header. +If the target location is undefined in the header, then the extraction region is the center of the IFU cube. For extended source data, anything specified in the reference file or step arguments will be ignored; the entire image will be extracted, and no background subtraction will be done. -For point sources a circular extraction aperture is used, along with an optional +For point sources, a circular extraction aperture is used, along with an optional circular annulus for background extraction and subtraction. The size of the extraction region and the background annulus size varies with wavelength. The extraction related vectors are found in the asdf extract1d reference file. -For each element in the ``wavelength`` vector there are three size components: ``radius``, ``inner_bkg``, and -``outer_bkg``. The radius vector sets the extraction size; while ``inner_bkg`` and ``outer_bkg`` specify the -limits of an annular background aperture. There are two additional vectors in the reference file, ``axis_ratio`` -and ``axis_pa``, which are placeholders for possible future functionality. +For each element in the `wavelength` vector there are three size components: `radius`, `inner_bkg`, and +`outer_bkg`. The radius vector sets the extraction size; while `inner_bkg` and `outer_bkg` specify the +limits of an annular background aperture. There are two additional vectors in the reference file, `axis_ratio` +and `axis_pa`, which are placeholders for possible future functionality. The extraction size parameters are given in units of arcseconds and converted to units of pixels in the extraction process. The region of overlap between an aperture and a pixel can be calculated by -one of three different methods, specified by the ``method`` parameter: "exact" +one of three different methods, specified by the `method` parameter: "exact" (default), limited only by finite precision arithmetic; "center", the full value in a pixel will be included if its center is within the aperture; or "subsample", which means pixels will be subsampled N x N and the "center" option will be used -for each sub-pixel. When ``method`` is "subsample", the parameter ``subpixels`` +for each sub-pixel. When `method` is "subsample", the parameter `subpixels` is used to set the resampling value. The default value is 10. For IFU cubes the error information is contained entirely in the ERR array, and is not broken out into the VAR_POISSON, VAR_RNOISE, and VAR_FLAT arrays. As such, ``extract_1d`` only propagates this non-differentiated error term. Since covariance is also extremely important for undersampled IFU data -(see discussion by Law et al. 2023; AJ, 166, 45) the optional parameter ``ifu_covar_scale`` +(see discussion by Law et al. 2023; AJ, 166, 45) the optional parameter `ifu_covar_scale` will multiply all ERR arrays in the extracted spectra by a constant prefactor to account for this covariance. As discussed by Law et al. 2023, this prefactor provides a reasonable first-order correction for the vast majority of use cases. Values for the prefactor @@ -338,11 +369,11 @@ is part of the :ref:`calwebb_spec2 ` pipeline, but currently it i information see :ref:`residual_fringe `. The pipeline also can apply a 1-D residual fringe correction. This correction is only relevant for MIRI MRS data and -can be turned on by setting the optional parameter ``extract_1d.ifu_rfcorr = True`` in the ``extract_1d`` step. +can be turned on by setting the optional parameter `ifu_rfcorr = True` in the ``extract_1d`` step. Empirically, the 1-D correction step has been found to work better than the 2-D correction step if it is applied to per-band spectra. -When using the ``ifu_rfcorr`` option in the ``extract_1d`` step to apply a 1-D residual fringe +When using the `ifu_rfcorr` option in the ``extract_1d`` step to apply a 1-D residual fringe correction, it is applied during the extraction of spectra from the IFU cube. The 1D residual fringe code can also be called outside the pipeline to correct an extracted spectrum. If running outside the pipeline, the correction works best on single-band cubes, and the channel of @@ -351,4 +382,4 @@ the data must be given. The steps to run this correction outside the pipeline ar from jwst.residual_fringe.utils import fit_residual_fringes_1d as rf1d flux_cor = rf1d(flux, wave, channel=4) -where ``flux`` is the extracted spectral data, and the data are from channel 4 for this example. +where `flux` is the extracted spectral data, and the data are from channel 4 for this example. diff --git a/docs/jwst/extract_1d/extract1d_api.rst b/docs/jwst/extract_1d/extract1d_api.rst new file mode 100644 index 0000000000..7b3b8e836c --- /dev/null +++ b/docs/jwst/extract_1d/extract1d_api.rst @@ -0,0 +1,6 @@ +Python interfaces for 1D Extraction +=================================== + +.. automodapi:: jwst.extract_1d.extract +.. automodapi:: jwst.extract_1d.extract1d +.. automodapi:: jwst.extract_1d.ifu diff --git a/docs/jwst/extract_1d/index.rst b/docs/jwst/extract_1d/index.rst index 3de6be1d40..2fa08593e8 100644 --- a/docs/jwst/extract_1d/index.rst +++ b/docs/jwst/extract_1d/index.rst @@ -10,6 +10,6 @@ Extract 1D Spectra description.rst arguments.rst reference_files.rst - reference_image.rst + extract1d_api.rst .. automodapi:: jwst.extract_1d diff --git a/docs/jwst/extract_1d/reference_image.rst b/docs/jwst/extract_1d/reference_image.rst deleted file mode 100644 index 3c9e61a997..0000000000 --- a/docs/jwst/extract_1d/reference_image.rst +++ /dev/null @@ -1,57 +0,0 @@ -Reference Image Format -====================== -An alternative EXTRACT1D reference format, an image, is also supported. -There are currently no files of this type in CRDS (there would be a conflict -with the current JSON-format reference files), but a user can create a file -in this format and specify that it be used as an override for the default -EXTRACT1D reference file. - -This format is a `~jwst.datamodels.MultiExtract1dImageModel`, which is -loosely based on `~jwst.datamodels.MultiSlitModel`. The file should -contain keyword DATAMODL, with value 'MultiExtract1dImageModel'; this is -not required, but it makes it possible to open the file simply with -`datamodels.open`. The reference image file contains one or more images, -which are of type `~jwst.datamodels.Extract1dImageModel`, and one can -iterate over the list of these images to find one that matches the -observing configuration. This iterable is the ``images`` attribute of -the model (``ref_model``, for purposes of discussion). Each element of -``ref_model.images`` can contain a ``name`` attribute (FITS keyword -SLTNAME) and a ``spectral_order`` attribute (FITS keyword SPORDER), which -can be compared with the slit name and spectral order respectively in the -science data model in order to select the matching reference image. The -wildcard for SLTNAME is "ANY", and any integer value for SPORDER greater -than or equal to 1000 is a wildcard for spectral order (SPORDER is an -integer, and an integer keyword may not be assigned a string value such as -"ANY"). For IFU data, the image to use is selected only on ``name``. - -For non-IFU data, the shape of the reference image should match the shape -of the science data, although the step can either trim the reference image -or pad it with zeros to match the size of the science data, pinned at -pixel [0, 0]. For IFU data, the shape of the reference image can be 3-D, -exactly matching the shape of the IFU data, or it can be 2-D, matching -the shape of one plane of the IFU data. If the reference image is 2-D, -it will be applied equally to each plane of the IFU data, i.e. it will be -broadcast over the dispersion direction. - -The data type of each image is float32, but the data values may only -be +1, 0, or -1. A value of +1 means that the matching pixel in the -science data will be included when accumulating data for the source -(target) region. A value of 0 means the pixel will not be used for -anything. A value of -1 means the pixel will be included for the -background; if there are no pixels with value -1, no background will be -subtracted. A pixel will either be included or not; there is no option -to include only a fraction of a pixel. - -For non-IFU data, values will be extracted column by column (if the -dispersion direction is horizontal, else row by row). The gross count -rate will be the sum of the source pixels in a column (or row). If -background region(s) were specified, the sum of those pixels will be -scaled by the ratio of the number of source pixels to the number of -background pixels (with possibly a different ratio for each column (row)) -before being subtracted from the gross count rate. The scaled background -is what will be saved in the output table. - -For IFU data, the values will be summed over each plane in the dispersion -direction, giving one value of flux and optionally one value of background -per plane. The background value will be scaled by the ratio of source -pixels to background pixels before being subtracted from the flux. diff --git a/docs/jwst/references_general/extract1d_reffile.inc b/docs/jwst/references_general/extract1d_reffile.inc index e0f35a569a..30efe8ef59 100644 --- a/docs/jwst/references_general/extract1d_reffile.inc +++ b/docs/jwst/references_general/extract1d_reffile.inc @@ -48,8 +48,7 @@ extraction locations within the image, so different elements of ``apertures`` are needed in order to specify those locations. If key ``dispaxis`` is specified, its value will be used to set the dispersion direction within the image. If ``dispaxis`` is -not specified, the dispersion direction will be taken to be the axis -along which the wavelengths change more rapidly. +not specified, the dispersion direction will be taken from the metadata for the exposure. Key ``region_type`` can be omitted, but if it is specified, its value must be "target". The remaining keys specify the characteristics of the source and background extraction regions. @@ -58,21 +57,21 @@ and background extraction regions. * spectral_order: the spectral order number (optional); this can be either positive or negative, but it should not be zero (int) * dispaxis: dispersion direction, 1 for X, 2 for Y (int) -* xstart: first pixel in the horizontal direction, X (int) (0-indexed) -* xstop: last pixel in the horizontal direction, X (int) (0-indexed) -* ystart: first pixel in the vertical direction, Y (int) (0-indexed) -* ystop: last pixel in the vertical direction, Y (int) (0-indexed) -* src_coeff: this takes priority for specifying the source extraction region - (list of lists of float) -* bkg_coeff: for specifying background subtraction regions - (list of lists of float) +* xstart: first pixel in the horizontal direction, X (float) (0-indexed) +* xstop: last pixel in the horizontal direction, X (float) (0-indexed) +* ystart: first pixel in the vertical direction, Y (float) (0-indexed) +* ystop: last pixel in the vertical direction, Y (float) (0-indexed) +* src_coeff: polynomial coefficients for source extraction region + limits (list of lists of float); takes precedence over start/stop values. +* bkg_coeff: polynomial coefficients for background extraction region + limits (list of lists of float) * independent_var: "wavelength" or "pixel" (string) * smoothing_length: width of boxcar for smoothing background regions along the dispersion direction (odd int) * bkg_fit: the type of background fit or computation (string) -* bkg_order: order of polynomial fit to background regions (int) +* bkg_order: order of polynomial fit to background data in the cross-dispersion direction (int) * extract_width: number of pixels in cross-dispersion direction (int) -* use_source_posn: adjust the extraction limits based on source RA/Dec (bool) +* use_source_posn: allow the extraction limits to be adjusted based on source RA/Dec (bool) .. note:: @@ -92,7 +91,7 @@ and use this modified file in ``extract_1d`` by specifying this modified referen (:ref:`override in python ` or :ref:`override in strun `). The format for JSON files has to be exact, for example, the format of a floating-point value with a fractional portion must include at least one decimal digit, so "1." is invalid, while "1.0" is valid. The best practice after editing a JSON reference -file is to run a JSON validator off-line, such as `https://jsonlint.com/`, and correct any format errors before using +file is to run a JSON validator off-line, such as `https://jsonlint.com/`, and correct any format errors before using the JSON reference file in the pipeline. @@ -126,54 +125,35 @@ are used in the extraction process. Example EXTRACT1D Reference File -------------------------------- The following JSON was taken as an example from reference file -jwst_niriss_extract1d_0003.json:: - - { - "REFTYPE": "EXTRACT1D", - "INSTRUME": "NIRISS", - "TELESCOP": "JWST", - "DETECTOR": "NIS", - "EXP_TYPE": "NIS_SOSS", - "PEDIGREE": "GROUND", - "DESCRIP": "NIRISS SOSS extraction params for ground testing", - "AUTHOR": "M.Wolfe, H.Bushouse", - "HISTORY": "Build 7.1 of the JWST Calibration pipeline. The regions are rectangular and do not follow the trace.", - "USEAFTER": "2015-11-01T00:00:00", - "apertures": [ - { - "id": "FULL", - "region_type": "target", - "bkg_coeff": [[2014.5],[2043.5]], - "xstart": 4, - "xstop": 2044, - "ystart": 1792, - "ystop": 1972, - "dispaxis": 1, - "extract_width": 181 - }, - - { - "id": "SUBSTRIP256", - "region_type": "target", - "bkg_coeff": [[221.5],[251.5]], - "xstart": 4, - "xstop": 2044, - "ystart": 20, - "ystop": 220, - "dispaxis": 1, - "extract_width": 201 - }, - - { - "id": "SUBSTRIP96", - "region_type": "target", - "bkg_coeff": [[1.5],[8.5],[92.5],[94.5]], - "xstart": 4, - "xstop": 2044, - "ystart": 10, - "ystop": 92, - "dispaxis": 1, - "extract_width": 83 - }] - } - +jwst_miri_extract1d_0006.json:: + + { + "reftype": "EXTRACT1D", + "instrument": "MIRI", + "telescope": "JWST", + "exp_type": "MIR_LRS-FIXEDSLIT|MIR_LRS-SLITLESS", + "pedigree": "INFLIGHT 2022-06-01 2022-12-01", + "description": "MIRI LRS extraction params for in-flight data", + "author": "K.Murray", + "history": "2024-Jun-25", + "useafter": "2022-01-01T00:00:00", + "apertures": [ + { + "id": "MIR_LRS-FIXEDSLIT", + "region_type": "target", + "bkg_order": 0, + "dispaxis": 2, + "xstart": 27, + "xstop": 34 + }, + { + "id": "MIR_LRS-SLITLESS", + "region_type": "target", + "bkg_order": 0, + "dispaxis": 2, + "xstart": 30, + "xstop": 41, + "use_source_posn": false + } + ] + } diff --git a/jwst/datamodels/container.py b/jwst/datamodels/container.py index 196c31ca5c..5b6170f870 100644 --- a/jwst/datamodels/container.py +++ b/jwst/datamodels/container.py @@ -330,7 +330,7 @@ def save(self, path : str or None - If None, the `meta.filename` is used for each model. - If a string, the string is used as a root and an index is - appended. + appended, along with the '.fits' extension. save_model_func: func or None Alternate function to save each model instead of @@ -353,8 +353,11 @@ def save(self, save_path = model.meta.filename else: if len(self) <= 1: - idx = None - save_path = path+str(idx)+".fits" + idx = '' + if path.endswith(".fits"): + save_path = path.replace(".fits", f"{idx}.fits") + else: + save_path = f"{path}{idx}.fits" output_paths.append( model.save(save_path, **kwargs) ) diff --git a/jwst/datamodels/tests/test_model_container.py b/jwst/datamodels/tests/test_model_container.py index 2266f7909d..8640b8302e 100644 --- a/jwst/datamodels/tests/test_model_container.py +++ b/jwst/datamodels/tests/test_model_container.py @@ -127,7 +127,8 @@ def test_group_id(tmp_path): assert asn_group_ids == model_droup_ids -def test_save(tmp_cwd, container): +@pytest.mark.parametrize("path", ["foo", "foo.fits"]) +def test_save(tmp_cwd, container, path): # container pushes us to data/ directory so need to go back to tmp_cwd # to avoid polluting the data/ directory @@ -142,8 +143,13 @@ def test_save(tmp_cwd, container): assert os.path.exists(fname) # test specifying path saves to custom path with indices - path = "foo" container.save(path) - expected_fnames = [path+str(i)+".fits" for i in range(len(container))] + expected_fnames = [path.replace(".fits", "")+str(i)+".fits" for i in range(len(container))] for fname in expected_fnames: assert os.path.exists(fname) + + # test saving path when the container has length 1: no index appended + container = ModelContainer([container[0]]) + container.save(path) + expected_fname = path.replace(".fits", "") + ".fits" + assert os.path.exists(expected_fname) diff --git a/jwst/extract_1d/apply_apcorr.py b/jwst/extract_1d/apply_apcorr.py index 3618719f89..1324508d36 100644 --- a/jwst/extract_1d/apply_apcorr.py +++ b/jwst/extract_1d/apply_apcorr.py @@ -53,7 +53,7 @@ class ApCorrBase(abc.ABC): } def __init__(self, input_model, apcorr_table, sizeunit, - location = None, slit_name = None, **match_kwargs): + location=None, slit_name=None, **match_kwargs): self.correction = None self.model = input_model @@ -155,7 +155,7 @@ def apply(self, spec_table): """ flux_cols_to_correct = ('flux', 'flux_error', 'surf_bright', 'sb_error') var_cols_to_correct = ('flux_var_poisson', 'flux_var_rnoise', 'flux_var_flat', - 'sb_var_poisson', 'sb_var_rnoise', 'sb_var_flat') + 'sb_var_poisson', 'sb_var_rnoise', 'sb_var_flat') for row in spec_table: correction = self.apcorr_func(row['npixels'], row['wavelength']) @@ -180,7 +180,7 @@ class ApCorrPhase(ApCorrBase): """ size_key = 'size' - def __init__(self, *args, pixphase = 0.5, **kwargs): + def __init__(self, *args, pixphase=0.5, **kwargs): self.phase = pixphase # In the future we'll attempt to measure the pixel phase from inputs. super().__init__(*args, **kwargs) @@ -297,7 +297,7 @@ class ApCorrRadial(ApCorrBase): """Aperture correction class used with spectral data produced from an extraction aperture radius.""" def __init__(self, input_model, apcorr_table, - location = None): + location=None): self.correction = None self.model = input_model @@ -339,7 +339,7 @@ def apply(self, spec_table): """ flux_cols_to_correct = ('flux', 'flux_error', 'surf_bright', 'sb_error') var_cols_to_correct = ('flux_var_poisson', 'flux_var_rnoise', 'flux_var_flat', - 'sb_var_poisson', 'sb_var_rnoise', 'sb_var_flat') + 'sb_var_poisson', 'sb_var_rnoise', 'sb_var_flat') for i, row in enumerate(spec_table): correction = self.apcorr_correction[i] @@ -377,7 +377,7 @@ def match_wavelengths(self, wavelength_ifu): def find_apcorr_func(self, iwave, radius_ifu): # at ifu wavelength plane (iwave), the extraction radius is radius_ifu # pull out the radius values (self.size) to use in the apcor ref file for this iwave - # self.size and self.apcorr have already been interpolated in wavelength to match the + # self.size and self.apcorr have already been interpolated in wavelength to match # the ifu wavelength range. radius_apcor = self.size[:, iwave] diff --git a/jwst/extract_1d/extract.py b/jwst/extract_1d/extract.py index fac9dc799f..13d24a6811 100644 --- a/jwst/extract_1d/extract.py +++ b/jwst/extract_1d/extract.py @@ -1,51 +1,34 @@ -import abc import logging -import copy import json -import math import numpy as np +from json.decoder import JSONDecodeError -from dataclasses import dataclass from astropy.modeling import polynomial -from astropy.io import fits from stdatamodels.jwst import datamodels -from stdatamodels.jwst.datamodels import dqflags from stdatamodels.jwst.datamodels.apcorr import ( MirLrsApcorrModel, MirMrsApcorrModel, NrcWfssApcorrModel, NrsFsApcorrModel, NrsMosApcorrModel, NrsIfuApcorrModel, NisWfssApcorrModel ) -from jwst.datamodels import SourceModelContainer - +from jwst.assign_wcs.util import wcs_bbox_from_shape +from jwst.datamodels import ModelContainer +from jwst.lib import pipe_utils +from jwst.lib.wcs_utils import get_wavelengths +from jwst.extract_1d import extract1d, spec_wcs +from jwst.extract_1d.apply_apcorr import select_apcorr -from ..assign_wcs import niriss # for specifying spectral order number -from ..assign_wcs.util import wcs_bbox_from_shape -from ..lib import pipe_utils -from ..lib.wcs_utils import get_wavelengths -from . import extract1d -from . import ifu -from . import spec_wcs -from .apply_apcorr import select_apcorr +__all__ = ['run_extract1d', 'read_extract1d_ref', 'read_apcorr_ref', + 'get_extract_parameters', 'box_profile', 'aperture_center', + 'location_from_wcs', 'shift_by_source_location', 'define_aperture', + 'extract_one_slit', 'create_extraction'] -from json.decoder import JSONDecodeError log = logging.getLogger(__name__) log.setLevel(logging.DEBUG) -# WFSS_EXPTYPES = ['NIS_WFSS', 'NRC_WFSS', 'NRC_GRISM', 'NRC_TSGRISM'] WFSS_EXPTYPES = ['NIS_WFSS', 'NRC_WFSS', 'NRC_GRISM'] """Exposure types to be regarded as wide-field slitless spectroscopy.""" -# These values are used to indicate whether the input extract1d reference file -# (if any) is JSON, IMAGE or ASDF (added for IFU data) -FILE_TYPE_JSON = "JSON" -FILE_TYPE_IMAGE = "IMAGE" -FILE_TYPE_ASDF = "ASDF" -FILE_TYPE_OTHER = "N/A" - -# This is to prevent calling offset_from_offset multiple times for multi-integration data. -OFFSET_NOT_ASSIGNED_YET = "not assigned yet" - ANY = "ANY" """Wildcard for slit name. @@ -59,137 +42,83 @@ name from the input data. """ -ANY_ORDER = 1000 -"""Wildcard for spectral order number in a reference image. - -Extended summary ----------------- -If the extract1d reference file contains images, keyword SPORDER gives the order -number of the spectrum that would be extracted using a given image in -the reference file. -""" - HORIZONTAL = 1 +"""Horizontal dispersion axis.""" VERTICAL = 2 -"""Dispersion direction, predominantly horizontal or vertical.""" - -# This is intended to be larger than any possible distance (in pixels) between the target and any point in the image; -# used by locn_from_wcs(). -HUGE_DIST = 1.e20 +"""Vertical dispersion axis.""" # These values are assigned in get_extract_parameters, using key "match". -# If there was an aperture in the reference file for which the "id" key matched, that's (at least) a partial match. +# If there was an aperture in the reference file for which the "id" key matched, +# that's (at least) a partial match. # If "spectral_order" also matched, that's an exact match. NO_MATCH = "no match" PARTIAL = "partial match" EXACT = "exact match" -@dataclass -class Aperture: - xstart: float - xstop: float - ystart: float - ystop: float - - -class Extract1dError(Exception): - pass - - -class InvalidSpectralOrderNumberError(Extract1dError): - """The spectral order number was invalid or off the detector.""" - pass - - -# Create custom error to pass continue from a function inside of a loop class ContinueError(Exception): + """Custom error to pass continue from a function inside a loop.""" pass -def open_extract1d_ref(refname, exptype): - """Open the extract1d reference file. +def read_extract1d_ref(refname): + """Read the extract1d reference file. Parameters ---------- refname : str - The name of the extract1d reference file. This file is expected to be - a JSON file or ASDF file giving extraction information, or a file - containing one or more images that are to be used as masks that - define the extraction region and optionally background regions. + The name of the extract1d reference file, or 'N/A'. + If specified, this file is expected to be a JSON file + giving extraction information. Returns ------- - ref_dict : dict - If the extract1d reference file is in JSON format, ref_dict will be the - dictionary returned by json.load(), except that the file type - ('JSON') will also be included with key 'ref_file_type'. - If the extract1d reference file is in asdf format, the ref_dict will - be a dictionary containing two keys: ref_dict['ref_file_type'] = 'ASDF' - and ref_dict['ref_model']. - If the reference file is an image, ref_dict will be a - dictionary with two keys: ref_dict['ref_file_type'] = 'IMAGE' - and ref_dict['ref_model']. The latter will be the open file - handle for the jwst.datamodels object for the extract1d file. + ref_dict : dict or None + If the extract1d reference file is specified, ref_dict will be the + dictionary returned by json.load(). """ - # the extract1d reference file can be 1 of three types: 'json', 'fits', or 'asdf' refname_type = refname[-4:].lower() if refname == "N/A": ref_dict = None else: + # If specified, the extract1d reference file can only be in json format if refname_type == 'json': fd = open(refname) try: ref_dict = json.load(fd) - ref_dict['ref_file_type'] = FILE_TYPE_JSON fd.close() except (UnicodeDecodeError, JSONDecodeError): # Input file does not load correctly as json file. # Probably an error in json file fd.close() - log.error("Extract1d json reference file has an error, run a json validator off line and fix the file") - raise RuntimeError("Invalid json extract 1d reference file, run json validator off line and fix file.") - elif refname_type == 'fits': - try: - fd = fits.open(refname) - extract_model = datamodels.MultiExtract1dImageModel(refname) - ref_dict = {'ref_file_type': FILE_TYPE_IMAGE, 'ref_model': extract_model} - fd.close() - except OSError: - log.error("Extract1d fits reference file has an error") - raise RuntimeError("Invalid fits extract 1d reference file- fix reference file.") - - elif refname_type == 'asdf': - extract_model = datamodels.Extract1dIFUModel(refname) - ref_dict = dict() - ref_dict['ref_file_type'] = FILE_TYPE_ASDF - ref_dict['ref_model'] = extract_model + log.error("Extract1d JSON reference file has an error. Run a json validator off line and fix the file.") + raise RuntimeError("Invalid JSON extract1d reference file.") else: - log.error("Invalid Extract 1d reference file, must be json, fits or asdf.") - raise RuntimeError("Invalid Extract 1d reference file, must be json, fits or asdf.") + log.error("Invalid Extract1d reference file: must be JSON.") + raise RuntimeError("Invalid extract1d reference file: must be JSON.") return ref_dict -def open_apcorr_ref(refname, exptype): - """Determine the appropriate DataModel class to use when opening the input APCORR reference file. +def read_apcorr_ref(refname, exptype): + """Read the apcorr reference file. + + Determine the appropriate DataModel class to use for an APCORR reference file + and read the file into it. Parameters ---------- refname : str - Path of the APCORR reference file + Path to the APCORR reference file. exptype : str - EXPTYPE of the input to the extract_1d step. + EXP_TYPE of the input to the extract_1d step. Returns ------- - Opened APCORR DataModel. - - Notes - ----- - This function should be removed after the DATAMODL keyword is required for the APCORR reference file. + DataModel + A datamodel containing the reference file input. """ apcorr_model_map = { @@ -209,45 +138,35 @@ def open_apcorr_ref(refname, exptype): return apcorr_model(refname) -def get_extract_parameters( - ref_dict, - input_model, - slitname, - sp_order, - meta, - smoothing_length, - bkg_fit, - bkg_order, - use_source_posn -): - """Get extract1d reference file values. +def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta, + smoothing_length=None, bkg_fit=None, bkg_order=None, + use_source_posn=None, subtract_background=None): + """Get extraction parameter values. Parameters ---------- ref_dict : dict or None - For an extract1d reference file in JSON format, `ref_dict` will be the entire - contents of the file. For an EXTRACT1D reference image, `ref_dict` will have - just two entries, 'ref_file_type' (a string) and 'ref_model', a - JWST data model for a collection of images. If there is no - extract1d reference file, `ref_dict` will be None. + For an extract1d reference file in JSON format, `ref_dict` will be + the entire contents of the file. If there is no extract1d reference + file, `ref_dict` will be None. - input_model : data model + input_model : JWSTDataModel This can be either the input science file or one SlitModel out of a list of slits. slitname : str - The name of the slit, or "ANY", or for NIRISS SOSS data this will - be the subarray name. + The name of the slit, or "ANY". sp_order : int The spectral order number. - meta : metadata for the actual input model, i.e. not just for the - current slit. + meta : ObjectNode + The metadata for the actual input model, i.e. not just for the + current slit, from input_model.meta. - smoothing_length : int or None + smoothing_length : int or None, optional Width of a boxcar function for smoothing the background regions. - If None, the smoothing length will be gotten from `ref_dict`, or + If None, the smoothing length will be retrieved from `ref_dict`, or it will be set to 0 (no background smoothing) if this key is not found in `ref_dict`. If `smoothing_length` is not None, that means that the user @@ -255,14 +174,14 @@ def get_extract_parameters( This argument is only used if background regions have been specified. - bkg_fit : str + bkg_fit : str or None, optional The type of fit to apply to background values in each column (or row, if the dispersion is vertical). The default `poly` results in a polynomial fit of order `bkg_order`. Other options are `mean` and `median`. If `mean` or `median` is selected, `bkg_order` is ignored. - bkg_order : int or None + bkg_order : int or None, optional Polynomial order for fitting to each column (or row, if the dispersion is vertical) of background. If None, the polynomial order will be gotten from `ref_dict`, or it will be set to 0 if @@ -274,26 +193,27 @@ def get_extract_parameters( This argument must be positive or zero, and it is only used if background regions have been specified. - use_source_posn : bool or None + use_source_posn : bool or None, optional If True, the target and background positions specified in `ref_dict` (or a default target position) will be shifted to account for the actual source location in the data. If None, the value specified in `ref_dict` will be used, or it will be set to True if not found in `ref_dict`. + subtract_background : bool or None, optional + If False, all background parameters will be ignored. + Returns ------- extract_params : dict - Information copied out of `ref_dict`. The items will be selected + Information copied from `ref_dict`. The items will be selected based on `slitname` and `sp_order`. Default values will be - assigned if `ref_dict` is None. For a reference image, the key - 'ref_image' gives the (open) image model. + assigned if `ref_dict` is None. """ extract_params = {'match': NO_MATCH} # initial value if ref_dict is None: # There is no extract1d reference file; use "reasonable" default values. - extract_params['ref_file_type'] = FILE_TYPE_OTHER extract_params['match'] = EXACT shape = input_model.data.shape extract_params['spectral_order'] = sp_order @@ -308,59 +228,60 @@ def get_extract_parameters( extract_params['bkg_fit'] = None # because no background sub. extract_params['bkg_order'] = 0 # because no background sub. extract_params['subtract_background'] = False - - if use_source_posn is None: - extract_params['use_source_posn'] = False - else: - extract_params['use_source_posn'] = use_source_posn - + extract_params['extraction_type'] = 'box' + extract_params['use_source_posn'] = False # no source position correction extract_params['position_correction'] = 0 extract_params['independent_var'] = 'pixel' - # Note that extract_params['dispaxis'] is not assigned. This will be done later, possibly slit by slit. - - elif ref_dict['ref_file_type'] == FILE_TYPE_JSON: - extract_params['ref_file_type'] = ref_dict['ref_file_type'] + # Note that extract_params['dispaxis'] is not assigned. + # This will be done later, possibly slit by slit. + else: for aper in ref_dict['apertures']: - if ('id' in aper and aper['id'] != "dummy" and - (aper['id'] == slitname or aper['id'] == "ANY" or - slitname == "ANY")): - extract_params['match'] = PARTIAL + if ('id' in aper and aper['id'] != "dummy" + and (aper['id'] == slitname + or aper['id'] == ANY + or slitname == ANY)): + # region_type is retained for backward compatibility; it is - # not required to be present. + # not required to be present, but if it is and is not set + # to target, then the aperture is not matched. + region_type = aper.get("region_type", "target") + if region_type != "target": + continue + + extract_params['match'] = PARTIAL + # spectral_order is a secondary selection criterion. The # default is the expected value, so if the key is not present # in the JSON file, the current aperture will be selected. # If the current aperture in the JSON file has # "spectral_order": "ANY", that aperture will be selected. - region_type = aper.get("region_type", "target") - - if region_type != "target": - continue - spectral_order = aper.get("spectral_order", sp_order) if spectral_order == sp_order or spectral_order == ANY: extract_params['match'] = EXACT extract_params['spectral_order'] = sp_order - # Note: extract_params['dispaxis'] is not assigned. This is done later, possibly slit by slit. - if meta.target.source_type == "EXTENDED": - shape = input_model.data.shape - extract_params['xstart'] = aper.get('xstart', 0) - extract_params['xstop'] = aper.get('xstop', shape[-1] - 1) - extract_params['ystart'] = aper.get('ystart', 0) - extract_params['ystop'] = aper.get('ystop', shape[-2] - 1) - else: - extract_params['xstart'] = aper.get('xstart') - extract_params['xstop'] = aper.get('xstop') - extract_params['ystart'] = aper.get('ystart') - extract_params['ystop'] = aper.get('ystop') + # Note: extract_params['dispaxis'] is not assigned. + # This is done later, possibly slit by slit. + + # Set default start/stop by shape if not specified + shape = input_model.data.shape + extract_params['xstart'] = aper.get('xstart', 0) + extract_params['xstop'] = aper.get('xstop', shape[-1] - 1) + extract_params['ystart'] = aper.get('ystart', 0) + extract_params['ystop'] = aper.get('ystop', shape[-2] - 1) extract_params['src_coeff'] = aper.get('src_coeff') extract_params['bkg_coeff'] = aper.get('bkg_coeff') - if extract_params['bkg_coeff'] is not None: + if (extract_params['bkg_coeff'] is not None + and subtract_background is not False): extract_params['subtract_background'] = True if bkg_fit is not None: + # Mean value for background fitting is equivalent + # to a polynomial fit with order 0. + if bkg_fit == 'mean': + bkg_fit = 'poly' + bkg_order = 0 extract_params['bkg_fit'] = bkg_fit else: extract_params['bkg_fit'] = aper.get('bkg_fit', 'poly') @@ -380,13 +301,14 @@ def get_extract_parameters( # parameter value on the command line is highest precedence, # then parameter value from the extract1d reference file, # and finally a default setting based on exposure type. - use_source_posn_aper = aper.get('use_source_posn', None) # value from the extract1d ref file + use_source_posn_aper = aper.get('use_source_posn', None) # value from the ref file if use_source_posn is None: # no value set on command line if use_source_posn_aper is None: # no value set in ref file # Use a suitable default - if meta.exposure.type in ['MIR_LRS-FIXEDSLIT', 'MIR_MRS', 'NRS_FIXEDSLIT', 'NRS_IFU', 'NRS_MSASPEC']: + if meta.exposure.type in ['MIR_LRS-FIXEDSLIT', 'NRS_FIXEDSLIT', 'NRS_MSASPEC']: use_source_posn = True - log.info(f"Turning on source position correction for exp_type = {meta.exposure.type}") + log.info(f"Turning on source position correction " + f"for exp_type = {meta.exposure.type}") else: use_source_posn = False else: # use the value from the ref file @@ -402,63 +324,45 @@ def get_extract_parameters( # If the user supplied a value, use that value. extract_params['smoothing_length'] = smoothing_length - break - - elif ref_dict['ref_file_type'] == FILE_TYPE_IMAGE: - # Note that we will use the supplied image-format extract1d reference file, - # without regard for the distinction between point source and - # extended source. - extract_params['ref_file_type'] = ref_dict['ref_file_type'] - im = None - - for im in ref_dict['ref_model'].images: - if im.name == slitname or im.name == ANY or slitname == ANY: - extract_params['match'] = PARTIAL + # Check that the smoothing length has a reasonable value + sm_length = extract_params['smoothing_length'] + if sm_length != int(sm_length): + sm_length = int(np.round(sm_length)) + log.warning(f'Smoothing length must be an integer. ' + f'Rounding to {sm_length}') + if sm_length != 0: + if sm_length < 3: + log.warning(f'Smoothing length {sm_length} ' + f'is not allowed. Setting it to 0.') + sm_length = 0 + elif sm_length % 2 == 0: + sm_length -= 1 + log.warning('Even smoothing lengths are not supported. ' + 'Rounding the smoothing length down to ' + f'{sm_length}.') + extract_params['smoothing_length'] = sm_length + + # Default the extraction type to 'box': 'optimal' + # is not yet supported. + extract_params['extraction_type'] = 'box' - if im.spectral_order == sp_order or im.spectral_order >= ANY_ORDER: - extract_params['match'] = EXACT - extract_params['spectral_order'] = sp_order break - if extract_params['match'] == EXACT: - extract_params['ref_image'] = im - # Note that extract_params['dispaxis'] is not assigned. This will be done later, possibly slit by slit. - if smoothing_length is None: - extract_params['smoothing_length'] = im.smoothing_length - else: - # The user-supplied value takes precedence. - extract_params['smoothing_length'] = smoothing_length - - if use_source_posn is None: - extract_params['use_source_posn'] = False - else: - extract_params['use_source_posn'] = use_source_posn - - extract_params['position_correction'] = 0 - - if -1 in extract_params['ref_image'].data: - extract_params['subtract_background'] = True - else: - extract_params['subtract_background'] = False - - else: - log.error("Reference file type {ref_dict['ref_file_type']} not recognized") - return extract_params def log_initial_parameters(extract_params): - """Log some of the initial extraction parameters. + """Log the initial extraction parameters. Parameters ---------- extract_params : dict Information read from the reference file. """ - if "xstart" not in extract_params: + if "dispaxis" not in extract_params: return - log.debug("Initial parameters:") + log.debug("Extraction parameters:") log.debug(f"dispaxis = {extract_params['dispaxis']}") log.debug(f"spectral order = {extract_params['spectral_order']}") log.debug(f"initial xstart = {extract_params['xstart']}") @@ -473,3159 +377,1146 @@ def log_initial_parameters(extract_params): log.debug(f"smoothing_length = {extract_params['smoothing_length']}") log.debug(f"independent_var = {extract_params['independent_var']}") log.debug(f"use_source_posn = {extract_params['use_source_posn']}") + log.debug(f"extraction_type = {extract_params['extraction_type']}") -def get_aperture(im_shape, wcs, extract_params): - """Get the extraction limits xstart, xstop, ystart, ystop. +def create_poly(coeff): + """Create a polynomial model from coefficients. Parameters ---------- - im_shape : tuple - The shape (2-D) of the input data. This will be for the current - integration, if the input contains more than one integration. - - wcs : a WCS object, or None - The wcs (if any) for the input data or slit. - - extract_params : dict - Parameters read from the reference file. + coeff : list of float + The coefficients of the polynomial, constant term first, highest + order term last. Returns ------- - ap_ref : Aperture NamedTuple or an empty dict - Keys are 'xstart', 'xstop', 'ystart', and 'ystop'. + `astropy.modeling.polynomial.Polynomial1D` or None + None is returned if `coeff` is empty. """ - if extract_params['ref_file_type'] == FILE_TYPE_IMAGE: - return {} - - ap_ref = aperture_from_ref(extract_params, im_shape) - ap_ref, truncated = update_from_shape(ap_ref, im_shape) - - if truncated: - log.debug("Extraction limits extended outside image borders; limits have been truncated.") - - if wcs is not None: - ap_wcs = aperture_from_wcs(wcs) - else: - ap_wcs = None + n = len(coeff) + if n < 1: + return None - # If the xstart, etc., values were not specified for the dispersion direction, the extraction region should be - # centered within the WCS bounding box (domain). - ap_ref = update_from_wcs(ap_ref, ap_wcs, extract_params["extract_width"], extract_params["dispaxis"]) - ap_ref = update_from_width(ap_ref, extract_params["extract_width"], extract_params["dispaxis"]) + coeff_dict = {f'c{i}': coeff[i] for i in range(n)} - return ap_ref + return polynomial.Polynomial1D(degree=n - 1, **coeff_dict) -def aperture_from_ref(extract_params, im_shape): - """Get extraction region from reference file or image shape. +def populate_time_keywords(input_model, output_model): + """Copy the integration times keywords to header keywords. Parameters ---------- - extract_params : dict - Parameters read from the reference file. + input_model : JWSTDataModel + The input science model. - im_shape : tuple of int - The last two elements are the height and width of the input image - (slit). + output_model : JWSTDataModel + The output science model. This may be modified in-place. - Returns - ------- - ap_ref : namedtuple - Keys are 'xstart', 'xstop', 'ystart', and 'ystop'. """ - nx = im_shape[-1] - ny = im_shape[-2] - _xstart = 0 - _xstop = nx - 1 - _ystart = 0 - _ystop = ny - 1 + nints = input_model.meta.exposure.nints + int_start = input_model.meta.exposure.integration_start - xstart = extract_params.get('xstart', 0) - xstop = extract_params.get('xstop', nx - 1) # limits are inclusive - ystart = extract_params.get('ystart', 0) - ystop = extract_params.get('ystop', ny - 1) + if hasattr(input_model, 'data'): + shape = input_model.data.shape - ap_ref = Aperture( - xstart=xstart if xstart is not None else _xstart, - xstop=xstop if xstop is not None else _xstop, - ystart=ystart if ystart is not None else _ystart, - ystop=ystop if ystop is not None else _ystop - ) + if len(shape) == 2: + num_integ = 1 + else: # len(shape) == 3 + num_integ = shape[0] + else: # e.g. MultiSlit data + num_integ = 1 - return ap_ref + # This assumes that the spec attribute of output_model has already been created, + # and spectra have been appended. + n_output_spec = len(output_model.spec) + # num_j is the number of spectra per integration, e.g. the number + # of fixed-slit spectra, MSA spectra, or different + # spectral orders; num_integ is the number of integrations. + # The total number of output spectra is n_output_spec = num_integ * num_j + num_j = n_output_spec // num_integ -def update_from_width(ap_ref, extract_width, direction): - """Update XD extraction limits based on extract_width. + if n_output_spec != num_j * num_integ: # sanity check + log.warning( + f"populate_time_keywords: Don't understand n_output_spec = {n_output_spec}, num_j = {num_j}, num_integ = " + f"{num_integ}" + ) + else: + log.debug( + f"Number of output spectra = {n_output_spec}; number of spectra for each integration = {num_j}; " + f"number of integrations = {num_integ}" + ) - If extract_width was specified, that value should override - ystop - ystart (or xstop - xstart, depending on dispersion direction). + if int_start is None: + log.warning("INTSTART not found; assuming a value of 1.") + int_start = 1 - Parameters - ---------- - ap_ref : Aperture NamedTuple - Contains xstart, xstop, ystart, ystop. These are the initial - values as read from the extract1d reference file, except that they may - have been truncated at the image borders. + int_start -= 1 # zero indexed + int_end = input_model.meta.exposure.integration_end - extract_width : int or None - The number of pixels in the cross-dispersion direction to add - together to make a 1-D spectrum from a 2-D image. + if int_end is None: + log.warning(f"INTEND not found; assuming a value of {nints}.") + int_end = nints - direction : int - HORIZONTAL (1) if the dispersion direction is predominantly - horizontal. VERTICAL (2) if the dispersion direction is - predominantly vertical. + int_end -= 1 # zero indexed - Returns - ------- - ap_width : namedtuple - Keys are 'xstart', 'xstop', 'ystart', and 'ystop'. - """ - if extract_width is None: - return ap_ref + if nints > 1: + num_integrations = int_end - int_start + 1 + else: + num_integrations = 1 - if direction == HORIZONTAL: - temp_width = ap_ref.ystop - ap_ref.ystart + 1 + if hasattr(input_model, 'int_times') and input_model.int_times is not None: + nrows = len(input_model.int_times) else: - temp_width = ap_ref.xstop - ap_ref.xstart + 1 + nrows = 0 - if extract_width == temp_width: - return ap_ref # OK as is + if nrows < 1: + log.warning("There is no INT_TIMES table in the input file - " + "Making best guess on integration numbers.") + for j in range(num_j): # for each spectrum or order + for k in range(num_integ): # for each integration + output_model.spec[(j * num_integ) + k].int_num = k + 1 # set int_num to (k+1) - 1-indexed integration + return - # An integral value corresponds to the center of a pixel. - # If the extraction limits were not specified via polynomial coefficients, assign_polynomial_limits will create - # polynomial functions using values from an Aperture, and these lower and upper limits will be expanded by 0.5 to - # give polynomials (constant functions) for the lower and upper edges of the bounding pixels. + # If we have a single plane (e.g. ImageModel or MultiSlitModel), + # we will only populate the keywords if the corresponding uncal file + # had one integration. + # If the data were or might have been segmented, we use the first and + # last integration numbers to determine whether the data were in fact + # averaged over integrations, and if so, we should not populate the + # int_times-related header keywords. + skip = False # initial value - width = float(extract_width) + if isinstance(input_model, (datamodels.MultiSlitModel, datamodels.ImageModel)): + if num_integrations > 1: + log.warning("Not using INT_TIMES table because the data have been averaged over integrations.") + skip = True + elif isinstance(input_model, (datamodels.CubeModel, datamodels.SlitModel)): + shape = input_model.data.shape - if direction == HORIZONTAL: - lower = float(ap_ref.ystart) - upper = float(ap_ref.ystop) - lower = (lower + upper) / 2. - (width - 1.) / 2. - upper = lower + (width - 1.) - ap_width = Aperture(xstart=ap_ref.xstart, xstop=ap_ref.xstop, ystart=lower, ystop=upper) - else: - lower = float(ap_ref.xstart) - upper = float(ap_ref.xstop) - lower = (lower + upper) / 2. - (width - 1.) / 2. - upper = lower + (width - 1.) - ap_width = Aperture(xstart=lower, xstop=upper, ystart=ap_ref.ystart, ystop=ap_ref.ystop) + if len(shape) == 2 and num_integrations > 1: + log.warning("Not using INT_TIMES table because the data have been averaged over integrations.") + skip = True + elif len(shape) != 3 or shape[0] > nrows: + # Later, we'll check that the integration_number column actually + # has a row corresponding to every integration in the input. + log.warning( + "Not using INT_TIMES table because the data shape is not consistent with the number of table rows." + ) + skip = True + elif isinstance(input_model, datamodels.IFUCubeModel): + log.warning("The INT_TIMES table will be ignored for IFU data.") + skip = True - return ap_width + if skip: + return + int_num = input_model.int_times['integration_number'] + start_time_mjd = input_model.int_times['int_start_MJD_UTC'] + mid_time_mjd = input_model.int_times['int_mid_MJD_UTC'] + end_time_mjd = input_model.int_times['int_end_MJD_UTC'] + start_tdb = input_model.int_times['int_start_BJD_TDB'] + mid_tdb = input_model.int_times['int_mid_BJD_TDB'] + end_tdb = input_model.int_times['int_end_BJD_TDB'] -def update_from_shape(ap, im_shape): - """Truncate extraction region based on input image shape. + data_range = (int_start, int_end) # Inclusive range of integration numbers in the input data, zero indexed. - Parameters - ---------- - ap : namedtuple - Extraction region. + # Inclusive range of integration numbers in the INT_TIMES table, zero indexed. + table_range = (int_num[0] - 1, int_num[-1] - 1) + offset = data_range[0] - table_range[0] + if ((data_range[0] < table_range[0] or data_range[1] > table_range[1]) + or offset > table_range[1]): + log.warning("Not using the INT_TIMES table because it does not include " + "rows for all integrations in the data.") + return - im_shape : tuple of int - The last two elements are the height and width of the input image. + log.debug("TSO data, so copying times from the INT_TIMES table.") - Returns - ------- - ap_shape : namedtuple - Keys are 'xstart', 'xstop', 'ystart', and 'ystop'. + n = 0 # Counter for spectra in output_model. - truncated : bool - True if any value was truncated at an image edge. - """ - nx = im_shape[-1] - ny = im_shape[-2] + for k in range(num_integ): # for each spectrum or order + for j in range(num_j): # for each integration + row = k + offset + spec = output_model.spec[n] # n is incremented below + spec.int_num = int_num[row] + spec.time_scale = "UTC" + spec.start_time_mjd = start_time_mjd[row] + spec.mid_time_mjd = mid_time_mjd[row] + spec.end_time_mjd = end_time_mjd[row] + spec.start_tdb = start_tdb[row] + spec.mid_tdb = mid_tdb[row] + spec.end_tdb = end_tdb[row] + n += 1 - xstart = ap.xstart - xstop = ap.xstop - ystart = ap.ystart - ystop = ap.ystop - truncated = False +def get_spectral_order(slit): + """Get the spectral order number. - if ap.xstart < 0: - xstart = 0 - truncated = True + Parameters + ---------- + slit : SlitModel + One slit from an input MultiSlitModel or similar. - if ap.xstop >= nx: - xstop = nx - 1 # limits are inclusive - truncated = True + Returns + ------- + int + Spectral order number for `slit`. If no information about spectral + order is available in `wcsinfo`, a default value of 1 will be + returned. - if ap.ystart < 0: - ystart = 0 - truncated = True + """ + sp_order = slit.meta.wcsinfo.spectral_order - if ap.ystop >= ny: - ystop = ny - 1 - truncated = True + if sp_order is None: + log.warning("spectral_order is None; using 1") + sp_order = 1 - ap_shape = Aperture(xstart=xstart, xstop=xstop, ystart=ystart, ystop=ystop) + return sp_order - return ap_shape, truncated +def is_prism(input_model): + """Determine whether the current observing mode used a prism. -def aperture_from_wcs(wcs): - """Get the limits over which the WCS is defined. + Extended summary + ---------------- + The reason for this test is so we can skip spectral extraction if the + spectral order is zero and the exposure was not made using a prism. + In this context, therefore, a grism is not considered to be a prism. Parameters ---------- - wcs : WCS - The world coordinate system interface. + input_model : JWSTDataModel + The input science model. Returns ------- - ap_wcs : Aperture or None - Keys are 'xstart', 'xstop', 'ystart', and 'ystop'. These are the - limits copied directly from wcs.bounding_box. + bool + True if the exposure used a prism; False otherwise. + """ - got_bounding_box = False + instrument = input_model.meta.instrument.name + + if instrument is None: + return False - try: - bounding_box = wcs.bounding_box - got_bounding_box = True - except AttributeError: - log.debug("wcs.bounding_box not found; using wcs.domain instead.") + instrument_filter = input_model.meta.instrument.filter - bounding_box = ( - (wcs.domain[0]['lower'], wcs.domain[0]['upper']), - (wcs.domain[1]['lower'], wcs.domain[1]['upper']) - ) + if instrument_filter is None: + instrument_filter = "NONE" + else: + instrument_filter = instrument_filter.upper() - if got_bounding_box and bounding_box is None: - log.warning("wcs.bounding_box is None") - return None + grating = input_model.meta.instrument.grating - # bounding_box should be a tuple of tuples, each of the latter consisting of (lower, upper) limits. - if len(bounding_box) < 2: - log.warning("wcs.bounding_box has the wrong shape") - return None + if grating is None: + grating = "NONE" + else: + grating = grating.upper() - # These limits are float, and they are inclusive. - xstart = bounding_box[0][0] - xstop = bounding_box[0][1] - ystart = bounding_box[1][0] - ystop = bounding_box[1][1] + prism_mode = False - ap_wcs = Aperture(xstart=xstart, xstop=xstop, ystart=ystart, ystop=ystop) + if ((instrument == "MIRI" and instrument_filter.find("P750L") >= 0) or + (instrument == "NIRSPEC" and grating.find("PRISM") >= 0)): + prism_mode = True - return ap_wcs + return prism_mode -def update_from_wcs(ap_ref, ap_wcs, extract_width, direction,): - """Limit the extraction region to the WCS bounding box. +def copy_keyword_info(slit, slitname, spec): + """Copy metadata from the input to the output spectrum. Parameters ---------- - ap_ref : namedtuple - Contains xstart, xstop, ystart, ystop. These are the values of - the extraction region as specified by the reference file or the - image size. - - ap_wcs : namedtuple or None - These are the bounding box limits. + slit : A SlitModel object + Metadata will be copied from the input `slit` to output `spec`. - extract_width : int - The number of pixels in the cross-dispersion direction to add - together to make a 1-D spectrum from a 2-D image. + slitname : str or None + The name of the slit. - direction : int - HORIZONTAL (1) if the dispersion direction is predominantly - horizontal. VERTICAL (2) if the dispersion direction is - predominantly vertical. + spec : One element of MultiSpecModel.spec + Metadata attributes will be updated in-place. - Returns - ------- - ap : namedtuple - Keys are 'xstart', 'xstop', 'ystart', and 'ystop'. """ - if ap_wcs is None: - return ap_ref - - # If the wcs limits don't pass the sanity test, ignore the bounding box. - if not sanity_check_limits(ap_ref, ap_wcs): - log.debug("Sanity check on WCS limits failed - using ap_ref") - return ap_ref - - # ap_wcs has the limits over which the WCS transformation is defined; take those as the outer limits over which we - # will extract. - xstart = compare_start(ap_ref.xstart, ap_wcs.xstart) - ystart = compare_start(ap_ref.ystart, ap_wcs.ystart) - xstop = compare_stop(ap_ref.xstop, ap_wcs.xstop) - ystop = compare_stop(ap_ref.ystop, ap_wcs.ystop) - - if extract_width is not None: - if direction == HORIZONTAL: - width = ystop - ystart + 1 - else: - width = xstop - xstart + 1 - - if width < extract_width: - log.debug(f"extract_width was truncated from {extract_width} to {width}") + if slitname is not None and slitname != "ANY": + spec.name = slitname - ap = Aperture(xstart=xstart, xstop=xstop, ystart=ystart, ystop=ystop) + # Copy over some attributes if present, even if they are None + copy_attributes = ["slitlet_id", "source_id", "source_xpos", "source_ypos", + "source_ra", "source_dec", "shutter_state"] + for key in copy_attributes: + if hasattr(slit, key): + setattr(spec, key, getattr(slit, key)) - return ap + # copy over some attributes only if they are present and not None + copy_populated_attributes = ["source_name", "source_alias", + "source_type", "stellarity"] + for key in copy_populated_attributes: + if getattr(slit, key, None) is not None: + setattr(spec, key, getattr(slit, key)) -def sanity_check_limits(ap_ref, ap_wcs,): - """Sanity check. +def _set_weight_from_limits(profile, idx, lower_limit, upper_limit, allow_partial=True): + """Set profile pixel weighting from a lower and upper limit. - Parameters - ---------- - ap_ref : namedtuple - Contains xstart, xstop, ystart, ystop. These are the values of - the extraction region as specified by the reference file or the - image size. + Pixels to be fully included in the aperture are set to 1.0. Pixels partially + included are set to fractional values. - ap_wcs : namedtuple - These are the bounding box limits. + The profile is updated in place, so aperture setting is cumulative. If + there are overlapping apertures specified, later ones will overwrite + earlier values. - Returns - ------- - flag : boolean - True if ap_ref and ap_wcs do overlap, i.e. if the sanity test passes. + Parameters + ---------- + profile : ndarray of float + The spatial profile to update. + idx : ndarray of int + Index values for the profile array, corresponding to the cross-dispersion + axis. Dimensions must match `profile` shape. + lower_limit : float or ndarray of float + Lower limit for the aperture. If not a single value, dimensions must + match `profile` shape. + upper_limit : float or ndarray of float + Upper limit for the aperture. If not a single value, dimensions must + match `profile` shape. + allow_partial : bool, optional + If True, partial pixel weights are set where the pixel index intersects + the limit values. If False, only whole integer weights are set. """ - if (ap_wcs.xstart >= ap_ref.xstop or ap_wcs.xstop <= ap_ref.xstart or - ap_wcs.ystart >= ap_ref.ystop or ap_wcs.ystop <= ap_ref.ystart): - log.warning("The WCS bounding box is outside the aperture:") - log.warning(f"\taperture: {ap_ref.xstart}, {ap_ref.xstop}, " - f"{ap_ref.ystart}, {ap_ref.ystop}") - log.warning(f"\twcs: {ap_wcs.xstart}, {ap_wcs.xstop}, " - f"{ap_wcs.ystart}, {ap_wcs.ystop}") - log.warning("so the wcs bounding box will be ignored.") - flag = False - else: - flag = True - - return flag - - -def compare_start(aperture_start, wcs_bb_lower_lim): - """Compare the start limit from the aperture with the WCS lower limit. - Extended summary - ---------------- - The more restrictive (i.e. larger) limit is the one upon which the - output value will be based. If the WCS limit is larger, the value will - be increased to an integer, on the assumption that WCS lower limits - correspond to the lower edge of the bounding pixel. If this value - will actually be used for an extraction limit (i.e. if the limits were - not already specified by polynomial coefficients), then - `assign_polynomial_limits` will create a polynomial function using this - value, except that it will be decreased by 0.5 to correspond to the - lower edge of the bounding pixels. + # Both limits are inclusive + profile[(idx >= lower_limit) & (idx <= upper_limit)] = 1.0 - Parameters - ---------- - aperture_start : int or float - xstart or ystart, as specified by the extract1d reference file or the image - size. + if allow_partial: - wcs_bb_lower_lim : int or float - The lower limit from the WCS bounding box. + # For each pixel, get the distance from the lower and upper limits, + # to set partial pixel weights to the fraction of the pixel included. + # For example, if the lower limit is 1.7, then the weight of pixel 1 + # should be set to 0.3. If the upper limit is 10.7, then the weight of + # pixel 11 should be set to 0.7. + lower_weight = idx - lower_limit + 1 + upper_weight = upper_limit - idx + 1 - Returns - ------- - value : int or float - The start limit, possibly constrained by the WCS start limit. - """ - if aperture_start >= wcs_bb_lower_lim: # ref is inside WCS limit - return aperture_start - else: # outside (below) WCS limit - return math.ceil(wcs_bb_lower_lim) + for partial_pixel_weight in [lower_weight, upper_weight]: + # Check for overlap values that are between 0 and 1, for which + # the profile does not already contain a higher weight + test = ((partial_pixel_weight > 0) + & (partial_pixel_weight < 1) + & (profile < partial_pixel_weight)) + # Set these values to the partial pixel weight + profile[test] = partial_pixel_weight[test] -def compare_stop(aperture_stop, wcs_bb_upper_lim): - """Compare the stop limit from the aperture with the WCS upper limit. - The more restrictive (i.e. smaller) limit is the one upon which the - output value will be based. If the WCS limit is smaller, the value - will be truncated to an integer, on the assumption that WCS upper - limits correspond to the upper edge of the bounding pixel. If this - value will actually be used for an extraction limit (i.e. if the - limits were not already specified by polynomial coefficients), then - `assign_polynomial_limits` will create a polynomial function using this - value, except that it will be increased by 0.5 to correspond to the - upper edge of the bounding pixels. +def box_profile(shape, extract_params, wl_array, coefficients='src_coeff', + label='aperture', return_limits=False): + """Create a spatial profile for box extraction. - Parameters - ---------- - aperture_stop : int or float - xstop or ystop, as specified by the extract1d reference file or the image - size. + The output profile is an image matching the input data shape, + containing weights for each pixel in the image. Pixels to be fully + included in an extraction aperture are set to 1.0. Pixels partially + included are set to values between 0 and 1, depending on the overlap + with the aperture limits. Pixels not included in the aperture are + set to 0.0. - wcs_bb_upper_lim : int or float - The upper limit from the WCS bounding box. + Upper and lower limits for the aperture are determined from the + `extract_params`, in this priority order: - Returns - ------- - value : int or float - The stop limit, possibly constrained by the WCS stop limit. - """ - if aperture_stop <= wcs_bb_upper_lim: # ref is inside WCS limit - return aperture_stop - else: # outside (above) WCS limit - return math.floor(wcs_bb_upper_lim) + 1. src_coeff upper and lower limits (or bkg_coeff, for a background profile) + 2. center of start/stop values +/- extraction width / 2 + 3. cross-dispersion start/stop values + 4. array limits. + Left and right limits are set from start/stop values only. -def create_poly(coeff): - """Create a polynomial model from coefficients. + Only a single aperture is supported at this time. The 'src_coeff' + and 'bkg_coeff' specifications allow multiple regions to be specified, + but all regions are considered part of the same aperture and are + added to the same profile, to be extracted at once. Parameters ---------- - coeff : list of float - The coefficients of the polynomial, constant term first, highest - order term last. + shape : tuple of int + Data shape for the output profile, to match the spectral image. + extract_params : dict + Extraction parameters, as returned from `get_extract_parameters`. + wl_array : ndarray + Array of wavelength values, matching `shape`, for each pixel in + the array. + coefficients : {'src_coeff', 'bkg_coeff'}, optional + The polynomial coefficients to look for in the `extract_params` + dictionary. If 'bkg_coeff', the output aperture contains background + regions; otherwise, it contains target source regions. + label : str, optional + A label to use for the aperture, while logging limit values. + return_limits : bool, optional + If True, an upper and lower limit value for the aperture are + returned along with the spatial profile. These are used for + recording the aperture extent in output metadata. For + apertures set from polynomial coefficients, the returned values + are averages of the true upper and lower limits, which may vary + by dispersion element. Returns ------- - `astropy.modeling.polynomial.Polynomial1D` object, or None if `coeff` - is empty. - """ - n = len(coeff) - - if n < 1: - return None - - coeff_dict = {f'c{i}': coeff[i] for i in range(n)} - - return polynomial.Polynomial1D(degree=n - 1, **coeff_dict) - - -class ExtractBase(abc.ABC): - """Base class for 1-D extraction info and methods. - - Attributes - ---------- - exp_type : str - Exposure type. - - ref_image : data model, or None - The reference image model. - - spectral_order : int - Spectral order number. - - dispaxis : int - Dispersion direction: 1 is horizontal, 2 is vertical. - - xstart : int or None - First pixel (zero indexed) in extraction region. - - xstop : int or None - Last pixel (zero indexed) in extraction region. - - ystart : int or None - First pixel (zero indexed) in extraction region. - - ystop : int or None - Last pixel (zero indexed) in extraction region. - - extract_width : int or None - Height (in the cross-dispersion direction) of the extraction - region. - - independent_var : str - The polynomial functions for computing the upper and lower - boundaries of extraction and background regions can be functions - of pixel number or wavelength (in microns). These options are - distinguished by `independent_var`. - - position_correction : float - If not zero, this will be added to the extraction region limits - for the cross-dispersion direction, both target and background. - - src_coeff : list of lists of float, or None - These are coefficients of polynomial functions that define the - cross-dispersion limits of one or more source extraction regions - (yes, there can be more than one extraction region). Note that - these are float values, and the limits can include fractions of - pixels. - If specified, this takes priority over `ystart`, `ystop`, and - `extract_width`, though `xstart` and `xstop` will still be used - for the limits in the dispersion direction. (Interchange "x" and - "y" if the dispersion is in the vertical direction.) - - For example, the value could be: - -- [[1, 2], [3, 4, 5], [6], [7, 8]] - - which means: - -- [1, 2] coefficients for the lower limit of the first region -- [3, 4, 5] coefficients for the upper limit of the first region -- [6] coefficient(s) for the lower limit of the second region -- [7, 8] coefficients for the upper limit of the second region - - The coefficients are listed with the constant term first, highest - order term last. For example, [3, 4, 5] means 3 + 4 * x + 5 * x^2. - - bkg_coeff : list of lists of float, or None - This has the same format as `src_coeff`, but the polynomials - define one or more background regions. - - p_src : list of astropy.modeling.polynomial.Polynomial1D - These are Astropy polynomial functions defining the limits in the - cross-dispersion direction of the source extraction region(s). - If `src_coeff` was specified, `p_src` will be created directly - from `src_coeff`; otherwise, a constant function based on - `ystart`, `ystop`, `extract_width` will be assigned to `p_src`. - - p_bkg : list of astropy.modeling.polynomial.Polynomial1D, or None - These are Astropy polynomial functions defining the limits in the - cross-dispersion direction of the background extraction regions. - This list will be populated from `bkg_coeff`, if that was specified. - - wcs : WCS object - For computing the right ascension, declination, and wavelength at - one or more pixels. + profile : ndarray of float + Aperture weights to use in box extraction from a spectral image. + lower_limit : float, optional + Average lower limit for the aperture. Returned only if `return_limits` + is set. + upper_limit : float, optional + Average upper limit for the aperture. Returned only if `return_limits` + is set. - smoothing_length : int - Width of a boxcar function for smoothing the background regions. - This argument must be an odd positive number or zero, and it is - only used if background regions have been specified. - - bkg_fit : str - Type of background fitting to perform in each column (or row, if - the dispersion is vertical). Allowed values are `poly` (default), - `mean`, and `median`. - - bkg_order : int - Polynomial order for fitting to each column (or row, if the - dispersion is vertical) of background. Only used if `bkg_fit` is - `poly`. - This argument must be positive or zero, and it is only used if - background regions have been specified. - - - subtract_background : bool or None - A flag that indicates whether the background should be subtracted. - If None, the value in the extract_1d reference file will be used. - If not None, this parameter overrides the value in the - extract_1d reference file. - - use_source_posn : bool or None - If True, the target and background positions specified in the - reference file (or the default position, if there is no - reference file) will be shifted to account for the actual - source position in the data. """ + # Get pixel index values for the array + yidx, xidx = np.mgrid[:shape[0], :shape[1]] + if extract_params['dispaxis'] == HORIZONTAL: + dval = yidx + else: + dval = xidx - def __init__( - self, - input_model, - slit = None, - ref_image = None, - dispaxis = HORIZONTAL, - spectral_order = 1, - xstart = None, - xstop = None, - ystart = None, - ystop = None, - extract_width = None, - src_coeff = None, - bkg_coeff = None, - independent_var = "pixel", - smoothing_length = 0, - bkg_fit = "poly", - bkg_order = 0, - position_correction = 0., - subtract_background = None, - use_source_posn = None, - match = None, - ref_file_type = None - ): - """ - Parameters - ---------- - input_model : data model - The input science data. - - slit : an input slit, or None if not used - For MultiSlit, `slit` is one slit from - a list of slits in the input. For other types of data, `slit` - will not be used. - - ref_image : data model, or None - The reference image. - - dispaxis : int - Dispersion direction: 1 is horizontal, 2 is vertical. - - spectral_order : int - Spectral order number. - - xstart : int - First pixel (zero indexed) in extraction region. - - xstop : int - Last pixel (zero indexed) in extraction region. - - ystart : int - First pixel (zero indexed) in extraction region. - - ystop : int - Last pixel (zero indexed) in extraction region. - - extract_width : int - Height (in the cross-dispersion direction) of the extraction - region. - - src_coeff : list of lists of float, or None - These are coefficients of polynomial functions that define - the cross-dispersion limits of one or more source extraction - regions. - - bkg_coeff : list of lists of float, or None - This has the same format as `src_coeff`, but the polynomials - define one or more background regions. - - independent_var : str - This can be either "pixel" or "wavelength" to specify the - independent variable for polynomial functions. - - smoothing_length : int - Width of a boxcar function for smoothing the background - regions. - - bkg_fit : str - Type of fitting to apply to background values in each column - (or row, if the dispersion is vertical). - - bkg_order : int - Polynomial order for fitting to each column (or row, if the - dispersion is vertical) of background. - - position_correction : float - If not zero, this will be added to the extraction region limits - for the cross-dispersion direction, both target and background. - - subtract_background : bool or None - A flag which indicates whether the background should be subtracted. - If None, the value in the extract_1d reference file will be used. - If not None, this parameter overrides the value in the - extract_1d reference file. - - use_source_posn : bool or None - If True, the target and background positions specified in the - reference file (or the default position, if there is no - reference file) will be shifted to account for the actual - source position in the data. - """ - self.exp_type = input_model.meta.exposure.type - - self.dispaxis = dispaxis - self.spectral_order = spectral_order - self.ref_image = ref_image - self.xstart = xstart - self.xstop = xstop - self.ystart = ystart - self.ystop = ystop - self.match = match - self.ref_file_type = ref_file_type - - # xstart, xstop, ystart, or ystop may be overridden with src_coeff, they may be limited by the input image size - # or by the WCS bounding box, or they may be modified if extract_width was specified (because extract_width - # takes precedence). - # If these values are specified, the limits in the cross-dispersion direction should be integers, but they may - # later be replaced with fractional values, depending on extract_width, in order to center the extraction window - # in the originally specified xstart to xstop (or ystart to ystop). - if self.dispaxis == VERTICAL: - if not isinstance(self.xstart, int) and self.xstart is not None: - self.xstart = round(self.xstart) - log.warning(f"xstart {xstart} should have been an integer; rounding to {self.xstart}") - - if not isinstance(self.xstop, int) and self.xstop is not None: - self.xstop = round(self.xstop) - log.warning(f"xstop {xstop} should have been an integer; rounding to {self.xstop}") - - if self.dispaxis == HORIZONTAL: - if not isinstance(self.ystart, int) and self.ystart is not None: - self.ystart = round(self.ystart) - log.warning(f"ystart {ystart} should have been an integer; rounding to {self.ystart}") - - if not isinstance(self.ystop, int) and self.ystop is not None: - self.ystop = round(self.ystop) - log.warning(f"ystop {ystop} should have been an integer; rounding to {self.ystop}") - - if extract_width is None: - self.extract_width = None - else: - self.extract_width = round(extract_width) - - self.independent_var = independent_var.lower() - - # Coefficients for source (i.e. target) and background limits and corresponding polynomial functions. - self.src_coeff = copy.deepcopy(src_coeff) - self.bkg_coeff = copy.deepcopy(bkg_coeff) - - # These functions will be assigned by assign_polynomial_limits. - # The "p" in the attribute name indicates a polynomial function. - self.p_src = None - self.p_bkg = None - - self.smoothing_length = smoothing_length if smoothing_length is not None else 0 - - if 0 < self.smoothing_length == self.smoothing_length // 2 * 2: - log.warning(f"smoothing_length was even ({self.smoothing_length}), so incremented by 1") - self.smoothing_length += 1 # must be odd - - self.bkg_fit = bkg_fit - self.bkg_order = bkg_order - self.use_source_posn = use_source_posn - self.position_correction = position_correction - self.subtract_background = subtract_background - - self.wcs = None # initial value - - if input_model.meta.exposure.type == "NIS_SOSS": - if hasattr(input_model.meta, 'wcs'): - try: - self.wcs = niriss.niriss_soss_set_input(input_model, self.spectral_order) - except ValueError: - raise InvalidSpectralOrderNumberError(f"Spectral order {self.spectral_order} is not valid") - elif slit is None: - if hasattr(input_model.meta, 'wcs'): - self.wcs = input_model.meta.wcs - elif hasattr(slit, 'meta') and hasattr(slit.meta, 'wcs'): - self.wcs = slit.meta.wcs - - if self.wcs is None: - log.warning("WCS function not found in input.") - - def update_extraction_limits(self, ap): - pass - - def assign_polynomial_limits(self): - pass - - @staticmethod - def get_target_coordinates(input_model, slit): - """Get the right ascension and declination of the target. - - For MultiSlitModel (or similar) data, each slit has the source - right ascension and declination as attributes, and this can vary - from one slit to another (e.g. for NIRSpec MOS, or for WFSS). In - this case, we want the celestial coordinates from the slit object. - For other models, however, the celestial coordinates of the source - are in input_model.meta.target. - - Parameters - ---------- - input_model : data model - The input science data model. - - slit : SlitModel or None - One slit from a MultiSlitModel (or similar), or None if - there are no slits. - - Returns - ------- - targ_ra : float or None - The right ascension of the target, or None - - targ_dec : float or None - The declination of the target, or None - """ - targ_ra = None - targ_dec = None - - if slit is not None: - # If we've been passed a slit object, get the RA/Dec - # from the slit source attributes - targ_ra = getattr(slit, 'source_ra', None) - targ_dec = getattr(slit, 'source_dec', None) - elif isinstance(input_model, datamodels.SlitModel): - # If the input model is a single SlitModel, again - # get the coords from the slit source attributes - targ_ra = getattr(input_model, 'source_ra', None) - targ_dec = getattr(input_model, 'source_dec', None) - - if targ_ra is None or targ_dec is None: - # Otherwise get it from the generic target coords - targ_ra = input_model.meta.target.ra - targ_dec = input_model.meta.target.dec - - # Issue a warning if none of the methods succeeded - if targ_ra is None or targ_dec is None: - log.warning("Target RA and Dec could not be determined") - targ_ra = targ_dec = None - - return targ_ra, targ_dec - - def offset_from_offset(self, input_model, slit): - """Get position offset from the target coordinates. - - Parameters - ---------- - input_model : data model - The input science data. - - slit : SlitModel or None - One slit from a MultiSlitModel (or similar), or None if - there are no slits. - - Returns - ------- - offset : float - The offset of the exposure from the nominal position, due to - source positioning. This is the component of the offset - perpendicular to the dispersion direction. A positive value - means that the spectrum is at a larger pixel number than the - nominal location. - - locn : float or None - The pixel coordinate of the target in the cross-dispersion - direction, at the middle of the spectrum in the dispersion - direction. - """ - targ_ra, targ_dec = self.get_target_coordinates(input_model, slit) - - if targ_ra is None or targ_dec is None: - return 0., None - - # Use the WCS function to find the cross-dispersion (XD) location that is closest to the target coordinates. - # This is the "actual" location of the spectrum, so the extraction region should be centered here. - locn_info = self.locn_from_wcs(input_model, slit, targ_ra, targ_dec) - - if locn_info is None: - middle = middle_wl = locn = locn_info - else: - middle, middle_wl, locn = locn_info - - if middle is not None: - log.debug(f"Spectrum location from WCS used column/row {middle}") - - # Find the nominal extraction location, i.e. the XD location specified in the reference file prior to adding any - # position offset. - # The difference is the position offset. - offset = 0. - - if middle is not None and locn is not None: - nominal_location = self.nominal_locn(middle, middle_wl) - - log.debug(f"Target spectrum is at {locn:.2f} in the cross-dispersion direction") - - if nominal_location is not None: - log.debug(f"and the nominal XD location of the spectrum is {nominal_location:.2f}") - - offset = locn - nominal_location - else: - log.debug("but couldn't determine the nominal XD location.") - - if np.isnan(offset): - log.warning("Source position offset is NaN; setting it to 0") - offset = 0. - - self.position_correction = offset - - return offset, locn - - def locn_from_wcs(self, input_model, slit, targ_ra, targ_dec): - """Get the location of the spectrum, based on the WCS. - - Parameters - ---------- - input_model : data model - The input science model. - - slit : one slit from a MultiSlitModel (or similar), or None - The WCS and target coordinates will be gotten from `slit` - unless `slit` is None, and in that case they will be gotten - from `input_model`. - - targ_ra : float or None - The right ascension of the target, or None - - targ_dec : float or None - The declination of the target, or None - - Returns - ------- - tuple (middle, middle_wl, locn) or None - middle : int - Pixel coordinate in the dispersion direction within the 2-D - cutout (or the entire input image) at the middle of the WCS - bounding box. This is the point at which to determine the - nominal extraction location, in case it varies along the - spectrum. The offset will then be the difference between - `locn` (below) and the nominal location. - - middle_wl : float - The wavelength at pixel `middle`. - - locn : float - Pixel coordinate in the cross-dispersion direction within the - 2-D cutout (or the entire input image) that has right ascension - and declination coordinates corresponding to the target location. - The spectral extraction region should be centered here. - - None will be returned if there was not sufficient information - available, e.g. if the wavelength attribute or wcs function is not - defined. - """ - # WFSS data are not currently supported by this function - if input_model.meta.exposure.type in WFSS_EXPTYPES: - log.warning("Can't use target coordinates to get location of spectrum " - f"for exp type {input_model.meta.exposure.type}") - return - - bb = self.wcs.bounding_box # ((x0, x1), (y0, y1)) - - if bb is None: - if slit is None: - shape = input_model.data.shape - else: - shape = slit.data.shape - - bb = wcs_bbox_from_shape(shape) - - if self.dispaxis == HORIZONTAL: - # Width (height) in the cross-dispersion direction, from the start of the 2-D cutout (or of the full image) - # to the upper limit of the bounding box. - # This may be smaller than the full width of the image, but it's all we need to consider. - xd_width = int(round(bb[1][1])) # must be an int - middle = int((bb[0][0] + bb[0][1]) / 2.) # Middle of the bounding_box in the dispersion direction. - x = np.empty(xd_width, dtype=np.float64) - x[:] = float(middle) - y = np.arange(xd_width, dtype=np.float64) - lower = bb[1][0] - upper = bb[1][1] - else: # dispaxis = VERTICAL - xd_width = int(round(bb[0][1])) # Cross-dispersion total width of bounding box; must be an int - middle = int((bb[1][0] + bb[1][1]) / 2.) # Mid-point of width along dispersion direction - x = np.arange(xd_width, dtype=np.float64) # 1-D vector of cross-dispersion (x) pixel indices - y = np.empty(xd_width, dtype=np.float64) # 1-D vector all set to middle y index - y[:] = float(middle) - - # lower and upper range in cross-dispersion direction - lower = bb[0][0] - upper = bb[0][1] - - # We need stuff[2], a 1-D array of wavelengths crossing the spectrum near its middle. - fwd_transform = self.wcs(x, y) - middle_wl = np.nanmean(fwd_transform[2]) - - if input_model.meta.exposure.type in ['NRS_FIXEDSLIT', 'NRS_MSASPEC', - 'NRS_BRIGHTOBJ']: - if slit is None: - xpos = input_model.source_xpos - ypos = input_model.source_ypos - else: - xpos = slit.source_xpos - ypos = slit.source_ypos - - slit2det = self.wcs.get_transform('slit_frame', 'detector') - x_y = slit2det(xpos, ypos, middle_wl) - log.info("Using source_xpos and source_ypos to center extraction.") - - elif input_model.meta.exposure.type == 'MIR_LRS-FIXEDSLIT': - try: - if slit is None: - dithra = input_model.meta.dither.dithered_ra - dithdec = input_model.meta.dither.dithered_dec - else: - dithra = slit.meta.dither.dithered_ra - dithdec = slit.meta.dither.dithered_dec - x_y = self.wcs.backward_transform(dithra, dithdec, middle_wl) - except AttributeError: - log.warning("Dithered pointing location not found in wcsinfo. " - "Defaulting to TARG_RA / TARG_DEC for centering.") - return - - # locn is the XD location of the spectrum: - if self.dispaxis == HORIZONTAL: - locn = x_y[1] - else: - locn = x_y[0] - - if locn < lower or locn > upper and targ_ra > 340.: - # Try this as a temporary workaround. - x_y = self.wcs.backward_transform(targ_ra - 360., targ_dec, middle_wl) - - if self.dispaxis == HORIZONTAL: - temp_locn = x_y[1] - else: - temp_locn = x_y[0] - - if lower <= temp_locn <= upper: - # Subtracting 360 from the right ascension worked! - locn = temp_locn - - log.debug(f"targ_ra changed from {targ_ra} to {targ_ra - 360.}") - - # If the target is at the edge of the image or at the edge of the non-NaN area, we can't use the WCS to find the - # location of the target spectrum. - if locn < lower or locn > upper: - log.warning(f"WCS implies the target is at {locn:.2f}, which is outside the bounding box,") - log.warning("so we can't get spectrum location using the WCS") - locn = None - - return middle, middle_wl, locn - - @abc.abstractmethod - def nominal_locn(self, middle, middle_wl): - # Implemented in the subclasses. - pass - - -class ExtractModel(ExtractBase): - """The extraction region was specified in a JSON file.""" - - def __init__(self, *base_args, **base_kwargs): - """Create a polynomial model from coefficients. - - Extended summary - ---------------- - If InvalidSpectralOrderNumberError is raised, processing of the - current slit or spectral order should be skipped. - - Parameters - ---------- - *base_args, **base_kwargs : - see ExtractBase parameters for more information. - - """ - super().__init__(*base_args, **base_kwargs) - - # The independent variable for functions for the lower and upper limits of target and background regions can be - # either 'pixel' or 'wavelength'. - if (self.independent_var != "wavelength" and - self.independent_var not in ["pixel", "pixels"]): - log.error(f"independent_var = {self.independent_var}'; specify 'wavelength' or 'pixel'") - raise RuntimeError("Invalid value for independent_var") - - # Do sanity checks between requested background subtraction and the - # existence of background region specifications - if self.subtract_background is not None: - if self.subtract_background: - # If background subtraction was requested, but no background region(s) - # were specified, turn it off - if self.bkg_coeff is None: - self.subtract_background = False - log.debug("Skipping background subtraction because background regions are not defined.") - else: - # If background subtraction was NOT requested, even though background region(s) - # were specified, blank out the bkg region info - if self.bkg_coeff is not None: - self.bkg_coeff = None - log.debug("Background subtraction was specified in the reference file,") - log.debug("but has been overridden by the step parameter.") - - def nominal_locn(self, middle, middle_wl): - """Find the nominal cross-dispersion location of the target spectrum. - - This version is for the case that the reference file is a JSON file, - or that there is no reference file. - - Parameters - ---------- - middle: int - The zero-indexed pixel number of the point in the dispersion - direction at which `locn_from_wcs` determined the actual - location (in the cross-dispersion direction) of the target - spectrum. This is used for evaluating the polynomial - functions if the independent variable is pixel. - - middle_wl: float - The wavelength at pixel `middle`. This is only used if the - independent variable for polynomial functions is wavelength. - - Returns - ------- - location: float or None - The nominal cross-dispersion location (i.e. unmodified by - position offset) of the target spectrum. - - """ - if self.src_coeff is None: - if self.dispaxis == HORIZONTAL: - location = float(self.ystart + self.ystop) / 2. - else: - location = float(self.xstart + self.xstop) / 2. - else: - if self.independent_var.startswith("wavelength"): - x = float(middle_wl) - else: - x = float(middle) - - # Create the polynomial functions. - # We'll do this again later, after adding the position offset to the coefficients, but we need to evaluate - # them at x now in order to get the nominal location of the spectrum. - self.assign_polynomial_limits() - - n_srclim = len(self.p_src) - sum_data = 0. - sum_weights = 0. - - for i in range(n_srclim): - lower = self.p_src[i][0](x) - upper = self.p_src[i][1](x) - weight = (upper - lower) - sum_data += weight * (lower + upper) / 2. - sum_weights += weight - - if sum_weights == 0.: - location = None - else: - location = sum_data / sum_weights - - return location - - def add_position_correction(self, shape): - """Add the position offset to the extraction location (in-place). - - Extended summary - ---------------- - If source extraction coefficients src_coeff were specified, this - method will add the source position correction to the first coefficient - of every coefficient list; otherwise, the source offset will be added - to xstart & xstop or to ystart & ystop. - If background extraction coefficients bkg_coeff were specified, - this method will add the source offset to the first coefficients. - Note that background coefficients are handled independently of - src_coeff. - - Parameters - ---------- - shape : tuple - The shape of the data array (may be just the last two axes). - This is used for truncating a shifted limit at the image edge. - """ - - if self.position_correction == 0.: - return - - if self.dispaxis == HORIZONTAL: - direction = "y" - self.ystart += self.position_correction - self.ystop += self.position_correction - # These values must not be negative. - self.ystart = max(self.ystart, 0) - self.ystop = max(self.ystop, 0) - self.ystart = min(self.ystart, shape[-2] - 1) - self.ystop = min(self.ystop, shape[-2] - 1) # inclusive limit - else: - direction = "x" - self.xstart += self.position_correction - self.xstop += self.position_correction - # These values must not be negative. - self.xstart = max(self.xstart, 0) - self.xstop = max(self.xstop, 0) - self.xstart = min(self.xstart, shape[-1] - 1) - self.xstop = min(self.xstop, shape[-1] - 1) # inclusive limit - - if self.src_coeff is None: - log.info( - f"Applying position offset of {self.position_correction:.2f} to {direction}start and {direction}stop") - - if self.src_coeff is not None or self.bkg_coeff is not None: - log.info(f"Applying position offset of {self.position_correction:.2f} to polynomial coefficients") - - if self.src_coeff is not None: - self._apply_position_corr(self.src_coeff) - - if self.bkg_coeff is not None: - self._apply_position_corr(self.bkg_coeff) - - def _apply_position_corr(self, coeffs): - for i in range(len(coeffs)): - coeff_list = coeffs[i] - coeff_list[0] += self.position_correction - coeffs[i] = copy.copy(coeff_list) - - def update_extraction_limits(self, ap): - """Update start and stop limits. - - Extended summary - ---------------- - Copy the values of xstart, etc., to the attributes. Note, however, - that if src_coeff was specified, that will override the values - given by xstart, etc. - The limits in the dispersion direction will be rounded to integer. - - Parameters - ---------- - ap : namedtuple - - """ - self.xstart = ap.xstart - self.xstop = ap.xstop - self.ystart = ap.ystart - self.ystop = ap.ystop - - if self.dispaxis == HORIZONTAL: - self.xstart = int(round(self.xstart)) - self.xstop = int(round(self.xstop)) - else: # vertical - self.ystart = int(round(self.ystart)) - self.ystop = int(round(self.ystop)) - - def log_extraction_parameters(self): - """Log the updated extraction parameters.""" - log.debug("Updated parameters:") - log.debug(f"position_correction = {self.position_correction}") - - if self.src_coeff is not None: - log.debug(f"src_coeff = {self.src_coeff}") - - # Since src_coeff was specified, that will be used instead of xstart & xstop (or ystart & ystop). - if self.dispaxis == HORIZONTAL: - # Only print xstart/xstop, because ystart/ystop are not used - log.debug(f"xstart = {self.xstart}") - log.debug(f"xstop = {self.xstop}") - else: - # Only print ystart/ystop, because xstart/xstop are not used - log.debug(f"ystart = {self.ystart}") - log.debug(f"ystop = {self.ystop}") - - if self.bkg_coeff is not None: - log.debug(f"bkg_coeff = {self.bkg_coeff}") - - def assign_polynomial_limits(self): - """Create polynomial functions for extraction limits. - - Extended summary - ---------------- - self.src_coeff and self.bkg_coeff contain lists of polynomial - coefficients. These will be used to create corresponding lists of - polynomial functions, self.p_src and self.p_bkg. Note, however, - that the structures of those two lists are not the same. - The coefficients lists have this form: - -- [[1, 2], [3, 4, 5], [6], [7, 8]] - - which means: - -- [1, 2] coefficients for the lower limit of the first region -- [3, 4, 5] coefficients for the upper limit of the first region -- [6] coefficient(s) for the lower limit of the second region -- [7, 8] coefficients for the upper limit of the second region - - The lists of coefficients must always be in pairs, for the lower - and upper limits respectively, but they're not explicitly in an - additional layer of two-element lists, - -- i.e., not like this: [[[1, 2], [3, 4, 5]], [[6], [7, 8]]] - - That seemed unnecessarily messy and harder for the user to specify. - For the lists of polynomial functions, on the other hand, that - additional layer of list is used: - -- [[fcn_lower1, fcn_upper1], [fcn_lower2, fcn_upper2]] - - where: - -- fcn_lower1 is 1 + 2 * x -- fcn_upper1 is 3 + 4 * x + 5 * x**2 -- fcn_lower2 is 6 -- fcn_upper2 is 7 + 8 * x - """ - if self.src_coeff is None: - # Create constant functions. - - if self.dispaxis == HORIZONTAL: - lower = float(self.ystart) - 0.5 - upper = float(self.ystop) + 0.5 - else: - lower = float(self.xstart) - 0.5 - upper = float(self.xstop) + 0.5 - - log.debug(f"Converting extraction limits to [[{lower}], [{upper}]]") + # Get start/stop values from parameters if present, + # or default to data shape + xstart = extract_params.get('xstart', 0) + xstop = extract_params.get('xstop', shape[1] - 1) + ystart = extract_params.get('ystart', 0) + ystop = extract_params.get('ystop', shape[0] - 1) - self.p_src = [[create_poly([lower]), create_poly([upper])]] + # Check if the profile should contain partial pixel weights + if coefficients == 'bkg_coeff' and extract_params['bkg_fit'] == 'median': + allow_partial = False + else: + allow_partial = True + + # Set aperture region, in this priority order: + # 1. src_coeff upper and lower limits (or bkg_coeff, for background profile) + # 2. center of start/stop values +/- extraction width + # 3. start/stop values + profile = np.full(shape, 0.0) + if extract_params[coefficients] is not None: + # Limits from source coefficients: ignore ystart/stop/width + if extract_params['independent_var'].startswith("wavelength"): + ival = wl_array + elif extract_params['dispaxis'] == HORIZONTAL: + ival = xidx.astype(np.float32) else: - # The source extraction can include more than one region. - n_src_coeff = len(self.src_coeff) - if n_src_coeff // 2 * 2 != n_src_coeff: - raise RuntimeError("src_coeff must contain alternating lists of lower and upper limits.") - - self.p_src = self._poly_per_region(self.src_coeff) - - if self.bkg_coeff is not None: - n_bkg_coeff = len(self.bkg_coeff) - if n_bkg_coeff // 2 * 2 != n_bkg_coeff: - raise RuntimeError("bkg_coeff must contain alternating lists of lower and upper limits.") - - self.p_bkg = self._poly_per_region(self.bkg_coeff) - - def _poly_per_region(self, coeffs): - result = [] - expect_lower = True # toggled in loop - - for coeff_list in coeffs: - if expect_lower: + ival = yidx.astype(np.float32) + + # The source extraction can include more than one region, + # but must contain pairs of lower and upper limits. + n_src_coeff = len(extract_params[coefficients]) + if n_src_coeff % 2 != 0: + raise RuntimeError(f"{coefficients} must contain alternating " + f"lists of lower and upper limits.") + + lower = None + lower_limit = None + upper_limit = None + for i, coeff_list in enumerate(extract_params[coefficients]): + if i % 2 == 0: lower = create_poly(coeff_list) else: upper = create_poly(coeff_list) - result.append([lower, upper]) - - expect_lower = not expect_lower - - return result - - def extract(self, data, var_poisson, var_rnoise, var_flat, wl_array): - """Do the extraction. - - Extended summary - ---------------- - This version is for the case that the reference file is a JSON - file, or that there is no reference file. - - Parameters - ---------- - data : ndarray, 2-D - Data array from which the spectrum will be extracted. - - var_poisson: ndarray, 2-D - Poisson noise variance array to be extracted following data extraction method. - - var_rnoise: ndarray, 2-D - Read noise variance array to be extracted following data extraction method. - - var_flat: ndarray, 2-D - Flat noise variance array to be extracted following data extraction method. - - wl_array : ndarray, 2-D, or None - Wavelengths corresponding to `data`, or None if no WAVELENGTH - extension was found in the input file. - - Returns - ------- - ra, dec : float - ra and dec are the right ascension and declination respectively - at the nominal center of the slit. - - wavelength : ndarray, 1-D, float64 - The wavelength in micrometers at each pixel. - - temp_flux : ndarray, 1-D - The sum of the data values in the extraction region minus the - sum of the data values in the background regions (scaled by the - ratio of the numbers of pixels), for each pixel. - The data values are in units of surface brightness, so this - value isn't really the flux, it's an intermediate value. - Dividing by `npixels` (to compute the average) will give the - array for the `surf_bright` (surface brightness) output column, - and multiplying by the solid angle of a pixel will give the - flux for a point source. - - f_var_poisson : ndarray, 1-D - The extracted poisson variance values to go along with the - temp_flux array. - - f_var_rnoise : ndarray, 1-D - The extracted read noise variance values to go along with the - temp_flux array. - - f_var_flat : ndarray, 1-D - The extracted flat field variance values to go along with the - temp_flux array. - - background : ndarray, 1-D, float64 - The background count rate that was subtracted from the sum of - the source data values to get `temp_flux`. - - b_var_poisson : ndarray, 1-D - The extracted poisson variance values to go along with the - background array. - - b_var_rnoise : ndarray, 1-D - The extracted read noise variance values to go along with the - background array. - - b_var_flat : ndarray, 1-D - The extracted flat field variance values to go along with the - background array. - - npixels : ndarray, 1-D, float64 - The number of pixels that were added together to get `temp_flux`. - - dq : ndarray, 1-D, uint32 - The data quality array. - - """ - # If the wavelength attribute exists and is populated, use it in preference to the wavelengths returned by the - # wcs function. - # But since we're now calling get_wavelengths from lib.wcs_utils, wl_array should be populated, and we should be - # able to remove some of this code. - got_wavelength = False if wl_array is None or len(wl_array) == 0 else True # may be reset later - - # The default value is 0, so all 0 values means that the wavelength attribute was not populated. - if not got_wavelength or (wl_array.min() == 0. and wl_array.max() == 0.): - got_wavelength = False - - if got_wavelength: - # We need a 1-D array of wavelengths, one element for each output table row. - # These are slice limits. - sx0 = int(round(self.xstart)) - sx1 = int(round(self.xstop)) + 1 - sy0 = int(round(self.ystart)) - sy1 = int(round(self.ystop)) + 1 - - # Convert non-positive values to NaN, to easily ignore them. - wl = wl_array.copy() # Don't modify wl_array - nan_flag = np.isnan(wl) - - # To avoid a warning about invalid value encountered in less_equal. - wl[nan_flag] = -1000. - wl = np.where(wl <= 0., np.nan, wl) - - if self.dispaxis == HORIZONTAL: - wavelength = np.nanmean(wl[sy0:sy1, sx0:sx1], axis=0) - else: - wavelength = np.nanmean(wl[sy0:sy1, sx0:sx1], axis=1) - - # Now call the wcs function to compute the celestial coordinates. - # Also use the returned wavelengths if we weren't able to get them from the wavelength attribute. - - # Used for computing the celestial coordinates. - if self.dispaxis == HORIZONTAL: - slice0 = int(round(self.xstart)) - slice1 = int(round(self.xstop)) + 1 - x_array = np.arange(slice0, slice1, dtype=np.float64) - y_array = np.empty(x_array.shape, dtype=np.float64) - y_array.fill((self.ystart + self.ystop) / 2.) - else: - slice0 = int(round(self.ystart)) - slice1 = int(round(self.ystop)) + 1 - y_array = np.arange(slice0, slice1, dtype=np.float64) - x_array = np.empty(y_array.shape, dtype=np.float64) - x_array.fill((self.xstart + self.xstop) / 2.) - - if self.wcs is not None: - coords = self.wcs(x_array, y_array) - ra = coords[0] - dec = coords[1] - wcs_wl = coords[2] - - # We need one right ascension and one declination, representing the direction of pointing. - mask = np.isnan(wcs_wl) - not_nan = np.logical_not(mask) - - if np.any(not_nan): - ra2 = ra[not_nan] - min_ra = ra2.min() - max_ra = ra2.max() - ra = (min_ra + max_ra) / 2. - dec2 = dec[not_nan] - min_dec = dec2.min() - max_dec = dec2.max() - dec = (min_dec + max_dec) / 2. - else: - log.warning("All wavelength values are NaN; assigning dummy value -999 to RA and Dec.") - ra = -999. - dec = -999. - else: - ra, dec, wcs_wl = None - - if not got_wavelength: - wavelength = wcs_wl # from wcs, or None - - if self.dispaxis == HORIZONTAL: - image = data - else: - image = np.transpose(data, (1, 0)) - var_poisson = np.transpose(var_poisson, (1, 0)) - var_rnoise = np.transpose(var_rnoise, (1, 0)) - var_flat = np.transpose(var_flat, (1, 0)) - - if wavelength is None: - log.warning("Wavelengths could not be determined.") - - if slice0 <= 0: - wavelength = np.arange(1, slice1 - slice0 + 1, dtype=np.float64) - else: - wavelength = np.arange(slice0, slice1, dtype=np.float64) - - temp_wl = wavelength.copy() - nan_mask = np.isnan(wavelength) - n_nan = nan_mask.sum(dtype=np.intp) - - if n_nan > 0: - log.debug(f"{n_nan} NaNs in wavelength array") - - temp_wl = np.nan_to_num(temp_wl, nan=0.01) # NaNs in the wavelength array cause problems; replace them. - - disp_range = [slice0, slice1] # Range (slice) of pixel numbers in the dispersion direction. - - temp_flux, f_var_poisson, f_var_rnoise, f_var_flat, background, \ - b_var_poisson, b_var_rnoise, b_var_flat, npixels = \ - extract1d.extract1d(image, var_poisson, var_rnoise, var_flat, - temp_wl, disp_range, self.p_src, self.p_bkg, - self.independent_var, self.smoothing_length, - self.bkg_fit, self.bkg_order, weights=None) - - del temp_wl - - dq = np.zeros(temp_flux.shape, dtype=np.uint32) - - if n_nan > 0: - wavelength, dq, nan_slc = nans_at_endpoints(wavelength, dq) - temp_flux = temp_flux[nan_slc] - background = background[nan_slc] - npixels = npixels[nan_slc] - f_var_poisson = f_var_poisson[nan_slc] - f_var_rnoise = f_var_rnoise[nan_slc] - f_var_flat = f_var_flat[nan_slc] - b_var_poisson = b_var_poisson[nan_slc] - b_var_rnoise = b_var_rnoise[nan_slc] - b_var_flat = b_var_flat[nan_slc] - - return (ra, dec, wavelength, - temp_flux, f_var_poisson, f_var_rnoise, f_var_flat, - background, b_var_poisson, b_var_rnoise, b_var_flat, npixels, dq) - -class ImageExtractModel(ExtractBase): - """This uses an image that specifies the extraction region. - - Extended summary - ---------------- - One of the requirements for this step is that for an extended target, - the entire aperture is supposed to be extracted (with no background - subtraction). It doesn't make any sense to use an image reference file - to extract the entire aperture; a trivially simple JSON reference file - would do. Therefore, we assume that if the user specified a reference - file in image format, the user actually wanted that reference file - to be used, so we will ignore the requirement and extract as specified - by the reference image. - """ - - def __init__(self, *base_args, **base_kwargs): - """Extract using a reference image to define the extraction and - background regions. - - Parameters - ---------- - *base_args, **base_kwargs : - See ExtractionBase for more. - - """ - super().__init__(*base_args, **base_kwargs) - - def nominal_locn(self, middle, middle_wl): - """Find the nominal cross-dispersion location of the target spectrum. - - This version is for the case that the reference file is an image. - - Parameters - ---------- - middle: int - The zero-indexed pixel number of the point in the dispersion - direction at which `locn_from_wcs` determined the actual - location (in the cross-dispersion direction) of the target - spectrum. - - middle_wl: float - The wavelength at pixel `middle`. This is not used in this - version. - - Returns - ------- - location: float or None - The nominal cross-dispersion location (i.e. unmodified by - source position offset) of the target spectrum. - The value will be None if `middle` is outside the reference - image or if the reference image does not specify any pixels - to extract at `middle`. - - """ - shape = self.ref_image.data.shape - middle_line = None - location = None - - if self.dispaxis == HORIZONTAL: - if 0 <= middle < shape[1]: - middle_line = self.ref_image.data[:, middle] + # NOTE: source coefficients currently have a different + # definition for pixel inclusion than the start/stop input. + # Source coefficients define 0 at the center of the pixel, + # start/stop defines 0 at the start of the lower pixel and + # includes the upper pixel. Here, we are setting limits + # in the spatial profile according to the more commonly used + # start/stop definition, so we need to modify the lower and + # upper limits from polynomial coefficients to match. + lower_limit_region = lower(ival) + 0.5 + upper_limit_region = upper(ival) - 0.5 + + _set_weight_from_limits(profile, dval, lower_limit_region, + upper_limit_region, + allow_partial=allow_partial) + mean_lower = np.mean(lower_limit_region) + mean_upper = np.mean(upper_limit_region) + log.info(f'Mean {label} start/stop from {coefficients}: ' + f'{mean_lower:.2f} -> {mean_upper:.2f} (inclusive)') + + if lower_limit is None: + lower_limit = mean_lower + upper_limit = mean_upper + else: + if mean_lower < lower_limit: + lower_limit = mean_lower + if mean_upper > upper_limit: + upper_limit = mean_upper + + elif extract_params['extract_width'] is not None: + # Limits from extraction width at center of ystart/stop if present, + # center of array if not + if extract_params['dispaxis'] == HORIZONTAL: + nominal_middle = (ystart + ystop) / 2.0 else: - if 0 <= middle < shape[0]: - middle_line = self.ref_image.data[middle, :] - - if middle_line is None: - log.warning( - f"Can't determine nominal location of spectrum because middle = {middle} is off the image." - ) - - return - - mask_target = np.where(middle_line > 0., 1., 0.) - x = np.arange(len(middle_line), dtype=np.float64) - - numerator = (x * mask_target).sum() - denominator = mask_target.sum() - - if denominator > 0.: - location = numerator / denominator - - return location - - def add_position_correction(self, shape): - """Shift the reference image (in-place). + nominal_middle = (xstart + xstop) / 2.0 - Parameters - ---------- - shape : tuple - Not sure if needed yet? + width = extract_params['extract_width'] + lower_limit = nominal_middle - (width - 1.0) / 2.0 + upper_limit = lower_limit + width - 1 - """ - if self.position_correction == 0: - return + _set_weight_from_limits(profile, dval, lower_limit, upper_limit) + log.info(f'{label.capitalize()} start/stop: ' + f'{lower_limit:.2f} -> {upper_limit:.2f} (inclusive)') - log.info(f"Applying source offset of {self.position_correction:.2f}") - - # Shift the image in the cross-dispersion direction. - ref = self.ref_image.data.copy() - shift = self.position_correction - ishift = round(shift) - - if ishift != shift: - log.info(f"Rounding source offset of {shift} to {ishift}") - - if self.dispaxis == HORIZONTAL: - if abs(ishift) >= ref.shape[0]: - log.warning(f"Nod offset {ishift} is too large, skipping ...") - - return - - self.ref_image.data[:, :] = 0. - - if ishift > 0: - self.ref_image.data[ishift:, :] = ref[:-ishift, :] - else: - ishift = -ishift - self.ref_image.data[:-ishift, :] = ref[ishift:, :] + else: + # Limits from start/stop only, defaulting to the full array + # if not specified + if extract_params['dispaxis'] == HORIZONTAL: + lower_limit = ystart + upper_limit = ystop else: - if abs(ishift) >= ref.shape[1]: - log.warning(f"Nod offset {ishift} is too large, skipping ...") + lower_limit = xstart + upper_limit = xstop - return + _set_weight_from_limits(profile, dval, lower_limit, upper_limit) + log.info(f'{label.capitalize()} start/stop: ' + f'{lower_limit:.2f} -> {upper_limit:.2f} (inclusive)') - self.ref_image.data[:, :] = 0. + # Set weights to zero outside left and right limits + if extract_params['dispaxis'] == HORIZONTAL: + profile[:, :int(np.ceil(xstart))] = 0 + profile[:, int(np.floor(xstop)) + 1:] = 0 + else: + profile[:int(np.ceil(ystart)), :] = 0 + profile[int(np.floor(ystop)) + 1:, :] = 0 - if ishift > 0: - self.ref_image.data[:, ishift:] = ref[:, :-ishift] - else: - ishift = -ishift - self.ref_image.data[:, :-ishift] = ref[:, ishift:] - - def log_extraction_parameters(self): - """Log the updated extraction parameters.""" - log.debug("Using a reference image that defines extraction regions.") - log.debug(f"dispaxis = {self.dispaxis}") - log.debug(f"spectral order = {self.spectral_order}") - log.debug(f"smoothing_length = {self.smoothing_length}") - log.debug(f"position_correction = {self.position_correction}") - - def extract(self, data, var_poisson, var_rnoise, var_flat, wl_array): - """ - Do the actual extraction, for the case that the extract1d reference file - is an image. - - Parameters - ---------- - data : ndarray, 2-D - Science data array. - - var_poisson : ndarray, 2-D - Poisson noise variance array to be extracted following data extraction method. - - var_rnoise : ndarray, 2-D - Read noise variance array to be extracted following data extraction method. - - var_flat : ndarray, 2-D - Flat noise variance array to be extracted following data extraction method. - - wl_array : ndarray, 2-D, or None - Wavelengths corresponding to `data`, or None if no WAVELENGTH - extension was found in the input file. - - Returns - ------- - ra, dec : float - ra and dec are the right ascension and declination respectively - at the nominal center of the slit. - - wavelength : ndarray, 1-D - The wavelength in micrometers at each pixel. - - temp_flux : ndarray, 1-D - The sum of the data values in the extraction region minus the - sum of the data values in the background regions (scaled by the - ratio of the numbers of pixels), for each pixel. - The data values are in units of surface brightness, so this - value isn't really the flux, it's an intermediate value. - Multiply `temp_flux` by the solid angle of a pixel to get the - flux for a point source (column "flux"). Divide `temp_flux` by - `npixels` (to compute the average) to get the array for the - "surf_bright" (surface brightness) output column. - - f_var_poisson : ndarray, 1-D - The extracted poisson variance values to go along with the - temp_flux array. - - f_var_rnoise : ndarray, 1-D - The extracted read noise variance values to go along with the - temp_flux array. - - f_var_flat : ndarray, 1-D - The extracted flat field variance values to go along with the - temp_flux array. - - background : ndarray, 1-D - The background count rate that was subtracted from the sum of - the source data values to get `temp_flux`. - - b_var_poisson : ndarray, 1-D - The extracted poisson variance values to go along with the - background array. - - b_var_rnoise : ndarray, 1-D - The extracted read noise variance values to go along with the - background array. - - b_var_flat : ndarray, 1-D - The extracted flat field variance values to go along with the - background array. - - npixels : ndarray, 1-D, float64 - The number of pixels that were added together to get `temp_flux`. - - dq : ndarray, 1-D, uint32 - """ - shape = data.shape - ref = self.match_shape(shape) # Truncate or expand reference image to match the science data. - - # This is the axis along which to add up the data. - if self.dispaxis == HORIZONTAL: - axis = 0 - else: - axis = 1 - - # The values of these arrays will be just 0 or 1. - # If ref did not define any background pixels, however, mask_bkg will be None. - (mask_target, mask_bkg) = self.separate_target_and_background(ref) - - # This is the number of pixels in the cross-dispersion direction, in the target extraction region. - n_target = mask_target.sum(axis=axis, dtype=float) - - # Extract the data. - gross = (data * mask_target).sum(axis=axis, dtype=float) - f_var_poisson = (var_poisson * mask_target).sum(axis=axis, dtype=float) - f_var_rnoise = (var_rnoise * mask_target).sum(axis=axis, dtype=float) - f_var_flat = (var_flat * mask_target).sum(axis=axis, dtype=float) - - # Compute the number of pixels that were added together to get gross. - temp = np.ones_like(data) - npixels = (temp * mask_target).sum(axis=axis, dtype=float) - - if self.subtract_background is not None: - if not self.subtract_background: - if mask_bkg is not None: - log.info("Background subtraction was turned off - skipping it.") - mask_bkg = None - else: - if mask_bkg is None: - log.info("Skipping background subtraction because background regions are not defined.") - - # Extract the background. - if mask_bkg is not None: - n_bkg = mask_bkg.sum(axis=axis, dtype=float) - n_bkg = np.where(n_bkg == 0., -1., n_bkg) # -1 is used as a flag, and also to avoid dividing by zero. - - background = (data * mask_bkg).sum(axis=axis, dtype=float) - b_var_poisson = (var_poisson * mask_bkg).sum(axis=axis, dtype=float) - b_var_rnoise = (var_rnoise * mask_bkg).sum(axis=axis, dtype=float) - b_var_flat = (var_flat * mask_bkg).sum(axis=axis, dtype=float) - - scalefactor = n_target / n_bkg - scalefactor = np.where(n_bkg > 0., scalefactor, 0.) - - background *= scalefactor - - if self.smoothing_length > 1: - background = extract1d.bxcar(background, self.smoothing_length) # Boxcar smoothing. - background = np.where(n_bkg > 0., background, 0.) - b_var_poisson = extract1d.bxcar(b_var_poisson, self.smoothing_length) - b_var_poisson = np.where(n_bkg > 0., b_var_poisson, 0.) - b_var_rnoise = extract1d.bxcar(b_var_rnoise, self.smoothing_length) - b_var_rnoise = np.where(n_bkg > 0., b_var_rnoise, 0.) - b_var_flat = extract1d.bxcar(b_var_flat, self.smoothing_length) - b_var_flat = np.where(n_bkg > 0., b_var_flat, 0.) - - temp_flux = gross - background - else: - background = np.zeros_like(gross) - b_var_poisson = np.zeros_like(gross) - b_var_rnoise = np.zeros_like(gross) - b_var_flat = np.zeros_like(gross) - temp_flux = gross.copy() - - del gross - - # Since we're now calling get_wavelengths from lib.wcs_utils, wl_array should be populated, and we should be - # able to remove some of this code. - if wl_array is None or len(wl_array) == 0: - got_wavelength = False - else: - got_wavelength = True # may be reset below - - # If wl_array has all 0 values, interpret that to mean that the wavelength attribute was not populated. - if not got_wavelength or wl_array.min() == 0. and wl_array.max() == 0.: - got_wavelength = False - - # Used for computing the celestial coordinates and the 1-D array of wavelengths. - flag = (mask_target > 0.) - grid = np.indices(shape) - masked_grid = flag.astype(float) * grid[axis] - g_sum = masked_grid.sum(axis=axis) - f_sum = flag.sum(axis=axis, dtype=float) - f_sum_zero = np.where(f_sum <= 0.) - f_sum[f_sum_zero] = 1. # to avoid dividing by zero - - spectral_trace = g_sum / f_sum - - del f_sum, g_sum, masked_grid, grid, flag - - # We want x_array and y_array to be 1-D arrays, with the X values initially running from 0 at the left edge of - # the input cutout to the right edge, and the Y values being near the middle of the spectral extraction region. - # So the locations (x_array[i], y_array[i]) should be the spectral trace. - # Near the left and right edges, there might not be any non-zero values in mask_target, so a slice will be - # extracted from both x_array and y_array in order to exclude pixels that are not within the extraction region. - if self.dispaxis == HORIZONTAL: - x_array = np.arange(shape[1], dtype=float) - y_array = spectral_trace - else: - x_array = spectral_trace - y_array = np.arange(shape[0], dtype=float) - - # Trim off the ends, if there's no data there. - # Save trim_slc. - mask = np.where(n_target > 0.) - - if len(mask[0]) > 0: - trim_slc = slice(mask[0][0], mask[0][-1] + 1) - temp_flux = temp_flux[trim_slc] - background = background[trim_slc] - f_var_poisson = f_var_poisson[trim_slc] - f_var_rnoise = f_var_rnoise[trim_slc] - f_var_flat = f_var_flat[trim_slc] - b_var_poisson = b_var_poisson[trim_slc] - b_var_rnoise = b_var_rnoise[trim_slc] - b_var_flat = b_var_flat[trim_slc] - npixels = npixels[trim_slc] - x_array = x_array[trim_slc] - y_array = y_array[trim_slc] - - if got_wavelength: - indx = np.around(x_array).astype(int) - indy = np.around(y_array).astype(int) - indx = np.where(indx < 0, 0, indx) - indx = np.where(indx >= shape[1], shape[1] - 1, indx) - indy = np.where(indy < 0, 0, indy) - indy = np.where(indy >= shape[0], shape[0] - 1, indy) - wavelength = wl_array[indy, indx] - - ra = dec = wcs_wl = None - - if self.wcs is not None: - stuff = self.wcs(x_array, y_array) - ra = stuff[0] - dec = stuff[1] - wcs_wl = stuff[2] - - # We need one right ascension and one declination, representing the direction of pointing. - middle = ra.shape[0] // 2 # ra and dec have same shape - mask = np.isnan(ra) - not_nan = np.logical_not(mask) - - if not_nan[middle]: - log.debug("Using midpoint of spectral trace for RA and Dec.") - - ra = ra[middle] - else: - if np.any(not_nan): - log.warning("Midpoint of coordinate array is NaN; " - "using the average of non-NaN min and max values.") + if return_limits: + return profile, lower_limit, upper_limit + return profile - ra = (np.nanmin(ra) + np.nanmax(ra)) / 2. - else: - log.warning("All right ascension values are NaN; assigning dummy value -999.") - ra = -999. - mask = np.isnan(dec) - not_nan = np.logical_not(mask) +def aperture_center(profile, dispaxis=1, middle_pix=None): + """Determine the nominal center of an aperture. - if not_nan[middle]: - dec = dec[middle] - else: - if np.any(not_nan): - dec = (np.nanmin(dec) + np.nanmax(dec)) / 2. - else: - log.warning("All declination values are NaN; assigning dummy value -999.") - dec = -999. + The center is determined from a weighted average of the pixel + coordinates, where the weights are set by the profile image. - if not got_wavelength: - wavelength = wcs_wl # from wcs, or None + If `middle_pix` is specified, it is expected to be the + dispersion element at which the cross-dispersion center + should be determined. In this case, the profile must contain + some non-zero elements at that dispersion coordinate. - if wavelength is None: - if self.dispaxis == HORIZONTAL: - wavelength = np.arange(shape[1], dtype=float) - else: - wavelength = np.arange(shape[0], dtype=float) - - wavelength = wavelength[trim_slc] - - dq = np.zeros(temp_flux.shape, dtype=np.uint32) - nan_mask = np.isnan(wavelength) - n_nan = nan_mask.sum(dtype=np.intp) - - if n_nan > 0: - log.debug(f"{n_nan} NaNs in wavelength array") - - wavelength, dq, nan_slc = nans_at_endpoints(wavelength, dq) - temp_flux = temp_flux[nan_slc] - background = background[nan_slc] - npixels = npixels[nan_slc] - f_var_poisson = f_var_poisson[nan_slc] - f_var_rnoise = f_var_rnoise[nan_slc] - f_var_flat = f_var_flat[nan_slc] - b_var_poisson = b_var_poisson[nan_slc] - b_var_rnoise = b_var_rnoise[nan_slc] - b_var_flat = b_var_flat[nan_slc] - return (ra, dec, wavelength, - temp_flux, f_var_poisson, f_var_rnoise, f_var_flat, - background, b_var_poisson, b_var_rnoise, b_var_flat, - npixels, dq) - - def match_shape(self, shape): - """Truncate or expand reference image to match the science data. - - Extended summary - ---------------- - The science data may be 2-D or 3-D, but the reference image only - needs to be 2-D. - - Parameters - ---------- - shape : tuple - The shape of the science data. - - Returns - ------- - ndarray, 2-D - This is either the reference image (the data array, not the - complete data model), or an array of the same type, either - larger or smaller than the actual reference image, but matching - the science data array both in size and location on the - detector. - - """ - ref = self.ref_image.data - ref_shape = ref.shape - - if shape == ref_shape: - return ref - - # This is the shape of the last two axes of the science data. - buf = np.zeros((shape[-2], shape[-1]), dtype=ref.dtype) - y_max = min(shape[-2], ref_shape[0]) - x_max = min(shape[-1], ref_shape[1]) - slc0 = slice(0, y_max) - slc1 = slice(0, x_max) - buf[slc0, slc1] = ref[slc0, slc1].copy() - - return buf - - @staticmethod - def separate_target_and_background(ref): - """Create masks for target and background. - - Parameters - ---------- - ref : ndarray, 2-D - This is the reference image as returned by `match_shape`, - i.e. it might be a subset of the original reference image. - - Returns - ------- - mask_target : ndarray, 2-D - This is an array of the same type and shape as the science - image, but with values of only 0 or 1. A value of 1 indicates - that the corresponding pixel in the science data array should - be included when adding up values to make the 1-D spectrum, - and a value of 0 means that it should not be included. - - mask_bkg : ndarray, 2-D, or None. - This is like `mask_target` but for background regions. - A negative value in the reference image flags a pixel that - should be included in the background region(s). If there is - no pixel in the reference image with a negative value, - `mask_bkg` will be set to None. - - """ - mask_target = np.where(ref > 0., 1., 0.) - mask_bkg = None - - if np.any(ref < 0.): - mask_bkg = np.where(ref < 0., 1., 0.) - - return mask_target, mask_bkg - - -def run_extract1d( - input_model, - extract_ref_name, - apcorr_ref_name, - smoothing_length, - bkg_fit, - bkg_order, - bkg_sigma_clip, - log_increment, - subtract_background, - use_source_posn, - center_xy, - ifu_autocen, - ifu_rfcorr, - ifu_set_srctype, - ifu_rscale, - ifu_covar_scale, - was_source_model, -): - """Extract 1-D spectra. - - This just reads the reference files (if any) and calls do_extract1d. + If `dispaxis` is 1 (the default), the return values are in (y,x) + order. Otherwise, the return values are in (x,y) order. Parameters ---------- - input_model : data model - The input science model. - - extract_ref_name : str - The name of the extract1d reference file, or "N/A". - - smoothing_length : int or None - Width of a boxcar function for smoothing the background regions. - - bkg_fit : str - Type of fitting to apply to background values in each column - (or row, if the dispersion is vertical). - - bkg_order : int or None - Polynomial order for fitting to each column (or row, if the - dispersion is vertical) of background. Only used if `bkg_fit` - is `poly`. - - bkg_sigma_clip : float - Sigma clipping value to use on background to remove noise/outliers - - log_increment : int - if `log_increment` is greater than 0 and the input data are - multi-integration, a message will be written to the log every - `log_increment` integrations. - - subtract_background : bool or None - User supplied flag indicating whether the background should be - subtracted. - If None, the value in the extract_1d reference file will be used. - If not None, this parameter overrides the value in the - extract_1d reference file. - - use_source_posn : bool or None - If True, the target and background positions specified in the - reference file (or the default position, if there is no reference - file) will be shifted to account for source position offset. - - center_xy : int or None - A list of 2 pixel coordinate values at which to place the center - of the extraction aperture for IFU data, overriding any centering - done by the step. Two values, in x,y order, are used for extraction - from IFU cubes. Default is None. - - ifu_autocen : bool - Switch to turn on auto-centering for point source spectral extraction - in IFU mode. Default is False. - - ifu_rfcorr : bool - Switch to select whether or not to apply a 1d residual fringe correction - for MIRI MRS IFU spectra. Default is False. - - ifu_set_srctype : str - For MIRI MRS IFU data override srctype and set it to either POINT or EXTENDED. - - ifu_rscale: float - For MRS IFU data a value for changing the extraction radius. The value provided is the number of PSF - FWHMs to use for the extraction radius. Values accepted are between 0.5 to 3.0. The - default extraction size is set to 2 * FWHM. Values below 2 will result in a smaller - radius, a value of 2 results in no change to the radius and a value above 2 results in a larger - extraction radius. - - ifu_covar_scale : float - Scaling factor by which to multiply the ERR values in extracted spectra to account - for covariance between adjacent spaxels in the IFU data cube. - - was_source_model : bool - True if and only if `input_model` is actually one SlitModel - obtained by iterating over a SourceModelContainer. The default - is False. - - apcorr_ref_name : str - Name of the APCORR reference file. Default is None + profile : ndarray + Pixel weights for the aperture. + dispaxis : int, optional + If 1, the dispersion axis is horizontal. Otherwise, + the dispersion axis is vertical. + middle_pix : int or None + Index value for the center of the slit along the + dispersion axis. If specified, it is returned as + `spec_center`. Returns ------- - output_model : data model - A new MultiSpecModel containing the extracted spectra. - + slit_center : float + Index value for the center of the aperture along the + cross-dispersion axis. + spec_center : float + Index value for the center of the aperture along the + dispersion axis. """ - # Read and interpret the extract1d reference file. - try: - ref_dict = open_extract1d_ref(extract_ref_name, input_model.meta.exposure.type) - except AttributeError: # Input is a ModelContainer of some type - ref_dict = open_extract1d_ref(extract_ref_name, input_model[0].meta.exposure.type) - - apcorr_ref_model = None + weights = profile.copy() + weights[weights <= 0] = 0.0 + + yidx, xidx = np.mgrid[:profile.shape[0], :profile.shape[1]] + if dispaxis == HORIZONTAL: + spec_idx = xidx + if middle_pix is not None: + slit_idx = yidx[:, middle_pix] + weights = weights[:, middle_pix] + else: + slit_idx = yidx + else: + spec_idx = yidx + if middle_pix is not None: + slit_idx = xidx[middle_pix, :] + weights = weights[middle_pix, :] + else: + slit_idx = xidx - if apcorr_ref_name is not None and apcorr_ref_name != 'N/A': - try: - apcorr_ref_model = open_apcorr_ref(apcorr_ref_name, input_model.meta.exposure.type) - except AttributeError: # SourceModelContainers don't have exposure nodes - apcorr_ref_model = open_apcorr_ref(apcorr_ref_name, input_model[0].meta.exposure.type) - - # This item is a flag to let us know that do_extract1d was called from run_extract1d; that is, we don't expect this - # key to be present in ref_dict if do_extract1d was called directly. - # If this key is not in ref_dict, or if it is but it's True, then we'll set S_EXTR1D to 'COMPLETE'. - if ref_dict is not None: - ref_dict['need_to_set_to_complete'] = False - - output_model = do_extract1d( - input_model, - ref_dict, - apcorr_ref_model, - smoothing_length, - bkg_fit, - bkg_order, - bkg_sigma_clip, - log_increment, - subtract_background, - use_source_posn, - center_xy, - ifu_autocen, - ifu_rfcorr, - ifu_set_srctype, - ifu_rscale, - ifu_covar_scale, - was_source_model, - ) + if np.sum(weights) == 0: + weights[:] = 1 - if apcorr_ref_model is not None: - apcorr_ref_model.close() + slit_center = np.average(slit_idx, weights=weights) + if middle_pix is not None: + spec_center = middle_pix + else: + spec_center = np.average(spec_idx, weights=weights) - # Remove target.source_type from the output model, so that it - # doesn't force creation of an empty SCI extension in the output - # x1d product just to hold this keyword. - output_model.meta.target.source_type = None + # if dispaxis == 1 (default), this returns center_x, center_y + return slit_center, spec_center - return output_model +def location_from_wcs(input_model, slit): + """Get the cross-dispersion location of the spectrum, based on the WCS. -def ref_dict_sanity_check(ref_dict): - """Check for required entries. + None values will be returned if there was insufficient information + available, e.g. if the wavelength attribute or wcs function is not + defined. Parameters ---------- - ref_dict : dict or None - The contents of the extract1d reference file. + input_model : DataModel + The input science model containing metadata information. + + slit : DataModel or None + One slit from a MultiSlitModel (or similar), or None. + The WCS and target coordinates will be retrieved from `slit` + unless `slit` is None. In that case, they will be retrieved + from `input_model`. Returns ------- - ref_dict : dict or None - + middle : int or None + Pixel coordinate in the dispersion direction within the 2-D + cutout (or the entire input image) at the middle of the WCS + bounding box. This is the point at which to determine the + nominal extraction location, in case it varies along the + spectrum. The offset will then be the difference between + `location` (below) and the nominal location. + + middle_wl : float or None + The wavelength at pixel `middle`. + + location : float or None + Pixel coordinate in the cross-dispersion direction within the + spectral image that is at the planned target location. + The spectral extraction region should be centered here. """ - if ref_dict is None: - return ref_dict - - if 'ref_file_type' not in ref_dict: # We can make an educated guess as to what this must be. - if 'ref_model' in ref_dict: - log.info("Assuming extract1d reference file type is image") - ref_dict['ref_file_type'] = FILE_TYPE_IMAGE - else: - log.info("Assuming extract1d reference file type is JSON") - ref_dict['ref_file_type'] = FILE_TYPE_JSON - - if 'apertures' not in ref_dict: - raise RuntimeError("Key 'apertures' must be present in the extract1d reference file") - - for aper in ref_dict['apertures']: - if 'id' not in aper: - log.warning(f"Key 'id' not found in aperture {aper} in extract1d reference file") - - return ref_dict - - -def do_extract1d( - input_model, - extract_ref_dict, - apcorr_ref_model = None, - smoothing_length = None, - bkg_fit = "poly", - bkg_order = None, - bkg_sigma_clip = 0, - log_increment = 50, - subtract_background = None, - use_source_posn = None, - center_xy = None, - ifu_autocen = None, - ifu_rfcorr = None, - ifu_set_srctype = None, - ifu_rscale = None, - ifu_covar_scale = 1.0, - was_source_model = False -): - """Extract 1-D spectra. - - In the pipeline, this function would be called by run_extract1d. - This exists as a separate function to allow a user to call this step - in a Python script, passing in a dictionary of parameters in order to - bypass reading reference files. - - Parameters - ---------- - input_model : data model - The input science model. - - extract_ref_dict : dict, or None - The contents of the extract1d reference file, or None in order to use - default values. If `ref_dict` is not None, use key 'ref_file_type' - to specify whether the parameters are those that could be read - from a JSON-format reference file - (i.e. ref_dict['ref_file_type'] = "JSON") - from a asdf-format reference file - (i.e. ref_dict['ref_file_type'] = "ASDF") - or parameters relevant for a reference image - (i.e. ref_dict['ref_file_type'] = "IMAGE"). - - smoothing_length : int or None - Width of a boxcar function for smoothing the background regions. - - bkg_fit : str - Type of fitting to perform on background values in each column - (or row, if the dispersion is vertical). - - bkg_order : int or None - Polynomial order for fitting to each column (or row, if the - dispersion is vertical) of background. - - bkg_sigma_clip : float - Sigma clipping value to use on background to remove noise/outliers - - log_increment : int - if `log_increment` is greater than 0 and the input data are - multi-integration, a message will be written to the log every - `log_increment` integrations. - - subtract_background : bool or None - User supplied flag indicating whether the background should be - subtracted. - If None, the value in the extract_1d reference file will be used. - If not None, this parameter overrides the value in the - extract_1d reference file. - - use_source_posn : bool or None - If True, the target and background positions specified in the - reference file (or the default position, if there is no reference - file) will be shifted to account for source position offset. - - center_xy : int or None - A list of 2 pixel coordinate values at which to place the center - of the IFU extraction aperture, overriding any centering done by the step. - Two values, in x,y order, are used for extraction from IFU cubes. - Default is None. - - ifu_autocen : bool - Switch to turn on auto-centering for point source spectral extraction - in IFU mode. Default is False. - - ifu_rfcorr : bool - Switch to select whether or not to apply a 1d residual fringe correction - for MIRI MRS IFU spectra. Default is False. - - ifu_set_srctype : str - For MIRI MRS IFU data override srctype and set it to either POINT or EXTENDED. - - ifu_rscale: float - For MRS IFU data a value for changing the extraction radius. The value provided is the number of PSF - FWHMs to use for the extraction radius. Values accepted are between 0.5 to 3.0. The - default extraction size is set to 2 * FWHM. Values below 2 will result in a smaller - radius, a value of 2 results in no change to the radius and a value above 2 results in a larger - extraction radius. - - ifu_covar_scale : float - Scaling factor by which to multiply the ERR values in extracted spectra to account - for covariance between adjacent spaxels in the IFU data cube. - - was_source_model : bool - True if and only if `input_model` is actually one SlitModel - obtained by iterating over a SourceModelContainer. The default - is False. - - apcorr_ref_model : `~fits.FITS_rec`, datamodel or None - Table of aperture correction values from the APCORR reference file. - - Returns - ------- - output_model : data model - A new MultiSpecModel containing the extracted spectra. - """ - - extract_ref_dict = ref_dict_sanity_check(extract_ref_dict) - - if isinstance(input_model, SourceModelContainer): - # log.debug('Input is a SourceModelContainer') - was_source_model = True - - # Set "meta_source" to either the first model in a container, or the individual input model, for convenience - # of retrieving meta attributes in subsequent statements - if was_source_model: - meta_source = input_model[0] + if slit is not None: + wcs_source = slit else: - meta_source = input_model - - # Setup the output model - output_model = datamodels.MultiSpecModel() - - if hasattr(meta_source, "int_times"): - output_model.int_times = meta_source.int_times.copy() - - output_model.update(meta_source, only='PRIMARY') - - # This will be relevant if we're asked to extract a spectrum and the spectral order is zero. - # That's only OK if the disperser is a prism. - prism_mode = is_prism(meta_source) - exp_type = meta_source.meta.exposure.type - - # use_source_posn doesn't apply to WFSS, so turn it off if it's currently on - if use_source_posn: - if exp_type in WFSS_EXPTYPES: - use_source_posn = False - log.warning( - f"Correcting for source position is not supported for exp_type = " - f"{meta_source.meta.exposure.type}, so use_source_posn will be set to False", - ) - - # Handle inputs that contain one or more slit models - if was_source_model or isinstance(input_model, datamodels.MultiSlitModel): - - is_multiple_slits = True - if was_source_model: # SourceContainer has a single list of SlitModels - log.debug("Input is a Source Model.") - if isinstance(input_model, datamodels.SlitModel): - # If input is a single SlitModel, as opposed to a list of SlitModels, - # put it into a list so that it's iterable later on - log.debug("Input SourceContainer holds one SlitModel.") - slits = [input_model] - else: - log.debug("Input SourceContainer holds a list of SlitModels.") - slits = input_model - - # The subsequent work on data uses the individual SlitModels, but there are many places where meta - # attributes are retrieved from input_model, so set this to allow that to work. - if not isinstance(input_model, datamodels.SlitModel): - input_model = input_model[0] - - elif isinstance(input_model, datamodels.MultiSlitModel): # A simple MultiSlitModel, not in a container - log.debug("Input is a MultiSlitModel.") - slits = input_model.slits - - # Save original use_source_posn value, because it can get - # toggled within the following loop over slits - save_use_source_posn = use_source_posn - - for slit in slits: # Loop over the slits in the input model - log.info(f'Working on slit {slit.name}') - log.debug(f'Slit is of type {type(slit)}') - - slitname = slit.name - prev_offset = OFFSET_NOT_ASSIGNED_YET - use_source_posn = save_use_source_posn # restore original value - - if np.size(slit.data) <= 0: - log.info(f'No data for slit {slit.name}, skipping ...') - continue - - sp_order = get_spectral_order(slit) - if sp_order == 0 and not prism_mode: - log.info("Spectral order 0 is a direct image, skipping ...") - continue - - try: - output_model = create_extraction( - extract_ref_dict, slit, slitname, sp_order, - smoothing_length, bkg_fit, bkg_order, use_source_posn, - prev_offset, exp_type, subtract_background, input_model, - output_model, apcorr_ref_model, log_increment, - is_multiple_slits - ) - except ContinueError: - continue - - else: - # Define source of metadata - slit = None - is_multiple_slits = False - - # These default values for slitname are not really slit names, and slitname may be assigned a better value - # below, in the sections for input_model being an ImageModel or a SlitModel. - slitname = input_model.meta.exposure.type - - if slitname is None: - slitname = ANY - - if slitname == 'NIS_SOSS': - slitname = input_model.meta.subarray.name - - # Loop over these spectral order numbers. - if input_model.meta.exposure.type == "NIS_SOSS": - # This list of spectral order numbers may need to be assigned differently for other exposure types. - spectral_order_list = [1, 2, 3] + wcs_source = input_model + wcs = wcs_source.meta.wcs + dispaxis = wcs_source.meta.wcsinfo.dispersion_direction + + bb = wcs.bounding_box # ((x0, x1), (y0, y1)) + if bb is None: + if slit is None: + shape = input_model.data.shape else: - spectral_order_list = ["not set yet"] # For this case, we'll call get_spectral_order to get the order. - - if isinstance(input_model, datamodels.ImageModel): - if hasattr(input_model, "name"): - slitname = input_model.name - - prev_offset = OFFSET_NOT_ASSIGNED_YET - - for sp_order in spectral_order_list: - if sp_order == "not set yet": - sp_order = get_spectral_order(input_model) - - if sp_order == 0 and not prism_mode: - log.info("Spectral order 0 is a direct image, skipping ...") - continue - - log.info(f'Processing spectral order {sp_order}') - - try: - output_model = create_extraction( - extract_ref_dict, slit, slitname, sp_order, - smoothing_length, bkg_fit, bkg_order, use_source_posn, - prev_offset, exp_type, subtract_background, input_model, - output_model, apcorr_ref_model, log_increment, - is_multiple_slits - ) - except ContinueError: - continue - - elif isinstance(input_model, (datamodels.CubeModel, datamodels.SlitModel)): - # SOSS uses this branch in both calspec2 and caltso3, where in both cases - # the input is a single CubeModel, because caltso3 loops over exposures - - # This branch will be invoked for inputs that are a CubeModel, which typically includes - # NIRSpec BrightObj (fixed slit) and NIRISS SOSS modes, as well as inputs that are a - # single SlitModel, which typically includes data from a single resampled/combined slit - # instance from level-3 processing of NIRSpec fixed slits and MOS modes. - - # Replace the default value for slitname with a more accurate value, if possible. - # For NRS_BRIGHTOBJ, the slit name comes from the slit model info - if input_model.meta.exposure.type == 'NRS_BRIGHTOBJ' and hasattr(input_model, "name"): - slitname = input_model.name - - # For NRS_FIXEDSLIT, the slit name comes from the FXD_SLIT keyword in the model meta - if input_model.meta.exposure.type == 'NRS_FIXEDSLIT': - if hasattr(input_model, "name") and input_model.name is not None: - slitname = input_model.name - else: - slitname = input_model.meta.instrument.fixed_slit - - # Loop over all spectral orders available for extraction - prev_offset = OFFSET_NOT_ASSIGNED_YET - for sp_order in spectral_order_list: - if sp_order == "not set yet": - sp_order = get_spectral_order(input_model) - - if sp_order == 0 and not prism_mode: - log.info("Spectral order 0 is a direct image, skipping ...") - continue - - log.info(f'Processing spectral order {sp_order}') - - try: - output_model = create_extraction( - extract_ref_dict, slit, slitname, sp_order, - smoothing_length, bkg_fit, bkg_order, use_source_posn, - prev_offset, exp_type, subtract_background, input_model, - output_model, apcorr_ref_model, log_increment, - is_multiple_slits - ) - except ContinueError: - continue - - elif isinstance(input_model, datamodels.IFUCubeModel): - try: - source_type = input_model.meta.target.source_type - except AttributeError: - source_type = "UNKNOWN" - - if source_type is None: - source_type = "UNKNOWN" - - if ifu_set_srctype is not None and input_model.meta.exposure.type == 'MIR_MRS': - source_type = ifu_set_srctype - log.info(f"Overriding source type and setting it to = {ifu_set_srctype}") - output_model = ifu.ifu_extract1d( - input_model, extract_ref_dict, source_type, subtract_background, - bkg_sigma_clip, apcorr_ref_model, center_xy, ifu_autocen, ifu_rfcorr, ifu_rscale, ifu_covar_scale - ) - + shape = slit.data.shape + + bb = wcs_bbox_from_shape(shape) + + if dispaxis == HORIZONTAL: + # Width (height) in the cross-dispersion direction, from the start of + # the 2-D cutout (or of the full image) to the upper limit of the bounding box. + # This may be smaller than the full width of the image, but it's all we + # need to consider. + xd_width = int(round(bb[1][1])) # must be an int + middle = int((bb[0][0] + bb[0][1]) / 2.) # Middle of the bounding_box in the dispersion direction. + x = np.empty(xd_width, dtype=np.float64) + x[:] = float(middle) + y = np.arange(xd_width, dtype=np.float64) + lower = bb[1][0] + upper = bb[1][1] + else: # dispaxis = VERTICAL + xd_width = int(round(bb[0][1])) # Cross-dispersion total width of bounding box; must be an int + middle = int((bb[1][0] + bb[1][1]) / 2.) # Mid-point of width along dispersion direction + x = np.arange(xd_width, dtype=np.float64) # 1-D vector of cross-dispersion (x) pixel indices + y = np.empty(xd_width, dtype=np.float64) # 1-D vector all set to middle y index + y[:] = float(middle) + + # lower and upper range in cross-dispersion direction + lower = bb[0][0] + upper = bb[0][1] + + # We need transform[2], a 1-D array of wavelengths crossing the spectrum + # near its middle. + fwd_transform = wcs(x, y) + middle_wl = np.nanmean(fwd_transform[2]) + + exp_type = input_model.meta.exposure.type + if exp_type in ['NRS_FIXEDSLIT', 'NRS_MSASPEC', 'NRS_BRIGHTOBJ']: + if slit is None: + xpos = input_model.source_xpos + ypos = input_model.source_ypos else: - log.error("The input file is not supported for this step.") - raise RuntimeError("Can't extract a spectrum from this file.") - - # Copy the integration time information from the INT_TIMES table to keywords in the output file. - if pipe_utils.is_tso(input_model): - populate_time_keywords(input_model, output_model) - else: - log.debug("Not copying from the INT_TIMES table because this is not a TSO exposure.") - if hasattr(output_model, "int_times"): - del output_model.int_times - - output_model.meta.wcs = None # See output_model.spec[i].meta.wcs instead. - - # If the extract1d reference file is an image, explicitly close it. - if extract_ref_dict is not None and 'ref_model' in extract_ref_dict: - extract_ref_dict['ref_model'].close() - - if (extract_ref_dict is None - or 'need_to_set_to_complete' not in extract_ref_dict - or extract_ref_dict['need_to_set_to_complete']): - output_model.meta.cal_step.extract_1d = 'COMPLETE' - - return output_model - - -def populate_time_keywords(input_model, output_model): - """Copy the integration times keywords to header keywords. - - Parameters - ---------- - input_model : data model - The input science model. - - output_model : data model - The output science model. This may be modified in-place. - - """ - nints = input_model.meta.exposure.nints - int_start = input_model.meta.exposure.integration_start - - if hasattr(input_model, 'data'): - shape = input_model.data.shape - - if len(shape) == 2: - num_integ = 1 - else: # len(shape) == 3 - num_integ = shape[0] - else: # e.g. MultiSlit data - num_integ = 1 - - # This assumes that the spec attribute of output_model has already been created, and spectra have been appended. - n_output_spec = len(output_model.spec) - - # num_j is the number of spectra per integration, e.g. the number of fixed-slit spectra, MSA spectra, or different - # spectral orders; num_integ is the number of integrations. - # The total number of output spectra is n_output_spec = num_integ * num_j - num_j = n_output_spec // num_integ - - if n_output_spec != num_j * num_integ: # sanity check - log.warning( - f"populate_time_keywords: Don't understand n_output_spec = {n_output_spec}, num_j = {num_j}, num_integ = " - f"{num_integ}" - ) - else: - log.debug( - f"Number of output spectra = {n_output_spec}; number of spectra for each integration = {num_j}; " - f"number of integrations = {num_integ}" - ) - - if int_start is None: - log.warning("INTSTART not found; assuming a value of 1.") - int_start = 1 - - int_start -= 1 # zero indexed - int_end = input_model.meta.exposure.integration_end - - if int_end is None: - log.warning(f"INTEND not found; assuming a value of {nints}.") - int_end = nints - - int_end -= 1 # zero indexed - - if nints > 1: - num_integrations = int_end - int_start + 1 - else: - num_integrations = 1 - - if hasattr(input_model, 'int_times') and input_model.int_times is not None: - nrows = len(input_model.int_times) - else: - nrows = 0 - - if nrows < 1: - log.warning("There is no INT_TIMES table in the input file - " - "Making best guess on integration numbers.") - for j in range(num_j): # for each spectrum or order - for k in range(num_integ): # for each integration - output_model.spec[(j * num_integ) + k].int_num = k + 1 # set int_num to (k+1) - 1-indexed integration - return - - # If we have a single plane (e.g. ImageModel or MultiSlitModel), we will only populate the keywords if the - # corresponding uncal file had one integration. - # If the data were or might have been segmented, we use the first and last integration numbers to determine whether - # the data were in fact averaged over integrations, and if so, we should not populate the int_times-related header - # keywords. - skip = False # initial value - - if isinstance(input_model, (datamodels.MultiSlitModel, datamodels.ImageModel)): - if num_integrations > 1: - log.warning("Not using INT_TIMES table because the data have been averaged over integrations.") - skip = True - elif isinstance(input_model, (datamodels.CubeModel, datamodels.SlitModel)): - shape = input_model.data.shape - - if len(shape) == 2 and num_integrations > 1: - log.warning("Not using INT_TIMES table because the data have been averaged over integrations.") - skip = True - elif len(shape) != 3 or shape[0] > nrows: - # Later, we'll check that the integration_number column actually - # has a row corresponding to every integration in the input. - log.warning( - "Not using INT_TIMES table because the data shape is not consistent with the number of table rows." - ) - skip = True - elif isinstance(input_model, datamodels.IFUCubeModel): - log.warning("The INT_TIMES table will be ignored for IFU data.") - skip = True - - if skip: - return - - int_num = input_model.int_times['integration_number'] - start_time_mjd = input_model.int_times['int_start_MJD_UTC'] - mid_time_mjd = input_model.int_times['int_mid_MJD_UTC'] - end_time_mjd = input_model.int_times['int_end_MJD_UTC'] - start_tdb = input_model.int_times['int_start_BJD_TDB'] - mid_tdb = input_model.int_times['int_mid_BJD_TDB'] - end_tdb = input_model.int_times['int_end_BJD_TDB'] - - data_range = (int_start, int_end) # Inclusive range of integration numbers in the input data, zero indexed. - - # Inclusive range of integration numbers in the INT_TIMES table, zero indexed. - table_range = (int_num[0] - 1, int_num[-1] - 1) - offset = data_range[0] - table_range[0] - - if data_range[0] < table_range[0] or data_range[1] > table_range[1]: - log.warning("Not using the INT_TIMES table because it does not include rows for all integrations in the data.") - return - - log.debug("TSO data, so copying times from the INT_TIMES table.") - - n = 0 # Counter for spectra in output_model. + xpos = slit.source_xpos + ypos = slit.source_ypos - for k in range(num_integ): # for each spectrum or order - for j in range(num_j): # for each integration - row = k + offset - spec = output_model.spec[n] # n is incremented below - spec.int_num = int_num[row] - spec.time_scale = "UTC" - spec.start_time_mjd = start_time_mjd[row] - spec.mid_time_mjd = mid_time_mjd[row] - spec.end_time_mjd = end_time_mjd[row] - spec.start_tdb = start_tdb[row] - spec.mid_tdb = mid_tdb[row] - spec.end_tdb = end_tdb[row] - n += 1 - - -def get_spectral_order(slit): - """Get the spectral order number. - - Parameters - ---------- - slit : SlitModel object - One slit from an input MultiSlitModel or similar. - - Returns - ------- - int - Spectral order number for `slit`. If no information about spectral - order is available in `wcsinfo`, a default value of 1 will be - returned. - - """ - if hasattr(slit.meta, 'wcsinfo'): - sp_order = slit.meta.wcsinfo.spectral_order - - if sp_order is None: - log.warning("spectral_order is None; using 1") - sp_order = 1 - else: - log.warning("slit.meta doesn't have attribute wcsinfo; setting spectral order to 1") - sp_order = 1 - - return sp_order - - -def is_prism(input_model): - """Determine whether the current observing mode used a prism. - - Extended summary - ---------------- - The reason for this test is so we can skip spectral extraction if the - spectral order is zero and the exposure was not made using a prism. - In this context, therefore, a grism is not considered to be a prism. - - Parameters - ---------- - input_model : data model - The input science model. - - Returns - ------- - bool - True if the exposure used a prism; False otherwise. - - """ - instrument = input_model.meta.instrument.name - - if instrument is None: - return False - - instrument_filter = input_model.meta.instrument.filter - - if instrument_filter is None: - instrument_filter = "NONE" - else: - instrument_filter = instrument_filter.upper() - - grating = input_model.meta.instrument.grating - - if grating is None: - grating = "NONE" - else: - grating = grating.upper() - - prism_mode = False - - if ((instrument == "MIRI" and instrument_filter.find("P750L") >= 0) or - (instrument == "NIRSPEC" and grating.find("PRISM") >= 0)): - prism_mode = True - - return prism_mode - - -def copy_keyword_info(slit, slitname, spec): - """Copy metadata from the input to the output spectrum. - - Parameters - ---------- - slit : A SlitModel object - Metadata will be copied from the input `slit` to output `spec`. - - slitname : str or None - The name of the slit. - - spec : One element of MultiSpecModel.spec - Metadata attributes will be updated in-place. + slit2det = wcs.get_transform('slit_frame', 'detector') + if 'gwa' in wcs.available_frames: + # Input is not resampled, wavelengths need to be meters + x_y = slit2det(xpos, ypos, middle_wl * 1e-6) + else: + x_y = slit2det(xpos, ypos, middle_wl) + log.info("Using source_xpos and source_ypos to center extraction.") - """ - if slitname is not None and slitname != "ANY": - spec.name = slitname + elif exp_type == 'MIR_LRS-FIXEDSLIT': + try: + if slit is None: + dithra = input_model.meta.dither.dithered_ra + dithdec = input_model.meta.dither.dithered_dec + else: + dithra = slit.meta.dither.dithered_ra + dithdec = slit.meta.dither.dithered_dec + x_y = wcs.backward_transform(dithra, dithdec, middle_wl) + except (AttributeError, TypeError): + log.warning("Dithered pointing location not found in wcsinfo.") + return None, None, None + else: + log.warning(f"Source position cannot be found for EXP_TYPE {exp_type}") + return None, None, None - if hasattr(slit, "slitlet_id"): - spec.slitlet_id = slit.slitlet_id + # location is the XD location of the spectrum: + if dispaxis == HORIZONTAL: + location = x_y[1] + else: + location = x_y[0] - if hasattr(slit, "source_id"): - spec.source_id = slit.source_id + if np.isnan(location): + log.warning('Source position could not be determined from WCS.') + return None, None, None - if hasattr(slit, "source_name") and slit.source_name is not None: - spec.source_name = slit.source_name + # If the target is at the edge of the image or at the edge of the + # non-NaN area, we can't use the WCS to find the + # location of the target spectrum. + if location < lower or location > upper: + log.warning(f"WCS implies the target is at {location:.2f}, which is outside the bounding box,") + log.warning("so we can't get spectrum location using the WCS") + return None, None, None - if hasattr(slit, "source_alias") and slit.source_alias is not None: - spec.source_alias = slit.source_alias + return middle, middle_wl, location - if hasattr(slit, "source_type") and slit.source_type is not None: - spec.source_type = slit.source_type - if hasattr(slit, "stellarity") and slit.stellarity is not None: - spec.stellarity = slit.stellarity +def shift_by_source_location(location, nominal_location, extract_params): + """Shift the nominal extraction parameters by the source location. - if hasattr(slit, "source_xpos"): - spec.source_xpos = slit.source_xpos + The offset applied is `location` - `nominal_location`, along + the cross-dispersion direction. - if hasattr(slit, "source_ypos"): - spec.source_ypos = slit.source_ypos + Start, stop, and polynomial coefficient values for source and + background are updated in place in the `extract_params` dictionary. - if hasattr(slit, "source_ra"): - spec.source_ra = slit.source_ra + Parameters + ---------- + location : float + The source location in the cross-dispersion direction + at which to center the extraction aperture. + nominal_location : float + The center of the nominal extraction aperture, in the + cross-dispersion direction, according to the extraction + parameters. + extract_params : dict + Extraction parameters to update, as created by + `get_extraction_parameters`, and corresponding to the + specified nominal location. - if hasattr(slit, "source_dec"): - spec.source_dec = slit.source_dec + """ - if hasattr(slit, "shutter_state"): - spec.shutter_state = slit.shutter_state + # Get the center of the nominal aperture + offset = location - nominal_location + log.info(f"Nominal location is {nominal_location:.2f}, " + f"so offset is {offset:.2f} pixels") + + # Shift aperture limits by the difference between the + # source location and the nominal center + coeff_params = ['src_coeff', 'bkg_coeff'] + for params in coeff_params: + if extract_params[params] is not None: + for coeff_list in extract_params[params]: + coeff_list[0] += offset + if extract_params['dispaxis'] == HORIZONTAL: + start_stop_params = ['ystart', 'ystop'] + else: + start_stop_params = ['xstart', 'xstop'] + for params in start_stop_params: + if extract_params[params] is not None: + extract_params[params] += offset -def extract_one_slit(input_model, slit, integ, prev_offset, extract_params): - """Extract data for one slit, or spectral order, or plane. +def define_aperture(input_model, slit, extract_params, exp_type): + """Define an extraction aperture from input parameters. Parameters ---------- - input_model : data model - The input science model. + input_model : DataModel + The input science model containing metadata information. + slit : DataModel or None + One slit from a MultiSlitModel (or similar), or None. + The spectral image and WCS information will be retrieved from `slit` + unless `slit` is None. In that case, they will be retrieved + from `input_model`. + extract_params : dict + Extraction parameters, as created by `get_extraction_parameters`. + exp_type : str + Exposure type for the input data. + + Returns + ------- + ra : float + A representative RA value for the source centered in the + aperture. + dec : float + A representative Dec value for the source centered in the + aperture. + wavelength : ndarray of float + The 1D wavelength array, matching the dispersion dimension of + the input spectral image, corresponding to the aperture. May + contain NaNs for invalid dispersion elements. + profile : ndarray of float + A 2D image containing pixel weights for the extraction aperture, + matching the dimensions of the input spectral image. Values + are between 0.0 (pixel not included in the extraction aperture) + and 1.0 (pixel fully included in the aperture). + bg_profile : ndarray of float or None + If background regions are specified in `extract_params['bkg_coeff']`, + and `extract_params['subtract_background']` is True, then + `bg_profile` is a 2D image containing pixel weights for background + regions, to be fit during extraction. Otherwise, `bg_profile` is + None. + limits : tuple of float + Index limit values for the aperture, returned as (lower_limit, upper_limit, + left_limit, right_limit). Upper/lower limits are along the + cross-dispersion axis. Left/right limits are along the dispersion axis. + All limits are inclusive and start at zero index value. + + """ + if slit is None: + data_model = input_model + else: + data_model = slit + data_shape = data_model.data.shape[-2:] + + # Get a wavelength array for the data + wl_array = get_wavelengths(data_model, exp_type, extract_params['spectral_order']) + + # Shift aperture definitions by source position if needed + # Extract parameters are updated in place + if extract_params['use_source_posn']: + # Source location from WCS + middle_pix, middle_wl, location = location_from_wcs(input_model, slit) + + if location is not None: + log.info(f"Computed source location is {location:.2f}, " + f"at pixel {middle_pix}, wavelength {middle_wl:.2f}") + + # Nominal location from extract params + located middle + nominal_profile = box_profile(data_shape, extract_params, wl_array, + label='nominal aperture') + nominal_location, _ = aperture_center( + nominal_profile, extract_params['dispaxis'], middle_pix=middle_pix) + + # Offet extract parameters by location - nominal + shift_by_source_location(location, nominal_location, extract_params) + + # Make a spatial profile, including source shifts if necessary + profile, lower_limit, upper_limit = box_profile( + data_shape, extract_params, wl_array, return_limits=True) + + # Make sure profile weights are zero where wavelengths are invalid + profile[~np.isfinite(wl_array)] = 0.0 + + # Get the effective left and right limits from the profile weights + nonzero_weight = np.where(np.sum(profile, axis=extract_params['dispaxis'] - 1) > 0)[0] + if len(nonzero_weight) > 0: + left_limit = nonzero_weight[0] + right_limit = nonzero_weight[-1] + else: + left_limit = None + right_limit = None + + # Make a background profile if necessary + # (will also include source shifts) + if (extract_params['subtract_background'] + and extract_params['bkg_coeff'] is not None): + bg_profile = box_profile(data_shape, extract_params, wl_array, + coefficients='bkg_coeff') + else: + bg_profile = None + + # Get 1D wavelength corresponding to the spatial profile + mask = np.isnan(wl_array) | (profile <= 0) + masked_wl = np.ma.masked_array(wl_array, mask=mask) + masked_weights = np.ma.masked_array(profile, mask=mask) + if extract_params['dispaxis'] == HORIZONTAL: + wavelength = np.average(masked_wl, weights=masked_weights, axis=0).filled(np.nan) + else: + wavelength = np.average(masked_wl, weights=masked_weights, axis=1).filled(np.nan) + + # Get RA and Dec corresponding to the center of the array, + # weighted by the spatial profile + center_y, center_x = aperture_center(profile, 1) + coords = data_model.meta.wcs(center_x, center_y) + if np.any(np.isnan(coords)): + ra = None + dec = None + else: + ra = float(coords[0]) + dec = float(coords[1]) - slit : one slit from a MultiSlitModel (or similar), or None - If slit is None, the data array is input_model.data; otherwise, - the data array is slit.data. - In the former case, if `integ` is zero or larger, the spectrum - will be extracted from the 2-D slice input_model.data[integ]. + # Return limits as a tuple with 4 elements: lower, upper, left, right + limits = (lower_limit, upper_limit, left_limit, right_limit) - integ : int - For the case that input_model is a SlitModel or a CubeModel, - `integ` is the integration number. If the integration number is - not relevant (i.e. the data array is 2-D), `integ` should be -1. + return ra, dec, wavelength, profile, bg_profile, limits - prev_offset : float or str - When extracting from multi-integration data, the source position - offset only needs to be determined once. `prev_offset` is either the - previously computed offset or a value (a string) indicating that - the offset hasn't been computed yet. In the latter case, method - `offset_from_offset` will be called to determine the offset. + +def extract_one_slit(data_model, integration, profile, bg_profile, extract_params): + """Extract data for one slit, or spectral order, or integration. + + Parameters + ---------- + data_model : JWSTDataModel + The input science model. May be a single slit from a MultiSlitModel + (or similar), or a single data type, like an ImageModel, SlitModel, + or CubeModel. + + integration : int + For the case that data_model is a SlitModel or a CubeModel, + `integration` is the integration number. If the integration number is + not relevant (i.e. the data array is 2-D), `integration` should be -1. + + profile : ndarray of float + Spatial profile indicating the aperture location. Must be a + 2D image matching the input, with floating point values between 0 + and 1 assigning a weight to each pixel. 0 means the pixel is not used, + 1 means the pixel is fully included in the aperture. + + bg_profile : ndarray of float or None + Background profile indicating any background regions to use, following + the same format as the spatial profile. Ignored if + extract_params['subtract_background'] is False. extract_params : dict - Parameters read from the extract1d reference file. + Parameters read from the extract1d reference file, as returned by + `get_extract_parameters`. Returns ------- - ra, dec : float - ra and dec are the right ascension and declination respectively - at the nominal center of the slit. - - wavelength : ndarray, 1-D, float64 - The wavelength in micrometers at each pixel. - - temp_flux : ndarray, 1-D, float64 + sum_flux : ndarray, 1-D, float64 The sum of the data values in the extraction region minus the sum of the data values in the background regions (scaled by the ratio of the numbers of pixels), for each pixel. - The data values are in units of surface brightness, so this value - isn't really the flux, it's an intermediate value. Multiply - `temp_flux` by the solid angle of a pixel to get the flux for a - point source (column "flux"). Divide `temp_flux` by `npixels` (to + The data values are usually in units of surface brightness, + so this value isn't the flux, it's an intermediate value. + Multiply `sum_flux` by the solid angle of a pixel to get the flux for a + point source (column "flux"). Divide `sum_flux` by `npixels` (to compute the average) to get the array for the "surf_bright" (surface brightness) output column. - f_var_poisson : ndarray, 1-D - The extracted poisson variance values to go along with the - temp_flux array. - f_var_rnoise : ndarray, 1-D The extracted read noise variance values to go along with the - temp_flux array. + sum_flux array. + + f_var_poisson : ndarray, 1-D + The extracted poisson variance values to go along with the + sum_flux array. f_var_flat : ndarray, 1-D The extracted flat field variance values to go along with the - temp_flux array. + sum_flux array. background : ndarray, 1-D The background count rate that was subtracted from the sum of - the source data values to get `temp_flux`. - - b_var_poisson : ndarray, 1-D - The extracted poisson variance values to go along with the - background array. + the source data values to get `sum_flux`. b_var_rnoise : ndarray, 1-D The extracted read noise variance values to go along with the background array. + b_var_poisson : ndarray, 1-D + The extracted poisson variance values to go along with the + background array. + b_var_flat : ndarray, 1-D The extracted flat field variance values to go along with the background array. npixels : ndarray, 1-D, float64 - The number of pixels that were added together to get `temp_flux`. + The number of pixels that were added together to get `sum_flux`, + including any fractional pixels included via non-integer weights + in the input profile. - dq : ndarray, 1-D, uint32 - The data quality array. - - offset : float - The source position offset in the cross-dispersion direction, either - computed by calling `offset_from_offset` in this function, or - copied from the input `prev_offset`. + scene_model : ndarray, 2-D, float64 + A 2D model of the flux in the spectral image, corresponding to + the extracted aperture. """ - - log_initial_parameters(extract_params) - - try: - exp_type = input_model.meta.exposure.type - except AttributeError: - exp_type = slit.meta.exposure.type - - if integ > -1: - log.info(f"Extracting integration {integ + 1}") - data = input_model.data[integ] - var_poisson = input_model.var_poisson[integ] - var_rnoise = input_model.var_rnoise[integ] - var_flat = input_model.var_flat[integ] - input_dq = input_model.dq[integ] - elif slit is None: - data = input_model.data - var_poisson = input_model.var_poisson - var_rnoise = input_model.var_rnoise - var_flat = input_model.var_flat - input_dq = input_model.dq + # Get the data and variance arrays + if integration > -1: + log.debug(f"Extracting integration {integration + 1}") + data = data_model.data[integration] + var_rnoise = data_model.var_rnoise[integration] + var_poisson = data_model.var_poisson[integration] + var_flat = data_model.var_flat[integration] else: - data = slit.data - var_poisson = slit.var_poisson - var_rnoise = slit.var_rnoise - var_flat = slit.var_flat - input_dq = slit.dq - - # Ensure variance arrays have been populated. If not, zero fill. - if np.shape(var_poisson) != np.shape(data): - var_poisson = np.zeros_like(data) - var_rnoise = np.zeros_like(data) + data = data_model.data + var_rnoise = data_model.var_rnoise + var_poisson = data_model.var_poisson + var_flat = data_model.var_flat - if np.shape(var_flat) != np.shape(data): + # Make sure variances match data + if var_rnoise.shape != data.shape: + var_rnoise = np.zeros_like(data) + if var_poisson.shape != data.shape: + var_poisson = np.zeros_like(data) + if var_flat.shape != data.shape: var_flat = np.zeros_like(data) - if input_dq.size == 0: - input_dq = None - - wl_array = get_wavelengths(input_model if slit is None else slit, exp_type, extract_params['spectral_order']) - data = replace_bad_values(data, input_dq, wl_array) - - if extract_params['ref_file_type'] == FILE_TYPE_IMAGE: # The reference file is an image. - extract_model = ImageExtractModel(input_model=input_model, slit=slit, **extract_params) - ap = None + # Transpose data for extraction + if extract_params['dispaxis'] == HORIZONTAL: + profile_view = profile + bg_profile_view = bg_profile else: - # If there is an extract1d reference file (there doesn't have to be), it's in JSON format. - extract_model = ExtractModel(input_model=input_model, slit=slit, **extract_params) - ap = get_aperture(data.shape, extract_model.wcs, extract_params) - extract_model.update_extraction_limits(ap) - - if extract_model.use_source_posn: - if prev_offset == OFFSET_NOT_ASSIGNED_YET: # Only call this method for the first integration. - offset, locn = extract_model.offset_from_offset(input_model, slit) - - if offset is not None and locn is not None: - log.debug(f"Computed source offset={offset:.2f}, source location={locn:.2f}") - - if not extract_model.use_source_posn: - offset = 0. + data = data.T + profile_view = profile.T + var_rnoise = var_rnoise.T + var_poisson = var_poisson.T + var_flat = var_flat.T + if bg_profile is not None: + bg_profile_view = bg_profile.T else: - offset = prev_offset + bg_profile_view = None + + # Extract spectra from the data + result = extract1d.extract1d(data, [profile_view], var_rnoise, var_poisson, var_flat, + profile_bg=bg_profile_view, + bg_smooth_length=extract_params['smoothing_length'], + fit_bkg=extract_params['subtract_background'], + bkg_fit_type=extract_params['bkg_fit'], + bkg_order=extract_params['bkg_order'], + extraction_type=extract_params['extraction_type']) + + # Extraction routine can return multiple spectra; + # here, we just want the first result + first_result = [] + for r in result[:-1]: + first_result.append(r[0]) + + # The last return value is the 2D model - there is only one, regardless + # of the number of input profiles. It may need to be transposed to match + # the input data. + scene_model = result[-1] + if extract_params['dispaxis'] == HORIZONTAL: + first_result.append(scene_model) else: - offset = 0. - - extract_model.position_correction = offset - - # Add the source position offset to the polynomial coefficients, or shift the reference image - # (depending on the type of reference file). - extract_model.add_position_correction(data.shape) - extract_model.log_extraction_parameters() - extract_model.assign_polynomial_limits() - - # store the extraction values we want to save in dictionary. - # define all the values to be None. If they are not valid for a mode - # they will not be written to fits header. - extraction_values = {} - extraction_values['xstart'] = None - extraction_values['xstop'] = None - extraction_values['ystart'] = None - extraction_values['ystop'] = None - - # Log the extraction limits being used - if integ < 1: - if extract_model.src_coeff is not None: - # Because src_coeff was specified, that will be used instead of xstart/xstop (or ystart/ystop). - if extract_model.dispaxis == HORIZONTAL: - # Only print xstart/xstop, because ystart/ystop are not used - log.info("Using extraction limits: " - f"xstart={extract_model.xstart}, " - f"xstop={extract_model.xstop}, and src_coeff") - extraction_values['xstart'] = extract_model.xstart + 1 - extraction_values['xstop'] = extract_model.xstop + 1 - else: - # Only print ystart/ystop, because xstart/xstop are not used - log.info("Using extraction limits: " - f"ystart={extract_model.ystart}, " - f"ystop={extract_model.ystop}, and src_coeff") - extraction_values['ystart'] = extract_model.ystart + 1 - extraction_values['ystop'] = extract_model.ystop + 1 - else: - # No src_coeff, so print all xstart/xstop and ystart/ystop values - log.info("Using extraction limits: " - f"xstart={extract_model.xstart}, xstop={extract_model.xstop}, " - f"ystart={extract_model.ystart}, ystop={extract_model.ystop}") - if extract_model.xstart is not None: - extraction_values['xstart'] = extract_model.xstart + 1 - if extract_model.xstop is not None: - extraction_values['xstop'] = extract_model.xstop + 1 - if extract_model.ystart is not None: - extraction_values['ystart'] = extract_model.ystart + 1 - if extract_model.ystop is not None: - extraction_values['ystop'] = extract_model.ystop + 1 - if extract_params['subtract_background']: - log.info("with background subtraction") - - ra, dec, wavelength, temp_flux, f_var_poisson, f_var_rnoise, f_var_flat, \ - background, b_var_poisson, b_var_rnoise, b_var_flat, npixels, dq = \ - extract_model.extract(data, var_poisson, var_rnoise, var_flat, - wl_array) - - return (ra, dec, wavelength, temp_flux, f_var_poisson, f_var_rnoise, f_var_flat, - background, b_var_poisson, b_var_rnoise, b_var_flat, npixels, dq, offset, - extraction_values) - - -def replace_bad_values(data, input_dq, wl_array): - """Replace values flagged with DO_NOT_USE or that have NaN wavelengths. + first_result.append(scene_model.T) + return first_result + + +def create_extraction(input_model, slit, output_model, + extract_ref_dict, slitname, sp_order, exp_type, + apcorr_ref_model=None, log_increment=50, + save_profile=False, save_scene_model=False, **kwargs): + """Extract spectra from an input model and append to an output model. + + Input data, specified in the `slit` or `input_model`, should contain data + with consistent wavelengths, target positions, and spectral image cutouts, + suitable for extraction with a shared aperture. Extraction parameters + are determined collectively, then multiple integrations, if present, + are each extracted separately. + + The output model must be a `MultiSpecModel`, created before calling this + function, and passed as `output_model`. It is updated in place, with + new spectral tables appended as they are created. + + The process is: + + 1. Retrieve extraction parameters from `extract_ref_dict` and + set defaults for any missing values as needed. + 2. Define an extraction aperture from the input parameters, as + well as the nominal RA, Dec, and 1D wavelength arrays for + the output spectrum. + 3. Set up an aperture correction to apply to each spectrum, + if `apcorr_ref_model` is not None. + 4. Loop over integrations to extract all spectra. + + For each integration, the extraction process is: + + 1. Extract summed flux and variance values from the aperture. + 2. Compute an average value (from flux / npixels), to be stored + as the surface brightness. + 3. Convert the summed flux to flux density (Jy). + 4. Compute an error spectrum from the square root of the sum + of the variance components. + 5. Set a DQ array, with DO_NOT_USE flags set where the + flux is NaN. + 6. Create a spectral table to contain all extracted values + and store it in a `SpecModel`. + 7. Apply the aperture correction to the spectral table, if + available. + 8. Append the new SpecModel to the `MultiSpecModel` provided + in `output_model`. Parameters ---------- - data : ndarray - The science data array. - - input_dq : ndarray or None - If not None, this will be checked for flag value DO_NOT_USE. The - science data will be set to NaN for every pixel that is flagged - with DO_NOT_USE in `input_dq`. - - wl_array : ndarray, 2-D - Wavelengths corresponding to `data`. For any element of this - array that is NaN, the corresponding element in `data` will be - set to NaN. + input_model : DataModel + Top-level datamodel containing metadata for the input data. + If slit is not specified, `input_model` must also contain the + spectral image(s) to extract, in the `data` attribute. + slit : SlitModel or None + One slit from an input MultiSlitModel or similar. If not None, + `slit `must contain the spectral image(s) to extract, in the `data` + attribute, along with appropriate WCS metadata. + output_model : MultiSpecModel + The output model to append spectral tables to. + extract_ref_dict : dict or None + Extraction parameters read in from the extract_1d reference file, + or None, if there was no reference file. + slitname : str + Slit name for the input data. + sp_order : int + Spectral order for the input data. + exp_type : str + Exposure type for the input data. + apcorr_ref_model : DataModel or None, optional + The aperture correction reference datamodel, containing the + APCORR reference file data. + log_increment : int, optional + If greater than 0 and the input data are multi-integration, a message + will be written to the log every `log_increment` integrations. + save_profile : bool, optional + If True, the spatial profile created for the aperture will be returned + as an ImageModel. If False, the return value is None. + save_scene_model : bool, optional + If True, the flux model created during extraction will be returned + as an ImageModel or CubeModel. If False, the return value is None. + kwargs : dict, optional + Additional options to pass to `get_extract_parameters`. Returns ------- - ndarray - A possibly modified copy of `data`. If no change was made, this - will be a view rather than a copy. Values that are set to NaN - should not be included when doing the 1-D spectral extraction. - + profile_model : ImageModel or None + If `save_profile` is True, the return value is an ImageModel containing + the spatial profile with aperture weights, used in extracting all + integrations. + scene_model : ImageModel, CubeModel, or None + If `save_scene_model` is True, the return value is an ImageModel or CubeModel + matching the input data, containing the flux model generated during + extraction. """ - mask = np.isnan(wl_array) - - if input_dq is not None: - bad_mask = np.bitwise_and(input_dq, dqflags.pixel['DO_NOT_USE']) > 0 - mask = np.logical_or(mask, bad_mask) - - if np.any(mask): - mod_data = data.copy() - mod_data[mask] = np.nan - return mod_data - - return data - - -def nans_at_endpoints(wavelength, dq): - """Flag NaNs in the wavelength array. - - Extended summary - ---------------- - Both input arrays should be 1-D and have the same shape. - If NaNs are present at endpoints of `wavelength`, the arrays will be - trimmed to remove the NaNs. NaNs at interior elements of `wavelength` - will be left in place, but they will be flagged with DO_NOT_USE in the - `dq` array. - - Parameters - ---------- - wavelength : ndarray - Array of wavelengths, possibly containing NaNs. - dq : ndarray - Data quality array. - - Returns - ------- - wavelength, dq : ndarray - The returned `dq` array may have NaNs flagged with DO_NOT_USE, - and both arrays may have been trimmed at either or both ends. - - slc : slice - The slice to be applied to other output arrays to match the modified - shape of the wavelength array. - """ - # The input arrays will not be modified in-place. - new_wl = wavelength.copy() - new_dq = dq.copy() - nelem = wavelength.shape[0] - slc = slice(nelem) - - nan_mask = np.isnan(wavelength) - new_dq[nan_mask] = np.bitwise_or(new_dq[nan_mask], dqflags.pixel['DO_NOT_USE']) - not_nan = np.logical_not(nan_mask) - flag = np.where(not_nan) - - if len(flag[0]) > 0: - n_trimmed = flag[0][0] + nelem - (flag[0][-1] + 1) - - if n_trimmed > 0: - log.debug(f"Output arrays have been trimmed by {n_trimmed} elements") - - slc = slice(flag[0][0], flag[0][-1] + 1) - new_wl = new_wl[slc] - new_dq = new_dq[slc] - else: - new_dq |= dqflags.pixel['DO_NOT_USE'] - - return new_wl, new_dq, slc - - -def create_extraction(extract_ref_dict, - slit, - slitname, - sp_order, - smoothing_length, - bkg_fit, - bkg_order, - use_source_posn, - prev_offset, - exp_type, - subtract_background, - input_model, - output_model, - apcorr_ref_model, - log_increment, - is_multiple_slits -): if slit is None: - meta_source = input_model + data_model = input_model else: - meta_source = slit + data_model = slit + + # Make sure NaNs and DQ flags match up in input + pipe_utils.match_nans_and_flags(data_model) if exp_type in WFSS_EXPTYPES: instrument = input_model.meta.instrument.name else: - instrument = meta_source.meta.instrument.name + instrument = data_model.meta.instrument.name if instrument is not None: instrument = instrument.upper() # We need a flag to indicate whether the photom step has been run. # If it hasn't, we'll copy the count rate to the flux column. - try: - s_photom = input_model.meta.cal_step.photom - except AttributeError: - s_photom = None - + s_photom = input_model.meta.cal_step.photom if s_photom is not None and s_photom.upper() == 'COMPLETE': photom_has_been_run = True flux_units = 'Jy' @@ -3640,9 +1531,9 @@ def create_extraction(extract_ref_dict, sb_var_units = 'DN^2 / s^2' log.warning("The photom step has not been run.") - # Turn off use_source_posn if the source is not POINT - if is_multiple_slits: - source_type = meta_source.source_type + # Get the source type for the data + if slit is not None: + source_type = slit.source_type else: if isinstance(input_model, datamodels.SlitModel): source_type = input_model.source_type @@ -3653,12 +1544,14 @@ def create_extraction(extract_ref_dict, source_type = input_model.meta.target.source_type # Turn off use_source_posn if the source is not POINT - if source_type != 'POINT': - use_source_posn = False - log.info(f"Setting use_source_posn to False for source type {source_type}") + if source_type != 'POINT' or exp_type in WFSS_EXPTYPES: + if kwargs.get('use_source_posn') is None: + kwargs['use_source_posn'] = False + log.info(f"Setting use_source_posn to False for exposure type {exp_type}, " + f"source type {source_type}") if photom_has_been_run: - pixel_solid_angle = meta_source.meta.photometry.pixelarea_steradians + pixel_solid_angle = data_model.meta.photometry.pixelarea_steradians if pixel_solid_angle is None: pixel_solid_angle = 1. log.warning("Pixel area (solid angle) is not populated; the flux will not be correct.") @@ -3666,19 +1559,7 @@ def create_extraction(extract_ref_dict, pixel_solid_angle = 1. # not needed extract_params = get_extract_parameters( - extract_ref_dict, - meta_source, - slitname, - sp_order, - input_model.meta, - smoothing_length, - bkg_fit, - bkg_order, - use_source_posn - ) - - if subtract_background is not None: - extract_params['subtract_background'] = subtract_background + extract_ref_dict, data_model, slitname, sp_order, input_model.meta, **kwargs) if extract_params['match'] == NO_MATCH: log.critical('Missing extraction parameters.') @@ -3687,15 +1568,61 @@ def create_extraction(extract_ref_dict, log.info(f'Spectral order {sp_order} not found, skipping ...') raise ContinueError() - extract_params['dispaxis'] = meta_source.meta.wcsinfo.dispersion_direction - + extract_params['dispaxis'] = data_model.meta.wcsinfo.dispersion_direction if extract_params['dispaxis'] is None: log.warning("The dispersion direction information is missing, so skipping ...") raise ContinueError() - # Loop over each integration in the input model - shape = meta_source.data.shape + # Set up spatial profiles and wavelength array, + # to be used for every integration + (ra, dec, wavelength, profile, bg_profile, limits) = define_aperture( + input_model, slit, extract_params, exp_type) + + valid = np.isfinite(wavelength) + wavelength = wavelength[valid] + if np.sum(valid) == 0: + log.error("Spectrum is empty; no valid data.") + raise ContinueError() + + # Save the profile if desired + if save_profile: + profile_model = datamodels.ImageModel(profile) + profile_model.update(input_model, only='PRIMARY') + profile_model.name = slitname + else: + profile_model = None + + # Set up aperture correction, to be used for every integration + apcorr_available = False + if source_type is not None and source_type.upper() == 'POINT' and apcorr_ref_model is not None: + log.info('Creating aperture correction.') + # NIRSpec needs to use a wavelength in the middle of the + # range rather than the beginning of the range + # for calculating the pixel scale since some wavelengths at the + # edges of the range won't map to the sky + if instrument == 'NIRSPEC': + wl = np.median(wavelength) + else: + wl = wavelength.min() + + kwargs = {'location': (ra, dec, wl)} + if not isinstance(input_model, datamodels.ImageModel): + kwargs['slit_name'] = slitname + if exp_type in ['NRS_FIXEDSLIT', 'NRS_BRIGHTOBJ']: + kwargs['slit'] = slitname + + apcorr_model = select_apcorr(input_model) + apcorr = apcorr_model(input_model, apcorr_ref_model.apcorr_table, + apcorr_ref_model.sizeunit, **kwargs) + else: + apcorr = None + + # Log the parameters before extracting + log_initial_parameters(extract_params) + # Set up integration iterations and progress messages + progress_msg_printed = True + shape = data_model.data.shape if len(shape) == 3 and shape[0] == 1: integrations = [0] elif len(shape) == 2: @@ -3703,27 +1630,36 @@ def create_extraction(extract_ref_dict, else: log.info(f"Beginning loop over {shape[0]} integrations ...") integrations = range(shape[0]) + progress_msg_printed = False + + # Set up a flux model to update if desired + if save_scene_model: + if len(integrations) > 1: + scene_model = datamodels.CubeModel(shape) + else: + scene_model = datamodels.ImageModel() + scene_model.update(input_model, only='PRIMARY') + scene_model.name = slitname + else: + scene_model = None - ra_last = dec_last = wl_last = apcorr = None - + # Extract each integration for integ in integrations: - try: - ra, dec, wavelength, temp_flux, f_var_poisson, f_var_rnoise, \ - f_var_flat, background, b_var_poisson, b_var_rnoise, \ - b_var_flat, npixels, dq, prev_offset, extraction_values = extract_one_slit( - input_model, - slit, - integ, - prev_offset, - extract_params - ) - except InvalidSpectralOrderNumberError as e: - log.info(f'{str(e)}, skipping ...') - raise ContinueError() + (sum_flux, f_var_rnoise, f_var_poisson, + f_var_flat, background, b_var_rnoise, b_var_poisson, + b_var_flat, npixels, scene_model_2d) = extract_one_slit( + data_model, integ, profile, bg_profile, extract_params) + + # Save the flux model + if save_scene_model: + if isinstance(scene_model, datamodels.CubeModel): + scene_model.data[integ] = scene_model_2d + else: + scene_model.data = scene_model_2d # Convert the sum to an average, for surface brightness. npixels_temp = np.where(npixels > 0., npixels, 1.) - surf_bright = temp_flux / npixels_temp # may be reset below + surf_bright = sum_flux / npixels_temp # may be reset below sb_var_poisson = f_var_poisson / npixels_temp / npixels_temp sb_var_rnoise = f_var_rnoise / npixels_temp / npixels_temp sb_var_flat = f_var_flat / npixels_temp / npixels_temp @@ -3735,18 +1671,18 @@ def create_extraction(extract_ref_dict, del npixels_temp # Convert to flux density. - # The input units will normally be MJy / sr, but for NIRSpec and NIRISS SOSS point-source spectra the units - # will be MJy. + # The input units will normally be MJy / sr, but for NIRSpec + # point-source spectra the units will be MJy. input_units_are_megajanskys = ( photom_has_been_run and source_type == 'POINT' - and (instrument == 'NIRSPEC' or exp_type == 'NIS_SOSS') + and instrument == 'NIRSPEC' ) if photom_has_been_run: - # for NIRSpec data and NIRISS SOSS, point source + # for NIRSpec point sources if input_units_are_megajanskys: - flux = temp_flux * 1.e6 # MJy --> Jy + flux = sum_flux * 1.e6 # MJy --> Jy f_var_poisson *= 1.e12 # MJy**2 --> Jy**2 f_var_rnoise *= 1.e12 # MJy**2 --> Jy**2 f_var_flat *= 1.e12 # MJy**2 --> Jy**2 @@ -3759,24 +1695,33 @@ def create_extraction(extract_ref_dict, b_var_rnoise = b_var_rnoise / pixel_solid_angle / pixel_solid_angle b_var_flat = b_var_flat / pixel_solid_angle / pixel_solid_angle else: - flux = temp_flux * pixel_solid_angle * 1.e6 # MJy / steradian --> Jy + flux = sum_flux * pixel_solid_angle * 1.e6 # MJy / steradian --> Jy f_var_poisson *= (pixel_solid_angle ** 2 * 1.e12) # (MJy / sr)**2 --> Jy**2 f_var_rnoise *= (pixel_solid_angle ** 2 * 1.e12) # (MJy / sr)**2 --> Jy**2 f_var_flat *= (pixel_solid_angle ** 2 * 1.e12) # (MJy / sr)**2 --> Jy**2 else: - flux = temp_flux # count rate + flux = sum_flux # count rate - del temp_flux + del sum_flux error = np.sqrt(f_var_poisson + f_var_rnoise + f_var_flat) sb_error = np.sqrt(sb_var_poisson + sb_var_rnoise + sb_var_flat) berror = np.sqrt(b_var_poisson + b_var_rnoise + b_var_flat) + # Set DQ from the flux value + dq = np.zeros(flux.shape, dtype=np.uint32) + dq[np.isnan(flux)] = datamodels.dqflags.pixel['DO_NOT_USE'] + + # Make a table of the values, trimming to points with valid wavelengths only otab = np.array( list( - zip(wavelength, flux, error, f_var_poisson, f_var_rnoise, f_var_flat, - surf_bright, sb_error, sb_var_poisson, sb_var_rnoise, sb_var_flat, - dq, background, berror, b_var_poisson, b_var_rnoise, b_var_flat, npixels) + zip(wavelength, flux[valid], error[valid], + f_var_poisson[valid], f_var_rnoise[valid], f_var_flat[valid], + surf_bright[valid], sb_error[valid], sb_var_poisson[valid], + sb_var_rnoise[valid], sb_var_flat[valid], + dq[valid], background[valid], berror[valid], + b_var_poisson[valid], b_var_rnoise[valid], b_var_flat[valid], + npixels[valid]) ), dtype=datamodels.SpecModel().spec_table.dtype ) @@ -3803,73 +1748,44 @@ def create_extraction(extract_ref_dict, spec.slit_dec = dec spec.spectral_order = sp_order spec.dispersion_direction = extract_params['dispaxis'] - spec.extraction_xstart = extraction_values['xstart'] - spec.extraction_xstop = extraction_values['xstop'] - spec.extraction_ystart = extraction_values['ystart'] - spec.extraction_ystop = extraction_values['ystop'] - - copy_keyword_info(meta_source, slitname, spec) - - if source_type is not None and source_type.upper() == 'POINT' and apcorr_ref_model is not None: - log.info('Applying Aperture correction.') - # NIRSpec needs to use a wavelength in the middle of the range rather then the beginning of the range - # for calculating the pixel scale since some wavelengths at the edges of the range won't map to the sky - if instrument == 'NIRSPEC': - wl = np.median(wavelength) - else: - wl = wavelength.min() - - # Determine whether we have a tabulated aperture correction - # available to save time. - - apcorr_available = False - if apcorr is not None: - if hasattr(apcorr, 'tabulated_correction'): - if apcorr.tabulated_correction is not None: - apcorr_available = True - + + # Record aperture limits as x/y start/stop values + lower_limit, upper_limit, left_limit, right_limit = limits + if spec.dispersion_direction == HORIZONTAL: + spec.extraction_xstart = left_limit + 1 + spec.extraction_xstop = right_limit + 1 + spec.extraction_ystart = lower_limit + 1 + spec.extraction_ystop = upper_limit + 1 + else: + spec.extraction_xstart = lower_limit + 1 + spec.extraction_xstop = upper_limit + 1 + spec.extraction_ystart = left_limit + 1 + spec.extraction_ystop = right_limit + 1 + + copy_keyword_info(data_model, slitname, spec) + + if apcorr is not None: + if hasattr(apcorr, 'tabulated_correction'): + if apcorr.tabulated_correction is not None: + apcorr_available = True + # See whether we can reuse the previous aperture correction # object. If so, just apply the pre-computed correction to # save a ton of time. - if ra == ra_last and dec == dec_last and wl == wl_last and apcorr_available: + if apcorr_available: # re-use the last aperture correction apcorr.apply(spec.spec_table, use_tabulated=True) - else: - if isinstance(input_model, datamodels.ImageModel): - apcorr = select_apcorr(input_model)( - input_model, - apcorr_ref_model.apcorr_table, - apcorr_ref_model.sizeunit, - location=(ra, dec, wl) - ) - else: - match_kwargs = {'location': (ra, dec, wl)} - if exp_type in ['NRS_FIXEDSLIT', 'NRS_BRIGHTOBJ']: - match_kwargs['slit'] = slitname - - apcorr = select_apcorr(input_model)( - input_model, - apcorr_ref_model.apcorr_table, - apcorr_ref_model.sizeunit, - slit_name=slitname, - **match_kwargs - ) # Attempt to tabulate the aperture correction for later use. # If this fails, fall back on the old method. try: apcorr.tabulate_correction(spec.spec_table) apcorr.apply(spec.spec_table, use_tabulated=True) - log.info("Tabulating aperture correction for use in multiple integrations.") + log.debug("Tabulating aperture correction for use in multiple integrations.") except AttributeError: - log.info("Computing aperture correction.") + log.debug("Computing aperture correction.") apcorr.apply(spec.spec_table) - # Save previous ra, dec, wavelength in case we can reuse - # the aperture correction object. - ra_last = ra - dec_last = dec - wl_last = wl output_model.spec.append(spec) if log_increment > 0 and (integ + 1) % log_increment == 0: @@ -3878,20 +1794,254 @@ def create_extraction(extract_ref_dict, elif integ == 0: if input_model.data.shape[0] == 1: log.info("1 integration done") + progress_msg_printed = True else: log.info("... 1 integration done") elif integ == input_model.data.shape[0] - 1: log.info(f"All {input_model.data.shape[0]} integrations done") + progress_msg_printed = True else: log.info(f"... {integ + 1} integrations done") - progress_msg_printed = True - else: - progress_msg_printed = False if not progress_msg_printed: - if input_model.data.shape[0] == 1: - log.info("1 integration done") + log.info(f"All {input_model.data.shape[0]} integrations done") + + return profile_model, scene_model + + +def run_extract1d(input_model, extract_ref_name="N/A", apcorr_ref_name=None, + smoothing_length=None, bkg_fit=None, bkg_order=None, + log_increment=50, subtract_background=None, + use_source_posn=None, save_profile=False, + save_scene_model=False): + """Extract all 1-D spectra from an input model. + + Parameters + ---------- + input_model : JWSTDataModel + The input science model. + extract_ref_name : str + The name of the extract1d reference file, or "N/A". + apcorr_ref_name : str or None + Name of the APCORR reference file. Default is None + smoothing_length : int or None + Width of a boxcar function for smoothing the background regions. + bkg_fit : str or None + Type of fitting to apply to background values in each column + (or row, if the dispersion is vertical). Allowed values are + 'mean', 'median', 'poly', or None. + bkg_order : int or None + Polynomial order for fitting to each column (or row, if the + dispersion is vertical) of background. Only used if `bkg_fit` + is `poly`. Allowed values are >= 0. + log_increment : int + if `log_increment` is greater than 0 and the input data are + multi-integration, a message will be written to the log every + `log_increment` integrations. + subtract_background : bool or None + User supplied flag indicating whether the background should be + subtracted. + If None, the value in the extract_1d reference file will be used. + If not None, this parameter overrides the value in the + extract_1d reference file. + use_source_posn : bool or None + If True, the target and background positions specified in the + reference file (or the default position, if there is no reference + file) will be shifted to account for source position offset. + save_profile : bool + If True, the spatial profiles created for the input model will be returned + as ImageModels. If False, the return value is None. + save_scene_model : bool + If True, a model of the 2D flux as defined by the extraction aperture + is returned as an ImageModel or CubeModel. If False, the return value + is None. + + Returns + ------- + output_model : MultiSpecModel + A new data model containing the extracted spectra. + profile_model : ModelContainer, ImageModel, or None + If `save_profile` is True, the return value is an ImageModel containing + the spatial profile with aperture weights, used in extracting a single + slit, or else a container of ImageModels, one for each slit extracted. + Otherwise, the return value is None. + scene_model : ModelContainer, ImageModel, CubeModel, or None + If `save_scene_model` is True, the return value is an ImageModel or CubeModel + matching the input data, containing a model of the flux as defined by the + aperture, created during extraction. Otherwise, the return value is None. + """ + # Set "meta_source" to either the first model in a container, + # or the individual input model, for convenience + # of retrieving meta attributes in subsequent statements + if isinstance(input_model, ModelContainer): + meta_source = input_model[0] + else: + meta_source = input_model + + # Get the exposure type + exp_type = meta_source.meta.exposure.type + + # Read in the extract1d reference file. + extract_ref_dict = read_extract1d_ref(extract_ref_name) + + # Read in the aperture correction reference file + apcorr_ref_model = None + if apcorr_ref_name is not None and apcorr_ref_name != 'N/A': + apcorr_ref_model = read_apcorr_ref(apcorr_ref_name, exp_type) + + # Set up the output model + output_model = datamodels.MultiSpecModel() + if hasattr(meta_source, "int_times"): + output_model.int_times = meta_source.int_times.copy() + output_model.update(meta_source, only='PRIMARY') + + # This will be relevant if we're asked to extract a spectrum + # and the spectral order is zero. + # That's only OK if the disperser is a prism. + prism_mode = is_prism(meta_source) + + # Handle inputs that contain one or more slit models + profile_model = None + scene_model = None + if isinstance(input_model, (ModelContainer, datamodels.MultiSlitModel)): + if isinstance(input_model, ModelContainer): + slits = input_model + else: + slits = input_model.slits + + # Save original use_source_posn value, because it can get + # toggled within the following loop over slits + save_use_source_posn = use_source_posn + + # Make a container for the profile models, if needed + if save_profile: + profile_model = ModelContainer() + if save_scene_model: + scene_model = ModelContainer() + + for slit in slits: # Loop over the slits in the input model + log.info(f'Working on slit {slit.name}') + log.debug(f'Slit is of type {type(slit)}') + + slitname = slit.name + use_source_posn = save_use_source_posn # restore original value + + if np.size(slit.data) <= 0: + log.info(f'No data for slit {slit.name}, skipping ...') + continue + + sp_order = get_spectral_order(slit) + if sp_order == 0 and not prism_mode: + log.info("Spectral order 0 is a direct image, skipping ...") + continue + + try: + profile, slit_scene_model = create_extraction( + meta_source, slit, output_model, + extract_ref_dict, slitname, sp_order, exp_type, + apcorr_ref_model=apcorr_ref_model, log_increment=log_increment, + save_profile=save_profile, save_scene_model=save_scene_model, + smoothing_length=smoothing_length, + bkg_fit=bkg_fit, bkg_order=bkg_order, + use_source_posn=use_source_posn, + subtract_background=subtract_background) + except ContinueError: + continue + + if save_profile: + profile_model.append(profile) + if save_scene_model: + scene_model.append(slit_scene_model) + + else: + # Define source of metadata + slit = None + + # This default value for slitname is not really a slit name. + # It may be assigned a better value below, in the sections for + # ImageModel or SlitModel. + slitname = exp_type + + if isinstance(input_model, datamodels.ImageModel): + if hasattr(input_model, "name") and input_model.name is not None: + slitname = input_model.name + + sp_order = get_spectral_order(input_model) + if sp_order == 0 and not prism_mode: + log.info("Spectral order 0 is a direct image, skipping ...") + else: + log.info(f'Processing spectral order {sp_order}') + try: + profile_model, scene_model = create_extraction( + input_model, slit, output_model, + extract_ref_dict, slitname, sp_order, exp_type, + apcorr_ref_model=apcorr_ref_model, log_increment=log_increment, + save_profile=save_profile, save_scene_model=save_scene_model, + smoothing_length=smoothing_length, + bkg_fit=bkg_fit, bkg_order=bkg_order, + use_source_posn=use_source_posn, + subtract_background=subtract_background) + except ContinueError: + pass + + elif isinstance(input_model, (datamodels.CubeModel, datamodels.SlitModel)): + # This branch will be invoked for inputs that are a CubeModel, which typically includes + # NIRSpec BrightObj (fixed slit) mode, as well as inputs that are a + # single SlitModel, which typically includes data from a single resampled/combined slit + # instance from level-3 processing of NIRSpec fixed slits and MOS modes. + + # Replace the default value for slitname with a more accurate value, if possible. + # For NRS_BRIGHTOBJ, the slit name comes from the slit model info + if exp_type == 'NRS_BRIGHTOBJ' and hasattr(input_model, "name"): + slitname = input_model.name + + # For NRS_FIXEDSLIT, the slit name comes from the FXD_SLIT keyword + # in the model meta if not present in the input model + if exp_type == 'NRS_FIXEDSLIT': + if hasattr(input_model, "name") and input_model.name is not None: + slitname = input_model.name + else: + slitname = input_model.meta.instrument.fixed_slit + + sp_order = get_spectral_order(input_model) + if sp_order == 0 and not prism_mode: + log.info("Spectral order 0 is a direct image, skipping ...") + else: + log.info(f'Processing spectral order {sp_order}') + + try: + profile_model, scene_model = create_extraction( + input_model, slit, output_model, + extract_ref_dict, slitname, sp_order, exp_type, + apcorr_ref_model=apcorr_ref_model, log_increment=log_increment, + save_profile=save_profile, save_scene_model=save_scene_model, + smoothing_length=smoothing_length, + bkg_fit=bkg_fit, bkg_order=bkg_order, + use_source_posn=use_source_posn, + subtract_background=subtract_background) + except ContinueError: + pass + else: - log.info(f"All {input_model.data.shape[0]} integrations done") + log.error("The input file is not supported for this step.") + raise RuntimeError("Can't extract a spectrum from this file.") + + # Copy the integration time information from the INT_TIMES table to keywords in the output file. + if pipe_utils.is_tso(input_model): + populate_time_keywords(input_model, output_model) + else: + log.debug("Not copying from the INT_TIMES table because this is not a TSO exposure.") + if hasattr(output_model, "int_times"): + del output_model.int_times + + output_model.meta.wcs = None # See output_model.spec[i].meta.wcs instead. + + if apcorr_ref_model is not None: + apcorr_ref_model.close() + + # Remove target.source_type from the output model, so that it + # doesn't force creation of an empty SCI extension in the output + # x1d product just to hold this keyword. + output_model.meta.target.source_type = None - return output_model + return output_model, profile_model, scene_model diff --git a/jwst/extract_1d/extract1d.py b/jwst/extract_1d/extract1d.py index 6f340b374b..46ceee056a 100644 --- a/jwst/extract_1d/extract1d.py +++ b/jwst/extract_1d/extract1d.py @@ -1,31 +1,14 @@ -""" -1-D spectral extraction +import warnings -:Authors: Mihai Cara (contact: help@stsci.edu) - -""" - -# STDLIB -import logging -import math -import copy - -# THIRD PARTY import numpy as np -from astropy.modeling import models, fitting +from astropy import convolution __all__ = ['extract1d'] -__taskname__ = 'extract1d' -__author__ = 'Mihai Cara' -log = logging.getLogger(__name__) -log.setLevel(logging.DEBUG) - -def extract1d(image, var_poisson, var_rnoise, var_flat, lambdas, disp_range, - p_src, p_bkg=None, independent_var="wavelength", - smoothing_length=0, bkg_fit="poly", bkg_order=0, weights=None): - """Extract the spectrum, optionally subtracting background. +def build_coef_matrix(image, profiles_2d=None, profile_bg=None, + weights=None, order=0): + """Build matrices and vectors to enable least-squares fits. Parameters: ----------- @@ -33,777 +16,468 @@ def extract1d(image, var_poisson, var_rnoise, var_flat, lambdas, disp_range, The array may have been transposed so that the dispersion direction is the second index. - var_poisson : 2-D ndarray - The array may have been transposed so that the dispersion direction - is the second index. + profiles_2d : list of 2-D ndarrays or None, optional + These arrays contain the weights for the extraction. These arrays + should be the same shape as image, with one array for each object + to extract. If set to None, the coefficient matrix values default + to unity. - var_rnoise : 2-D ndarray - The array may have been transposed so that the dispersion direction - is the second index. + profile_bg : boolean 2-D ndarray or None, optional + Array of the same shape as image, with nonzero elements where the + background is to be estimated. If not specified, no additional + pixels are included for background calculations. - var_flat : 2-D ndarray - The array may have been transposed so that the dispersion direction - is the second index. + weights : 2-D ndarray or None, optional + Array of (float) weights for the extraction. If using inverse + variance weighting, these should be the square root of the inverse + variance. If not supplied, unit weights will be used. - lambdas : 1-D array - Wavelength at each pixel within `disp_range`. For example, - lambdas[0] is the wavelength for image[:, disp_range[0]]. - - disp_range : two-element list - Limits of a slice for extracting the spectrum from `image`. - - p_src : list of two-element lists of functions - These are Astropy polynomial functions defining the limits in the - cross-dispersion direction of the source extraction region(s). - - p_bkg : list of two-element lists of functions, or None - These are Astropy polynomial functions defining the limits in the - cross-dispersion direction of the background extraction regions. - - independent_var : string - The value may be "wavelength" or "pixel", indicating whether the - independent variable for the source and background polynomial - functions is wavelength or pixel number. - - smoothing_length : int - If this is greater than one and background regions have been - specified, the background regions will be boxcar smoothed by this - length (must be zero or an odd integer). - This argument is only used if background regions have been - specified. - - bkg_fit : string - Type of fitting to apply to background values in each column (or - row, if the dispersion is vertical). - - bkg_order : int - Polynomial order for fitting to each column of background. A value - of 0 means that a simple average of the background regions, column - by column, will be used. - This argument must be positive or zero, and it is only used if - background regions have been specified and if `bkg_fit` is `poly`. - - weights : function or None - If not None, this computes the weights for the source extraction - region as a function of the wavelength (a single float) for the - current column and an array of Y pixel coordinates. + order : int, optional + Polynomial order for fitting to each column of background. + Default 0 (uniform background). Returns: -------- - temp_flux : ndarray, 1-D, float64 - The extracted spectrum, typically in units of MJy/sr, except for observations - with NIRSpec and NIRISS SOSS, where temp_flux is in units of MJy. - - f_var_poisson : ndarray, 1-D, float64 - The extracted variance due to Poisson noise, units of (counts/s)^2 - - f_var_rnoise : ndarray, 1-D, float64 - The extracted variance due to read noise, units of (counts/s)^2 - - f_var_flat : ndarray, 1-D, float64 - The extracted flat-field variance, units of (counts/s)^2 - - background : ndarray, 1-D, float64 - The background that was subtracted from the source. + matrix : ndarray, 3-D, float64 + Design matrix for each pixel, shape (npixels, npar, npar) - b_var_poisson : ndarray, 1-D, float64 - The extracted background variance due to Poisson noise, units of (counts/s)^2 + vec : ndarray, 2-D, float64 + Target vectors for the design matrix, shape (npixels, npar) - b_var_rnoise : ndarray, 1-D, float64 - The extracted background variance due to read noise, units of (counts/s)^2 + coefmatrix : ndarray, 3-D, float64 + Matrix of coefficients for each parameter for each pixel, + used to reconstruct the model fit. Shape (npixels, npixels_y, npar) - b_var_flat : ndarray, 1-D, float64 - The extracted background flat-field variance, units of (counts/s)^2 - - npixels : ndarray, 1-D, float64 - For each column, this is the number of pixels that were added - together to get `temp_flux`. """ - nl = lambdas.shape[0] - - # Evaluate the functions for source and (optionally) background limits, - # saving the resulting arrays of lower and upper limits in srclim and - # bkglim. - # If independent_var is pixel, the zero point is the first column in - # the input image, rather than the first pixel in the extracted - # spectrum. The spectrum will be extracted from image columns - # disp_range[0] to disp_range[1], so disp_range[0] is the value of - # the independent variable (if not wavelength) at the first pixel of - # the extracted spectrum. - - if not (independent_var.startswith("pixel") or - independent_var.startswith("wavelength")): - log.warning("independent_var was '%s'; using 'pixel' instead.", - independent_var) - independent_var = "pixel" - if independent_var.startswith("pixel"): - # Temporary array for the independent variable. - pixels = np.arange(disp_range[0], disp_range[1], dtype=np.float64) - - srclim = [] # this will be a list of lists, like p_src - n_srclim = len(p_src) - for i in range(n_srclim): - lower = p_src[i][0] - upper = p_src[i][1] - if independent_var.startswith("wavelength"): # OK if 'wavelengths' - srclim.append([lower(lambdas), upper(lambdas)]) - else: - srclim.append([lower(pixels), upper(pixels)]) - - if p_bkg is None: - nbkglim = 0 + if profiles_2d is None: + profiles_2d = [] + + # Independent variable values for the polynomial fit. + y = np.linspace(-1, 1, image.shape[0]) + + # Build the matrix of terms that multiply the coefficients. + # Polynomial terms first, then source terms if those arrays + # are supplied. + coefmatrix = np.ones((image.shape[1], image.shape[0], order + 1 + len(profiles_2d))) + for i in range(1, order + 1): + coefmatrix[..., i] = coefmatrix[..., i - 1] * y[np.newaxis, :] + + # Here are the source terms. + for i in range(len(profiles_2d)): + coefmatrix[..., i + order + 1] = profiles_2d[i].T + + # Construct a boolean array for the pixels that are nonzero + # in any of our profiles. + pixels_used = np.zeros(image.T.shape, dtype=bool) + for profile_2d in profiles_2d: + pixels_used = pixels_used | (profile_2d.T != 0) + if profile_bg is not None: + pixels_used = pixels_used | (profile_bg.T != 0) + pixels_used = pixels_used & np.isfinite(image.T) + + # Target vector and coefficient vector for the least squares fit. + targetvector = image.T.copy() + + # We don't want to be ruined by NaNs in regions we are not fitting anyway. + targetvector[~pixels_used] = 0 + coefmatrix_masked = coefmatrix * pixels_used[:, :, np.newaxis] + + # Weighting goes here. If we are using inverse variance weighting, + # weight here is the square root of the inverse variance. + if weights is not None: + coefmatrix_masked *= weights.T[:, :, np.newaxis] + targetvector *= weights.T + + # Products of the coefficient matrices suitable for passing to + # linalg.solve. These are matrices of size (npixels, npar, npar) + # and (npixels, npar). + matrix = np.einsum('lji,ljk->lik', coefmatrix_masked, coefmatrix_masked) + vec = np.einsum('lji,lj->li', coefmatrix_masked, targetvector) + + return matrix, vec, coefmatrix, coefmatrix_masked + + +def _fit_background_for_box_extraction( + image, profiles_2d, variance_rn, variance_phnoise, variance_flat, + profile_bg, weights, bg_smooth_length, bkg_fit_type, bkg_order): + """Fit a background level for box extraction.""" + + # Start by copying the input image + input_background = image.copy() + + # Smooth the image, if desired, for computing a background. + # Astropy's convolve routine will replace NaNs. The convolution + # is done along the dispersion direction. + if bg_smooth_length > 1: + if not bg_smooth_length % 2 == 1: + raise ValueError("bg_smooth_length should be an odd integer >= 1.") + kernel = np.ones((1, bg_smooth_length)) / bg_smooth_length + input_background = convolution.convolve(input_background, kernel, boundary='extend') + + if bkg_fit_type == 'median': + input_background[profile_bg == 0] = np.nan + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=RuntimeWarning, message="All-NaN") + bkg_1d = np.nanmedian(input_background, axis=0) + bkg_2d = bkg_1d[np.newaxis, :] + + # Putting an uncertainty on the median is a bit harder. + # It is typically about 1.2 times the uncertainty on the mean. + # The details depend on the number of pixels averaged... + # bkg_npix is the total weight that we will need to apply + # to the background value when removing it from the 2D image. + # pixwgt normalizes the total weights of all pixels used to 1. + bkg_npix = np.sum(profiles_2d[0], axis=0) + variance = variance_rn + variance_phnoise + variance_flat + wgt = np.isfinite(image) * np.isfinite(variance) * (profile_bg != 0) + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=RuntimeWarning, message="invalid value") + pixwgt = wgt / np.sum(wgt, axis=0)[np.newaxis, :] + + var_bkg_rn = 1.2 ** 2 * bkg_npix ** 2 * np.array([np.nansum(variance_rn * pixwgt ** 2, axis=0)]) + var_bkg_phnoise = 1.2 ** 2 * bkg_npix ** 2 * np.array([np.nansum(variance_phnoise * pixwgt ** 2, axis=0)]) + var_bkg_flat = 1.2 ** 2 * bkg_npix ** 2 * np.array([np.nansum(variance_flat * pixwgt ** 2, axis=0)]) + + elif bkg_fit_type == 'poly' and bkg_order >= 0: + + if not bkg_order == int(bkg_order): + raise ValueError("If bkg_fit_type is 'poly', bkg_order must be an integer >= 0.") + + # Build the matrices to fit a polynomial column-by-column. + result = build_coef_matrix(input_background, profile_bg=profile_bg, + weights=weights, order=bkg_order) + matrix, vec, coefmatrix, coefmatrix_masked = result + + # Don't try to solve singular matrices. Background will be + # zero in these cases. Could make them NaN if you want. + ok = np.linalg.cond(matrix) < 1e10 + cov_bg_coefs = np.zeros(matrix.shape) + cov_bg_coefs[ok] = np.linalg.inv(matrix[ok]) + + # These are the pixel-dependent weights to compute our coefficients. + # We will use them to propagate errors. + pixwgt = weights.T[:, :, np.newaxis] * np.einsum('ijk,ilj->ilk', cov_bg_coefs, coefmatrix_masked) + bkg_mat = np.sum(np.swapaxes(coefmatrix, 0, 1) * profiles_2d[0][:, :, np.newaxis], axis=0) + + # Sum of all the contributions to the background at the pixels + # where we will do the extraction. Used to propagate errors. + pixwgt_tot = np.sum(bkg_mat[:, np.newaxis, :] * pixwgt, axis=-1) + + var_bkg_rn = np.array([np.nansum(variance_rn.T * pixwgt_tot ** 2, axis=1)]) + var_bkg_phnoise = np.array([np.nansum(variance_phnoise.T * pixwgt_tot ** 2, axis=1)]) + var_bkg_flat = np.array([np.nansum(variance_flat.T * pixwgt_tot ** 2, axis=1)]) + + coefs = np.einsum('ijk,ij->ik', cov_bg_coefs, vec) + + # Reconstruct the 2D background. + bkg_2d = np.sum(coefs[:, np.newaxis, :] * coefmatrix, axis=-1).T + else: - nbkglim = len(p_bkg) - bkglim = [] # this will be a list of lists, like p_bkg - for i in range(nbkglim): - lower = p_bkg[i][0] - upper = p_bkg[i][1] - if independent_var.startswith("wavelength"): - bkglim.append([lower(lambdas), upper(lambdas)]) - else: - bkglim.append([lower(pixels), upper(pixels)]) - - # Sanity check: check for extraction limits that are out of bounds, - # or a lower limit that's above the upper limit (limit curves just - # swapped, or crossing each other). - # Truncate extraction limits that are out of bounds, but log a warning. - shape = image.shape - for i in range(n_srclim): - lower = srclim[i][0] - upper = srclim[i][1] - diff = upper - lower - if diff.min() < 0.: - if diff.max() < 0.: - log.error("Lower and upper source extraction limits" - " appear to be swapped.") - raise ValueError("Lower and upper source extraction limits" - " appear to be swapped.") - else: - log.error("Lower and upper source extraction limits" - " cross each other.") - raise ValueError("Lower and upper source extraction limits" - " cross each other.") - del diff - if np.any(lower < -0.5) or np.any(upper < -0.5): - log.warning("Source extraction limit extends below -0.5") - srclim[i][0][:] = np.where(lower < -0.5, -0.5, lower) - srclim[i][1][:] = np.where(upper < -0.5, -0.5, upper) - upper_limit = float(shape[0]) - 0.5 - if np.any(lower > upper_limit) or np.any(upper > upper_limit): - log.warning("Source extraction limit extends above %g", upper_limit) - srclim[i][0][:] = np.where(lower > upper_limit, upper_limit, lower) - srclim[i][1][:] = np.where(upper > upper_limit, upper_limit, upper) - for i in range(nbkglim): - lower = bkglim[i][0] - upper = bkglim[i][1] - diff = upper - lower - if diff.min() < 0.: - if diff.max() < 0.: - log.error("Lower and upper background extraction limits" - " appear to be swapped.") - raise ValueError("Lower and upper background extraction limits" - " appear to be swapped.") - else: - log.error("Lower and upper background extraction limits" - " cross each other.") - raise ValueError("Lower and upper background extraction limits" - " cross each other.") - del diff - if np.any(lower < -0.5) or np.any(upper < -0.5): - log.warning("Background limit extends below -0.5") - bkglim[i][0][:] = np.where(lower < -0.5, -0.5, lower) - bkglim[i][1][:] = np.where(upper < -0.5, -0.5, upper) - upper_limit = float(shape[0]) - 0.5 - if np.any(lower > upper_limit) or np.any(upper > upper_limit): - log.warning("Background limit extends above %g", upper_limit) - bkglim[i][0][:] = np.where(lower > upper_limit, upper_limit, lower) - bkglim[i][1][:] = np.where(upper > upper_limit, upper_limit, upper) - - # Smooth the input image, and use the smoothed image for extracting - # the background. temp_image is only needed for background data. - if nbkglim > 0 and smoothing_length > 1: - temp_image = bxcar(image, smoothing_length) + raise ValueError("bkg_fit_type should be 'median' or 'poly'. " + "If 'poly', bkg_order must be an integer >= 0.") + + return bkg_2d, var_bkg_rn, var_bkg_phnoise, var_bkg_flat + + +def _box_extract( + image, profiles_2d, variance_rn, variance_phnoise, variance_flat, + bkg_2d, var_bkg_rn, var_bkg_phnoise, var_bkg_flat, model): + """Perform box extraction.""" + # This only makes sense with a single profile, i.e., pulling out + # a single spectrum. + nobjects = 1 + profile_2d = profiles_2d[0].copy() + image_masked = image.copy() + + # Mask NaNs and infs for the extraction + profile_2d[~np.isfinite(image_masked)] = 0 + image_masked[profile_2d == 0] = 0 + + # Return array of shape (1, npixels) for generality + fluxes = np.array([np.sum(image_masked * profile_2d, axis=0)]) + + # Number of contributing pixels at each wavelength. + npixels = np.array([np.sum(profile_2d, axis=0)]) + + # Add average flux over the aperture to the model, so that + # a sum over the cross-dispersion direction reproduces the summed flux + valid = npixels[0] > 0 + model[:, valid] += (fluxes[0][valid] / npixels[0][valid]) * profile_2d[:, valid] + + # Compute the variance on the sum, same shape as f. + # Need to decompose this into read noise, photon noise, and flat noise. + var_rn = np.array([np.nansum(variance_rn * profile_2d ** 2, axis=0)]) + var_phnoise = np.array([np.nansum(variance_phnoise * profile_2d ** 2, axis=0)]) + var_flat = np.array([np.nansum(variance_flat * profile_2d ** 2, axis=0)]) + + if bkg_2d is not None: + var_rn += var_bkg_rn + var_phnoise += var_bkg_phnoise + var_flat += var_bkg_flat + bkg = np.array([np.sum(bkg_2d * profile_2d, axis=0)]) + bkg[~np.isfinite(bkg)] = 0.0 else: - temp_image = image - - ################################################# - # Perform spectral extraction: # - ################################################# - - bkg_model = None - b_var_poisson_model = None - b_var_rnoise_model = None - b_var_flat_model = None - - temp_flux = np.zeros(nl, dtype=np.float64) - f_var_poisson = np.zeros(nl, dtype=np.float64) - f_var_rnoise = np.zeros(nl, dtype=np.float64) - f_var_flat = np.zeros(nl, dtype=np.float64) - background = np.zeros(nl, dtype=np.float64) - b_var_poisson = np.zeros(nl, dtype=np.float64) - b_var_rnoise = np.zeros(nl, dtype=np.float64) - b_var_flat = np.zeros(nl, dtype=np.float64) - npixels = np.zeros(nl, dtype=np.float64) - # x is an index (column number) within `image`, while j is an index in - # lambdas, temp_flux, background, npixels, and the arrays in - # srclim and bkglim. - x = disp_range[0] - for j in range(nl): - lam = lambdas[j] - - if nbkglim > 0: - - # Compute the background for the current column, - # using the (optionally) smoothed background. - (bkg_model, b_var_poisson_model, b_var_rnoise_model, - b_var_flat_model, bkg_npts) = _fit_background_model( - temp_image, var_poisson, var_rnoise, var_flat, x, - j, bkglim, bkg_fit, bkg_order - ) - - if bkg_npts == 0: - bkg_model = None - log.debug(f"Not enough valid pixels to determine background " - f"for lambda={lam:.6f} (column {x:d})") - - elif len(bkg_model) < bkg_order: - log.debug(f"Not enough valid pixels to determine background " - f"with the required order for lambda={lam:.6f} " - f"(column {x:d})") - log.debug(f"Lowering background order to {len(bkg_model)}") - - # Extract the source, and optionally subtract background using the - # fit to the background for this column. Even if - # background smoothing was done, we must extract the source from - # the original, unsmoothed image. - # source total flux, background total flux, area, total weight - (temp_flux[j], f_var_poisson[j], f_var_rnoise[j], f_var_flat[j], - bkg_flux, b_var_poisson_val, b_var_rnoise_val, b_var_flat_val, - npixels[j], twht) = _extract_src_flux( - image, var_poisson, var_rnoise, var_flat, x, j, lam, srclim, - weights=weights, bkgmodels=[bkg_model, b_var_poisson_model, - b_var_rnoise_model, b_var_flat_model] - ) - if nbkglim > 0: - background[j] = bkg_flux - b_var_poisson[j] = b_var_poisson_val - b_var_rnoise[j] = b_var_rnoise_val - b_var_flat[j] = b_var_flat_val - - x += 1 - continue - - return (temp_flux, f_var_poisson, f_var_rnoise, f_var_flat, - background, b_var_poisson, b_var_rnoise, b_var_flat, npixels) - - -def bxcar(image, smoothing_length): - """Smooth with a 1-D interval, along the last axis. - - Extended summary - ---------------- - Note that the entire input array will be smoothed, including the - region containing the source. The source extraction must therefore - be done from the original, unsmoothed array. + bkg = np.zeros((nobjects, image.shape[1])) - Parameters: - ----------- - image : 2-D ndarray - The input data array. + return (fluxes, var_rn, var_phnoise, var_flat, + bkg, var_bkg_rn, var_bkg_phnoise, var_bkg_flat, npixels, model) - Returns: - -------- - ndarray, 1-D - The smoothed input array. - """ - half = smoothing_length // 2 +def _optimal_extract( + image, profiles_2d, variance_rn, variance_phnoise, variance_flat, + weights, profile_bg, fit_bkg, bkg_order, + bkg_2d, var_bkg_rn, var_bkg_phnoise, var_bkg_flat, model): + """Perform optimal extraction.""" - shape0 = image.shape - width = shape0[-1] - shape = shape0[0:-1] + (width + 2 * half,) - temp_im = np.zeros(shape, dtype=np.float64) + # Background fitting needs to be done simultaneously with the + # fitting of the spectra in this case. If we are not fitting a + # background, pass -1 for the order of the polynomial correction. + if fit_bkg: + order = bkg_order + if order != int(order) or order < 0: + raise ValueError("For optimal extraction, bkg_order must be an integer >= 0.") + else: + order = -1 + + result = build_coef_matrix(image, profiles_2d=profiles_2d, weights=weights, + profile_bg=profile_bg, order=order) + matrix, vec, coefmatrix, coefmatrix_masked = result + + # Don't try to solve equations with singular matrices. + # Fluxes will be zero in these cases. We will make them NaN later. + ok = np.linalg.cond(matrix) < 1e10 + + # These are the covariance matrices for all parameters if inverse + # variance weights are passed to the build_coef_matrix above. For + # generality, variances are actually computed using the weights on + # each pixel and the associated variance in the input image. + covariances = np.zeros(matrix.shape) + covariances[ok] = np.linalg.inv(matrix[ok]) + + # These are the pixel-dependent weights to compute our coefficients. + # We will use them to propagate errors. + pixwgt = weights.T[:, :, np.newaxis] * np.einsum('ijk,ilj->ilk', covariances, coefmatrix_masked) + + # Don't use NaN pixels in the sum. These will already be zero in + # pixwgt. coefs are the best-fit coefficients of the source and + # background components. + coefs = np.nansum(pixwgt * image.T[:, :, np.newaxis], axis=1) + + # Effective number of contributing pixels at each wavelength for each source. + nobjects = len(profiles_2d) + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=RuntimeWarning, message="invalid value") + wgt_src_pix = [profiles_2d[i] * (weights > 0) / np.sum(profiles_2d[i] ** 2, axis=0) + for i in range(nobjects)] + npixels = np.sum(wgt_src_pix, axis=1) + + if order > -1: + bkg_2d = np.sum(coefs[:, np.newaxis, :order + 1] * coefmatrix[..., :order + 1], axis=-1).T + + # Variances for each object (discard variances for background here) + var_rn = np.nansum(pixwgt[..., -nobjects:] ** 2 * variance_rn.T[:, :, np.newaxis], axis=1).T + var_phnoise = np.nansum(pixwgt[..., -nobjects:] ** 2 * variance_phnoise.T[:, :, np.newaxis], axis=1).T + var_flat = np.nansum(pixwgt[..., -nobjects:] ** 2 * variance_flat.T[:, :, np.newaxis], axis=1).T + + # Computing a background contribution to the noise is harder in a joint fit. + # Here, I am computing the weighting coefficients I would have without a background. + # I then compute those variances, and subtract them from the actual variances. + if order > -1: + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=RuntimeWarning, message="invalid value") + + wgt_nobkg = [profiles_2d[i] * weights / np.sum(profiles_2d[i] ** 2 * weights, axis=0) + for i in range(nobjects)] + + bkg = np.array([np.sum(wgt_nobkg[i] * bkg_2d, axis=0) for i in range(nobjects)]) + + var_bkg_rn = np.array([var_rn[i] - np.sum(wgt_nobkg[i] ** 2 * variance_rn, axis=0) + for i in range(nobjects)]) + var_bkg_phnoise = np.array([var_phnoise[i] - np.sum(wgt_nobkg[i] ** 2 * variance_phnoise, axis=0) + for i in range(nobjects)]) + var_bkg_flat = np.array([var_flat[i] - np.sum(wgt_nobkg[i] ** 2 * variance_flat, axis=0) + for i in range(nobjects)]) + + # Make sure background values are finite + bkg[~np.isfinite(bkg)] = 0.0 + + # We did our best to estimate the background contribution to the variance + # in this case. Don't let it go negative. + var_bkg_rn *= var_bkg_rn > 0 + var_bkg_phnoise *= var_bkg_phnoise > 0 + var_bkg_flat *= var_bkg_flat > 0 + else: + bkg = np.zeros((nobjects, image.shape[1])) - i = 0 - for k in range(smoothing_length): - temp_im[..., i:i + width] += image - i += 1 - temp_im /= float(smoothing_length) + # Reshape to (nobjects, npixels) + fluxes = coefs[:, -nobjects:].T + model += np.sum(coefs[:, np.newaxis, :] * coefmatrix, axis=-1).T - return temp_im[..., half:half + width].astype(image.dtype) + return (fluxes, var_rn, var_phnoise, var_flat, + bkg, var_bkg_rn, var_bkg_phnoise, var_bkg_flat, npixels, model) -def _extract_src_flux(image, var_poisson, var_rnoise, var_flat, x, j, lam, srclim, weights, bkgmodels): - """Extract the source and subtract background. +def extract1d(image, profiles_2d, variance_rn, variance_phnoise, variance_flat, + weights=None, profile_bg=None, extraction_type='box', + bg_smooth_length=0, fit_bkg=False, bkg_fit_type='poly', bkg_order=0): + """Extract the spectrum, optionally subtracting background. Parameters: ----------- image : 2-D ndarray - The input data array. - - var_poisson : 2-D ndarray - The input poisson variance array. + The array may have been transposed so that the dispersion direction + is the second index. - var_rnoise : 2-D ndarray - The input read noise variance array. + profiles_2d : list of 2-D ndarray + These arrays contain the weights for the extraction. A box + extraction will add up the flux multiplied by these weights; an + optimal extraction will fit an amplitude to the weight map at each + column in the dispersion direction. These arrays should be the + same shape as image, with one array for each object to extract. + Box extraction only works if exactly one profile is supplied + (i.e. this is a one-element list). - var_flat : 2-D ndarray - The input flat field variance array. + variance_rn : 2-D ndarray + Read noise component of the variance. - x : int - This is an index (column number) within `image`. + variance_phnoise : 2-D ndarray + Photon noise component of the variance. - j : int - This is an index starting with 0 at the first pixel of the - extracted spectrum. See `disp_range` in function `extract1d`. - j = 0 when x = disp_range[0]. + variance_flat : 2-D ndarray + Flat component of the variance. - lam : float + weights : 2-D ndarray or None + Weights for the individual pixels in fitting a profile. If None + (default), use uniform weights (ones for all valid pixels). - srclim : list of lists of ndarrays - For each i, srclim[i] is a two-element list. Those two elements - are arrays of the lower and upper limits of one of the extraction - regions. Of course, there may only be one extraction region. + profile_bg : 2-D ndarray or None + Array of the same shape as image, with nonzero elements where the + background is to be estimated. - weights : function or None - If not None, this function gives the weights for pixels within - an extraction region. + extraction_type : string + Type of spectral extraction. Currently must be either "box" + or "optimal". - bkgmodels : function + bg_smooth_length : int + Smoothing length for box smoothing of the background along the + dispersion direction. Should be odd, >=1. - Returns: - -------- - total_flux : float or NaN - Sum of counts within the source extraction region for the current - column. This will be NaN if there is no data in the source - extraction region for the current column. - - f_var_poisson : float or NaN - Sum of variance within source extraction region for current column. - - f_var_rnoise : float or NaN - Sum of variance within source extraction region for current column. - - f_var_flat : float or NaN - Sum of variance within source extraction region for current column. - - bkg_flux : float - Sum of counts within the background extraction regions for the - current column. - - b_var_poisson : float or NaN - Sum of variance within background extraction region for current column. - - b_var_rnoise : float or NaN - Sum of variance within background extraction region for current column. - - b_var_flat : float or NaN - Sum of variance within background extraction region for current column. - - tarea : float - The sum of the number of pixels in the source extraction region for - the current column. If only a fraction of a pixel is included at - an endpoint, that fraction is what would be included in the sum. - For example, if the source limits (in `srclim`) are (3, 7), then - the extraction region extends from the middle of pixel 3 to the - middle of pixel 7, so the pixel areas in the extraction region - would be: 0.5, 1.0, 1.0, 1.0, 0.5, resulting in `tarea` = 4. - - twht : float - Two different weights are applied to the pixels. One is the - fraction of a pixel that is included in the extraction (this will - be 1.0 except possibly at the endpoints); see also `tarea`. The - other weight depends on the `weights` argument. If `weights` is - None, then this weight will be 1.0 for each pixel, and `twht` will - be the sum of these values. If the source limits are (3, 7) as in - the example for `tarea`, `twht` would be 5.0. - """ - - # extract pixel values along the column that are within - # source limits: - y, val, area = _extract_colpix(image, x, j, srclim) - - # find indices of "good" (finite) values: - good = np.isfinite(val) - npts = good.sum() - - if npts == 0: - return np.nan, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 - - # filter-out bad values: - # TODO: in the future we may need to develop a way of interpolating - # over missing values either from a model or from adjacent columns - val = val[good] - area = area[good] - y = y[good] - if bkgmodels[0] is None: - bkg = np.zeros_like(val, dtype=np.float64) - b_var_poisson = np.zeros_like(val, dtype=np.float64) - b_var_rnoise = np.zeros_like(val, dtype=np.float64) - b_var_flat = np.zeros_like(val, dtype=np.float64) - else: - bkg = bkgmodels[0](y) - b_var_poisson = bkgmodels[1](y) * area - b_var_rnoise = bkgmodels[2](y) * area - b_var_flat = bkgmodels[3](y) * area - - dummy_y, f_var_poisson_val, dummy_area = _extract_colpix(var_poisson, x, j, srclim) - f_var_poisson = f_var_poisson_val[good] * area - dummy_y, f_var_rnoise_val, dummy_area = _extract_colpix(var_rnoise, x, j, srclim) - f_var_rnoise = f_var_rnoise_val[good] * area - dummy_y, f_var_flat_val, dummy_area = _extract_colpix(var_flat, x, j, srclim) - f_var_flat = f_var_flat_val[good] * area - - # subtract background per pixel: - val -= bkg - - # add background variance to variance in source extraction region - f_var_poisson += b_var_poisson - f_var_rnoise += b_var_rnoise - f_var_flat += b_var_flat - - # scale per pixel values by pixel area included in extraction - val *= area - bkg *= area - - # compute weights: - if weights is None: - wht = np.ones_like(y, dtype=np.float64) - else: - wht = weights(lam, y) - - # compute weighted total flux - # NOTE: the correct formulae must be derived depending on - # final interpretation of weights [not available at - # initial release v0.0.1] - tarea = area.sum(dtype=np.float64) - twht = wht.sum(dtype=np.float64) - mwht = twht / wht.shape[0] - total_flux = (val * wht).sum(dtype=np.float64) / mwht - f_var_poisson = (f_var_poisson * wht).sum(dtype=np.float64) / mwht - f_var_rnoise = (f_var_rnoise * wht).sum(dtype=np.float64) / mwht - f_var_flat = (f_var_flat * wht).sum(dtype=np.float64) / mwht - bkg_flux = bkg.sum(dtype=np.float64) - b_var_poisson = b_var_poisson.sum(dtype=np.float64) - b_var_rnoise = b_var_rnoise.sum(dtype=np.float64) - b_var_flat = b_var_flat.sum(dtype=np.float64) - - return (total_flux, f_var_poisson, f_var_rnoise, f_var_flat, - bkg_flux, b_var_poisson, b_var_rnoise, b_var_flat, tarea, twht) - - -def _fit_background_model(image, var_poisson, var_rnoise, var_flat, - x, j, bkglim, bkg_fit, bkg_order): - """Extract background pixels and fit a polynomial. If the number of good data - points is <= 1, the fit model will be forced to 0 to avoid divergence. + fit_bkg : bool + Fit a background? Default False - Parameters: - ----------- - image : 2-D ndarray - The input data array. - - var_poisson : 2-D ndarray - The input poisson variance array. + bkg_fit_type : string + Type of fitting to apply to background values in each column (or + row, if the dispersion is vertical). - var_rnoise : 2-D ndarray - The input read noise variance array. + bkg_order : int + Polynomial order for fitting to each column of background. A value + of 0 means that a simple average of the background regions, column + by column, will be used. + This argument must be positive or zero, and it is only used if + background regions have been specified and if `bkg_fit` is `poly`. - var_flat : 2-D ndarray - The input flat field variance array. + Returns: + -------- + fluxes : ndarray, n-D, float64 + The extracted spectrum/spectra. Units are currently arbitrary. + The first dimension is the same as the length of profiles_2d - x : int - This is an index (column number) within `image`. + var_rn : ndarray, n-D, float64 + The variances of the extracted spectrum/spectra due to read noise. + Units are the same as flux^2, shape is the same as flux. - j : int - This is an index starting with 0 at the first pixel of the - extracted spectrum. See `disp_range` in function `extract1d`. - j = 0 when x = disp_range[0]. + var_phnoise : ndarray, n-D, float64 + The variances of the extracted spectrum/spectra due to photon noise. + Units are the same as flux^2, shape is the same as flux. - bkglim : list of lists of arrays - For each i, bkglim[i] is a two-element list. Those two elements - are arrays of the lower and upper limits of one of the background - extraction regions. + var_flat : ndarray, n-D, float64 + The variances of the extracted spectrum/spectra due to flatfield + uncertainty. Units are the same as flux^2, shape is the same as flux. - bkg_fit : str - Type of "fitting" to perform: "poly" = polynomial, "mean", or - "median". Note that mathematically the result for "mean" is - identical to "poly" with `bkg_order`=0. + bkg : ndarray, n-D, float64 + Background level that would be obtained for each source if performing + a 1-D extraction on the 2D background. - bkg_order : int - Polynomial order for fitting to the background regions of the - current column. + var_bkg_rn : ndarray, n-D, float64 + As above, for read noise. Nonzero because read noise adds an error term + to the derived background level. - Returns: - -------- - bkg_model : function - Polynomial fit to the background regions for the current column. - When the requested fit is either "mean" or "median", the returned - model is a 0th-order polynomial with c0 equal to the mean/median. + var_bkg_phnoise : ndarray, n-D, float64 + The variances of the extracted spectrum/spectra due to background photon + noise. Units are the same as flux^2, shape is the same as flux. This + background contribution is already included in var_phnoise. - b_var_poisson_model : function - Polynomial fit to the background regions for var_poisson values. + var_bkg_flat : ndarray, n-D, float64 + As above, for the flatfield. - b_var_rnoise_model : function - Polynomial fit to the background regions for var_rnoise values. + npixels : ndarray, n-D, int64 + Number of pixels that contribute to the flux measurement for each source - b_var_flat_model : function - Polynomial fit to the background regions for var_flat values. + model : ndarray, 2-D, float64 + The model of the scene, the same shape as the input image (and + hopefully also similar in value). - npts : int - This is intended to be the number of good values in the background - regions. If the background limits are at pixel edges, however, - `npts` can include a pixel with zero weight; that is, `npts` can be - 1 larger than one might expect. """ - # extract pixel values along the column that are within - # background limits: - y, val, wht = _extract_colpix(image, x, j, bkglim) - - # find indices of "good" (finite) values: - good = np.isfinite(val) - npts = good.sum() - - if npts <= 1 or not np.any(good): - return (models.Polynomial1D(0), models.Polynomial1D(0), - models.Polynomial1D(0), models.Polynomial1D(0), 0) - - # filter-out bad values: - val = val[good] - wht = wht[good] - y = y[good] - - if wht.sum() == 0: - return (models.Polynomial1D(0), models.Polynomial1D(0), - models.Polynomial1D(0), models.Polynomial1D(0), 0) - - # Find values for each variance array according to locations of - # "good" image values - dummy_y, var_poisson_val, dummy_area = _extract_colpix(var_poisson, x, j, bkglim) - var_poisson_val = var_poisson_val[good] - dummy_y, var_rnoise_val, dummy_area = _extract_colpix(var_rnoise, x, j, bkglim) - var_rnoise_val = var_rnoise_val[good] - dummy_y, var_flat_val, dummy_area = _extract_colpix(var_flat, x, j, bkglim) - var_flat_val = var_flat_val[good] - - # Compute the fit - if bkg_fit == 'poly': - - # Fit the background values with a polynomial of the requested order - lsqfitter = fitting.LinearLSQFitter() - bkg_model = lsqfitter(models.Polynomial1D(min(bkg_order, npts - 1)), - y, val, weights=wht) - b_var_poisson_model = lsqfitter(models.Polynomial1D(min(bkg_order, npts - 1)), - y, var_poisson_val, weights=wht) - b_var_rnoise_model = lsqfitter(models.Polynomial1D(min(bkg_order, npts - 1)), - y, var_rnoise_val, weights=wht) - b_var_flat_model = lsqfitter(models.Polynomial1D(min(bkg_order, npts - 1)), - y, var_flat_val, weights=wht) - - elif bkg_fit == 'mean': - - # Compute the mean of the (good) background values - # only use values with weight=1 - bkg_model = models.Polynomial1D(degree=0, c0=np.mean(val[wht == 1])) - b_var_poisson_model = models.Polynomial1D(degree=0, c0=np.mean(var_poisson_val[wht == 1])) - b_var_rnoise_model = models.Polynomial1D(degree=0, c0=np.mean(var_rnoise_val[wht == 1])) - b_var_flat_model = models.Polynomial1D(degree=0, c0=np.mean(var_flat_val[wht == 1])) - - elif bkg_fit == 'median': - - # Compute the median of the (good) background values - # only use values with weight=1 - bkg_model = models.Polynomial1D(degree=0, c0=np.median(val[wht == 1])) - b_var_poisson_model = models.Polynomial1D(degree=0, c0=np.median(var_poisson_val[wht == 1])) - b_var_rnoise_model = models.Polynomial1D(degree=0, c0=np.median(var_rnoise_val[wht == 1])) - b_var_flat_model = models.Polynomial1D(degree=0, c0=np.median(var_flat_val[wht == 1])) - - return bkg_model, b_var_poisson_model, b_var_rnoise_model, b_var_flat_model, npts - - -def _extract_colpix(image_data, x, j, limits): - """Extract either the source or background data. + nobjects = len(profiles_2d) # hopefully at least one! + model = np.zeros(image.shape) - Parameters: - ----------- - image_data : 2-D ndarray - The input data array. + # If weights are not supplied, equally weight all valid pixels. + variance = variance_rn + variance_phnoise + variance_flat + if weights is None: + weights = np.isfinite(variance) * np.isfinite(image) + + # This is the case of a background fit independent of a flux fit. + # This is done only if we have a background region to fit, the + # boolean variable to fit is set, and we are using box extraction. + # Inverse variance weights should be used with care, as they have the + # potential to introduce biases. + if profile_bg is not None and fit_bkg and extraction_type == 'box': + (bkg_2d, var_bkg_rn, + var_bkg_phnoise, var_bkg_flat) = _fit_background_for_box_extraction( + image, profiles_2d, variance_rn, variance_phnoise, variance_flat, + profile_bg, weights, bg_smooth_length, bkg_fit_type, bkg_order) + + model += bkg_2d + image_sub = image - bkg_2d + else: + image_sub = image.copy() - x : int - This is an index (column number) within `image_data`. + # Set background image to None + bkg_2d = None - j : int - This is an index starting with 0 at the first pixel of the - extracted spectrum. See `disp_range` in function `extract1d`. - j = 0 when x = disp_range[0]. + # Set background uncertainties to zero. + var_bkg_rn = np.zeros((nobjects, image.shape[1])) + var_bkg_phnoise = np.zeros((nobjects, image.shape[1])) + var_bkg_flat = np.zeros((nobjects, image.shape[1])) - limits : list of lists of ndarrays - For each i, limits[i] is a two-element list. Those two elements - are 1-D arrays of the pixel coordinates for the lower and upper - limits of one of the source or background extraction regions. The - number of elements in each of these arrays is the number of pixels - in the domain from disp_range[0] to and including disp_range[1]. + # This is the case of box extraction. + if extraction_type == 'box' and len(profiles_2d) == 1: - Returns: - -------- - y : ndarray, float64 - Y pixel coordinates within the current column, for every pixel that - is included in any of the intervals in `limits`. - - val : ndarray, float64 - The image values at the pixels given by `y`. That is, - val[i] = image_data[y[i], x]. - - wht : ndarray, float64 - The weight associated with each element in `val`. The weight - ranges from 0 to 1, giving the fraction of a pixel that is included - within an interval. For example, suppose one of the elements in - `limits` has values 3.0 and 9.0 for the lower and upper limits for - pixel number `x`. That corresponds to this list of Y pixel values: - [3., 4., 5., 6., 7., 8., 9.]. (You would then see this list as a - section in `y`.) Because the integer value is the center of the - pixel, the lower and upper limits are in the middle of those - pixels, so the corresponding weights would be: - [0.5, 1., 1., 1., 1., 1., 0.5]. - """ + (fluxes, var_rn, var_phnoise, var_flat, + bkg, var_bkg_rn, var_bkg_phnoise, + var_bkg_flat, npixels, model) = _box_extract( + image_sub, profiles_2d, variance_rn, variance_phnoise, variance_flat, + bkg_2d, var_bkg_rn, var_bkg_phnoise, var_bkg_flat, model) - # These are the extraction limits in image pixel coordinates: - intervals = [] - for lim in limits: - intervals.append([lim[0][j], lim[1][j]]) - - if len(intervals) == 0: - return [], [], [] - - # optimize limits: - intervals = _coalesce_bounds(intervals) - - # compute number of data points: - ns = image_data.shape[0] - 1 - ns12 = ns + 0.5 - npts = 0 - for interval in intervals: - if interval[0] == interval[1]: - # no data between limits - continue - maxval = min(ns, int(math.floor(interval[1] + 0.5))) - minval = max(0, int(math.floor(interval[0] + 0.5))) - if maxval - minval + 1 > 0: - npts += maxval - minval + 1 - if npts == 0: - return [], [], [] - - # pre-allocate data arrays: - y = np.empty(npts, dtype=np.float64) - val = np.empty(npts, dtype=np.float64) - wht = np.ones(npts, dtype=np.float64) - - # populate data and weights: - k = 0 - for i in intervals: - if i[0] == i[1]: - continue - i1 = i[0] if i[0] >= -0.5 else -0.5 - i2 = i[1] if i[1] <= ns12 else ns12 - - ii1 = max(0, int(math.floor(i1 + 0.5))) - ii1 = min(ii1, ns) - ii2 = min(ns, int(math.floor(i2 + 0.5))) - - # special case: ii1 == ii2: - # take the value at the pixel - if ii1 == ii2: - v = image_data[ii1, x] - val[k] = v - wht[k] = i2 - i1 - k += 1 - continue - - # bounds in different pixels: - # a. lower bound: - v = image_data[ii1, x] - val[k] = v - wht[k] = 1.0 - divmod(i1 - 0.5, 1)[1] if i1 >= -0.5 else 1.0 - - # b. upper bound: - v = image_data[ii2, x] - kn = k + ii2 - ii1 - val[kn] = v - wht[kn] = divmod(i2 + 0.5, 1)[1] if i2 < ns12 else 1.0 - - # c. all other intermediate pixels: - val[k + 1:kn] = image_data[ii1 + 1:ii2, x] - y[k:kn + 1] = np.arange(ii1, ii2 + 1, 1, dtype=np.float64) - - k += ii2 - ii1 + 1 - - return y, val, wht # pixel coordinate, value, wheight=fractional pixel area - - -def _coalesce_bounds(segments): - """Optimize limits. + # This is optimal extraction (well, profile-based extraction). It + # is "optimal extraction" if the weights are the inverse variances, + # but you must be careful about biases in this case. + elif extraction_type == 'optimal': - Parameters: - ----------- - segments : list of two-element lists of float - Each element of `segments` is a list containing the lower and upper - limits of a source or background extraction region for one of the - columns in the input image. + (fluxes, var_rn, var_phnoise, var_flat, + bkg, var_bkg_rn, var_bkg_phnoise, + var_bkg_flat, npixels, model) = _optimal_extract( + image_sub, profiles_2d, variance_rn, variance_phnoise, variance_flat, + weights, profile_bg, fit_bkg, bkg_order, + bkg_2d, var_bkg_rn, var_bkg_phnoise, var_bkg_flat, model) - Returns: - -------- - list of two-element lists of float - A copy of `segments`, but sorted, and with overlapping intervals - merged into a smaller number of equivalent intervals. - """ - if not isinstance(segments, list): - raise TypeError("'segments' must be a list") - - intervals = copy.deepcopy(segments) - - if all([isinstance(x, list) for x in intervals]): - # make sure each nested list is a list of two numbers: - for x in intervals: - if len(x) != 2: - raise ValueError("Each list in the 'segments' list must have " - "exactly two elements") - try: - map(float, x) - except TypeError: - raise TypeError("Each segment in 'segments' must be a list of " - "two numbers") - else: - # we have a single "interval" (or "segment"): - if len(intervals) != 2: - raise ValueError("'segments' list must be a list of lists of two " - "numbers or 'segments' must be a list of exactly " - "two numbers") - try: - map(float, intervals) - except TypeError: - raise TypeError("Each bound in 'segments' must be a number") - - intervals.sort() - return intervals - - # sort each "segment" in an increasing order, then sort the entire list - # in the increasing order of lower limit: - for segment in intervals: - segment.sort() - intervals.sort(key=lambda x: x[0]) - - # coalesce intervals/segments: - if len(intervals) > 0: - cint = [intervals.pop(0)] else: - return [[]] + raise ValueError(f"Extraction method {extraction_type} not " + f"supported with {nobjects} input profiles.") - while len(intervals) > 0: - pint = cint[-1] - interval = intervals.pop(0) - if interval[0] <= pint[1]: - pint[1] = max(interval[1], pint[1]) - continue - cint.append(interval) + no_data = np.isclose(npixels, 0) + fluxes[no_data] = np.nan - return cint + return (fluxes, var_rn, var_phnoise, var_flat, + bkg, var_bkg_rn, var_bkg_phnoise, var_bkg_flat, npixels, model) diff --git a/jwst/extract_1d/extract_1d_step.py b/jwst/extract_1d/extract_1d_step.py index 3ff620e6c5..a613bedf26 100644 --- a/jwst/extract_1d/extract_1d_step.py +++ b/jwst/extract_1d/extract_1d_step.py @@ -5,6 +5,7 @@ from ..stpipe import Step from . import extract from .soss_extract import soss_extract +from .ifu import ifu_extract1d __all__ = ["Extract1dStep"] @@ -14,6 +15,26 @@ class Extract1dStep(Step): Attributes ---------- + subtract_background : bool or None + A flag which indicates whether the background should be subtracted. + If None, the value in the extract_1d reference file will be used. + If not None, this parameter overrides the value in the + extract_1d reference file. + + apply_apcorr : bool + Switch to select whether to apply an APERTURE correction during + the Extract1dStep. Default is True. + + use_source_posn : bool or None + If True, the source and background extraction positions specified in + the extract1d reference file (or the default position, if there is no + reference file) will be shifted to account for the computed position + of the source in the data. If None (the default), the values in the + reference file will be used. Aperture offset is determined by computing + the pixel location of the source based on its RA and Dec. It does not + make sense to apply aperture offsets for extended sources, so this + parameter can be overridden (set to False) internally by the step. + smoothing_length : int or None If not None, the background regions (if any) will be smoothed with a boxcar function of this width along the dispersion @@ -36,10 +57,6 @@ class Extract1dStep(Step): If both `smoothing_length` and `bkg_order` are not None, the boxcar smoothing will be done first. - bkg_sigma_clip : float - Background sigma clipping value to use on background to remove outliers - and maximize the quality of the 1d spectrum - log_increment : int if `log_increment` is greater than 0 (the default is 50) and the input data are multi-integration (which can be CubeModel or @@ -47,21 +64,13 @@ class Extract1dStep(Step): INFO every `log_increment` integrations. This is intended to provide progress information when invoking the step interactively. - subtract_background : bool or None - A flag which indicates whether the background should be subtracted. - If None, the value in the extract_1d reference file will be used. - If not None, this parameter overrides the value in the - extract_1d reference file. + save_profile : bool + If True, the spatial profile containing the extraction aperture + is saved to disk. Ignored for IFU and NIRISS SOSS extractions. - use_source_posn : bool or None - If True, the source and background extraction positions specified in - the extract1d reference file (or the default position, if there is no - reference file) will be shifted to account for the computed position - of the source in the data. If None (the default), the values in the - reference file will be used. Aperture offset is determined by computing - the pixel location of the source based on its RA and Dec. It does not - make sense to apply aperture offsets for extended sources, so this - parameter can be overridden (set to False) internally by the step. + save_scene_model : bool + If True, a model of the 2D flux as defined by the extraction aperture + is saved to disk. Ignored for IFU and NIRISS SOSS extractions. center_xy : int or None A list of 2 pixel coordinate values at which to place the center @@ -69,14 +78,14 @@ class Extract1dStep(Step): Two values, in x,y order, are used for extraction from IFU cubes. Default is None. - apply_apcorr : bool - Switch to select whether or not to apply an APERTURE correction during - the Extract1dStep. Default is True - ifu_autocen : bool Switch to turn on auto-centering for point source spectral extraction in IFU mode. Default is False. + bkg_sigma_clip : float + Background sigma clipping value to use on background to remove outliers + and maximize the quality of the 1d spectrum. Used for IFU mode only. + ifu_rfcorr : bool Switch to select whether or not to apply a 1d residual fringe correction for MIRI MRS IFU spectra. Default is False. @@ -108,29 +117,6 @@ class Extract1dStep(Step): Oversampling factor of the underlying wavelength grid when modeling the SOSS trace in ATOCA. Default is 2. - soss_transform : list[float] - Rotation applied to the reference files to match the observation orientation. - Default is None. - - soss_tikfac : float - The regularization factor used for extraction in ATOCA. If left to default - value of None, ATOCA will find an optimized value. - - soss_width : float - Aperture width used to extract the SOSS spectrum from the decontaminated - trace in ATOCA. Default is 40. - - soss_bad_pix : str - Method used to handle bad pixels, accepts either "model" or "masking". Default - method is "model". - - soss_modelname : str - Filename for optional model output of ATOCA traces and pixel weights. - - soss_estimate : str or SpecModel or None - Filename or SpecModel of the estimate of the target flux. The estimate must - be a SpecModel with wavelength and flux values. - soss_wave_grid_in : str or SossWaveGrid or None Filename or SossWaveGrid containing the wavelength grid used by ATOCA to model each pixel valid pixel of the detector. If not given, the grid is determined @@ -140,6 +126,10 @@ class Extract1dStep(Step): soss_wave_grid_out : str or None Filename to hold the wavelength grid calculated by ATOCA. + soss_estimate : str or SpecModel or None + Filename or SpecModel of the estimate of the target flux. The estimate must + be a SpecModel with wavelength and flux values. + soss_rtol : float The relative tolerance needed on a pixel model. It is used to determine the sampling of the soss_wave_grid when not directly given. @@ -147,25 +137,46 @@ class Extract1dStep(Step): soss_max_grid_size: int Maximum grid size allowed. It is used when soss_wave_grid is not provided to make sure the computation time or the memory used stays reasonable. + + soss_tikfac : float + The regularization factor used for extraction in ATOCA. If left to default + value of None, ATOCA will find an optimized value. + + soss_width : float + Aperture width used to extract the SOSS spectrum from the decontaminated + trace in ATOCA. Default is 40. + + soss_bad_pix : str + Method used to handle bad pixels, accepts either "model" or "masking". Default + method is "model". + + soss_modelname : str + Filename for optional model output of ATOCA traces and pixel weights. + """ class_alias = "extract_1d" spec = """ + subtract_background = boolean(default=None) # subtract background? + apply_apcorr = boolean(default=True) # apply aperture corrections? + + use_source_posn = boolean(default=None) # use source coords to center extractions? smoothing_length = integer(default=None) # background smoothing size bkg_fit = option("poly", "mean", "median", None, default=None) # background fitting type bkg_order = integer(default=None, min=0) # order of background polynomial fit - bkg_sigma_clip = float(default=3.0) # background sigma clipping threshold log_increment = integer(default=50) # increment for multi-integration log messages - subtract_background = boolean(default=None) # subtract background? - use_source_posn = boolean(default=None) # use source coords to center extractions? + save_profile = boolean(default=False) # save spatial profile to disk + save_scene_model = boolean(default=False) # save flux model to disk + center_xy = float_list(min=2, max=2, default=None) # IFU extraction x/y center - apply_apcorr = boolean(default=True) # apply aperture corrections? ifu_autocen = boolean(default=False) # Auto source centering for IFU point source data. + bkg_sigma_clip = float(default=3.0) # background sigma clipping threshold for IFU ifu_rfcorr = boolean(default=False) # Apply 1d residual fringe correction ifu_set_srctype = option("POINT", "EXTENDED", None, default=None) # user-supplied source type ifu_rscale = float(default=None, min=0.5, max=3) # Radius in terms of PSF FWHM to scale extraction radii ifu_covar_scale = float(default=1.0) # Scaling factor to apply to errors to account for IFU cube covariance + soss_atoca = boolean(default=True) # use ATOCA algorithm soss_threshold = float(default=1e-2) # TODO: threshold could be removed from inputs. Its use is too specific now. soss_n_os = integer(default=2) # minimum oversampling factor of the underlying wavelength grid used when modeling trace. @@ -182,6 +193,146 @@ class Extract1dStep(Step): reference_file_types = ['extract1d', 'apcorr', 'pastasoss', 'specprofile', 'speckernel'] + def _get_extract_reference_files_by_mode(self, model, exp_type): + """Get extraction reference files with special handling by exposure type.""" + if isinstance(model, ModelContainer): + model = model[0] + + if exp_type not in extract.WFSS_EXPTYPES: + extract_ref = self.get_reference_file(model, 'extract1d') + else: + extract_ref = 'N/A' + if extract_ref != 'N/A': + self.log.info(f'Using EXTRACT1D reference file {extract_ref}') + + if self.apply_apcorr: + apcorr_ref = self.get_reference_file(model, 'apcorr') + else: + apcorr_ref = 'N/A' + if apcorr_ref != 'N/A': + self.log.info(f'Using APCORR file {apcorr_ref}') + + return extract_ref, apcorr_ref + + def _extract_soss(self, model): + """Extract NIRISS SOSS spectra.""" + # Set the filter configuration + if model.meta.instrument.filter == 'CLEAR': + self.log.info('Exposure is through the GR700XD + CLEAR (science).') + soss_filter = 'CLEAR' + else: + self.log.error('The SOSS extraction is implemented for the CLEAR filter only. ' + f'Requested filter is {model.meta.instrument.filter}.') + self.log.error('extract_1d will be skipped.') + model.meta.cal_step.extract_1d = 'SKIPPED' + return model + + # Set the subarray mode being processed + if model.meta.subarray.name == 'SUBSTRIP256': + self.log.info('Exposure is in the SUBSTRIP256 subarray.') + self.log.info('Traces 1 and 2 will be modelled and decontaminated before extraction.') + subarray = 'SUBSTRIP256' + elif model.meta.subarray.name == 'SUBSTRIP96': + self.log.info('Exposure is in the SUBSTRIP96 subarray.') + self.log.info('Traces of orders 1 and 2 will be modelled but only order 1 ' + 'will be decontaminated before extraction.') + subarray = 'SUBSTRIP96' + else: + self.log.error('The SOSS extraction is implemented for the SUBSTRIP256 ' + 'and SUBSTRIP96 subarrays only. Subarray is currently ' + f'{model.meta.subarray.name}.') + self.log.error('Extract1dStep will be skipped.') + model.meta.cal_step.extract_1d = 'SKIPPED' + return model + + # Load reference files. + pastasoss_ref_name = self.get_reference_file(model, 'pastasoss') + specprofile_ref_name = self.get_reference_file(model, 'specprofile') + speckernel_ref_name = self.get_reference_file(model, 'speckernel') + + # Build SOSS kwargs dictionary. + soss_kwargs = dict() + soss_kwargs['threshold'] = self.soss_threshold + soss_kwargs['n_os'] = self.soss_n_os + soss_kwargs['tikfac'] = self.soss_tikfac + soss_kwargs['width'] = self.soss_width + soss_kwargs['bad_pix'] = self.soss_bad_pix + soss_kwargs['subtract_background'] = self.subtract_background + soss_kwargs['rtol'] = self.soss_rtol + soss_kwargs['max_grid_size'] = self.soss_max_grid_size + soss_kwargs['wave_grid_in'] = self.soss_wave_grid_in + soss_kwargs['wave_grid_out'] = self.soss_wave_grid_out + soss_kwargs['estimate'] = self.soss_estimate + soss_kwargs['atoca'] = self.soss_atoca + # Set flag to output the model and the tikhonov tests + soss_kwargs['model'] = True if self.soss_modelname else False + + # Run the extraction. + result, ref_outputs, atoca_outputs = soss_extract.run_extract1d( + model, + pastasoss_ref_name, + specprofile_ref_name, + speckernel_ref_name, + subarray, + soss_filter, + soss_kwargs) + + # Set the step flag to complete + if result is None: + return None + else: + result.meta.cal_step.extract_1d = 'COMPLETE' + result.meta.target.source_type = None + + model.close() + + if self.soss_modelname: + soss_modelname = self.make_output_path( + basepath=self.soss_modelname, + suffix='SossExtractModel' + ) + ref_outputs.save(soss_modelname) + + if self.soss_modelname: + soss_modelname = self.make_output_path( + basepath=self.soss_modelname, + suffix='AtocaSpectra' + ) + atoca_outputs.save(soss_modelname) + + return result + + def _extract_ifu(self, model, exp_type, extract_ref, apcorr_ref): + """Extract IFU spectra from a single datamodel.""" + source_type = model.meta.target.source_type + if self.ifu_set_srctype is not None and exp_type == 'MIR_MRS': + source_type = self.ifu_set_srctype + self.log.info(f"Overriding source type and setting it to {self.ifu_set_srctype}") + + result = ifu_extract1d( + model, extract_ref, source_type, self.subtract_background, + self.bkg_sigma_clip, apcorr_ref, self.center_xy, + self.ifu_autocen, self.ifu_rfcorr, self.ifu_rscale, + self.ifu_covar_scale + ) + return result + + def _save_intermediate(self, intermediate_model, suffix): + """Save an intermediate output file.""" + if isinstance(intermediate_model, ModelContainer): + # Save the profile with the slit name + suffix 'profile' + for model in intermediate_model: + slit = str(model.name).lower() + output_path = self.make_output_path(suffix=f'{slit}_{suffix}') + self.log.info(f"Saving {suffix} {output_path}") + model.save(output_path) + else: + # Only one profile - just use the suffix 'profile' + output_path = self.make_output_path(suffix=suffix) + self.log.info(f"Saving {suffix} {output_path}") + intermediate_model.save(output_path) + intermediate_model.close() + def process(self, input): """Execute the step. @@ -202,331 +353,94 @@ def process(self, input): else: input_model = datamodels.open(input) - was_source_model = False # default value - if isinstance(input_model, datamodels.CubeModel): - # It's a 3-D multi-integration model - self.log.debug('Input is a CubeModel for a multiple integ. file') - elif isinstance(input_model, datamodels.ImageModel): - # It's a single 2-D image. This could be a resampled 2-D image - self.log.debug('Input is an ImageModel') - elif isinstance(input_model, SourceModelContainer): - self.log.debug('Input is a SourceModelContainer') - was_source_model = True - elif isinstance(input_model, ModelContainer): - self.log.debug('Input is a ModelContainer') + if isinstance(input_model, (datamodels.CubeModel, datamodels.ImageModel, + datamodels.SlitModel, datamodels.IFUCubeModel, + ModelContainer, SourceModelContainer)): + # Acceptable input type, just log it + self.log.debug(f'Input is a {str(type(input_model))}.') elif isinstance(input_model, datamodels.MultiSlitModel): - # If input is a 3D rateints (which is unsupported) skip the step + # If input is multislit, with 3D calints, skip the step + self.log.debug('Input is a MultiSlitModel') if len((input_model[0]).shape) == 3: self.log.warning('3D input is unsupported; step will be skipped') input_model.meta.cal_step.extract_1d = 'SKIPPED' return input_model - self.log.debug('Input is a MultiSlitModel') - elif isinstance(input_model, datamodels.MultiExposureModel): - self.log.warning('Input is a MultiExposureModel, ' - 'which is not currently supported') - elif isinstance(input_model, datamodels.IFUCubeModel): - self.log.debug('Input is an IFUCubeModel') - elif isinstance(input_model, datamodels.SlitModel): - # NRS_BRIGHTOBJ and MIRI LRS fixed-slit (resampled) modes - self.log.debug('Input is a SlitModel') else: self.log.error(f'Input is a {str(type(input_model))}, ') - self.log.error('which was not expected for extract_1d') - self.log.error('extract_1d will be skipped.') + self.log.error('which was not expected for extract_1d.') + self.log.error('The extract_1d step will be skipped.') input_model.meta.cal_step.extract_1d = 'SKIPPED' return input_model - if isinstance(input_model, datamodels.IFUCubeModel): + if not isinstance(input_model, ModelContainer): exp_type = input_model.meta.exposure.type - elif isinstance(input_model, ModelContainer): - exp_type = input_model[0].meta.exposure.type + + # Make the input iterable + input_model = [input_model] else: - exp_type = None - - if self.ifu_rfcorr: - if exp_type != "MIR_MRS": - self.log.warning("The option to apply a residual refringe correction is" - f" not supported for {input_model.meta.exposure.type} data.") - - if self.ifu_rscale is not None: - if exp_type != "MIR_MRS": - self.log.warning("The option to change the extraction radius is" - f" not supported for {input_model.meta.exposure.type} data.") - - if self.ifu_set_srctype is not None: - if exp_type != "MIR_MRS": - self.log.warning("The option to change the source type is" - f" not supported for {input_model.meta.exposure.type} data.") - - # ______________________________________________________________________ - # Do the extraction for ModelContainer - this might only be WFSS data - if isinstance(input_model, ModelContainer): - - # This is the branch WFSS data take - if len(input_model) > 1: - self.log.debug(f"Input contains {len(input_model)} items") - - # -------------------------------------------------------------- - # Data is WFSS - if input_model[0].meta.exposure.type in extract.WFSS_EXPTYPES: - - # For WFSS level-3, the input is a single entry of a - # SourceContainer, which contains a list of multiple - # SlitModels for a single source. Send the whole list - # into extract1d and put all results in a single product. - apcorr_ref = ( - self.get_reference_file(input_model[0], 'apcorr') if self.apply_apcorr is True else 'N/A' - ) + exp_type = input_model[0].meta.exposure.type + self.log.debug(f"Input for EXP_TYPE {exp_type} contains {len(input_model)} items") - if apcorr_ref == 'N/A': - self.log.info('APCORR reference file name is "N/A"') - self.log.info('APCORR will NOT be applied') - else: - self.log.info(f'Using APCORR file {apcorr_ref}') + if len(input_model) > 1 and exp_type in extract.WFSS_EXPTYPES: + # For WFSS level-3, the input is a single entry of a + # SourceContainer, which contains a list of multiple + # SlitModels for a single source. Send the whole container + # into extract1d and put all results in a single product. + input_model = [input_model] - extract_ref = 'N/A' - self.log.info('No EXTRACT1D reference file will be used') + if exp_type == 'NIS_SOSS': + # Data is NIRISS SOSS observation, use its own extraction routines + self.log.info( + 'Input is a NIRISS SOSS observation, the specialized SOSS ' + 'extraction (ATOCA) will be used.') - result = extract.run_extract1d( - input_model, + # There is only one input model for this mode + model = input_model[0] + result = self._extract_soss(model) + + else: + result = ModelContainer() + for model in input_model: + # Get the reference file names + extract_ref, apcorr_ref = self._get_extract_reference_files_by_mode( + model, exp_type) + + profile = None + scene_model = None + if isinstance(model, datamodels.IFUCubeModel): + # Call the IFU specific extraction routine + extracted = self._extract_ifu(model, exp_type, extract_ref, apcorr_ref) + else: + # Call the general extraction routine + extracted, profile, scene_model = extract.run_extract1d( + model, extract_ref, apcorr_ref, self.smoothing_length, self.bkg_fit, self.bkg_order, - self.bkg_sigma_clip, self.log_increment, self.subtract_background, self.use_source_posn, - self.center_xy, - self.ifu_autocen, - self.ifu_rfcorr, - self.ifu_set_srctype, - self.ifu_rscale, - self.ifu_covar_scale, - was_source_model=was_source_model - ) - # Set the step flag to complete - result.meta.cal_step.extract_1d = 'COMPLETE' - - # -------------------------------------------------------------- - # Data is a ModelContainer but is not WFSS - else: - result = ModelContainer() - for model in input_model: - # Get the reference file names - extract_ref = self.get_reference_file(model, 'extract1d') - self.log.info(f'Using EXTRACT1D reference file {extract_ref}') - - apcorr_ref = self.get_reference_file(model, 'apcorr') if self.apply_apcorr is True else 'N/A' - - if apcorr_ref == 'N/A': - self.log.info('APCORR reference file name is "N/A"') - self.log.info('APCORR will NOT be applied') - else: - self.log.info(f'Using APCORR file {apcorr_ref}') - - temp = extract.run_extract1d( - model, - extract_ref, - apcorr_ref, - self.smoothing_length, - self.bkg_fit, - self.bkg_order, - self.bkg_sigma_clip, - self.log_increment, - self.subtract_background, - self.use_source_posn, - self.center_xy, - self.ifu_autocen, - self.ifu_rfcorr, - self.ifu_set_srctype, - self.ifu_rscale, - self.ifu_covar_scale, - was_source_model=was_source_model, - ) - # Set the step flag to complete in each MultiSpecModel - temp.meta.cal_step.extract_1d = 'COMPLETE' - result.append(temp) - del temp - # ------------------------------------------------------------------------ - # Still in ModelContainer type, but only 1 model - elif len(input_model) == 1: - if input_model[0].meta.exposure.type in extract.WFSS_EXPTYPES: - extract_ref = 'N/A' - self.log.info('No EXTRACT1D reference file will be used') - else: - # Get the extract1d reference file name for the one model in input - extract_ref = self.get_reference_file(input_model[0], 'extract1d') - self.log.info(f'Using EXTRACT1D reference file {extract_ref}') - - apcorr_ref = self.get_reference_file(input_model[0], 'apcorr') if self.apply_apcorr is True else 'N/A' - - if apcorr_ref == 'N/A': - self.log.info('APCORR reference file name is "N/A"') - self.log.info('APCORR will NOT be applied') - else: - self.log.info(f'Using APCORR file {apcorr_ref}') - - result = extract.run_extract1d( - input_model[0], - extract_ref, - apcorr_ref, - self.smoothing_length, - self.bkg_fit, - self.bkg_order, - self.bkg_sigma_clip, - self.log_increment, - self.subtract_background, - self.use_source_posn, - self.center_xy, - self.ifu_autocen, - self.ifu_rfcorr, - self.ifu_set_srctype, - self.ifu_rscale, - self.ifu_covar_scale, - was_source_model=was_source_model, - ) - - # Set the step flag to complete - result.meta.cal_step.extract_1d = 'COMPLETE' - else: - self.log.error('Input model is empty;') - self.log.error('extract_1d will be skipped.') - return input_model - - # ______________________________________________________________________ - # Data that is not a ModelContainer (IFUCube and other single models) - else: - # Data is NRISS SOSS observation. - if input_model.meta.exposure.type == 'NIS_SOSS': - - self.log.info( - 'Input is a NIRISS SOSS observation, the specialized SOSS extraction (ATOCA) will be used.') - - # Set the filter configuration - if input_model.meta.instrument.filter == 'CLEAR': - self.log.info('Exposure is through the GR700XD + CLEAR (science).') - soss_filter = 'CLEAR' - else: - self.log.error('The SOSS extraction is implemented for the CLEAR filter only.' - f'Requested filter is {input_model.meta.instrument.filter}.') - self.log.error('extract_1d will be skipped.') - input_model.meta.cal_step.extract_1d = 'SKIPPED' - return input_model - - # Set the subarray mode being processed - if input_model.meta.subarray.name == 'SUBSTRIP256': - self.log.info('Exposure is in the SUBSTRIP256 subarray.') - self.log.info('Traces 1 and 2 will be modelled and decontaminated before extraction.') - subarray = 'SUBSTRIP256' - elif input_model.meta.subarray.name == 'SUBSTRIP96': - self.log.info('Exposure is in the SUBSTRIP96 subarray.') - self.log.info('Traces of orders 1 and 2 will be modelled but only order 1' - ' will be decontaminated before extraction.') - subarray = 'SUBSTRIP96' - else: - self.log.error('The SOSS extraction is implemented for the SUBSTRIP256' - 'and SUBSTRIP96 subarrays only. Subarray is currently ' - f'{input_model.meta.subarray.name}.') - self.log.error('Extract1dStep will be skipped.') - input_model.meta.cal_step.extract_1d = 'SKIPPED' - return input_model - - # Load reference files. - pastasoss_ref_name = self.get_reference_file(input_model, 'pastasoss') - specprofile_ref_name = self.get_reference_file(input_model, 'specprofile') - speckernel_ref_name = self.get_reference_file(input_model, 'speckernel') - - # Build SOSS kwargs dictionary. - soss_kwargs = dict() - soss_kwargs['threshold'] = self.soss_threshold - soss_kwargs['n_os'] = self.soss_n_os - soss_kwargs['tikfac'] = self.soss_tikfac - soss_kwargs['width'] = self.soss_width - soss_kwargs['bad_pix'] = self.soss_bad_pix - soss_kwargs['subtract_background'] = self.subtract_background - soss_kwargs['rtol'] = self.soss_rtol - soss_kwargs['max_grid_size'] = self.soss_max_grid_size - soss_kwargs['wave_grid_in'] = self.soss_wave_grid_in - soss_kwargs['wave_grid_out'] = self.soss_wave_grid_out - soss_kwargs['estimate'] = self.soss_estimate - soss_kwargs['atoca'] = self.soss_atoca - # Set flag to output the model and the tikhonov tests - soss_kwargs['model'] = True if self.soss_modelname else False - - # Run the extraction. - result, ref_outputs, atoca_outputs = soss_extract.run_extract1d( - input_model, - pastasoss_ref_name, - specprofile_ref_name, - speckernel_ref_name, - subarray, - soss_filter, - soss_kwargs) - - # Set the step flag to complete - if result is None: - return None - else: - result.meta.cal_step.extract_1d = 'COMPLETE' - result.meta.target.source_type = None - - input_model.close() - - if self.soss_modelname: - soss_modelname = self.make_output_path( - basepath=self.soss_modelname, - suffix='SossExtractModel' - ) - ref_outputs.save(soss_modelname) - - if self.soss_modelname: - soss_modelname = self.make_output_path( - basepath=self.soss_modelname, - suffix='AtocaSpectra' + self.save_profile, + self.save_scene_model ) - atoca_outputs.save(soss_modelname) - else: - # Get the reference file names - if input_model.meta.exposure.type in extract.WFSS_EXPTYPES: - extract_ref = 'N/A' - self.log.info('No EXTRACT1D reference file will be used') - else: - extract_ref = self.get_reference_file(input_model, 'extract1d') - self.log.info(f'Using EXTRACT1D reference file {extract_ref}') - apcorr_ref = self.get_reference_file(input_model, 'apcorr') if self.apply_apcorr is True else 'N/A' + # Set the step flag to complete in each model + extracted.meta.cal_step.extract_1d = 'COMPLETE' + result.append(extracted) + del extracted - if apcorr_ref == 'N/A': - self.log.info('APCORR reference file name is "N/A"') - self.log.info('APCORR will NOT be applied') - else: - self.log.info(f'Using APCORR file {apcorr_ref}') - - result = extract.run_extract1d( - input_model, - extract_ref, - apcorr_ref, - self.smoothing_length, - self.bkg_fit, - self.bkg_order, - self.bkg_sigma_clip, - self.log_increment, - self.subtract_background, - self.use_source_posn, - self.center_xy, - self.ifu_autocen, - self.ifu_rfcorr, - self.ifu_set_srctype, - self.ifu_rscale, - self.ifu_covar_scale, - was_source_model=False, - ) + # Save profile if needed + if self.save_profile and profile is not None: + self._save_intermediate(profile, 'profile') - # Set the step flag to complete - result.meta.cal_step.extract_1d = 'COMPLETE' + # Save model if needed + if self.save_scene_model and scene_model is not None: + self._save_intermediate(scene_model, 'scene_model') - input_model.close() + # If only one result, return the model instead of the container + if len(result) == 1: + result = result[0] return result diff --git a/jwst/extract_1d/ifu.py b/jwst/extract_1d/ifu.py index d2d7334509..3a68efd150 100644 --- a/jwst/extract_1d/ifu.py +++ b/jwst/extract_1d/ifu.py @@ -1,31 +1,27 @@ import logging + import numpy as np +from astropy import stats +from astropy.stats import sigma_clipped_stats as sigclip from photutils.aperture import (CircularAperture, CircularAnnulus, RectangularAperture, aperture_photometry) +from photutils.detection import DAOStarFinder +from scipy.interpolate import interp1d from stdatamodels.jwst import datamodels from stdatamodels.jwst.datamodels import dqflags -from .apply_apcorr import select_apcorr -from ..assign_wcs.util import compute_scale -from astropy import stats - -from . import spec_wcs -from scipy.interpolate import interp1d +from jwst.assign_wcs.util import compute_scale +from jwst.extract_1d import spec_wcs +from jwst.extract_1d.apply_apcorr import select_apcorr +from jwst.extract_1d.extract import read_apcorr_ref +from jwst.residual_fringe import utils as rfutils -from astropy.stats import sigma_clipped_stats as sigclip -from photutils.detection import DAOStarFinder -from ..residual_fringe import utils as rfutils +__all__ = ['ifu_extract1d'] log = logging.getLogger(__name__) log.setLevel(logging.DEBUG) -# These values are used to indicate whether the input extract1d reference file -# (if any) is ASDF (default) or IMAGE - -FILE_TYPE_ASDF = "ASDF" -FILE_TYPE_IMAGE = "IMAGE" - # This is to prevent calling offset_from_offset multiple times for # multi-integration data. OFFSET_NOT_ASSIGNED_YET = "not assigned yet" @@ -35,8 +31,8 @@ HUGE_DIST = 1.e10 -def ifu_extract1d(input_model, ref_dict, source_type, subtract_background, - bkg_sigma_clip, apcorr_ref_model=None, center_xy=None, +def ifu_extract1d(input_model, ref_file, source_type, subtract_background, + bkg_sigma_clip, apcorr_ref_file=None, center_xy=None, ifu_autocen=False, ifu_rfcorr=False, ifu_rscale=None, ifu_covar_scale=1.0): """Extract a 1-D spectrum from an IFU cube. @@ -45,8 +41,8 @@ def ifu_extract1d(input_model, ref_dict, source_type, subtract_background, input_model : JWST data model for an IFU cube (IFUCubeModel) The input model. - ref_dict : dict - The contents of the extract1d reference file. + ref_file : str + File name for the extract1d reference file, in ASDF format. source_type : string "POINT" or "EXTENDED" @@ -60,8 +56,8 @@ def ifu_extract1d(input_model, ref_dict, source_type, subtract_background, bkg_sigma_clip : float Background sigma clipping value to use to remove noise/outliers in background - apcorr_ref_model : apcorr datamodel or None - Aperture correction table. + apcorr_ref_file : str or None + File name for aperture correction reference file. center_xy : float or None A list of 2 pixel coordinate values at which to place the center @@ -117,7 +113,7 @@ def ifu_extract1d(input_model, ref_dict, source_type, subtract_background, if slitname is None: slitname = "ANY" - extract_params = get_extract_parameters(ref_dict, bkg_sigma_clip, slitname) + extract_params = get_extract_parameters(ref_file, bkg_sigma_clip) # Add info about IFU auto-centroiding, residual fringe correction and extraction radius scale # to extract_params for use later @@ -146,20 +142,10 @@ def ifu_extract1d(input_model, ref_dict, source_type, subtract_background, "the source is extended.") extract_params['subtract_background'] = subtract_background - if extract_params: - if extract_params['ref_file_type'] == FILE_TYPE_ASDF: - (ra, dec, wavelength, temp_flux, f_var_poisson, f_var_rnoise, f_var_flat, - background, b_var_poisson, b_var_rnoise, b_var_flat, npixels, dq, npixels_bkg, - radius_match, x_center, y_center) = \ - extract_ifu(input_model, source_type, extract_params) - else: # FILE_TYPE_IMAGE - (ra, dec, wavelength, temp_flux, f_var_poisson, f_var_rnoise, f_var_flat, - background, b_var_poisson, b_var_rnoise, b_var_flat, npixels, dq, npixels_bkg, - x_center, y_center) = \ - image_extract_ifu(input_model, source_type, extract_params) - else: - log.critical('Missing extraction parameters.') - raise ValueError('Missing extraction parameters.') + (ra, dec, wavelength, temp_flux, f_var_poisson, f_var_rnoise, f_var_flat, + background, b_var_poisson, b_var_rnoise, b_var_flat, npixels, dq, npixels_bkg, + radius_match, x_center, y_center) = \ + extract_ifu(input_model, source_type, extract_params) npixels_temp = np.where(npixels > 0., npixels, 1.) npixels_bkg_temp = np.where(npixels_bkg > 0., npixels_bkg, 1.) @@ -296,9 +282,10 @@ def ifu_extract1d(input_model, ref_dict, source_type, subtract_background, spec.extraction_x = x_center spec.extraction_y = y_center - if source_type == 'POINT' and apcorr_ref_model is not None: - log.info('Applying Aperture correction.') + if source_type == 'POINT' and apcorr_ref_file is not None and apcorr_ref_file != 'N/A': + apcorr_ref_model = read_apcorr_ref(apcorr_ref_file, input_model.meta.exposure.type) + log.info('Applying Aperture correction.') if instrument == 'NIRSPEC': wl = np.median(wavelength) else: @@ -325,69 +312,50 @@ def ifu_extract1d(input_model, ref_dict, source_type, subtract_background, # See output_model.spec[0].meta.wcs instead. output_model.meta.wcs = None + output_model.meta.wcs = None # See output_model.spec[i].meta.wcs instead. + return output_model -def get_extract_parameters(ref_dict, bkg_sigma_clip, slitname): +def get_extract_parameters(ref_file, bkg_sigma_clip): """Read extraction parameters for an IFU. Parameters ---------- - ref_dict : dict - The contents of the extract1d reference file. + ref_file : dict + File name for the extract1d reference file, in ASDF format - slitname : str - The name of the slit, or "ANY". + bkg_sigma_clip : float + Background sigma clipping value to use to remove noise/outliers in background. Returns ------- dict The extraction parameters. """ - extract_params = {} + # for consistency put the bkg_sigma_clip in dictionary: extract_params extract_params['bkg_sigma_clip'] = bkg_sigma_clip - if ref_dict['ref_file_type'] == FILE_TYPE_ASDF: - extract_params['ref_file_type'] = FILE_TYPE_ASDF - refmodel = ref_dict['ref_model'] - subtract_background = refmodel.meta.subtract_background - method = refmodel.meta.method - subpixels = refmodel.meta.subpixels - - data = refmodel.data - wavelength = data.wavelength - radius = data.radius - inner_bkg = data.inner_bkg - outer_bkg = data.outer_bkg - - extract_params['subtract_background'] = bool(subtract_background) - extract_params['method'] = method - extract_params['subpixels'] = subpixels - extract_params['wavelength'] = wavelength - extract_params['radius'] = radius - extract_params['inner_bkg'] = inner_bkg - extract_params['outer_bkg'] = outer_bkg - - elif ref_dict['ref_file_type'] == FILE_TYPE_IMAGE: - extract_params['ref_file_type'] = FILE_TYPE_IMAGE - foundit = False - for im in ref_dict['ref_model'].images: - if (im.name is None or im.name == "ANY" or slitname == "ANY" or - im.name == slitname): - extract_params['ref_image'] = im - foundit = True - break - - if not foundit: - log.error("No match for slit name %s in reference image", slitname) - raise RuntimeError("Specify slit name or use 'any' in ref image.") - - else: - log.error("Reference file type %s not recognized", - ref_dict['ref_file_type']) - raise RuntimeError("extract1d reference file must be ASDF, JSON or FITS image.") + refmodel = datamodels.Extract1dIFUModel(ref_file) + subtract_background = refmodel.meta.subtract_background + method = refmodel.meta.method + subpixels = refmodel.meta.subpixels + + data = refmodel.data + wavelength = data.wavelength + radius = data.radius + inner_bkg = data.inner_bkg + outer_bkg = data.outer_bkg + + extract_params['subtract_background'] = bool(subtract_background) + extract_params['method'] = method + extract_params['subpixels'] = subpixels + extract_params['wavelength'] = wavelength + extract_params['radius'] = radius + extract_params['inner_bkg'] = inner_bkg + extract_params['outer_bkg'] = outer_bkg return extract_params @@ -568,9 +536,6 @@ def extract_ifu(input_model, source_type, extract_params): subpixels = extract_params['subpixels'] subtract_background = extract_params['subtract_background'] - radius = None - inner_bkg = None - outer_bkg = None width = None height = None theta = None @@ -650,7 +615,6 @@ def extract_ifu(input_model, source_type, extract_params): # get aperture for extended it will not change with wavelength if source_type == 'EXTENDED': aperture = RectangularAperture(position, width, height, theta) - annulus = None for k in range(shape[0]): # looping over wavelength inner_bkg = None @@ -677,7 +641,6 @@ def extract_ifu(input_model, source_type, extract_params): normalization = 1. temp_weightmap = weightmap[k, :, :] temp_weightmap[temp_weightmap > 1] = 1 - aperture_area = 0 annulus_area = 0 # Make a boolean mask to ignore voxels with no valid data @@ -829,7 +792,7 @@ def locn_from_wcs(input_model, ra_targ, dec_targ): Parameters ---------- - input_model : data model + input_model : JWSTDataModel The input science model. ra_targ, dec_targ : float or None @@ -914,256 +877,6 @@ def celestial_to_cartesian(ra, dec): return cart -def image_extract_ifu(input_model, source_type, extract_params): - """Extraction using a extract1d reference image. - - Extended summary - ---------------- - The IMAGE extract1d reference file should have pixels with avalue of 1 for the - source extraction region, 0 for pixels not to include in source or background, - and -1 for the background region. - For SRCTYPE=POINT the source extraction region is defined by pixels in the ref - image = 1 and the background region is defined by pixels in the ref image with - -1. - For SRCTYPE=EXTENDED the extraction region is defined by pixels in the ref image - = 1 (only the source region is used). The default procedure of using the extract 1d - asdf reference files extracts the entire region for EXTENDED source data. However, - if the user supplies the reference image it is assumed they have defined a specific - region to be extracted instead of the entire field. At each wavelength bin sigma - clipping is performed on the extraction region and is store in the background column of - spec table to be used in masterbackground subtraction. In the extended source case - pixels flagged as background (-1) in the reference image are ignored. - - Parameters - ---------- - input_model : IFUCubeModel - The input model. - - source_type : string - "POINT" or "EXTENDED" - - extract_params : dict - The extraction parameters. One of these is an open file handle - for an image that specifies which pixels should be included as - source and which (if any) should be used as background. - - Returns - ------- - ra, dec : float - ra and dec are the right ascension and declination respectively - at the centroid of the target region in the reference image. - - wavelength : ndarray, 1-D - The wavelength in micrometers at each plane of the IFU cube. - - temp_flux : ndarray, 1-D - The sum of the data values in the extraction aperture minus the - sum of the data values in the background region (scaled by the - ratio of areas), for each plane. - The data values are in units of surface brightness, so this value - isn't really the flux, it's an intermediate value. Dividing by - `npixels` (to compute the average) will give the value for the - `surf_bright` (surface brightness) column, and multiplying by - the solid angle of a pixel will give the flux for a point source. - - f_var_poisson : ndarray, 1-D - The extracted poisson variance values to go along with the - temp_flux array. - - f_var_rnoise : ndarray, 1-D - The extracted read noise variance values to go along with the - temp_flux array. - - f_var_flat : ndarray, 1-D - The extracted flat field variance values to go along with the - temp_flux array. - - background : ndarray, 1-D - The background count rate that was subtracted from the total - source data values to get `temp_flux`. - - b_var_poisson : ndarray, 1-D - The extracted poisson variance values to go along with the - background array. - - b_var_rnoise : ndarray, 1-D - The extracted read noise variance values to go along with the - background array. - - b_var_flat : ndarray, 1-D - The extracted flat field variance values to go along with the - background array. - - npixels : ndarray, 1-D, float64 - For each slice, this is the number of pixels that were added - together to get `temp_flux`. - - dq : ndarray, 1-D, uint32 - The data quality array. - - n_bkg : ndarray, 1-D, float64 - For each slice, this is the number of pixels that were added - together to get background. - - x_center, y_center : float - The x and y center of the extraction region - """ - - data = input_model.data - try: - var_poisson = input_model.var_poisson - var_rnoise = input_model.var_rnoise - var_flat = input_model.var_flat - except AttributeError: - log.info("Input model does not break out variance information. Passing only generalized errors.") - var_poisson = input_model.err * input_model.err - var_rnoise = np.zeros_like(input_model.data) - var_flat = np.zeros_like(input_model.data) - # The axes are (z, y, x) in the sense that (x, y) are the ordinary - # axes of the image for one plane, i.e. at one wavelength. The - # wavelengths vary along the z-axis. - shape = data.shape - - # The dispersion direction is the first axis. - npixels = np.ones(shape[0], dtype=np.float64) - n_bkg = np.ones(shape[0], dtype=np.float64) - - dq = np.zeros(shape[0], dtype=np.uint32) - - ref_image = extract_params['ref_image'] - ref = ref_image.data - subtract_background = extract_params['subtract_background'] - - (mask_target, mask_bkg) = separate_target_and_background(ref) - - # subtracting the background is only allowed for source_type=POINT - # subtract_background = False for EXTENDED data (set in ifu_extract1d) - - if subtract_background: - if mask_bkg is None: - log.info("Skipping background subtraction because " - "background regions are not defined.") - subtract_background = False - - ra_targ = input_model.meta.target.ra - dec_targ = input_model.meta.target.dec - locn = locn_from_wcs(input_model, ra_targ, dec_targ) - x_center = None - y_center = None - if locn is not None: - log.info("Target location is x_center = %g, y_center = %g, " - "based on TARG_RA and TARG_DEC.", locn[0], locn[1]) - - # Use the centroid of mask_target as the point where the target - # would be without any source position correction. - (y0, x0) = im_centroid(data, mask_target) - log.debug("Target location based on reference image is X = %g, Y = %g", - x0, y0) - - # TODO - check if shifting location should be done for reference_image option - if locn is None or np.isnan(locn[0]): - log.warning("Couldn't determine pixel location from WCS, so " - "source position correction will not be applied.") - else: - (x_center, y_center) = locn - # Shift the reference image so it will be centered at locn. - # Only shift by a whole number of pixels. - delta_x = int(round(x_center - x0)) # must be integer - delta_y = int(round(y_center - y0)) - log.debug("Shifting reference image by %g in X and %g in Y", - delta_x, delta_y) - temp = shift_ref_image(mask_target, delta_y, delta_x) - if temp is not None: - mask_target = temp - del temp - if mask_bkg is not None: - mask_bkg = shift_ref_image(mask_bkg, delta_y, delta_x) - # Since we have shifted mask_target and mask_bkg to - # x_center and y_center, update x0 and y0. - x0 = x_center - y0 = y_center - - # Extract the data. - # First add up the values along the x direction, then add up the - # values along the y direction. - gross = np.nansum(np.nansum(data * mask_target, axis=2, dtype=np.float64), axis=1) - f_var_poisson = np.nansum(np.nansum(var_poisson * mask_target, axis=2, dtype=np.float64), axis=1) - f_var_rnoise = np.nansum(np.nansum(var_rnoise * mask_target, axis=2, dtype=np.float64), axis=1) - f_var_flat = np.nansum(np.nansum(var_flat * mask_target, axis=2, dtype=np.float64), axis=1) - - # Compute the number of pixels that were added together to get gross. - normalization = 1. - - weightmap = input_model.weightmap - temp_weightmap = weightmap - temp_weightmap[temp_weightmap > 1] = 1 - npixels[:] = np.nansum(np.nansum(temp_weightmap * mask_target, axis=2, dtype=np.float64), axis=1) - bkg_sigma_clip = extract_params['bkg_sigma_clip'] - - # Point Source data 1. extract background and subtract 2. do not - if source_type == 'POINT': - if subtract_background and mask_bkg is not None: - n_bkg[:] = np.nansum(np.nansum(temp_weightmap * mask_bkg, axis=2, dtype=np.float64), axis=1) - n_bkg[:] = np.where(n_bkg <= 0., 1., n_bkg) - normalization = npixels / n_bkg - background = np.nansum(np.nansum(data * mask_bkg, axis=2, dtype=np.float64), axis=1) - b_var_poisson = np.nansum(np.nansum(var_poisson * mask_bkg, axis=2, dtype=np.float64), axis=1) - b_var_rnoise = np.nansum(np.nansum(var_rnoise * mask_bkg, axis=2, dtype=np.float64), axis=1) - b_var_flat = np.nansum(np.nansum(var_flat * mask_bkg, axis=2, dtype=np.float64), axis=1) - temp_flux = gross - background * normalization - else: - background = np.zeros_like(gross) - b_var_poisson = np.zeros_like(gross) - b_var_rnoise = np.zeros_like(gross) - b_var_flat = np.zeros_like(gross) - temp_flux = gross.copy() - else: - temp_flux = np.nansum(np.nansum(data * mask_target, axis=2, dtype=np.float64), axis=1) - - # Extended source data, sigma clip outliers of extraction region is performed - # at each wavelength plane. - (background, b_var_poisson, b_var_rnoise, - b_var_flat, n_bkg) = sigma_clip_extended_region(data, var_poisson, var_rnoise, - var_flat, mask_target, - temp_weightmap, bkg_sigma_clip) - - del temp_weightmap - - # Compute the ra, dec, and wavelength at the pixels of a column through - # the IFU cube at the target location. ra and dec should be constant - # (so they're scalars), but wavelength will vary from plane to plane. - - log.debug("IFU 1-D extraction parameters (using reference image):") - log.debug(" x_center = %s", str(x0)) - log.debug(" y_center = %s", str(y0)) - log.debug(" subtract_background parameter = %s", str(subtract_background)) - if mask_bkg is not None: - log.debug(" background will be subtracted") - else: - log.debug(" background will not be subtracted") - - (ra, dec, wavelength) = get_coordinates(input_model, x0, y0) - - # Check for NaNs in the wavelength array, flag them in the dq array, - # and truncate the arrays if NaNs are found at endpoints (unless the - # entire array is NaN). - wavelength, dq, nan_slc = nans_in_wavelength(wavelength, dq) - temp_flux = temp_flux[nan_slc] - background = background[nan_slc] - npixels = npixels[nan_slc] - n_bkg = n_bkg[nan_slc] - f_var_poisson = f_var_poisson[nan_slc] - f_var_rnoise = f_var_rnoise[nan_slc] - f_var_flat = f_var_flat[nan_slc] - b_var_poisson = b_var_poisson[nan_slc] - b_var_rnoise = b_var_rnoise[nan_slc] - b_var_flat = b_var_flat[nan_slc] - - return (ra, dec, wavelength, temp_flux, f_var_poisson, f_var_rnoise, f_var_flat, - background, b_var_poisson, b_var_rnoise, b_var_flat, - npixels, dq, n_bkg, x_center, y_center) - - def get_coordinates(input_model, x0, y0): """Get celestial coordinates and wavelengths. diff --git a/jwst/extract_1d/tests/conftest.py b/jwst/extract_1d/tests/conftest.py new file mode 100644 index 0000000000..467f03f385 --- /dev/null +++ b/jwst/extract_1d/tests/conftest.py @@ -0,0 +1,444 @@ +import numpy as np +import pytest +import stdatamodels.jwst.datamodels as dm +from astropy.io import fits +from astropy.table import Table + +from jwst.assign_wcs.util import wcs_bbox_from_shape +from jwst.exp_to_source import multislit_to_container + + +@pytest.fixture() +def simple_wcs(): + shape = (50, 50) + xcenter = shape[1] // 2.0 + + def simple_wcs_function(x, y): + """ Simple WCS for testing """ + crpix1 = xcenter + crpix3 = 1.0 + cdelt1 = 0.1 + cdelt2 = 0.1 + cdelt3 = 0.01 + + crval1 = 45.0 + crval2 = 45.0 + crval3 = 7.5 + + wave = (x + 1 - crpix3) * cdelt3 + crval3 + ra = (x + 1 - crpix1) * cdelt1 + crval1 + dec = np.full_like(ra, crval2 + 1 * cdelt2) + + return ra, dec, wave + + # Add a bounding box + simple_wcs_function.bounding_box = wcs_bbox_from_shape(shape) + + # Add a few expected attributes, so they can be monkeypatched as needed + simple_wcs_function.get_transform = None + simple_wcs_function.backward_transform = None + simple_wcs_function.available_frames = [] + + return simple_wcs_function + + +@pytest.fixture() +def simple_wcs_transpose(): + shape = (50, 50) + ycenter = shape[0] // 2.0 + + def simple_wcs_function(x, y): + """ Simple WCS for testing """ + crpix2 = ycenter + crpix3 = 1.0 + cdelt1 = 0.1 + cdelt2 = 0.1 + cdelt3 = -0.01 + + crval1 = 45.0 + crval2 = 45.0 + crval3 = 7.5 + + wave = (y + 1 - crpix3) * cdelt3 + crval3 + ra = (y + 1 - crpix2) * cdelt1 + crval1 + dec = np.full_like(ra, crval2 + 1 * cdelt2) + + return ra, dec, wave + + # Add a bounding box + simple_wcs_function.bounding_box = wcs_bbox_from_shape(shape) + + # Add a few expected attributes, so they can be monkeypatched as needed + simple_wcs_function.get_transform = None + simple_wcs_function.backward_transform = None + simple_wcs_function.available_frames = [] + + return simple_wcs_function + + +@pytest.fixture() +def simple_wcs_ifu(): + shape = (10, 50, 50) + xcenter = shape[1] // 2.0 + + def simple_wcs_function(x, y, z): + """ Simple WCS for testing """ + crpix1 = xcenter + crpix3 = 1.0 + cdelt1 = 0.1 + cdelt2 = 0.1 + cdelt3 = 0.01 + + crval1 = 45.0 + crval2 = 45.0 + crval3 = 7.5 + + wave = (z + 1 - crpix3) * cdelt3 + crval3 + ra = (z + 1 - crpix1) * cdelt1 + crval1 + dec = np.full_like(ra, crval2 + 1 * cdelt2) + + return ra, dec, wave[::-1] + + simple_wcs_function.bounding_box = wcs_bbox_from_shape(shape) + + return simple_wcs_function + + +@pytest.fixture() +def mock_nirspec_fs_one_slit(simple_wcs): + model = dm.SlitModel() + model.meta.instrument.name = 'NIRSPEC' + model.meta.instrument.detector = 'NRS1' + model.meta.instrument.filter = 'F290LP' + model.meta.instrument.grating = 'G395H' + model.meta.observation.date = '2023-07-22' + model.meta.observation.time = '06:24:45.569' + model.meta.instrument.fixed_slit = 'S200A1' + model.meta.exposure.nints = 1 + model.meta.exposure.type = 'NRS_FIXEDSLIT' + model.meta.subarray.name = 'ALLSLITS' + model.source_type = 'EXTENDED' + + model.meta.wcsinfo.dispersion_direction = 1 + model.meta.wcs = simple_wcs + + model.data = np.arange(50 * 50, dtype=float).reshape((50, 50)) + model.var_poisson = model.data * 0.02 + model.var_rnoise = model.data * 0.02 + model.var_flat = model.data * 0.05 + yield model + model.close() + + +@pytest.fixture() +def mock_nirspec_mos(mock_nirspec_fs_one_slit): + model = dm.MultiSlitModel() + model.meta.instrument.name = 'NIRSPEC' + model.meta.instrument.detector = 'NRS1' + model.meta.observation.date = '2023-07-22' + model.meta.observation.time = '06:24:45.569' + model.meta.exposure.type = 'NRS_MSASPEC' + model.meta.exposure.nints = 1 + + nslit = 3 + for i in range(nslit): + slit = mock_nirspec_fs_one_slit.copy() + slit.name = str(i + 1) + model.slits.append(slit) + + yield model + model.close() + + +@pytest.fixture() +def mock_nirspec_bots(simple_wcs): + model = dm.CubeModel() + model.meta.instrument.name = 'NIRSPEC' + model.meta.instrument.detector = 'NRS1' + model.meta.instrument.filter = 'F290LP' + model.meta.instrument.grating = 'G395H' + model.meta.observation.date = '2022-05-30' + model.meta.observation.time = '01:03:16.369' + model.meta.instrument.fixed_slit = 'S1600A1' + model.meta.exposure.type = 'NRS_BRIGHTOBJ' + model.meta.subarray.name = 'SUB2048' + model.meta.exposure.nints = 10 + model.meta.visit.tsovisit = True + + model.name = 'S1600A1' + model.meta.target.source_type = 'POINT' + model.meta.wcsinfo.dispersion_direction = 1 + model.meta.wcs = simple_wcs + + model.data = np.arange(10 * 50 * 50, dtype=float).reshape((10, 50, 50)) + model.var_poisson = model.data * 0.02 + model.var_rnoise = model.data * 0.02 + model.var_flat = model.data * 0.05 + + # Add an int_times table + integrations = [ + (1, 59729.04367729, 59729.04378181, 59729.04388632, 59729.04731706, 59729.04742158, 59729.04752609), + (2, 59729.04389677, 59729.04400128, 59729.04410579, 59729.04753654, 59729.04764105, 59729.04774557), + (3, 59729.04411625, 59729.04422076, 59729.04432527, 59729.04775602, 59729.04786053, 59729.04796504), + (4, 59729.04433572, 59729.04444023, 59729.04454475, 59729.04797549, 59729.04808001, 59729.04818452), + (5, 59729.0445552, 59729.04465971, 59729.04476422, 59729.04819497, 59729.04829948, 59729.048404), + (6, 59729.04477467, 59729.04487918, 59729.0449837, 59729.04841445, 59729.04851896, 59729.04862347), + (7, 59729.04499415, 59729.04509866, 59729.04520317, 59729.04863392, 59729.04873844, 59729.04884295), + (8, 59729.04521362, 59729.04531813, 59729.04542265, 59729.0488534 , 59729.04895791, 59729.04906242), + (9, 59729.0454331, 59729.04553761, 59729.04564212, 59729.04907288, 59729.04917739, 59729.0492819), + (10, 59729.04565257, 59729.04575709, 59729.0458616, 59729.04929235, 59729.04939686, 59729.04950138), + ] + + integration_table = np.array(integrations, dtype=[('integration_number', 'i4'), + ('int_start_MJD_UTC', 'f8'), + ('int_mid_MJD_UTC', 'f8'), + ('int_end_MJD_UTC', 'f8'), + ('int_start_BJD_TDB', 'f8'), + ('int_mid_BJD_TDB', 'f8'), + ('int_end_BJD_TDB', 'f8')]) + model.int_times = integration_table + + yield model + model.close() + + +@pytest.fixture() +def mock_miri_lrs_fs(simple_wcs_transpose): + model = dm.ImageModel() + model.meta.instrument.name = 'MIRI' + model.meta.instrument.detector = 'MIRIMAGE' + model.meta.observation.date = '2023-07-22' + model.meta.observation.time = '06:24:45.569' + model.meta.exposure.nints = 1 + model.meta.exposure.type = 'MIR_LRS-FIXEDSLIT' + model.meta.subarray.name = 'FULL' + model.meta.target.source_type = 'EXTENDED' + model.meta.dither.dithered_ra = 45.0 + model.meta.dither.dithered_ra = 45.0 + + model.meta.wcsinfo.dispersion_direction = 2 + model.meta.wcs = simple_wcs_transpose + + model.data = np.arange(50 * 50, dtype=float).reshape((50, 50)) + model.var_poisson = model.data * 0.02 + model.var_rnoise = model.data * 0.02 + model.var_flat = model.data * 0.05 + yield model + model.close() + + +@pytest.fixture() +def mock_miri_ifu(simple_wcs_ifu): + model = dm.IFUCubeModel() + model.meta.instrument.name = 'MIRI' + model.meta.instrument.detector = 'MIRIFULONG' + model.meta.observation.date = '2023-07-22' + model.meta.observation.time = '06:24:45.569' + model.meta.exposure.type = 'MIR_MRS' + + model.meta.wcsinfo.dispersion_direction = 2 + model.meta.photometry.pixelarea_steradians = 1.0 + model.meta.wcs = simple_wcs_ifu + + model.data = np.arange(10 * 50 * 50, dtype=float).reshape((10, 50, 50)) + model.var_poisson = model.data * 0.02 + model.var_rnoise = model.data * 0.02 + model.var_flat = model.data * 0.05 + model.weightmap = np.full_like(model.data, 1.0) + yield model + model.close() + + +@pytest.fixture() +def mock_niriss_wfss_l3(mock_nirspec_fs_one_slit): + model = dm.MultiSlitModel() + model.meta.instrument.name = 'NIRISS' + model.meta.instrument.detector = 'NIS' + model.meta.observation.date = '2023-07-22' + model.meta.observation.time = '06:24:45.569' + model.meta.exposure.type = 'NIS_WFSS' + + nslit = 3 + for i in range(nslit): + slit = mock_nirspec_fs_one_slit.copy() + slit.name = str(i + 1) + slit.meta.exposure.type = 'NIS_WFSS' + model.slits.append(slit) + + container = multislit_to_container([model])['0'] + + yield container + container.close() + + +@pytest.fixture() +def mock_niriss_soss(simple_wcs): + model = dm.CubeModel() + model.meta.instrument.name = 'NIRISS' + model.meta.instrument.detector = 'NIS' + model.meta.instrument.filter = 'CLEAR' + model.meta.instrument.pupil_position = 245.79 + model.meta.observation.date = '2023-07-22' + model.meta.observation.time = '06:24:45.569' + model.meta.exposure.type = 'NIS_SOSS' + model.meta.exposure.nints = 3 + + model.meta.target.source_type = 'POINT' + model.meta.wcsinfo.dispersion_direction = 1 + model.meta.wcs = simple_wcs + + yield model + model.close() + + +@pytest.fixture() +def mock_niriss_soss_256(mock_niriss_soss): + model = mock_niriss_soss + model.meta.subarray.name = 'SUBSTRIP256' + + shape = (3, 256, 2048) + model.data = np.ones(shape, dtype=np.float32) + model.dq = np.zeros(shape, dtype=np.uint32) + model.err = model.data * 0.02 + model.var_poisson = model.data * 0.001 + model.var_rnoise = model.data * 0.001 + model.var_flat = model.data * 0.001 + return model + + +@pytest.fixture() +def mock_niriss_soss_96(mock_niriss_soss): + model = mock_niriss_soss + model.meta.subarray.name = 'SUBSTRIP96' + + shape = (3, 96, 2048) + model.data = np.ones(shape, dtype=np.float32) + model.dq = np.zeros(shape, dtype=np.uint32) + model.err = model.data * 0.02 + model.var_poisson = model.data * 0.001 + model.var_rnoise = model.data * 0.001 + model.var_flat = model.data * 0.001 + return model + + +def make_spec_model(name='slit1', value=1.0): + wavelength = np.arange(20, dtype=np.float32) + flux = np.full(10, value) + error = 0.05 * flux + f_var_poisson = error ** 2 + f_var_rnoise = np.zeros_like(flux) + f_var_flat = np.zeros_like(flux) + surf_bright = flux / 10 + sb_error = error / 10 + sb_var_poisson = f_var_poisson / 10**2 + sb_var_rnoise = f_var_rnoise / 10**2 + sb_var_flat = f_var_flat / 10**2 + dq = np.zeros(20, dtype=np.uint32) + background = np.zeros_like(flux) + berror = np.zeros_like(flux) + b_var_poisson = np.zeros_like(flux) + b_var_rnoise = np.zeros_like(flux) + b_var_flat = np.zeros_like(flux) + npixels = np.full(20, 10) + + spec_dtype = dm.SpecModel().spec_table.dtype + otab = np.array( + list( + zip( + wavelength, flux, error, f_var_poisson, f_var_rnoise, f_var_flat, + surf_bright, sb_error, sb_var_poisson, sb_var_rnoise, sb_var_flat, + dq, background, berror, b_var_poisson, b_var_rnoise, b_var_flat, + npixels + ), + ), dtype=spec_dtype + ) + + spec_model = dm.SpecModel(spec_table=otab) + spec_model.name = name + + return spec_model + + +@pytest.fixture() +def mock_one_spec(): + model = dm.MultiSpecModel() + spec_model = make_spec_model() + model.spec.append(spec_model) + + yield model + model.close() + + +@pytest.fixture() +def mock_10_spec(mock_one_spec): + model = dm.MultiSpecModel() + + for i in range(10): + spec_model = make_spec_model(name=f'slit{i + 1}', value=i + 1) + model.spec.append(spec_model) + + yield model + model.close() + + +@pytest.fixture() +def miri_lrs_apcorr(): + table = Table( + { + 'subarray': ['FULL', 'SLITLESSPRISM'], + 'wavelength': [[1, 2, 3], [1, 2, 3]], + 'nelem_wl': [3, 3], + 'size': [[1, 2, 3], [1, 2, 3]], + 'nelem_size': [3, 3], + 'apcorr': np.full((2, 3, 3), 0.5), + 'apcorr_err': np.full((2, 3, 3), 0.01) + } + ) + table = fits.table_to_hdu(table) + table.header['EXTNAME'] = 'APCORR' + table.header['SIZEUNIT'] = 'pixels' + hdul = fits.HDUList([fits.PrimaryHDU(), table]) + + apcorr_model = dm.MirLrsApcorrModel(hdul) + yield apcorr_model + apcorr_model.close() + + +@pytest.fixture() +def miri_lrs_apcorr_file(tmp_path, miri_lrs_apcorr): + filename = str(tmp_path / 'miri_lrs_apcorr.fits') + miri_lrs_apcorr.save(filename) + return filename + + +@pytest.fixture() +def nirspec_fs_apcorr(): + table = Table( + { + 'filter': ['clear', 'f290lp'], + 'grating': ['prism', 'g395h'], + 'slit': ['S200A1', 'S200A1'], + 'wavelength': [[1, 2, 3], [1, 2, 3]], + 'nelem_wl': [3, 3], + 'size': np.full((2, 3, 3), [0.3, 0.5, 1]), + 'nelem_size': [3, 3], + 'pixphase': [[0.0, 0.25, 0.5], [0.0, 0.25, 0.5]], + 'apcorr': np.full((2, 3, 3, 3), 0.5), + 'apcorr_err': np.full((2, 3, 3, 3), 0.01) + } + ) + table = fits.table_to_hdu(table) + table.header['EXTNAME'] = 'APCORR' + table.header['SIZEUNIT'] = 'pixels' + hdul = fits.HDUList([fits.PrimaryHDU(), table]) + + apcorr_model = dm.NrsFsApcorrModel(hdul) + yield apcorr_model + apcorr_model.close() + + +@pytest.fixture() +def nirspec_fs_apcorr_file(tmp_path, nirspec_fs_apcorr): + filename = str(tmp_path / 'nirspec_fs_apcorr.fits') + nirspec_fs_apcorr.save(filename) + return filename diff --git a/jwst/extract_1d/tests/test_expected_skips.py b/jwst/extract_1d/tests/test_expected_skips.py index 2f761107b8..d3e3bbe239 100644 --- a/jwst/extract_1d/tests/test_expected_skips.py +++ b/jwst/extract_1d/tests/test_expected_skips.py @@ -55,3 +55,21 @@ def test_expected_skip_niriss_soss_f277w(mock_niriss_f277w): assert result2.meta.cal_step.photom == 'SKIPPED' assert result2.meta.cal_step.extract_1d == 'SKIPPED' assert np.all(result2.data == model.data) + + +def test_expected_skip_multi_int_multi_slit(): + model = dm.MultiSlitModel() + model.slits.append(dm.SlitModel(np.zeros((10, 10, 10)))) + result = Extract1dStep().process(model) + assert result.meta.cal_step.extract_1d == 'SKIPPED' + assert np.all(result.slits[0].data == model.slits[0].data) + model.close() + result.close() + + +def test_expected_skip_unexpected_model(): + model = dm.MultiExposureModel() + result = Extract1dStep().process(model) + assert result.meta.cal_step.extract_1d == 'SKIPPED' + model.close() + result.close() diff --git a/jwst/extract_1d/tests/test_extract.py b/jwst/extract_1d/tests/test_extract.py new file mode 100644 index 0000000000..a66b19c55c --- /dev/null +++ b/jwst/extract_1d/tests/test_extract.py @@ -0,0 +1,1570 @@ +import json +import logging +import numpy as np +import pytest +import stdatamodels.jwst.datamodels as dm +from astropy.modeling import polynomial + +from jwst.datamodels import ModelContainer +from jwst.extract_1d import extract as ex +from jwst.tests.helpers import LogWatcher + + +@pytest.fixture +def log_watcher(monkeypatch): + # Set a log watcher to check for a log message at any level + # in the extract_1d.extract module + watcher = LogWatcher('') + logger = logging.getLogger('jwst.extract_1d.extract') + for level in ['debug', 'info', 'warning', 'error']: + monkeypatch.setattr(logger, level, watcher) + return watcher + + +@pytest.fixture() +def extract1d_ref_dict(): + apertures = [{'id': 'slit1'}, + {'id': 'slit2', 'region_type': 'other'}, + {'id': 'slit3', 'xstart': 10, 'xstop': 20, 'ystart': 10, 'ystop': 20}, + {'id': 'slit4', 'bkg_coeff': [[10], [20]]}, + {'id': 'slit5', 'bkg_coeff': None}, + {'id': 'slit6', 'use_source_posn': True}, + {'id': 'slit7', 'spectral_order': 20}, + {'id': 'S200A1'} + ] + ref_dict = {'apertures': apertures} + return ref_dict + + +@pytest.fixture() +def extract1d_ref_file(tmp_path, extract1d_ref_dict): + filename = str(tmp_path / 'extract1d_ref.json') + with open(filename, 'w') as fh: + json.dump(extract1d_ref_dict, fh) + return filename + + +@pytest.fixture() +def extract_defaults(): + default = {'bkg_coeff': None, + 'bkg_fit': None, + 'bkg_order': 0, + 'extract_width': None, + 'extraction_type': 'box', + 'independent_var': 'pixel', + 'match': 'exact match', + 'position_correction': 0, + 'smoothing_length': 0, + 'spectral_order': 1, + 'src_coeff': None, + 'subtract_background': False, + 'use_source_posn': False, + 'xstart': 0, + 'xstop': 49, + 'ystart': 0, + 'ystop': 49} + return default + + +@pytest.fixture() +def simple_profile(): + profile = np.zeros((50, 50), dtype=np.float32) + profile[20:30, :] = 1.0 + return profile + + +@pytest.fixture() +def background_profile(): + profile = np.zeros((50, 50), dtype=np.float32) + profile[:10, :] = 1.0 + profile[40:, :] = 1.0 + return profile + + +@pytest.fixture() +def create_extraction_inputs(mock_nirspec_fs_one_slit, extract1d_ref_dict): + input_model = mock_nirspec_fs_one_slit + slit = None + output_model = dm.MultiSpecModel() + ref_dict = extract1d_ref_dict + slitname = 'S200A1' + sp_order = 1 + exp_type = 'NRS_FIXEDSLIT' + yield [input_model, slit, output_model, ref_dict, + slitname, sp_order, exp_type] + output_model.close() + + +def test_read_extract1d_ref(extract1d_ref_dict, extract1d_ref_file): + ref_dict = ex.read_extract1d_ref(extract1d_ref_file) + assert ref_dict == extract1d_ref_dict + + +def test_read_extract1d_ref_bad_json(tmp_path): + filename = str(tmp_path / 'bad_ref.json') + with open(filename, 'w') as fh: + fh.write('apertures: [bad,]\n') + + with pytest.raises(RuntimeError, match='Invalid JSON extract1d reference'): + ex.read_extract1d_ref(filename) + + +def test_read_extract1d_ref_bad_type(tmp_path): + filename = str(tmp_path / 'bad_ref.fits') + with open(filename, 'w') as fh: + fh.write('bad file\n') + + with pytest.raises(RuntimeError, match='must be JSON'): + ex.read_extract1d_ref(filename) + + +def test_read_extract1d_ref_na(): + ref_dict = ex.read_extract1d_ref('N/A') + assert ref_dict is None + + +def test_read_apcorr_ref(): + apcorr_model = ex.read_apcorr_ref(None, 'MIR_LRS-FIXEDSLIT') + assert isinstance(apcorr_model, dm.MirLrsApcorrModel) + + +def test_get_extract_parameters_default( + mock_nirspec_fs_one_slit, extract1d_ref_dict, extract_defaults): + input_model = mock_nirspec_fs_one_slit + + # match a bare entry + params = ex.get_extract_parameters( + extract1d_ref_dict, input_model, 'slit1', 1, input_model.meta) + + # returned value has defaults except that use_source_posn + # is switched on for NRS_FIXEDSLIT + expected = extract_defaults + expected['use_source_posn'] = True + + assert params == expected + + +def test_get_extract_parameters_na(mock_nirspec_fs_one_slit, extract_defaults): + input_model = mock_nirspec_fs_one_slit + + # no reference input: defaults returned + params = ex.get_extract_parameters(None, input_model, 'slit1', 1, input_model.meta) + assert params == extract_defaults + + +@pytest.mark.parametrize('bgsub', [None, True]) +@pytest.mark.parametrize('bgfit', ['poly', 'mean', None]) +def test_get_extract_parameters_background( + mock_nirspec_fs_one_slit, extract1d_ref_dict, bgsub, bgfit): + input_model = mock_nirspec_fs_one_slit + + # match a slit with background defined + params = ex.get_extract_parameters( + extract1d_ref_dict, input_model, 'slit4', 1, input_model.meta, + subtract_background=bgsub, bkg_fit=bgfit) + + # returned value has background switched on + assert params['subtract_background'] is True + assert params['bkg_coeff'] is not None + assert params['bkg_fit'] == 'poly' + assert params['bkg_order'] == 0 + + +def test_get_extract_parameters_bg_ignored(mock_nirspec_fs_one_slit, extract1d_ref_dict): + input_model = mock_nirspec_fs_one_slit + + # match a slit with background defined + params = ex.get_extract_parameters( + extract1d_ref_dict, input_model, 'slit4', 1, input_model.meta, + subtract_background=False) + + # background parameters are ignored + assert params['subtract_background'] is False + assert params['bkg_fit'] is None + + +@pytest.mark.parametrize('slit', ['slit2', 'no_match']) +def test_get_extract_parameters_no_match( + mock_nirspec_fs_one_slit, extract1d_ref_dict, slit): + input_model = mock_nirspec_fs_one_slit + + # no slit with an appropriate region_type matched + params = ex.get_extract_parameters( + extract1d_ref_dict, input_model, slit, 1, input_model.meta) + assert params == {'match': ex.NO_MATCH} + + +def test_get_extract_parameters_source_posn_exptype( + mock_nirspec_bots, extract1d_ref_dict, extract_defaults): + input_model = mock_nirspec_bots + + # match a bare entry + params = ex.get_extract_parameters( + extract1d_ref_dict, input_model, 'slit1', 1, input_model.meta, + use_source_posn=None) + + # use_source_posn is switched off for NRS_BRIGHTOBJ + assert params['use_source_posn'] is False + + +def test_get_extract_parameters_source_posn_from_ref( + mock_nirspec_bots, extract1d_ref_dict, extract_defaults): + input_model = mock_nirspec_bots + + # match an entry that explicity sets use_source_posn + params = ex.get_extract_parameters( + extract1d_ref_dict, input_model, 'slit6', 1, input_model.meta, + use_source_posn=None) + + # returned value has use_source_posn switched off by default + # for NRS_BRIGHTOBJ, but ref file overrides + assert params['use_source_posn'] is True + + +@pytest.mark.parametrize('length', [3, 4, 2.8, 3.5]) +def test_get_extract_parameters_smoothing( + mock_nirspec_fs_one_slit, extract1d_ref_dict, + extract_defaults, length): + input_model = mock_nirspec_fs_one_slit + + params = ex.get_extract_parameters( + extract1d_ref_dict, input_model, 'slit1', 1, input_model.meta, + smoothing_length=length) + + # returned value has input smoothing length, rounded to an + # odd integer if necessary + assert params['smoothing_length'] == 3 + + +@pytest.mark.parametrize('length', [-1, 1, 2, 1.3]) +def test_get_extract_parameters_smoothing_bad_value( + mock_nirspec_fs_one_slit, extract1d_ref_dict, + extract_defaults, length): + input_model = mock_nirspec_fs_one_slit + + params = ex.get_extract_parameters( + extract1d_ref_dict, input_model, 'slit1', 1, input_model.meta, + smoothing_length=length) + + # returned value has smoothing length 0 + assert params['smoothing_length'] == 0 + + +def test_log_params(extract_defaults, log_watcher): + log_watcher.message = 'Extraction parameters' + + # Defaults don't have dispaxis assigned yet - parameters are not logged + ex.log_initial_parameters(extract_defaults) + assert not log_watcher.seen + + # Add dispaxis: parameters are now logged + extract_defaults['dispaxis'] = 1 + ex.log_initial_parameters(extract_defaults) + log_watcher.assert_seen() + + +def test_create_poly(): + coeff = [1, 2, 3] + poly = ex.create_poly(coeff) + assert isinstance(poly, polynomial.Polynomial1D) + assert poly.degree == 2 + assert poly(2) == 1 + 2 * 2 + 3 * 2**2 + + +def test_create_poly_empty(): + coeff = [] + assert ex.create_poly(coeff) is None + + +def test_populate_time_keywords(mock_nirspec_bots, mock_10_spec): + ex.populate_time_keywords(mock_nirspec_bots, mock_10_spec) + + # time keywords now added to output spectra + for i, spec in enumerate(mock_10_spec.spec): + assert spec.int_num == i + 1 + assert spec.start_time_mjd == mock_nirspec_bots.int_times['int_start_MJD_UTC'][i] + assert spec.end_tdb == mock_nirspec_bots.int_times['int_end_BJD_TDB'][i] + + +def test_populate_time_keywords_no_table(mock_nirspec_fs_one_slit, mock_one_spec): + ex.populate_time_keywords(mock_nirspec_fs_one_slit, mock_one_spec) + + # only int_num is added to spec + assert mock_one_spec.spec[0].int_num == 1 + + +def test_populate_time_keywords_multislit(mock_nirspec_mos, mock_10_spec): + mock_nirspec_mos.meta.exposure.nints = 10 + ex.populate_time_keywords(mock_nirspec_mos, mock_10_spec) + + # no int_times - only int_num is added to spec + # It is set to 1 for all spectra - no integrations in multislit data. + assert mock_10_spec.spec[0].int_num == 1 + assert mock_10_spec.spec[9].int_num == 1 + + +def test_populate_time_keywords_multislit_table( + mock_nirspec_mos, mock_nirspec_bots, mock_10_spec, log_watcher): + mock_nirspec_mos.meta.exposure.nints = 10 + mock_nirspec_mos.int_times = mock_nirspec_bots.int_times + + log_watcher.message = 'Not using INT_TIMES table' + ex.populate_time_keywords(mock_nirspec_mos, mock_10_spec) + log_watcher.assert_seen() + + # int_times present but not used - no update + assert mock_10_spec.spec[0].int_num is None + + +def test_populate_time_keywords_averaged( + mock_nirspec_fs_one_slit, mock_nirspec_bots, mock_10_spec, log_watcher): + mock_nirspec_fs_one_slit.meta.exposure.nints = 10 + mock_nirspec_fs_one_slit.int_times = mock_nirspec_bots.int_times + + log_watcher.message = 'Not using INT_TIMES table' + ex.populate_time_keywords(mock_nirspec_fs_one_slit, mock_10_spec) + log_watcher.assert_seen() + + # int_times not used - no update + assert mock_10_spec.spec[0].int_num is None + + +def test_populate_time_keywords_mismatched_table(mock_nirspec_bots, mock_10_spec, log_watcher): + # mock 20 integrations - table has 10 + mock_nirspec_bots.data = np.vstack([mock_nirspec_bots.data, mock_nirspec_bots.data]) + log_watcher.message = 'Not using INT_TIMES table' + ex.populate_time_keywords(mock_nirspec_bots, mock_10_spec) + log_watcher.assert_seen() + + # int_times not used - no update + assert mock_10_spec.spec[0].int_num is None + + +def test_populate_time_keywords_missing_ints(mock_nirspec_bots, mock_10_spec, log_watcher): + mock_nirspec_bots.meta.exposure.integration_start = 20 + log_watcher.message = 'does not include rows' + ex.populate_time_keywords(mock_nirspec_bots, mock_10_spec) + log_watcher.assert_seen() + + # int_times not used - no update + assert mock_10_spec.spec[0].int_num is None + + +def test_populate_time_keywords_ifu_table( + mock_miri_ifu, mock_nirspec_bots, mock_10_spec, log_watcher): + mock_miri_ifu.meta.exposure.nints = 10 + mock_miri_ifu.int_times = mock_nirspec_bots.int_times + + log_watcher.message = 'ignored for IFU' + ex.populate_time_keywords(mock_miri_ifu, mock_10_spec) + log_watcher.assert_seen() + + # int_times present but not used - no update + assert mock_10_spec.spec[0].int_num is None + + +def test_populate_time_keywords_mismatched_spec( + mock_nirspec_bots, mock_one_spec, log_watcher): + log_watcher.message = "Don't understand n_output_spec" + ex.populate_time_keywords(mock_nirspec_bots, mock_one_spec) + log_watcher.assert_seen() + + +def test_get_spectral_order(mock_nirspec_fs_one_slit): + slit = mock_nirspec_fs_one_slit + + slit.meta.wcsinfo.spectral_order = 2 + assert ex.get_spectral_order(slit) == 2 + + slit.meta.wcsinfo.spectral_order = None + assert ex.get_spectral_order(slit) == 1 + + slit.meta.wcsinfo = None + assert ex.get_spectral_order(slit) == 1 + + del slit.meta.wcsinfo + assert ex.get_spectral_order(slit) == 1 + + +def test_is_prism_false_conditions(mock_nirspec_fs_one_slit): + model = mock_nirspec_fs_one_slit + assert ex.is_prism(model) is False + + model.meta.instrument.filter = None + model.meta.instrument.grating = None + assert ex.is_prism(model) is False + + model.meta.instrument.name = None + assert ex.is_prism(model) is False + + +def test_is_prism_nirspec(mock_nirspec_fs_one_slit): + mock_nirspec_fs_one_slit.meta.instrument.grating = 'PRISM' + assert ex.is_prism(mock_nirspec_fs_one_slit) is True + + +def test_is_prism_miri(mock_miri_ifu): + mock_miri_ifu.meta.instrument.filter = 'P750L' + assert ex.is_prism(mock_miri_ifu) is True + +def test_copy_keyword_info(mock_nirspec_fs_one_slit, mock_one_spec): + expected = {'slitlet_id': 2, + 'source_id': 3, + 'source_name': '4', + 'source_alias': '5', + 'source_type': 'POINT', + 'stellarity': 0.5, + 'source_xpos': -0.5, + 'source_ypos': -0.5, + 'source_ra': 10.0, + 'source_dec': 10.0, + 'shutter_state': 'x'} + for key, value in expected.items(): + setattr(mock_nirspec_fs_one_slit, key, value) + assert not hasattr(mock_one_spec, key) + + ex.copy_keyword_info(mock_nirspec_fs_one_slit, 'slit_name', mock_one_spec) + assert mock_one_spec.name == 'slit_name' + for key, value in expected.items(): + assert getattr(mock_one_spec, key) == value + + +@pytest.mark.parametrize("partial", [True, False]) +@pytest.mark.parametrize("lower,upper", + [(0, 19), (-1, 21), + (np.full(10, 0.0), np.full(10, 19.0)), + (np.linspace(-1, 0, 10), np.linspace(19, 20, 10)),]) +def test_set_weights_from_limits_whole_array(lower, upper, partial): + shape = (20, 10) + profile = np.zeros(shape, dtype=float) + yidx, _ = np.mgrid[:shape[0], :shape[1]] + + ex._set_weight_from_limits(profile, yidx, lower, upper, allow_partial=partial) + assert np.all(profile == 1.0) + + +@pytest.mark.parametrize("lower,upper", + [(10, 12), (9.5, 12.5), + (np.linspace(9.5, 10, 10), np.linspace(12, 12.5, 10)),]) +def test_set_weights_from_limits_whole_pixel(lower, upper): + shape = (20, 10) + profile = np.zeros(shape, dtype=np.float32) + yidx, _ = np.mgrid[:shape[0], :shape[1]] + + ex._set_weight_from_limits(profile, yidx, lower, upper, allow_partial=False) + assert np.all(profile[10:13] == 1.0) + + +@pytest.mark.parametrize("lower,upper", + [(9.5, 12.5), + (np.linspace(9.5, 10, 10), np.linspace(12, 12.5, 10)),]) +def test_set_weights_from_limits_partial_pixel(lower, upper): + shape = (20, 10) + profile = np.zeros(shape, dtype=np.float32) + yidx, _ = np.mgrid[:shape[0], :shape[1]] + + ex._set_weight_from_limits(profile, yidx, lower, upper, allow_partial=True) + assert np.allclose(profile[10:13], 1.0) + assert np.allclose(profile[9], 10 - lower) + assert np.allclose(profile[13], upper - 12) + + +def test_set_weights_from_limits_overlap(): + shape = (20, 10) + profile = np.zeros(shape, dtype=np.float32) + yidx, _ = np.mgrid[:shape[0], :shape[1]] + + # Set an aperture with partial pixel edges + ex._set_weight_from_limits(profile, yidx, 9.5, 10.5, allow_partial=True) + assert np.allclose(profile[9], 0.5) + assert np.allclose(profile[11], 0.5) + assert np.allclose(profile[12], 0.0) + + # Set an overlapping region in the same profile + ex._set_weight_from_limits(profile, yidx, 9.8, 11.5, allow_partial=True) + + # Higher weight from previous profile remains + assert np.allclose(profile[9], 0.5) + + # Previous partial pixel is now fully included + assert np.allclose(profile[11], 1.0) + + # New partial weight set on upper limit + assert np.allclose(profile[12], 0.5) + + +def test_box_profile_horizontal(extract_defaults): + shape = (10, 10) + wl_array = np.empty(shape) + wl_array[:] = np.linspace(3, 5, 10) + + params = extract_defaults + params['dispaxis'] = 1 + + # Exclude 2 pixels at left and right + params['xstart'] = 1.5 + params['xstop'] = 7.5 + + # Exclude 2 pixels at top and bottom, set half pixel weights + # for another pixel at top and bottom edges + params['ystart'] = 2.5 + params['ystop'] = 6.5 + profile = ex.box_profile(shape, extract_defaults, wl_array) + + # ystart/stop sets partial weights, xstart/stop sets whole pixels only + assert np.all(profile[2:3, 3:8] == 0.5) + assert np.all(profile[7:8, 3:8] == 0.5) + assert np.all(profile[3:7, 3:8] == 1.0) + assert np.all(profile[:2] == 0.0) + assert np.all(profile[8:] == 0.0) + assert np.all(profile[:, :2] == 0.0) + assert np.all(profile[:, 8:] == 0.0) + + +def test_box_profile_vertical(extract_defaults): + shape = (10, 10) + wl_array = np.empty(shape) + wl_array[:] = np.linspace(3, 5, 10) + + params = extract_defaults + params['dispaxis'] = 2 + + # Exclude 2 pixels at "left" and "right" - in transposed aperture + params['ystart'] = 1.5 + params['ystop'] = 7.5 + + # Exclude 2 pixels at "top" and "bottom", set half pixel weights + # for another pixel at top and bottom edges + params['xstart'] = 2.5 + params['xstop'] = 6.5 + profile = ex.box_profile(shape, extract_defaults, wl_array) + + # xstart/stop sets partial weights, ystart/stop sets whole pixels only + assert np.all(profile[3:8, 2:3] == 0.5) + assert np.all(profile[3:8, 7:8] == 0.5) + assert np.all(profile[3:8, 3:7] == 1.0) + + assert np.all(profile[:2] == 0.0) + assert np.all(profile[8:] == 0.0) + assert np.all(profile[:, :2] == 0.0) + assert np.all(profile[:, 8:] == 0.0) + + +@pytest.mark.parametrize('dispaxis', [1, 2]) +def test_box_profile_bkg_coeff(extract_defaults, dispaxis): + shape = (10, 10) + wl_array = np.empty(shape) + wl_array[:] = np.linspace(3, 5, 10) + + params = extract_defaults + params['dispaxis'] = dispaxis + + # the definition for bkg_coeff is half a pixel off from start/stop definitions - + # this will set the equivalent of start/stop 0-2, 7-9 - + # 3 pixels at top and bottom of the array + params['bkg_coeff'] = [[-0.5], [2.5], [6.5], [9.5]] + + profile, lower, upper = ( + ex.box_profile(shape, extract_defaults, wl_array, + coefficients='bkg_coeff', return_limits=True)) + if dispaxis == 2: + profile = profile.T + + assert np.all(profile[:3] == 1.0) + assert np.all(profile[7:] == 1.0) + assert np.all(profile[3:7] == 0.0) + assert lower == 0.0 + assert upper == 9.0 + + +def test_box_profile_bkg_coeff_median(extract_defaults): + shape = (10, 10) + wl_array = np.empty(shape) + wl_array[:] = np.linspace(3, 5, 10) + + params = extract_defaults + params['dispaxis'] = 1 + params['bkg_fit'] = 'median' + + # Attempt to set partial pixels at middle edges + params['bkg_coeff'] = [[-0.5], [3.0], [6.0], [9.5]] + + profile, lower, upper = ( + ex.box_profile(shape, extract_defaults, wl_array, + coefficients='bkg_coeff', return_limits=True)) + + # partial pixels are not allowed for fit type median - the profile is + # set for whole pixels only + assert np.all(profile[:3] == 1.0) + assert np.all(profile[7:] == 1.0) + assert np.all(profile[3:7] == 0.0) + assert lower == 0.0 + assert upper == 9.0 + + +@pytest.mark.parametrize('swap_order', [False, True]) +def test_box_profile_bkg_coeff_poly(extract_defaults, swap_order): + shape = (10, 10) + wl_array = np.empty(shape) + wl_array[:] = np.linspace(3, 5, 10) + + params = extract_defaults + params['dispaxis'] = 1 + params['bkg_fit'] = 'poly' + + # Attempt to set partial pixels at middle edges + if swap_order: + # upper region first - this should make no difference. + params['bkg_coeff'] = [[6.0], [9.5], [-0.5], [3.0]] + else: + params['bkg_coeff'] = [[-0.5], [3.0], [6.0], [9.5]] + + profile, lower, upper = ( + ex.box_profile(shape, extract_defaults, wl_array, + coefficients='bkg_coeff', return_limits=True)) + + # partial pixels are allowed for fit type poly + assert np.all(profile[:3] == 1.0) + assert np.all(profile[7:] == 1.0) + assert np.all(profile[3] == 0.5) + assert np.all(profile[6] == 0.5) + assert np.all(profile[4:6] == 0.0) + assert lower == 0.0 + assert upper == 9.0 + + +@pytest.mark.parametrize('independent_var', ['pixel', 'wavelength']) +def test_box_profile_src_coeff_constant(extract_defaults, independent_var): + shape = (10, 10) + wl_array = np.empty(shape) + wl_array[:] = np.linspace(3, 5, 10) + + params = extract_defaults + params['dispaxis'] = 1 + params['independent_var'] = independent_var + + # the definition for src_coeff is half a pixel off from start/stop definitions - + # this will set the equivalent of start/stop 3-6, excluding + # 3 pixels at top and bottom of the array + params['src_coeff'] = [[2.5], [6.5]] + + profile, lower, upper = ( + ex.box_profile(shape, extract_defaults, wl_array, + coefficients='src_coeff', return_limits=True)) + assert np.all(profile[3:7] == 1.0) + assert np.all(profile[:3] == 0.0) + assert np.all(profile[7:] == 0.0) + assert lower == 3.0 + assert upper == 6.0 + + +@pytest.mark.parametrize('independent_var', ['pixel', 'wavelength']) +def test_box_profile_src_coeff_linear(extract_defaults, independent_var): + shape = (10, 10) + wl_array = np.empty(shape) + wl_array[:] = np.linspace(3, 5, 10) + + params = extract_defaults + params['dispaxis'] = 1 + params['independent_var'] = independent_var + + if independent_var == 'wavelength': + slope = 1 / (wl_array[0, 1] - wl_array[0, 0]) + start = -0.5 - wl_array[0, 0] * slope + stop = start + 4 + else: + slope = 1.0 + start = -0.5 + stop = 3.5 + + # Set linearly increasing upper and lower edges, + # starting at the bottom of the array, with a width of 4 pixels + params['src_coeff'] = [[start, slope], [stop, slope]] + expected = [[1, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 0, 0, 0, 0, 0, 0], + [0, 1, 1, 1, 1, 0, 0, 0, 0, 0], + [0, 0, 1, 1, 1, 1, 0, 0, 0, 0], + [0, 0, 0, 1, 1, 1, 1, 0, 0, 0], + [0, 0, 0, 0, 1, 1, 1, 1, 0, 0], + [0, 0, 0, 0, 0, 1, 1, 1, 1, 0], + [0, 0, 0, 0, 0, 0, 1, 1, 1, 1]] + + profile, lower, upper = ( + ex.box_profile(shape, extract_defaults, wl_array, + coefficients='src_coeff', return_limits=True)) + assert np.allclose(profile, expected) + + # upper and lower limits are averages + assert np.isclose(lower, 4.5) + assert np.isclose(upper, 7.5) + + +def test_box_profile_mismatched_coeff(extract_defaults): + shape = (10, 10) + wl_array = np.empty(shape) + wl_array[:] = np.linspace(3, 5, 10) + + params = extract_defaults + params['dispaxis'] = 1 + + # set some mismatched coefficients limits + params['src_coeff'] = [[2.5], [6.5], [9.5]] + + with pytest.raises(RuntimeError, match='must contain alternating lists'): + ex.box_profile(shape, extract_defaults, wl_array) + + +@pytest.mark.parametrize('dispaxis', [1, 2]) +def test_box_profile_from_width(extract_defaults, dispaxis): + shape = (10, 10) + wl_array = np.empty(shape) + wl_array[:] = np.linspace(3, 5, 10) + + params = extract_defaults + params['dispaxis'] = dispaxis + + if dispaxis == 1: + # Set ystart and ystop to center on pixel 4 + params['ystart'] = 2.0 + params['ystop'] = 6.0 + else: + # Set xstart and xstop to center on pixel 4 + params['xstart'] = 2.0 + params['xstop'] = 6.0 + + # Set a width to 6 pixels + params['extract_width'] = 6.0 + + profile = ex.box_profile(shape, extract_defaults, wl_array) + if dispaxis == 2: + profile = profile.T + + # Aperture is centered at pixel 4, to start at 1.5, end at 6.5 + assert np.all(profile[2:7] == 1.0) + assert np.all(profile[1] == 0.5) + assert np.all(profile[7] == 0.5) + assert np.all(profile[0] == 0.0) + assert np.all(profile[8:] == 0.0) + + +@pytest.mark.parametrize('middle', [None, 7]) +@pytest.mark.parametrize('dispaxis', [1, 2]) +def test_aperture_center(middle, dispaxis): + profile = np.zeros((10, 10), dtype=np.float32) + profile[1:4] = 1.0 + if dispaxis != 1: + profile = profile.T + slit_center, spec_center = ex.aperture_center( + profile, dispaxis=dispaxis, middle_pix=middle) + assert slit_center == 2.0 + if middle is None: + assert spec_center == 4.5 + else: + assert spec_center == middle + + +@pytest.mark.parametrize('middle', [None, 7]) +@pytest.mark.parametrize('dispaxis', [1, 2]) +def test_aperture_center_zero_weight(middle, dispaxis): + profile = np.zeros((10, 10), dtype=np.float32) + slit_center, spec_center = ex.aperture_center( + profile, dispaxis=dispaxis, middle_pix=middle) + assert slit_center == 4.5 + if middle is None: + assert spec_center == 4.5 + else: + assert spec_center == middle + + +@pytest.mark.parametrize('middle', [None, 7]) +@pytest.mark.parametrize('dispaxis', [1, 2]) +def test_aperture_center_variable_weight_by_slit(middle, dispaxis): + profile = np.zeros((10, 10), dtype=np.float32) + profile[1:4] = np.arange(10) + if dispaxis != 1: + profile = profile.T + slit_center, spec_center = ex.aperture_center( + profile, dispaxis=dispaxis, middle_pix=middle) + assert slit_center == 2.0 + if middle is None: + assert np.isclose(spec_center, 6.3333333) + else: + assert spec_center == middle + + +@pytest.mark.parametrize('middle', [None, 7]) +@pytest.mark.parametrize('dispaxis', [1, 2]) +def test_aperture_center_variable_weight_by_spec(middle, dispaxis): + profile = np.zeros((10, 10), dtype=np.float32) + profile[:, 1:4] = np.arange(10)[:, None] + if dispaxis != 1: + profile = profile.T + slit_center, spec_center = ex.aperture_center( + profile, dispaxis=dispaxis, middle_pix=middle) + if middle is None: + assert np.isclose(slit_center, 6.3333333) + assert np.isclose(spec_center, 2.0) + else: + assert np.isclose(slit_center, 4.5) + assert spec_center == middle + + +@pytest.mark.parametrize('resampled', [True, False]) +@pytest.mark.parametrize('is_slit', [True, False]) +@pytest.mark.parametrize('missing_bbox', [True, False]) +def test_location_from_wcs_nirspec( + monkeypatch, mock_nirspec_fs_one_slit, resampled, is_slit, missing_bbox): + model = mock_nirspec_fs_one_slit + + # monkey patch in a transform for the wcs + def slit2det(*args, **kwargs): + def return_one(*args, **kwargs): + return 0.0, 1.0 + return return_one + + monkeypatch.setattr(model.meta.wcs, 'get_transform', slit2det) + + if not resampled: + # also mock available frames, so it looks like unresampled cal data + monkeypatch.setattr(model.meta.wcs, 'available_frames', ['gwa']) + + if missing_bbox: + # also mock a missing bounding box - should have same results + # for the test data + monkeypatch.setattr(model.meta.wcs, 'bounding_box', None) + + if is_slit: + middle, middle_wl, location = ex.location_from_wcs(model, model) + else: + middle, middle_wl, location = ex.location_from_wcs(model, None) + + # middle pixel is center of dispersion axis + assert middle == int((model.data.shape[1] - 1) / 2) + + # middle wavelength is the wavelength at that point, from the mock wcs + assert np.isclose(middle_wl, 7.74) + + # location is 1.0 - from the mocked transform function + assert location == 1.0 + + +@pytest.mark.parametrize('is_slit', [True, False]) +def test_location_from_wcs_miri(monkeypatch, mock_miri_lrs_fs, is_slit): + model = mock_miri_lrs_fs + + # monkey patch in a transform for the wcs + def radec2det(*args, **kwargs): + def return_one(*args, **kwargs): + return 1.0, 0.0 + return return_one + + monkeypatch.setattr(model.meta.wcs, 'backward_transform', radec2det()) + + # Get the slit center from the WCS + if is_slit: + middle, middle_wl, location = ex.location_from_wcs(model, model) + else: + middle, middle_wl, location = ex.location_from_wcs(model, None) + + # middle pixel is center of dispersion axis + assert middle == int((model.data.shape[0] - 1) / 2) + + # middle wavelength is the wavelength at that point, from the mock wcs + assert np.isclose(middle_wl, 7.26) + + # location is 1.0 - from the mocked transform function + assert location == 1.0 + + +def test_location_from_wcs_missing_data(mock_miri_lrs_fs, log_watcher): + # model is missing WCS information - None values are returned + log_watcher.message = "Dithered pointing location not found" + result = ex.location_from_wcs(mock_miri_lrs_fs, None) + assert result == (None, None, None) + log_watcher.assert_seen() + + +def test_location_from_wcs_wrong_exptype(mock_niriss_soss, log_watcher): + # model is not a handled exposure type + log_watcher.message = "Source position cannot be found for EXP_TYPE" + result = ex.location_from_wcs(mock_niriss_soss, None) + assert result == (None, None, None) + log_watcher.assert_seen() + + +def test_location_from_wcs_bad_location( + monkeypatch, mock_nirspec_fs_one_slit, log_watcher): + model = mock_nirspec_fs_one_slit + + # monkey patch in a transform for the wcs + def slit2det(*args, **kwargs): + def return_one(*args, **kwargs): + return 0.0, np.nan + return return_one + + monkeypatch.setattr(model.meta.wcs, 'get_transform', slit2det) + + # WCS transform returns NaN for the location + log_watcher.message = "Source position could not be determined" + result = ex.location_from_wcs(model, None) + assert result == (None, None, None) + log_watcher.assert_seen() + + +def test_location_from_wcs_location_out_of_range( + monkeypatch, mock_nirspec_fs_one_slit, log_watcher): + model = mock_nirspec_fs_one_slit + + # monkey patch in a transform for the wcs + def slit2det(*args, **kwargs): + def return_one(*args, **kwargs): + return 0.0, 2000 + return return_one + + monkeypatch.setattr(model.meta.wcs, 'get_transform', slit2det) + + # WCS transform a value outside the bounding box + log_watcher.message = "outside the bounding box" + result = ex.location_from_wcs(model, None) + assert result == (None, None, None) + log_watcher.assert_seen() + + +def test_shift_by_source_location_horizontal(extract_defaults): + location = 12.5 + nominal_location = 15.0 + offset = location - nominal_location + + extract_params = extract_defaults.copy() + extract_params['dispaxis'] = 1 + + ex.shift_by_source_location(location, nominal_location, extract_params) + assert extract_params['xstart'] == extract_defaults['xstart'] + assert extract_params['xstop'] == extract_defaults['xstop'] + assert extract_params['ystart'] == extract_defaults['ystart'] + offset + assert extract_params['ystop'] == extract_defaults['ystop'] + offset + + +def test_shift_by_source_location_vertical(extract_defaults): + location = 12.5 + nominal_location = 15.0 + offset = location - nominal_location + + extract_params = extract_defaults.copy() + extract_params['dispaxis'] = 2 + + ex.shift_by_source_location(location, nominal_location, extract_params) + assert extract_params['xstart'] == extract_defaults['xstart'] + offset + assert extract_params['xstop'] == extract_defaults['xstop'] + offset + assert extract_params['ystart'] == extract_defaults['ystart'] + assert extract_params['ystop'] == extract_defaults['ystop'] + + +def test_shift_by_source_location_coeff(extract_defaults): + location = 6.5 + nominal_location = 4.0 + offset = location - nominal_location + + extract_params = extract_defaults.copy() + extract_params['dispaxis'] = 1 + extract_params['src_coeff'] = [[2.5, 1.0], [6.5, 1.0]] + extract_params['bkg_coeff'] = [[-0.5], [3.0], [6.0], [9.5]] + + ex.shift_by_source_location(location, nominal_location, extract_params) + assert extract_params['src_coeff'] == [[2.5 + offset, 1.0], [6.5 + offset, 1.0]] + assert extract_params['bkg_coeff'] == [[-0.5 + offset], [3.0 + offset], + [6.0 + offset], [9.5 + offset]] + + +@pytest.mark.parametrize('is_slit', [True, False]) +def test_define_aperture_nirspec(mock_nirspec_fs_one_slit, extract_defaults, is_slit): + model = mock_nirspec_fs_one_slit + extract_defaults['dispaxis'] = 1 + if is_slit: + slit = model + else: + slit = None + exptype = 'NRS_FIXEDSLIT' + result = ex.define_aperture(model, slit, extract_defaults, exptype) + ra, dec, wavelength, profile, bg_profile, limits = result + assert np.isclose(ra, 45.05) + assert np.isclose(dec, 45.1) + assert wavelength.shape == (model.data.shape[1],) + assert profile.shape == model.data.shape + + # Default profile is the full array + assert np.all(profile == 1.0) + assert limits == (0, model.data.shape[0] - 1, 0, model.data.shape[1] - 1) + + # Default bg profile is None + assert bg_profile is None + + +@pytest.mark.parametrize('is_slit', [True, False]) +def test_define_aperture_miri(mock_miri_lrs_fs, extract_defaults, is_slit): + model = mock_miri_lrs_fs + extract_defaults['dispaxis'] = 2 + if is_slit: + slit = model + else: + slit = None + exptype = 'MIR_LRS-FIXEDSLIT' + result = ex.define_aperture(model, slit, extract_defaults, exptype) + ra, dec, wavelength, profile, bg_profile, limits = result + assert np.isclose(ra, 45.05) + assert np.isclose(dec, 45.1) + assert wavelength.shape == (model.data.shape[1],) + assert profile.shape == model.data.shape + + # Default profile is the full array + assert np.all(profile == 1.0) + assert limits == (0, model.data.shape[0] - 1, 0, model.data.shape[1] - 1) + + # Default bg profile is None + assert bg_profile is None + + +def test_define_aperture_with_bg(mock_nirspec_fs_one_slit, extract_defaults): + model = mock_nirspec_fs_one_slit + extract_defaults['dispaxis'] = 1 + slit = None + exptype = 'NRS_FIXEDSLIT' + + extract_defaults['subtract_background'] = True + extract_defaults['bkg_coeff'] = [[-0.5], [2.5]] + + result = ex.define_aperture(model, slit, extract_defaults, exptype) + bg_profile = result[-2] + + # Bg profile has 1s in the first 3 rows + assert bg_profile.shape == model.data.shape + assert np.all(bg_profile[:3] == 1.0) + assert np.all(bg_profile[3:] == 0.0) + + +def test_define_aperture_empty_aperture(mock_nirspec_fs_one_slit, extract_defaults): + model = mock_nirspec_fs_one_slit + extract_defaults['dispaxis'] = 1 + slit = None + exptype = 'NRS_FIXEDSLIT' + + # Set the extraction limits out of range + extract_defaults['ystart'] = 2000 + extract_defaults['ystop'] = 3000 + + result = ex.define_aperture(model, slit, extract_defaults, exptype) + _, _, _, profile, _, limits = result + + assert np.all(profile == 0.0) + assert limits == (2000, 3000, None, None) + + +def test_define_aperture_bad_wcs(monkeypatch, mock_nirspec_fs_one_slit, extract_defaults): + model = mock_nirspec_fs_one_slit + extract_defaults['dispaxis'] = 1 + slit = None + exptype = 'NRS_FIXEDSLIT' + + # Set a wavelength so wcs is not called to retrieve it + model.wavelength = np.empty_like(model.data) + model.wavelength[:] = np.linspace(3, 5, model.data.shape[1]) + + # mock a bad wcs + def return_nan(*args): + return np.nan, np.nan, np.nan + + monkeypatch.setattr(model.meta, 'wcs', return_nan) + + result = ex.define_aperture(model, slit, extract_defaults, exptype) + ra, dec = result[:2] + + # RA and Dec returned are none + assert ra is None + assert dec is None + + +def test_define_aperture_use_source(monkeypatch, mock_nirspec_fs_one_slit, extract_defaults): + model = mock_nirspec_fs_one_slit + extract_defaults['dispaxis'] = 1 + slit = None + exptype = 'NRS_FIXEDSLIT' + + # mock the source location function + def mock_source_location(*args): + return 24, 7.74, 9.5 + + monkeypatch.setattr(ex, 'location_from_wcs', mock_source_location) + + # set parameters to extract a 6 pixel aperture, centered on source location + extract_defaults['use_source_posn'] = True + extract_defaults['extract_width'] = 6.0 + + result = ex.define_aperture(model, slit, extract_defaults, exptype) + _, _, _, profile, _, limits = result + + assert np.all(profile[:7] == 0.0) + assert np.all(profile[7:13] == 1.0) + assert np.all(profile[13:] == 0.0) + + +def test_extract_one_slit_horizontal(mock_nirspec_fs_one_slit, extract_defaults, + simple_profile, background_profile): + # update parameters to subtract background + extract_defaults['dispaxis'] = 1 + extract_defaults['subtract_background'] = True + extract_defaults['bkg_fit'] = 'poly' + extract_defaults['bkg_order'] = 1 + + # set a source in the profile region + mock_nirspec_fs_one_slit.data[simple_profile != 0] += 1.0 + + result = ex.extract_one_slit(mock_nirspec_fs_one_slit, -1, simple_profile, + background_profile, extract_defaults) + + for data in result[:-1]: + assert np.all(data > 0) + assert data.shape == (mock_nirspec_fs_one_slit.data.shape[1],) + + # residuals from the 2D scene model should be zero - this simple case + # is exactly modeled with a box profile + scene_model = result[-1] + assert scene_model.shape == mock_nirspec_fs_one_slit.data.shape + assert np.allclose(np.abs(mock_nirspec_fs_one_slit.data - scene_model), 0) + + # flux should be 1.0 * npixels + flux = result[0] + npixels = result[-2] + assert np.allclose(flux, npixels) + + # npixels is sum of profile + assert np.all(npixels == np.sum(simple_profile, axis=0)) + + +def test_extract_one_slit_vertical(mock_miri_lrs_fs, extract_defaults, + simple_profile, background_profile): + model = mock_miri_lrs_fs + profile = simple_profile.T + profile_bg = background_profile.T + + # update parameters to subtract background + extract_defaults['dispaxis'] = 2 + extract_defaults['subtract_background'] = True + extract_defaults['bkg_fit'] = 'poly' + extract_defaults['bkg_order'] = 1 + + # set a source in the profile region + model.data[profile != 0] += 1.0 + + result = ex.extract_one_slit(model, -1, profile, profile_bg, extract_defaults) + + for data in result[:-1]: + assert np.all(data > 0) + assert data.shape == (model.data.shape[0],) + + # residuals from the 2D scene model should be zero - this simple case + # is exactly modeled with a box profile + scene_model = result[-1] + assert scene_model.shape == model.data.shape + assert np.allclose(np.abs(model.data - scene_model), 0) + + # flux should be 1.0 * npixels + flux = result[0] + npixels = result[-2] + assert np.allclose(flux, npixels) + + # npixels is sum of profile + assert np.all(npixels == np.sum(profile, axis=1)) + + +def test_extract_one_slit_vertical_no_bg(mock_miri_lrs_fs, extract_defaults, + simple_profile): + model = mock_miri_lrs_fs + profile = simple_profile.T + extract_defaults['dispaxis'] = 2 + + result = ex.extract_one_slit(model, -1, profile, None, extract_defaults) + + # flux and variances are nonzero + for data in result[:4]: + assert np.all(data > 0) + assert data.shape == (model.data.shape[0],) + + # background and variances are zero + for data in result[4:8]: + assert np.all(data == 0) + assert data.shape == (model.data.shape[0],) + + # npixels is the sum of the profile + assert np.allclose(result[8], np.sum(simple_profile, axis=0)) + + # scene model has 2D shape + assert result[-1].shape == model.data.shape + + +def test_extract_one_slit_multi_int(mock_nirspec_bots, extract_defaults, + simple_profile, log_watcher): + model = mock_nirspec_bots + extract_defaults['dispaxis'] = 1 + + log_watcher.message = "Extracting integration 2" + result = ex.extract_one_slit(model, 1, simple_profile, None, extract_defaults) + log_watcher.assert_seen() + + # flux and variances are nonzero + for data in result[:4]: + assert np.all(data > 0) + assert data.shape == (model.data.shape[2],) + + # background and variances are zero + for data in result[4:8]: + assert np.all(data == 0) + assert data.shape == (model.data.shape[2],) + + # npixels is the sum of the profile + assert np.allclose(result[8], np.sum(simple_profile, axis=0)) + + # scene model has 2D shape + assert result[-1].shape == model.data.shape[-2:] + + +def test_extract_one_slit_missing_var(mock_nirspec_fs_one_slit, extract_defaults, + simple_profile): + model = mock_nirspec_fs_one_slit + extract_defaults['dispaxis'] = 1 + + # Test that mismatched variances still extract okay. + # This is probably only possible for var_flat, which is optional and + # uninitialized if flat fielding is skipped, but the code has handling + # for all 3 variance arrays. + model.var_rnoise = np.zeros((10, 10)) + model.var_poisson = np.zeros((10, 10)) + model.var_flat = np.zeros((10, 10)) + + result = ex.extract_one_slit(model, -1, simple_profile, None, extract_defaults) + + # flux is nonzero + assert np.all(result[0] > 0) + assert result[0].shape == (model.data.shape[1],) + + # variances are zero + for data in result[1:4]: + assert np.all(data == 0) + assert data.shape == (model.data.shape[1],) + + +def test_create_extraction_with_photom(create_extraction_inputs): + model = create_extraction_inputs[0] + model.meta.cal_step.photom = 'COMPLETE' + + ex.create_extraction(*create_extraction_inputs) + + output_model = create_extraction_inputs[2] + assert output_model.spec[0].spec_table.columns['flux'].unit == 'Jy' + + +def test_create_extraction_without_photom(create_extraction_inputs): + model = create_extraction_inputs[0] + model.meta.cal_step.photom = 'SKIPPED' + + ex.create_extraction(*create_extraction_inputs) + + output_model = create_extraction_inputs[2] + assert output_model.spec[0].spec_table.columns['flux'].unit == 'DN/s' + + +def test_create_extraction_missing_src_type(create_extraction_inputs): + model = create_extraction_inputs[0] + model.source_type = None + model.meta.target.source_type = 'EXTENDED' + + ex.create_extraction(*create_extraction_inputs) + + output_model = create_extraction_inputs[2] + assert output_model.spec[0].source_type == 'EXTENDED' + + +def test_create_extraction_no_match(create_extraction_inputs): + create_extraction_inputs[4] = 'bad slitname' + with pytest.raises(ValueError, match="Missing extraction parameters"): + ex.create_extraction(*create_extraction_inputs) + + +def test_create_extraction_partial_match(create_extraction_inputs, log_watcher): + # match a slit that has a mismatched spectral order specified + create_extraction_inputs[4] = 'slit7' + + log_watcher.message = 'Spectral order 1 not found' + with pytest.raises(ex.ContinueError): + ex.create_extraction(*create_extraction_inputs) + log_watcher.assert_seen() + + +def test_create_extraction_missing_dispaxis(create_extraction_inputs, log_watcher): + create_extraction_inputs[0].meta.wcsinfo.dispersion_direction = None + log_watcher.message = 'dispersion direction information is missing' + with pytest.raises(ex.ContinueError): + ex.create_extraction(*create_extraction_inputs) + log_watcher.assert_seen() + + +def test_create_extraction_missing_wavelengths(create_extraction_inputs, log_watcher): + model = create_extraction_inputs[0] + model.wavelength = np.full_like(model.data, np.nan) + log_watcher.message = 'Spectrum is empty; no valid data' + with pytest.raises(ex.ContinueError): + ex.create_extraction(*create_extraction_inputs) + log_watcher.assert_seen() + + +def test_create_extraction_nrs_apcorr(create_extraction_inputs, nirspec_fs_apcorr, + mock_nirspec_bots, log_watcher): + model = mock_nirspec_bots + model.source_type = 'POINT' + model.meta.cal_step.photom = 'COMPLETE' + create_extraction_inputs[0] = model + + log_watcher.message = 'Tabulating aperture correction' + ex.create_extraction(*create_extraction_inputs, apcorr_ref_model=nirspec_fs_apcorr, + use_source_posn=False) + log_watcher.assert_seen() + + +def test_create_extraction_one_int(create_extraction_inputs, mock_nirspec_bots, log_watcher): + # Input model is a cube, but with only one integration + model = mock_nirspec_bots + model.data = model.data[0].reshape(1, *model.data.shape[-2:]) + create_extraction_inputs[0] = model + + log_watcher.message = '1 integration done' + ex.create_extraction(*create_extraction_inputs, log_increment=1) + output_model = create_extraction_inputs[2] + assert len(output_model.spec) == 1 + log_watcher.assert_seen() + + +def test_create_extraction_log_increment( + create_extraction_inputs, mock_nirspec_bots, log_watcher): + create_extraction_inputs[0] = mock_nirspec_bots + + # all integrations are logged + log_watcher.message = '... 9 integrations done' + ex.create_extraction(*create_extraction_inputs, log_increment=1) + log_watcher.assert_seen() + + +@pytest.mark.parametrize('use_source', [True, False, None]) +@pytest.mark.parametrize('source_type', ['POINT', 'EXTENDED']) +def test_create_extraction_use_source( + monkeypatch, create_extraction_inputs, mock_nirspec_fs_one_slit, + use_source, source_type, log_watcher): + model = mock_nirspec_fs_one_slit + model.source_type = source_type + create_extraction_inputs[0] = model + + # mock the source location function + def mock_source_location(*args): + return 24, 7.74, 9.5 + + monkeypatch.setattr(ex, 'location_from_wcs', mock_source_location) + + if source_type != 'POINT' and use_source is None: + # If not specified, source position should be used only if POINT + log_watcher.message = 'Setting use_source_posn to False' + elif use_source is True or (source_type == 'POINT' and use_source is None): + # If explicitly set to True, or unspecified + source type is POINT, + # source position is used + log_watcher.message = 'Aperture start/stop: -15' + else: + # If False, source position is not used + log_watcher.message = 'Aperture start/stop: 0' + ex.create_extraction(*create_extraction_inputs, use_source_posn=use_source) + log_watcher.assert_seen() + + +def test_run_extract1d(mock_nirspec_mos): + model = mock_nirspec_mos + output_model, profile_model, scene_model = ex.run_extract1d(model) + assert isinstance(output_model, dm.MultiSpecModel) + assert profile_model is None + assert scene_model is None + output_model.close() + + +def test_run_extract1d_save_models(mock_niriss_wfss_l3): + model = mock_niriss_wfss_l3 + output_model, profile_model, scene_model = ex.run_extract1d( + model, save_profile=True, save_scene_model=True) + assert isinstance(output_model, dm.MultiSpecModel) + assert isinstance(profile_model, ModelContainer) + assert isinstance(scene_model, ModelContainer) + + assert len(profile_model) == len(model) + assert len(scene_model) == len(model) + + for pmodel in profile_model: + assert isinstance(pmodel, dm.ImageModel) + for smodel in scene_model: + assert isinstance(smodel, dm.ImageModel) + + output_model.close() + profile_model.close() + scene_model.close() + + +def test_run_extract1d_save_cube_scene(mock_nirspec_bots): + model = mock_nirspec_bots + output_model, profile_model, scene_model = ex.run_extract1d( + model, save_profile=True, save_scene_model=True) + assert isinstance(output_model, dm.MultiSpecModel) + assert isinstance(profile_model, dm.ImageModel) + assert isinstance(scene_model, dm.CubeModel) + + assert profile_model.data.shape == model.data.shape[-2:] + assert scene_model.data.shape == model.data.shape + + output_model.close() + profile_model.close() + scene_model.close() + + + +def test_run_extract1d_tso(mock_nirspec_bots): + model = mock_nirspec_bots + output_model, _, _ = ex.run_extract1d(model) + + # time and integration keywords are populated + for i, spec in enumerate(output_model.spec): + assert spec.int_num == i + 1 + + output_model.close() + + +@pytest.mark.parametrize('from_name_attr', [True, False]) +def test_run_extract1d_slitmodel_name(mock_nirspec_fs_one_slit, from_name_attr): + model = mock_nirspec_fs_one_slit + slit_name = 'S200A1' + if from_name_attr: + model.name = slit_name + model.meta.instrument.fixed_slit = None + else: + model.name = None + model.meta.instrument.fixed_slit = 'S200A1' + + output_model, _, _ = ex.run_extract1d(model) + assert output_model.spec[0].name == 'S200A1' + + output_model.close() + + +@pytest.mark.parametrize('from_name_attr', [True, False]) +def test_run_extract1d_imagemodel_name(mock_miri_lrs_fs, from_name_attr): + model = mock_miri_lrs_fs + slit_name = 'test_slit_name' + if from_name_attr: + model.name = slit_name + else: + model.name = None + + output_model, _, _ = ex.run_extract1d(model) + if from_name_attr: + assert output_model.spec[0].name == 'test_slit_name' + else: + assert output_model.spec[0].name == 'MIR_LRS-FIXEDSLIT' + output_model.close() + + +def test_run_extract1d_apcorr(mock_miri_lrs_fs, miri_lrs_apcorr_file, log_watcher): + model = mock_miri_lrs_fs + model.meta.target.source_type = 'POINT' + + log_watcher.message = 'Creating aperture correction' + output_model, _, _ = ex.run_extract1d(model, apcorr_ref_name=miri_lrs_apcorr_file) + log_watcher.assert_seen() + + output_model.close() + + +def test_run_extract1d_invalid(): + model = dm.MultiSpecModel() + with pytest.raises(RuntimeError, match="Can't extract a spectrum"): + ex.run_extract1d(model) + + +def test_run_extract1d_zeroth_order_slit(mock_nirspec_fs_one_slit): + model = mock_nirspec_fs_one_slit + model.meta.wcsinfo.spectral_order = 0 + output_model, _, _ = ex.run_extract1d(model) + + # no spectra extracted for zeroth order + assert len(output_model.spec) == 0 + output_model.close() + + +def test_run_extract1d_zeroth_order_image(mock_miri_lrs_fs): + model = mock_miri_lrs_fs + model.meta.wcsinfo.spectral_order = 0 + output_model, _, _ = ex.run_extract1d(model) + + # no spectra extracted for zeroth order + assert len(output_model.spec) == 0 + output_model.close() + + +def test_run_extract1d_zeroth_order_multispec(mock_nirspec_mos): + model = mock_nirspec_mos + for slit in model.slits: + slit.meta.wcsinfo.spectral_order = 0 + output_model, _, _ = ex.run_extract1d(model) + + # no spectra extracted for zeroth order + assert len(output_model.spec) == 0 + output_model.close() + + +def test_run_extract1d_no_data(mock_niriss_wfss_l3): + container = mock_niriss_wfss_l3 + for model in container: + model.data = np.array([]) + output_model, _, _ = ex.run_extract1d(container) + + # no spectra extracted + assert len(output_model.spec) == 0 + output_model.close() + + +def test_run_extract1d_continue_error_slit(monkeypatch, mock_nirspec_fs_one_slit): + def raise_continue_error(*args, **kwargs): + raise ex.ContinueError('Test error') + + monkeypatch.setattr(ex, 'create_extraction', raise_continue_error) + output_model, _, _ = ex.run_extract1d(mock_nirspec_fs_one_slit) + + # no spectra extracted + assert len(output_model.spec) == 0 + output_model.close() + + +def test_run_extract1d_continue_error_image(monkeypatch, mock_miri_lrs_fs): + def raise_continue_error(*args, **kwargs): + raise ex.ContinueError('Test error') + + monkeypatch.setattr(ex, 'create_extraction', raise_continue_error) + output_model, _, _ = ex.run_extract1d(mock_miri_lrs_fs) + + # no spectra extracted + assert len(output_model.spec) == 0 + output_model.close() + + +def test_run_extract1d_continue_error_multislit(monkeypatch, mock_nirspec_mos): + def raise_continue_error(*args, **kwargs): + raise ex.ContinueError('Test error') + + monkeypatch.setattr(ex, 'create_extraction', raise_continue_error) + output_model, _, _ = ex.run_extract1d(mock_nirspec_mos) + + # no spectra extracted + assert len(output_model.spec) == 0 + output_model.close() diff --git a/jwst/extract_1d/tests/test_extract_1d_step.py b/jwst/extract_1d/tests/test_extract_1d_step.py index 3e096b38b0..98e0aaeb63 100644 --- a/jwst/extract_1d/tests/test_extract_1d_step.py +++ b/jwst/extract_1d/tests/test_extract_1d_step.py @@ -1,79 +1,12 @@ +import os + import numpy as np import pytest import stdatamodels.jwst.datamodels as dm -from jwst.assign_wcs.util import wcs_bbox_from_shape +from jwst.datamodels import ModelContainer from jwst.extract_1d.extract_1d_step import Extract1dStep - - -@pytest.fixture -def simple_wcs(): - shape = (50, 50) - xcenter = shape[1] // 2.0 - - def simple_wcs_function(x, y): - """ Simple WCS for testing """ - crpix1 = xcenter - crpix3 = 1.0 - cdelt1 = 0.1 - cdelt2 = 0.1 - cdelt3 = 0.01 - - crval1 = 45.0 - crval2 = 45.0 - crval3 = 7.5 - - wave = (x + 1 - crpix3) * cdelt3 + crval3 - ra = (x + 1 - crpix1) * cdelt1 + crval1 - dec = np.full_like(ra, crval2 + 1 * cdelt2) - - return ra, dec, wave - - simple_wcs_function.bounding_box = wcs_bbox_from_shape(shape) - - return simple_wcs_function - - -@pytest.fixture() -def mock_nirspec_fs_one_slit(simple_wcs): - model = dm.SlitModel() - model.meta.instrument.name = 'NIRSPEC' - model.meta.instrument.detector = 'NRS1' - model.meta.observation.date = '2023-07-22' - model.meta.observation.time = '06:24:45.569' - model.meta.instrument.fixed_slit = 'S200A1' - model.meta.exposure.type = 'NRS_FIXEDSLIT' - model.meta.subarray.name = 'ALLSLITS' - - model.source_type = 'EXTENDED' - model.meta.wcsinfo.dispersion_direction = 1 - model.meta.wcs = simple_wcs - - model.data = np.arange(50 * 50, dtype=float).reshape((50, 50)) - model.var_poisson = model.data * 0.02 - model.var_rnoise = model.data * 0.02 - model.var_flat = model.data * 0.05 - yield model - model.close() - - -@pytest.fixture() -def mock_nirspec_mos(mock_nirspec_fs_one_slit): - model = dm.MultiSlitModel() - model.meta.instrument.name = 'NIRSPEC' - model.meta.instrument.detector = 'NRS1' - model.meta.observation.date = '2023-07-22' - model.meta.observation.time = '06:24:45.569' - model.meta.exposure.type = 'NRS_MSASPEC' - - nslit = 3 - for i in range(nslit): - slit = mock_nirspec_fs_one_slit.copy() - slit.name = str(i + 1) - model.slits.append(slit) - - yield model - model.close() +from jwst.extract_1d.soss_extract import soss_extract @pytest.mark.parametrize('slit_name', [None, 'S200A1', 'S1600A1']) @@ -116,3 +49,170 @@ def test_extract_nirspec_mos_multi_slit(mock_nirspec_mos, simple_wcs): assert np.all(spec.spec_table['FLUX_ERROR'] > 0) result.close() + + +def test_extract_nirspec_bots(mock_nirspec_bots, simple_wcs): + result = Extract1dStep.call(mock_nirspec_bots, apply_apcorr=False) + assert result.meta.cal_step.extract_1d == 'COMPLETE' + assert (result.spec[0].name == 'S1600A1') + + # output wavelength is the same as input + _, _, expected_wave = simple_wcs(np.arange(50), np.arange(50)) + assert np.allclose(result.spec[0].spec_table['WAVELENGTH'], expected_wave) + + # output flux and errors are non-zero, exact values will depend + # on extraction parameters + assert np.all(result.spec[0].spec_table['FLUX'] > 0) + assert np.all(result.spec[0].spec_table['FLUX_ERROR'] > 0) + result.close() + + +@pytest.mark.parametrize('ifu_set_srctype', [None, 'EXTENDED']) +def test_extract_miri_ifu(mock_miri_ifu, simple_wcs_ifu, ifu_set_srctype): + # Source type defaults to extended, results should be the + # same with and without override + result = Extract1dStep.call(mock_miri_ifu, ifu_covar_scale=1.0, + ifu_set_srctype=ifu_set_srctype) + assert result.meta.cal_step.extract_1d == 'COMPLETE' + + # output wavelength is the same as input + _, _, expected_wave = simple_wcs_ifu(np.arange(50), np.arange(50), np.arange(10)) + assert np.allclose(result.spec[0].spec_table['WAVELENGTH'], expected_wave) + + # output flux for extended data is a simple sum over all data + # with a conversion factor for Jy + assert np.allclose(result.spec[0].spec_table['FLUX'], + 1e6 * np.sum(mock_miri_ifu.data, axis=(1, 2))) + + # output error comes from the sum of the variance components + variance_sum = (np.sum(mock_miri_ifu.var_rnoise, axis=(1, 2)) + + np.sum(mock_miri_ifu.var_poisson, axis=(1, 2)) + + np.sum(mock_miri_ifu.var_flat, axis=(1, 2))) + assert np.allclose(result.spec[0].spec_table['FLUX_ERROR'], + 1e6 * np.sqrt(variance_sum)) + result.close() + + +def test_extract_container(mock_nirspec_mos, mock_nirspec_fs_one_slit, simple_wcs): + # if not WFSS, the container is looped over, so contents need not match + container = ModelContainer([mock_nirspec_mos, mock_nirspec_fs_one_slit]) + result = Extract1dStep.call(container) + assert isinstance(result, ModelContainer) + + for model in result: + for spec in model.spec: + # output wavelength is the same as input + _, _, expected_wave = simple_wcs(np.arange(50), np.arange(50)) + assert np.allclose(spec.spec_table['WAVELENGTH'], expected_wave) + + # output flux and errors are non-zero, exact values will depend + # on extraction parameters + assert np.all(spec.spec_table['FLUX'] > 0) + assert np.all(spec.spec_table['FLUX_ERROR'] > 0) + + result.close() + + +def test_extract_niriss_wfss(mock_niriss_wfss_l3, simple_wcs): + # input is a SourceModelContainer + result = Extract1dStep.call(mock_niriss_wfss_l3) + + # output is a single spectral model (not a container) + assert isinstance(result, dm.MultiSpecModel) + assert result.meta.cal_step.extract_1d == 'COMPLETE' + + for i, spec in enumerate(result.spec): + assert spec.name == str(i + 1) + + # output wavelength is the same as input + _, _, expected_wave = simple_wcs(np.arange(50), np.arange(50)) + assert np.allclose(spec.spec_table['WAVELENGTH'], expected_wave) + + # output flux and errors are non-zero, exact values will depend + # on extraction parameters + assert np.all(spec.spec_table['FLUX'] > 0) + assert np.all(spec.spec_table['FLUX_ERROR'] > 0) + + result.close() + + +@pytest.mark.slow +def test_extract_niriss_soss_256(tmp_path, mock_niriss_soss_256): + result = Extract1dStep.call(mock_niriss_soss_256, soss_rtol=0.1, + soss_modelname='soss_model.fits', + output_dir=str(tmp_path)) + assert result.meta.cal_step.extract_1d == 'COMPLETE' + + # output flux and errors are non-zero, exact values will depend + # on extraction parameters + assert np.all(result.spec[0].spec_table['FLUX'] > 0) + assert np.all(result.spec[0].spec_table['FLUX_ERROR'] > 0) + result.close() + + # soss output files are saved + assert os.path.isfile(tmp_path / 'soss_model_SossExtractModel.fits') + assert os.path.isfile(tmp_path / 'soss_model_AtocaSpectra.fits') + + +@pytest.mark.slow +def test_extract_niriss_soss_96(tmp_path, mock_niriss_soss_96): + result = Extract1dStep.call(mock_niriss_soss_96, soss_rtol=0.1, + soss_modelname='soss_model.fits', + output_dir=str(tmp_path)) + assert result.meta.cal_step.extract_1d == 'COMPLETE' + + # output flux and errors are non-zero, exact values will depend + # on extraction parameters + assert np.all(result.spec[0].spec_table['FLUX'] > 0) + assert np.all(result.spec[0].spec_table['FLUX_ERROR'] > 0) + result.close() + + # soss output files are saved + assert os.path.isfile(tmp_path / 'soss_model_SossExtractModel.fits') + assert os.path.isfile(tmp_path / 'soss_model_AtocaSpectra.fits') + + +def test_extract_niriss_soss_fail(tmp_path, monkeypatch, mock_niriss_soss_96): + # Mock an error condition in the soss extraction + def mock(*args, **kwargs): + return None, None, None + monkeypatch.setattr(soss_extract, 'run_extract1d', mock) + + # None is returned + result = Extract1dStep.call(mock_niriss_soss_96) + assert result is None + + +def test_save_output_single(tmp_path, mock_nirspec_fs_one_slit): + mock_nirspec_fs_one_slit.meta.filename = 'test_s2d.fits' + result = Extract1dStep.call(mock_nirspec_fs_one_slit, + save_results=True, save_profile=True, + save_scene_model=True, output_dir=str(tmp_path), + suffix='x1d') + + output_path = str(tmp_path / 'test_x1d.fits') + + assert os.path.isfile(output_path) + assert os.path.isfile(output_path.replace('x1d', 'profile')) + assert os.path.isfile(output_path.replace('x1d', 'scene_model')) + + result.close() + + +def test_save_output_multislit(tmp_path, mock_nirspec_mos): + mock_nirspec_mos.meta.filename = 'test_s2d.fits' + result = Extract1dStep.call(mock_nirspec_mos, + save_results=True, save_profile=True, + save_scene_model=True, output_dir=str(tmp_path), + suffix='x1d') + + output_path = str(tmp_path / 'test_x1d.fits') + + assert os.path.isfile(output_path) + + # intermediate files for multislit data contain the slit name + for slit in mock_nirspec_mos.slits: + assert os.path.isfile(output_path.replace('x1d', f'{slit.name}_profile')) + assert os.path.isfile(output_path.replace('x1d', f'{slit.name}_scene_model')) + + result.close() diff --git a/jwst/extract_1d/tests/test_extract_src_flux.py b/jwst/extract_1d/tests/test_extract_src_flux.py index cfab06f553..b6632a1cfc 100644 --- a/jwst/extract_1d/tests/test_extract_src_flux.py +++ b/jwst/extract_1d/tests/test_extract_src_flux.py @@ -13,105 +13,156 @@ def inputs_constant(): shape = (9, 5) image = np.arange(shape[0] * shape[1], dtype=np.float32).reshape(shape) - var_poisson = image.copy() var_rnoise = image.copy() + var_poisson = image.copy() var_rflat = image.copy() - x = 2 - j = 2 - lam = 1.234 # an arbitrary value (not actually used) weights = None - bkgmodels = [None, None, None, None] - lower = np.zeros(shape[1], dtype=np.float64) + 3. # middle of pixel 3 - upper = np.zeros(shape[1], dtype=np.float64) + 7. # middle of pixel 7 - srclim = [[lower, upper]] + profile_bg = None - return (image, var_poisson, var_rnoise, var_rflat, - x, j, lam, srclim, weights, bkgmodels) + profile = np.zeros_like(image) + profile[3] = 0.5 # lower limit: middle of pixel 3 + profile[4:7] = 1.0 # center of aperture + profile[7] = 0.5 # upper limit: middle of pixel 7 + return (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg) -def test_extract_src_flux(inputs_constant): - (image, var_poisson, var_rnoise, var_rflat, - x, j, lam, srclim, weights, bkgmodels) = inputs_constant - (total_flux, f_var_poisson, f_var_rnoise, f_var_flat, - bkg_flux, b_var_poisson, b_var_rnoise, b_var_flat, - tarea, twht) = extract1d._extract_src_flux( - image, var_poisson, var_rnoise, var_rflat, - x, j, lam, srclim, weights, bkgmodels) +@pytest.fixture +def inputs_with_source(): + shape = (9, 5) + image = np.full(shape, 0.0) + image[3] = 5.0 + image[4] = 10.0 + image[5] = 5.0 + var_rnoise = np.full(shape, 0.1) + var_poisson = image * 0.05 + var_rflat = image * 0.05 + weights = None + profile_bg = None + + profile = np.zeros_like(image) + profile[3] = 0.25 + profile[4] = 0.5 + profile[5] = 0.25 - # 0.5 * 17. + 22. + 27. + 32. + 0.5 * 37. - assert math.isclose(total_flux, 108., rel_tol=1.e-8, abs_tol=1.e-8) + return (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg) - assert bkg_flux == 0. - assert math.isclose(tarea, 4., rel_tol=1.e-8, abs_tol=1.e-8) +def test_extract_src_flux(inputs_constant): + (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg) = inputs_constant + + (total_flux, f_var_rnoise, f_var_poisson, f_var_flat, + bkg_flux, b_var_rnoise, b_var_poisson, b_var_flat, + npixels, model) = extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg) + + # check the value at column 2 + # 0.5 * 17. + 22. + 27. + 32. + 0.5 * 37. + assert math.isclose(total_flux[0][2], 108., rel_tol=1.e-8, abs_tol=1.e-8) + assert bkg_flux[0][2] == 0. + assert math.isclose(npixels[0][2], 4., rel_tol=1.e-8, abs_tol=1.e-8) + # set a NaN value in the column of interest image[5, 2] = np.nan - (total_flux, f_var_poisson, f_var_rnoise, f_var_flat, - bkg_flux, b_var_poisson, b_var_rnoise, b_var_flat, - tarea, twht) = extract1d._extract_src_flux( - image, var_poisson, var_rnoise, var_rflat, - x, j, lam, srclim, weights, bkgmodels) + (total_flux, f_var_rnoise, f_var_poisson, f_var_flat, + bkg_flux, b_var_rnoise, b_var_poisson, b_var_flat, + npixels, model) = extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg) # 0.5 * 17. + 22. + 32. + 0.5 * 37. - assert math.isclose(total_flux, 81., rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(tarea, 3., rel_tol=1.e-8, abs_tol=1.e-8) + assert math.isclose(total_flux[0][2], 81., rel_tol=1.e-8, abs_tol=1.e-8) + assert math.isclose(npixels[0][2], 3., rel_tol=1.e-8, abs_tol=1.e-8) + # set the whole column to NaN image[:, 2] = np.nan - (total_flux, f_var_poisson, f_var_rnoise, f_var_flat, - bkg_flux, b_var_poisson, b_var_rnoise, b_var_flat, - tarea, twht) = extract1d._extract_src_flux( - image, var_poisson, var_rnoise, var_rflat, - x, j, lam, srclim, weights, bkgmodels) - - assert np.isnan(total_flux) + (total_flux, f_var_rnoise, f_var_poisson, f_var_flat, + bkg_flux, b_var_rnoise, b_var_poisson, b_var_flat, + npixels, model) = extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg) - assert tarea == 0. + assert np.isnan(total_flux[0][2]) + assert npixels[0][2] == 0. -@pytest.mark.parametrize('test_type', ['all_empty', 'all_equal']) -def test_extract_src_flux_empty_interval(inputs_constant, test_type): - (image, var_poisson, var_rnoise, var_rflat, - x, j, lam, srclim, weights, bkgmodels) = inputs_constant +def test_extract_src_flux_empty_interval(inputs_constant): + (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg) = inputs_constant - if test_type == 'all_empty': - # no limits provided - srclim = [] - else: - # empty extraction range: upper equals lower - srclim[0][1] = srclim[0][0].copy() + # empty extraction range + profile[:] = 0.0 - (total_flux, f_var_poisson, f_var_rnoise, f_var_flat, - bkg_flux, b_var_poisson, b_var_rnoise, b_var_flat, - tarea, twht) = extract1d._extract_src_flux( - image, var_poisson, var_rnoise, var_rflat, - x, j, lam, srclim, weights, bkgmodels) + (total_flux, f_var_rnoise, f_var_poisson, f_var_flat, + bkg_flux, b_var_rnoise, b_var_poisson, b_var_flat, + npixels, model) = extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg) # empty interval, so no flux returned - assert np.isnan(total_flux) - assert bkg_flux == 0. - assert tarea == 0. + assert np.all(np.isnan(total_flux)) + assert np.all(bkg_flux == 0.) + assert np.all(npixels == 0.) + + +@pytest.mark.parametrize('use_weights', [True, False]) +def test_extract_optimal(inputs_with_source, use_weights): + (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg) = inputs_with_source + + if use_weights: + weights = 1 / var_rnoise + + (total_flux, f_var_rnoise, f_var_poisson, f_var_flat, + bkg_flux, b_var_rnoise, b_var_poisson, b_var_flat, + npixels, model) = extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg, extraction_type='optimal') + + # Total flux should be well modeled + assert np.allclose(total_flux[0], 20) + assert np.allclose(bkg_flux[0], 0) + assert np.allclose(npixels[0], 2.66667) + + # set a NaN value in a column of interest + image[4, 2] = np.nan + if use_weights: + weights[4, 2] = 0 + + (total_flux, f_var_rnoise, f_var_poisson, f_var_flat, + bkg_flux, b_var_rnoise, b_var_poisson, b_var_flat, + npixels, model) = extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg, extraction_type='optimal') + + # Total flux is still well modeled from 2 pixels + assert np.allclose(total_flux[0], 20) + assert np.isclose(npixels[0, 2], 1.33333) + + # set the whole column to NaN + image[:, 2] = np.nan + (total_flux, f_var_rnoise, f_var_poisson, f_var_flat, + bkg_flux, b_var_rnoise, b_var_poisson, b_var_flat, + npixels, model) = extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg) -@pytest.mark.parametrize('offset', [-100, 100]) -def test_extract_src_flux_interval_out_of_range(inputs_constant, offset): - (image, var_poisson, var_rnoise, var_rflat, - x, j, lam, srclim, weights, bkgmodels) = inputs_constant + # Now the flux can no longer be estimated in that column + assert np.isnan(total_flux[0][2]) + assert npixels[0][2] == 0. - # extraction limits out of range - srclim[0][0] += offset - srclim[0][1] += offset - (total_flux, f_var_poisson, f_var_rnoise, f_var_flat, - bkg_flux, b_var_poisson, b_var_rnoise, b_var_flat, - tarea, twht) = extract1d._extract_src_flux( - image, var_poisson, var_rnoise, var_rflat, - x, j, lam, srclim, weights, bkgmodels) +def test_too_many_profiles(inputs_constant): + (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg) = inputs_constant - # empty interval, so no flux returned - assert np.isnan(total_flux) - assert bkg_flux == 0. - assert tarea == 0. + with pytest.raises(ValueError, match="not supported with 2 input profiles"): + extract1d.extract1d(image, [profile, profile.copy()], + var_rnoise, var_poisson, var_rflat) diff --git a/jwst/extract_1d/tests/test_fit_background_model.py b/jwst/extract_1d/tests/test_fit_background_model.py index 5aeb596954..7fe302cfb5 100644 --- a/jwst/extract_1d/tests/test_fit_background_model.py +++ b/jwst/extract_1d/tests/test_fit_background_model.py @@ -1,9 +1,6 @@ """ -Test for extract_1d._fit_background_model +Tests for extract_1d background fitting """ -import math -from copy import deepcopy - import numpy as np import pytest @@ -17,212 +14,215 @@ def inputs_constant(): var_poisson = image.copy() var_rnoise = image.copy() var_rflat = image.copy() - x = 2 - j = 2 - b_lower = np.zeros(shape[1], dtype=np.float64) + 3.5 # 4, inclusive - b_upper = np.zeros(shape[1], dtype=np.float64) + 4.5 # 4, inclusive - bkglim = [[b_lower, b_upper]] + profile = np.zeros_like(image) + profile[1] = 1.0 # one pixel aperture + weights = None + + profile_bg = np.zeros_like(image) + profile_bg[4:8] = 1.0 # background region bkg_order = 0 bkg_fit = "poly" - return image, var_poisson, var_rnoise, var_rflat, x, j, bkglim, bkg_fit, bkg_order - - -def test_fit_background_model(inputs_constant): - - (bkg_model, b_var_poisson_model, b_var_rnoise_model, b_var_flat_model, npts) = \ - extract1d._fit_background_model(*inputs_constant) - - assert math.isclose(bkg_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(bkg_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_poisson_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_poisson_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) + return (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg, bkg_fit, bkg_order) - assert math.isclose(b_var_rnoise_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_rnoise_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_flat_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_flat_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert npts == 2 +@pytest.fixture +def inputs_with_source(): + shape = (9, 5) + image = np.full(shape, 1.0) + image[3] += 5.0 + image[4] += 10.0 + image[5] += 5.0 + var_rnoise = np.full(shape, 0.1) + var_poisson = image * 0.05 + var_rflat = image * 0.05 + weights = 1 / var_rnoise -def test_fit_background_mean(inputs_constant): - image, var_poisson, var_rnoise, var_rflat, x, j, bkglim, bkg_fit, bkg_order = inputs_constant - bkg_fit = "mean" + # Most of the image is set to a low but non-zero weight + # (contribution to PSF is small) + profile = np.full_like(image, 1e-6) - (bkg_model, b_var_poisson_model, b_var_rnoise_model, b_var_flat_model, npts) = \ - extract1d._fit_background_model(image, var_poisson, var_rnoise, var_rflat, - x, j, bkglim, bkg_fit, bkg_order) + # The source crosses 3 pixels, with the central pixel + # twice as high as the other two + profile[3] = 1 + profile[4] = 2 + profile[5] = 1 - assert math.isclose(bkg_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(bkg_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) + # Normalize across the spatial dimension + profile /= np.sum(profile, axis=0) - assert math.isclose(b_var_poisson_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_poisson_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) + profile_bg = None - assert math.isclose(b_var_rnoise_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_rnoise_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) + return (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg) - assert math.isclose(b_var_flat_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_flat_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert npts == 2 +@pytest.mark.parametrize('bkg_fit_type', ['poly', 'median']) +def test_fit_background(inputs_constant, bkg_fit_type): + (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg, bkg_fit, bkg_order) = inputs_constant + (total_flux, f_var_rnoise, f_var_poisson, f_var_flat, + bkg_flux, b_var_rnoise, b_var_poisson, b_var_flat, + npixels, model) = extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg, fit_bkg=True, + bkg_fit_type=bkg_fit_type, bkg_order=bkg_order) -def test_fit_background_median(inputs_constant): - image, var_poisson, var_rnoise, var_rflat, x, j, bkglim, bkg_fit, bkg_order = inputs_constant - bkg_fit = "median" + if bkg_fit_type == 'median': + extra_factor = 1.2 ** 2 + else: + extra_factor = 1.0 - (bkg_model, b_var_poisson_model, b_var_rnoise_model, b_var_flat_model, npts) = \ - extract1d._fit_background_model(image, var_poisson, var_rnoise, var_rflat, - x, j, bkglim, bkg_fit, bkg_order) + assert np.allclose(bkg_flux[0], np.sum(image[4:8], axis=0) / 4) + assert np.allclose(b_var_rnoise[0], extra_factor * np.sum(image[4:8], axis=0) / 16) + assert np.allclose(b_var_poisson[0], extra_factor * np.sum(image[4:8], axis=0) / 16) + assert np.allclose(b_var_flat[0], extra_factor * np.sum(image[4:8], axis=0) / 16) - assert math.isclose(bkg_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(bkg_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_poisson_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_poisson_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) +@pytest.mark.parametrize('bkg_fit_type', ['poly', 'median']) +def test_fit_background_with_smoothing(inputs_constant, bkg_fit_type): + (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg, bkg_fit, bkg_order) = inputs_constant - assert math.isclose(b_var_rnoise_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_rnoise_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) + (total_flux, f_var_rnoise, f_var_poisson, f_var_flat, + bkg_flux, b_var_rnoise, b_var_poisson, b_var_flat, + npixels, model) = extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg, fit_bkg=True, + bkg_fit_type=bkg_fit_type, bkg_order=bkg_order, bg_smooth_length=3) - assert math.isclose(b_var_flat_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_flat_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) + if bkg_fit_type == 'median': + extra_factor = 1.2 ** 2 + else: + extra_factor = 1.0 - assert npts == 2 + # Should be the same as background fit without smoothing, + # except for edge effects + expected = np.sum(image[4:8], axis=0) / 4 + assert np.allclose(bkg_flux[0][1:-1], expected[1:-1]) + assert np.allclose(b_var_rnoise[0][1:-1], extra_factor * expected[1:-1] / 4) + assert np.allclose(b_var_poisson[0][1:-1], extra_factor * expected[1:-1] / 4) + assert np.allclose(b_var_flat[0][1:-1], extra_factor * expected[1:-1] / 4) def test_handles_nan(inputs_constant): - image, var_poisson, var_rnoise, var_rflat, x, j, bkglim, bkg_fit, bkg_order = inputs_constant + (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg, bkg_fit, bkg_order) = inputs_constant image[:, 2] = np.nan - (bkg_model, b_var_poisson_model, b_var_rnoise_model, b_var_flat_model, npts) = \ - extract1d._fit_background_model(image, var_poisson, var_rnoise, var_rflat, - x, j, bkglim, bkg_fit, bkg_order) - - assert math.isclose(bkg_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(bkg_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_poisson_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_poisson_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_rnoise_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_rnoise_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_flat_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_flat_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert npts == 0 - - -def test_handles_one_value(inputs_constant): - image, var_poisson, var_rnoise, var_rflat, x, j, bkglim, bkg_fit, bkg_order = inputs_constant - image[np.where(image == 22)] = np.nan # During extraction, only two pixels are returned "normally"; set one to Nan - - # If only one data point is available, the polynomial fit is forced to 0 - (bkg_model, b_var_poisson_model, b_var_rnoise_model, b_var_flat_model, npts) = \ - extract1d._fit_background_model(image, var_poisson, var_rnoise, var_rflat, - x, j, bkglim, bkg_fit, bkg_order) - - assert math.isclose(bkg_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(bkg_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_poisson_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_poisson_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_rnoise_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_rnoise_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_flat_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_flat_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert npts == 0 - - -@pytest.mark.parametrize('test_type', ['all_empty', 'all_equal']) -def test_handles_empty_interval(inputs_constant, test_type): - image, var_poisson, var_rnoise, var_rflat, x, j, bkglim, bkg_fit, bkg_order = inputs_constant - - if test_type == 'all_empty': - # no limits provided - bkglim = [] + (total_flux, f_var_rnoise, f_var_poisson, f_var_flat, + bkg_flux, b_var_rnoise, b_var_poisson, b_var_flat, + npixels, model) = extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg, fit_bkg=True, + bkg_fit_type=bkg_fit, bkg_order=bkg_order) + + assert np.allclose(bkg_flux[0], np.nansum(image[4:8], axis=0) / 4) + assert np.allclose(b_var_rnoise[0], np.nansum(image[4:8], axis=0) / 16) + assert np.allclose(b_var_poisson[0], np.nansum(image[4:8], axis=0) / 16) + assert np.allclose(b_var_flat[0], np.nansum(image[4:8], axis=0) / 16) + + +@pytest.mark.parametrize('bkg_order_val', [0, 1, 2]) +def test_handles_one_value(inputs_constant, bkg_order_val): + (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg, bkg_fit, bkg_order) = inputs_constant + profile_bg[:] = 0.0 + profile_bg[4:6] = 1.0 + image[4] = np.nan + + (total_flux, f_var_rnoise, f_var_poisson, f_var_flat, + bkg_flux, b_var_rnoise, b_var_poisson, b_var_flat, + npixels, model) = extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg, fit_bkg=True, + bkg_fit_type=bkg_fit, bkg_order=bkg_order_val) + + if bkg_order_val == 0: + expected = image[5] else: - # empty extraction range: upper equals lower - bkglim[0][1] = bkglim[0][0].copy() - - # No data available: background model is 0 - (bkg_model, b_var_poisson_model, b_var_rnoise_model, b_var_flat_model, npts) = \ - extract1d._fit_background_model(image, var_poisson, var_rnoise, var_rflat, - x, j, bkglim, bkg_fit, bkg_order) - - assert math.isclose(bkg_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(bkg_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_poisson_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_poisson_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_rnoise_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_rnoise_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_flat_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_flat_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert npts == 0 - - -@pytest.mark.parametrize('offset', [-100, 100]) -def test_handles_interval_out_of_range(inputs_constant, offset): - image, var_poisson, var_rnoise, var_rflat, x, j, bkglim, bkg_fit, bkg_order = inputs_constant - bkglim[0][0] += offset - bkglim[0][1] += offset - - # No data available: background model is 0 - (bkg_model, b_var_poisson_model, b_var_rnoise_model, b_var_flat_model, npts) = \ - extract1d._fit_background_model(image, var_poisson, var_rnoise, var_rflat, - x, j, bkglim, bkg_fit, bkg_order) - - assert math.isclose(bkg_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(bkg_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_poisson_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_poisson_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_rnoise_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_rnoise_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_flat_model(0.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_flat_model(8.), 0.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert npts == 0 - - -def test_handles_one_empty_interval(inputs_constant): - image, var_poisson, var_rnoise, var_rflat, x, j, bkglim, bkg_fit, bkg_order = inputs_constant - - # add an extra interval that is empty - bkglim.append(deepcopy(bkglim[0])) - bkglim[1][0] += 2.0 - bkglim[1][1] = bkglim[1][0].copy() - print(bkglim) - - # should ignore the second interval and return a valid answer for the first - (bkg_model, b_var_poisson_model, b_var_rnoise_model, b_var_flat_model, npts) = \ - extract1d._fit_background_model(image, var_poisson, var_rnoise, var_rflat, - x, j, bkglim, bkg_fit, bkg_order) - - assert math.isclose(bkg_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(bkg_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_poisson_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_poisson_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_rnoise_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_rnoise_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert math.isclose(b_var_flat_model(0.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - assert math.isclose(b_var_flat_model(8.), 22.0, rel_tol=1.e-8, abs_tol=1.e-8) - - assert npts == 2 + # not enough input for fit: value is set to 0.0 + expected = 0.0 + + assert np.allclose(bkg_flux[0], expected) + assert np.allclose(b_var_rnoise[0], expected) + assert np.allclose(b_var_poisson[0], expected) + assert np.allclose(b_var_flat[0], expected) + + +def test_handles_empty_interval(inputs_constant): + (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg, bkg_fit, bkg_order) = inputs_constant + profile_bg[:] = 0.0 + + (total_flux, f_var_rnoise, f_var_poisson, f_var_flat, + bkg_flux, b_var_rnoise, b_var_poisson, b_var_flat, + npixels, model) = extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg, fit_bkg=True, + bkg_fit_type=bkg_fit, bkg_order=bkg_order) + + assert np.allclose(bkg_flux[0], 0.0) + assert np.allclose(b_var_rnoise[0], 0.0) + assert np.allclose(b_var_poisson[0], 0.0) + assert np.allclose(b_var_flat[0], 0.0) + + +def test_bad_fit_type(inputs_constant): + (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg, bkg_fit, bkg_order) = inputs_constant + with pytest.raises(ValueError, match="bkg_fit_type should be 'median' or 'poly'"): + extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg, fit_bkg=True, + bkg_fit_type='mean', bkg_order=bkg_order) + + +@pytest.mark.parametrize('smooth', [1.5, 2]) +def test_smooth_length(inputs_constant, smooth): + (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg, bkg_fit, bkg_order) = inputs_constant + with pytest.raises(ValueError, match="should be an odd integer >= 1"): + extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg, fit_bkg=True, + bkg_fit_type=bkg_fit, bkg_order=bkg_order, bg_smooth_length=smooth) + + +@pytest.mark.parametrize('extraction_type', ['box', 'optimal']) +@pytest.mark.parametrize('bkg_order_val', [-1, 2.3]) +def test_bad_fit_order(inputs_constant, extraction_type, bkg_order_val): + (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg, bkg_fit, bkg_order) = inputs_constant + with pytest.raises(ValueError, match="bkg_order must be an integer >= 0"): + extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg, fit_bkg=True, + bkg_fit_type='poly', bkg_order=bkg_order_val, + extraction_type=extraction_type) + + +@pytest.mark.parametrize('use_weights', [True, False]) +@pytest.mark.parametrize('bkg_order_val', [0, 1, 2]) +def test_fit_background_optimal(inputs_with_source, use_weights, bkg_order_val): + (image, var_rnoise, var_poisson, var_rflat, + profile, weights, profile_bg) = inputs_with_source + + if not use_weights: + weights = None + + result = extract1d.extract1d( + image, [profile], var_rnoise, var_poisson, var_rflat, + weights=weights, profile_bg=profile_bg, fit_bkg=True, + bkg_fit_type='poly', bkg_order=bkg_order_val, + extraction_type='optimal') + + flux = result[0][0] + background = result[4][0] + + assert np.allclose(flux, 20.0) + assert np.allclose(background, 2.66667) diff --git a/jwst/extract_1d/tests/test_ifu_image_ref.py b/jwst/extract_1d/tests/test_ifu_image_ref.py deleted file mode 100644 index b46392b450..0000000000 --- a/jwst/extract_1d/tests/test_ifu_image_ref.py +++ /dev/null @@ -1,183 +0,0 @@ -""" -Test using a reference image in extract_1d for IFU data -""" -import numpy as np - -from stdatamodels.jwst import datamodels - -from jwst.extract_1d import extract - - -data_shape = (941, 48, 46) -x_center = 23. -y_center = 26. # offset +2 from image center -radius = 11.5 -inner_bkg = 11.5 -outer_bkg = 16.5 -method = "exact" - - -def test_ifu_3d(): - """Test 1""" - - input = make_ifu_cube(data_shape, source=5., background=3.7, - x_center=x_center, y_center=y_center, - radius=radius, - inner_bkg=inner_bkg, outer_bkg=outer_bkg) - - ref_image_2d = make_ref_image(data_shape[-2:], # 2-D ref image - x_center=x_center, y_center=y_center, - radius=radius, - inner_bkg=inner_bkg, outer_bkg=outer_bkg) - - ref_image_3d = make_ref_image(data_shape, # 3-D ref image - x_center=x_center, y_center=y_center, - radius=radius, - inner_bkg=inner_bkg, outer_bkg=outer_bkg) - - ref_dict_2d = {"ref_file_type": extract.FILE_TYPE_IMAGE, - "ref_model": ref_image_2d} - truth = extract.do_extract1d(input, ref_dict_2d, smoothing_length=0, - bkg_order=0, log_increment=50, - subtract_background=True) - - ref_dict_3d = {"ref_file_type": extract.FILE_TYPE_IMAGE, - "ref_model": ref_image_3d} - output = extract.do_extract1d(input, ref_dict_3d, smoothing_length=0, - bkg_order=0, log_increment=50, - subtract_background=True) - - true_wl = truth.spec[0].spec_table['wavelength'] - true_flux = truth.spec[0].spec_table['flux'] - true_bkg = truth.spec[0].spec_table['background'] - - wavelength = output.spec[0].spec_table['wavelength'] - flux = output.spec[0].spec_table['flux'] - background = output.spec[0].spec_table['background'] - - # These should all be the same because the reference image is the - # same in every plane. - - assert np.allclose(wavelength, true_wl, rtol=1.e-14, atol=1.e-14) - - assert np.allclose(flux, true_flux, atol=1.e-14) - - assert np.allclose(background, true_bkg, atol=1.e-14) - - input.close() - truth.close() - output.close() - del ref_dict_2d, ref_dict_3d - ref_image_2d.close() - ref_image_3d.close() - - -def make_ifu_cube(data_shape, source=None, background=None, - x_center=None, y_center=None, - radius=None, inner_bkg=None, outer_bkg=None): - """Create "science" data for testing. - - Returns - ------- - input_model : `~jwst.datamodels.ifucube.IFUCubeModel` - """ - - data = np.zeros(data_shape, dtype=np.float32) - weightmap = np.zeros(data_shape, dtype=np.float32) - r_2 = radius**2 - if inner_bkg is not None and outer_bkg is not None: - create_background = True - bkg = background - inner_2 = inner_bkg**2 - outer_2 = outer_bkg**2 - elif inner_bkg is not None or outer_bkg is not None: - raise RuntimeError("Specify both inner_bkg and outer_bkg or neither.") - else: - create_background = False - bkg = 0. - - for j in range(data_shape[-2]): - for i in range(data_shape[-1]): - dist_2 = (float(i) - x_center)**2 + (float(j) - y_center)**2 - if dist_2 <= r_2: - data[:, j, i] = source + bkg - weightmap[:, j, i] = 1 - if create_background: - if dist_2 > inner_2 and dist_2 <= outer_2: - data[:, j, i] = bkg - weightmap[:, j, i] = 1 - - dq = np.zeros(data_shape, dtype=np.uint32) - input_model = datamodels.IFUCubeModel(data=data, dq=dq, weightmap=weightmap) - # Populate the BUNIT keyword so that in ifu.py the net will be moved - # to the flux column. - input_model.meta.bunit_data = 'MJy/sr' - - def mock_wcs(x, y, z): - """Fake wcs method.""" - - wavelength = np.linspace(0.5975, 5.2975, 941, endpoint=True, - retstep=False, dtype=np.float64) - - if hasattr(z, 'dtype'): - iz = np.around(z).astype(np.int64) - else: - iz = round(z) - - wl = wavelength[iz] - ra = wl.copy() # dummy values - dec = wl.copy() # dummy values - return ra, dec, wl - - input_model.meta.wcs = mock_wcs - input_model.meta.target.source_type = 'POINT' - - return input_model - - -# Functions user_bkg_spec_a, user_bkg_spec_b, and user_bkg_spec_c create -# "user-supplied" background data for the tests above. - -def make_ref_image(shape, - x_center=None, y_center=None, - radius=None, inner_bkg=None, outer_bkg=None): - """Create an image reference file for testing. - - Returns - ------- - ref_image : `~jwst.datamodels.MultiExtract1dImageModel` - """ - - ref_image = datamodels.MultiExtract1dImageModel() - - if len(shape) < 2 or len(shape) > 3: - raise RuntimeError("shape must be either 2-D or 3-D") - - mask = np.zeros(shape, dtype=np.float32) - - r_2 = radius**2 - if inner_bkg is not None and outer_bkg is not None: - create_background = True - inner_2 = inner_bkg**2 - outer_2 = outer_bkg**2 - elif inner_bkg is not None or outer_bkg is not None: - raise RuntimeError("Specify both inner_bkg and outer_bkg or neither.") - else: - create_background = False - - for j in range(shape[-2]): - for i in range(shape[-1]): - dist_2 = (float(i) - x_center)**2 + (float(j) - y_center)**2 - if dist_2 <= r_2: - mask[..., j, i] = 1. # source - if create_background: - if dist_2 > inner_2 and dist_2 <= outer_2: - mask[..., j, i] = -1. # background - - ref_image = datamodels.MultiExtract1dImageModel() - - im = datamodels.Extract1dImageModel(data=mask) - im.name = "ANY" - ref_image.images.append(im) - - return ref_image diff --git a/jwst/pixel_replace/pixel_replace_step.py b/jwst/pixel_replace/pixel_replace_step.py index c7695440b8..5191ac3930 100644 --- a/jwst/pixel_replace/pixel_replace_step.py +++ b/jwst/pixel_replace/pixel_replace_step.py @@ -57,7 +57,7 @@ def process(self, input): datamodels.ImageModel, datamodels.IFUImageModel, datamodels.CubeModel)): - self.log.debug('Input is a {input_model.meta.model_type}.') + self.log.debug(f'Input is a {input_model.meta.model_type}.') elif isinstance(input_model, datamodels.ModelContainer): self.log.debug('Input is a ModelContainer.') else: diff --git a/jwst/regtest/test_miri_lrs_slit_spec2.py b/jwst/regtest/test_miri_lrs_slit_spec2.py index bd9d2a3020..65765fd86a 100644 --- a/jwst/regtest/test_miri_lrs_slit_spec2.py +++ b/jwst/regtest/test_miri_lrs_slit_spec2.py @@ -64,24 +64,6 @@ def test_miri_lrs_extract1d_from_cal(run_pipeline, rtdata_module, fitsdiff_defau assert diff.identical, diff.report() -@pytest.mark.bigdata -def test_miri_lrs_extract1d_image_ref(run_pipeline, rtdata_module, fitsdiff_default_kwargs): - rtdata = rtdata_module - - rtdata.get_data("miri/lrs/jw01530005001_03103_00001_mirimage_image_ref.fits") - rtdata.input = "jw01530005001_03103_00001_mirimage_cal.fits" - Extract1dStep.call(rtdata.input, - override_extract1d="jw01530005001_03103_00001_mirimage_image_ref.fits", - suffix='x1dfromrefimage', - save_results=True) - output = "jw01530005001_03103_00001_mirimage_x1dfromrefimage.fits" - rtdata.output = output - rtdata.get_truth(f"truth/test_miri_lrs_slit_spec2/{output}") - - diff = FITSDiff(rtdata.output, rtdata.truth, **fitsdiff_default_kwargs) - assert diff.identical, diff.report() - - @pytest.mark.bigdata def test_miri_lrs_slit_wcs(run_pipeline, rtdata_module, fitsdiff_default_kwargs): rtdata = rtdata_module diff --git a/jwst/regtest/test_nirspec_bots_extract1d.py b/jwst/regtest/test_nirspec_bots_extract1d.py new file mode 100644 index 0000000000..148b60e41f --- /dev/null +++ b/jwst/regtest/test_nirspec_bots_extract1d.py @@ -0,0 +1,44 @@ +import os +import pytest + +from astropy.io.fits.diff import FITSDiff + +from jwst.stpipe import Step + + +@pytest.fixture(scope="module") +def run_extract(rtdata_module, request): + """Run the extract1d step with custom parameters on BOTS data (S1600A1 slit).""" + + rtdata = rtdata_module + + # Get the custom reference file and input exposure + ref_file = 'jwst_nirspec_extract1d_custom_g395h.json' + rtdata.get_data(f'nirspec/tso/{ref_file}') + rtdata.get_data('nirspec/tso/jw01118005001_04101_00001-first20_nrs1_calints.fits') + + # Run the calwebb_spec2 pipeline; + args = ["extract_1d", rtdata.input, + f"--override_extract1d={ref_file}", + "--suffix=x1dints"] + Step.from_cmdline(args) + + return rtdata + + +@pytest.mark.bigdata +def test_nirspec_bots_custom_extraction(run_extract, fitsdiff_default_kwargs): + """ + Regression test of calwebb_spec2 pipeline performed on NIRSpec + fixed-slit data that uses the NRS_BRIGHTOBJ mode (S1600A1 slit). + """ + rtdata = run_extract + output = 'jw01118005001_04101_00001-first20_nrs1_x1dints.fits' + rtdata.output = output + + # Get the truth files + rtdata.get_truth(os.path.join("truth/test_nirspec_bots_extract1d", output)) + + # Compare the results + diff = FITSDiff(rtdata.output, rtdata.truth, **fitsdiff_default_kwargs) + assert diff.identical, diff.report()