from __future__ import division
import matplotlib
# =============================================================================
# IMPORTS
# =============================================================================
import os
import astropy.io.fits as pyfits
from tqdm import trange
from jwst import datamodels
from jwst.associations.load_as_asn import LoadAsLevel2Asn
from jwst.outlier_detection.outlier_detection_step import OutlierDetectionStep
from jwst.pipeline import Image2Pipeline
import logging
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)
# =============================================================================
# MAIN
# =============================================================================
[docs]
class Coron2Pipeline_spaceKLIP(Image2Pipeline):
"""
The spaceKLIP JWST stage 2 pipeline class.
"""
class_alias = "calwebb_coron2"
spec = """
save_intermediates = boolean(default=False) # Save all intermediate step results
"""
def __init__(self,
**kwargs):
"""
Initialize the spaceKLIP JWST stage 2 pipeline class.
Parameters
----------
**kwargs : keyword arguments
Default JWST stage 2 image pipeline keyword arguments.
Returns
-------
None.
"""
# Add outlier detection.
self.step_defs['outlier_detection'] = OutlierDetectionStep
# Initialize Image2Pipeline class.
super(Coron2Pipeline_spaceKLIP, self).__init__(**kwargs)
# Set additional step parameters.
self.outlier_detection.skip = False
[docs]
def process(self,
input):
"""
Process an input JWST datamodel with the spaceKLIP JWST stage 2
pipeline.
Parameters
----------
input : jwst.datamodel
Input JWST datamodel to be processed.
Returns
-------
all_res : list of jwst.datamodel
List of output JWST datamodels.
"""
# Open input as asn model.
asn = LoadAsLevel2Asn.load(input, basename=self.output_file)
# Loop through all products.
all_res = []
for product in asn['products']:
if self.save_results:
self.output_file = product['name']
try:
getattr(asn, 'filename')
except AttributeError:
asn.filename = 'singleton'
# Process exposure.
filebase = os.path.basename(asn.filename)
res = self.process_exposure_product(product, asn['asn_pool'], filebase)
# Run outlier detection on CubeModel
try:
if isinstance(res, datamodels.CubeModel):
res = self.outlier_detection.run(res)
except Exception as e:
log.warning('Error in outlier_detection step. Skipping outlier detection.')
log.error(e)
pass
# Save results.
suffix = 'calints' if isinstance(res, datamodels.CubeModel) else 'cal'
res.meta.filename = self.make_output_path(suffix=suffix)
all_res.append(res)
# If outlier detection was run but intermediates were not request
# to be saved, remove the intermediate files.
if not self.save_intermediates and not self.outlier_detection.skip:
file_median = res.meta.filename.replace('calints', 'median')
file_blot = res.meta.filename.replace('calints', 'blot')
file_outlier = res.meta.filename.replace('calints.fits', 'outlier_i2d.fits')
foutlier_local = os.path.basename(file_outlier)
for f in [file_median, file_blot, file_outlier, foutlier_local]:
if os.path.exists(f):
os.remove(f)
# Setup output file.
self.output_use_model = True
self.suffix = False
return all_res
[docs]
def run_single_file(fitspath, output_dir, steps={}, verbose=False, **kwargs):
""" Run the JWST stage 2 image pipeline on a single file.
This customized implementation will also run the 'outlier_detection' step
if not skipped.
Parameters
----------
database : spaceKLIP.Database
SpaceKLIP database on which the JWST stage 2 image pipeline shall be
run.
steps : dict, optional
See here for how to use the steps parameter:
https://jwst-pipeline.readthedocs.io/en/latest/jwst/user_documentation/running_pipeline_python.html#configuring-a-pipeline-step-in-python
Custom step parameters are:
- n/a
The default is {}.
subdir : str, optional
Name of the directory where the data products shall be saved. The
default is 'stage2'.
do_rates : bool, optional
In addition to processing rateints files, also process rate files
if they exist? The default is False.
overwrite : bool, optional
Overwrite existing files? Default is False.
quiet : bool, optional
Use progress bar to track progress instead of messages.
Overrides verbose and sets it to False. Default is False.
verbose : bool, optional
Print all info messages? Default is False.
Keyword Args
------------
save_results : bool, optional
Save the JWST pipeline products? The default is True.
skip_bg : bool, optional
Skip the background subtraction step? The default is False.
skip_photom : bool, optional
Skip the photometric correction step? The default is False.
skip_resample : bool, optional
Skip the resampling (drizzle) step? While the default is set
to False, this step only applies to normal imaging modes and
only on rate/cal files (not rateints/calints).
For coronagraphic observations, resampling occurs in Stage 3.
skip_wcs : bool, optional
Skip the WCS assignment step? The default is False.
skip_flat : bool, optional
Skip the flat field correction step? The default is False.
skip_outlier : bool, optional
Skip the outlier detection step? The default is False except
for target acquisition subarray data, which will always be True.
Returns
-------
None.
"""
# Print all info message if verbose, otherwise only errors or critical.
from .logging_tools import all_logging_disabled
log_level = logging.INFO if verbose else logging.ERROR
# Create output directory if it doesn't exist.
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# Initialize Coron1Pipeline.
with all_logging_disabled(log_level):
pipeline = Coron2Pipeline_spaceKLIP(output_dir=output_dir)
# Options for saving results
pipeline.save_results = kwargs.get('save_results', True)
pipeline.save_intermediates = kwargs.get('save_intermediates', False)
# Skip certain steps?
pipeline.bkg_subtract.skip = kwargs.get('skip_bg', True)
pipeline.photom.skip = kwargs.get('skip_photom', False)
pipeline.resample.skip = kwargs.get('skip_resample', False)
pipeline.assign_wcs.skip = kwargs.get('skip_wcs', False)
pipeline.flat_field.skip = kwargs.get('skip_flat', False)
# Don't perform outlier step for TA data.
hdr0 = pyfits.getheader(fitspath)
if 'NRC_TA' in hdr0['EXP_TYPE']:
pipeline.outlier_detection.skip = True
else:
pipeline.outlier_detection.skip = kwargs.get('skip_outlier', False)
# Set step parameters.
for key1 in steps.keys():
for key2 in steps[key1].keys():
setattr(getattr(pipeline, key1), key2, steps[key1][key2])
# Run Coron2Pipeline. Raise exception on error.
# Ensure that pipeline is closed out.
try:
with all_logging_disabled(log_level):
res = pipeline.run(fitspath)
except Exception as e:
raise RuntimeError(
'Caught exception during pipeline processing.'
'\nException: {}'.format(e)
)
finally:
try:
pipeline.closeout()
except AttributeError:
# Method deprecated as of stpipe version 0.6.0
pass
if isinstance(res, list):
res = res[0]
return res
[docs]
def run_obs(database,
steps={},
subdir='stage2',
do_rates=False,
overwrite=True,
quiet=False,
verbose=False,
**kwargs):
"""
Run the JWST stage 2 image pipeline on the input observations database.
This customized implementation will also run the 'outlier_detection' step
if not skipped.
Parameters
----------
database : spaceKLIP.Database
SpaceKLIP database on which the JWST stage 2 image pipeline shall be
run.
steps : dict, optional
See here for how to use the steps parameter:
https://jwst-pipeline.readthedocs.io/en/latest/jwst/user_documentation/running_pipeline_python.html#configuring-a-pipeline-step-in-python
Custom step parameters are:
- n/a
The default is {}.
subdir : str, optional
Name of the directory where the data products shall be saved. The
default is 'stage2'.
do_rates : bool, optional
In addition to processing rateints files, also process rate files
if they exist? The default is False.
overwrite : bool, optional
Overwrite existing files? Default is False.
quiet : bool, optional
Use progress bar to track progress instead of messages.
Overrides verbose and sets it to False. Default is False.
verbose : bool, optional
Print all info messages? Default is False.
Keyword Args
------------
save_results : bool, optional
Save the JWST pipeline products? The default is True.
skip_bg : bool, optional
Skip the background subtraction step? The default is False.
skip_photom : bool, optional
Skip the photometric correction step? The default is False.
skip_resample : bool, optional
Skip the resampling (drizzle) step? While the default is set
to False, this step only applies to normal imaging modes and
only on rate/cal files (not rateints/calints).
For coronagraphic observations, resampling occurs in Stage 3.
skip_wcs : bool, optional
Skip the WCS assignment step? The default is False.
skip_flat : bool, optional
Skip the flat field correction step? The default is False.
skip_outlier : bool, optional
Skip the outlier detection step? The default is False except
for target acquisition subarray data, which will always be True.
Returns
-------
None.
"""
# Set output directory.
output_dir = os.path.join(database.output_dir, subdir)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# Get list of concatenation keys.
keys = list(database.obs.keys())
nkeys = len(keys)
if quiet:
verbose = False
itervals = trange(nkeys, desc='Concatenations')
else:
itervals = range(nkeys)
# Loop through concatenations.
for i in itervals:
key = keys[i]
if not quiet: log.info('--> Concatenation ' + key)
# Loop through FITS files.
nfitsfiles = len(database.obs[key])
jtervals = trange(nfitsfiles, desc='FITS files', leave=False) if quiet else range(nfitsfiles)
for j in jtervals:
# Skip non-stage 1 files.
head, tail = os.path.split(database.obs[key]['FITSFILE'][j])
fitspath = os.path.abspath(database.obs[key]['FITSFILE'][j])
if database.obs[key]['DATAMODL'][j] != 'STAGE1':
if not quiet: log.info(' --> Coron2Pipeline: skipping non-stage 1 file ' + tail)
else:
# Get expected output file name
outfile_name = tail.replace('rateints.fits', 'calints.fits')
fitsout_path = os.path.join(output_dir, outfile_name)
# Skip if file already exists and overwrite is False.
if os.path.isfile(fitsout_path) and not overwrite:
if not quiet: log.info(' --> Coron2Pipeline: skipping already processed file ' + tail)
else:
if not quiet: log.info(' --> Coron2Pipeline: processing ' + tail)
res = run_single_file(fitspath, output_dir, steps=steps,
verbose=verbose, **kwargs)
# Update spaceKLIP database.
database.update_obs(key, j, fitsout_path, update_pxar=True)
# Also process rate files?
if do_rates:
fitspath = fitspath.replace('rateints', 'rate')
fitsout_path = fitsout_path.replace('calints', 'cal')
if os.path.isfile(fitsout_path) and not overwrite:
if not quiet: log.info(' --> Coron2Pipeline: skipping already processed file ' + tail)
else:
if not quiet: log.info(' --> Coron2Pipeline: processing rate.fits file')
res = run_single_file(fitspath, output_dir, steps=steps,
verbose=verbose, **kwargs)