Skip to content

Commit

Permalink
First version of PSF extraction, contains hacks
Browse files Browse the repository at this point in the history 
Co-authored-by: jemorrison <[email protected]>
  • Loading branch information
@melanieclarke @jemorrison
melanieclarke and jemorrison committed Nov 19, 2024
1 parent 6ea94a3 commit e31fcb5
Show file tree
Hide file tree
Showing 3 changed files with 352 additions and 21 deletions.
101 changes: 84 additions & 17 deletions jwst/extract_1d/extract.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,17 @@
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 jwst.extract_1d.psf_profile import psf_profile

log = logging.getLogger(__name__)
log.setLevel(logging.DEBUG)

WFSS_EXPTYPES = ['NIS_WFSS', 'NRC_WFSS', 'NRC_GRISM']
"""Exposure types to be regarded as wide-field slitless spectroscopy."""

OPTIMAL_EXPTYPES = ['MIR_LRS-FIXEDSLIT']
"""Exposure types for which optimal extraction is available."""

ANY = "ANY"
"""Wildcard for slit name.
Expand Down Expand Up @@ -139,7 +143,8 @@ def read_apcorr_ref(refname, exptype):

def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
smoothing_length, bkg_fit, bkg_order, use_source_posn,
subtract_background):
subtract_background, extraction_type, specwcs_ref_name,
psf_ref_name):
"""Get extraction parameter values.
Parameters
Expand Down Expand Up @@ -202,6 +207,17 @@ def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
subtract_background : bool or None
If False, all background parameters will be ignored.
extraction_type : str
Extraction type ('box' or 'optimal'). Optimal extraction is
only available if `specwcs_ref_name` and `psf_ref_name` are
not 'N/A'.
specwcs_ref_name : str
The name of the specwcs reference file, or "N/A".
psf_ref_name : str
The name of the PSF reference file, or "N/A".
Returns
-------
extract_params : dict
Expand All @@ -228,6 +244,8 @@ def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
extract_params['bkg_order'] = 0 # because no background sub.
extract_params['subtract_background'] = False
extract_params['extraction_type'] = 'box'
extract_params['specwcs_ref_name'] = 'N/A'
extract_params['psf_ref_name'] = 'N/A'

if use_source_posn is None:
extract_params['use_source_posn'] = False
Expand Down Expand Up @@ -326,9 +344,10 @@ def get_extract_parameters(ref_dict, input_model, slitname, sp_order, meta,
# If the user supplied a value, use that value.
extract_params['smoothing_length'] = smoothing_length

# Default the extraction type to 'box': 'optimal'
# is not yet supported.
extract_params['extraction_type'] = 'box'
# Set the extraction type to 'box' or 'optimal'
extract_params['extraction_type'] = extraction_type
extract_params['specwcs'] = specwcs_ref_name
extract_params['psf'] = psf_ref_name

break

Expand Down Expand Up @@ -1206,10 +1225,27 @@ def define_aperture(input_model, slit, extract_params, exp_type):

# Offet extract parameters by location - nominal
shift_by_source_location(location, nominal_location, extract_params)
else:
middle_pix, middle_wl, location = None, None, None

# 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)
if extract_params['extraction_type'] == 'optimal':
if location is None:
nominal_profile = box_profile(data_shape, extract_params, wl_array,
label='nominal aperture')
middle_pix, location = aperture_center(nominal_profile)
if extract_params['dispaxis'] == HORIZONTAL:
middle_wl = wl_array[middle_pix, location]
else:
middle_wl = wl_array[location, middle_pix]

profile, lower_limit, upper_limit = psf_profile(
data_model, extract_params['psf'], extract_params['specwcs'],
middle_wl, location)

else:
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
Expand Down Expand Up @@ -1383,8 +1419,9 @@ def extract_one_slit(data_model, integ, profile, bg_profile, extract_params):
def create_extraction(input_model, slit, output_model,
extract_ref_dict, slitname, sp_order, smoothing_length,
bkg_fit, bkg_order, use_source_posn, exp_type,
subtract_background, apcorr_ref_model, log_increment,
save_profile):
subtract_background, apcorr_ref_model,
extraction_type, specwcs_ref_name, psf_ref_name,
log_increment, save_profile):
"""Extract spectra from an input model and append to an output model.
Input data, specified in the `slit` or `input_model`, should contain data
Expand Down Expand Up @@ -1463,6 +1500,14 @@ def create_extraction(input_model, slit, output_model,
apcorr_ref_model : DataModel or None
The aperture correction reference datamodel, containing the
APCORR reference file data.
extraction_type : str
Extraction type ('box' or 'optimal'). Optimal extraction is
only available if `specwcs_ref_name` and `psf_ref_name` are
not 'N/A'.
specwcs_ref_name : str
The name of the specwcs reference file, or "N/A".
psf_ref_name : str
The name of the PSF reference file, or "N/A".
log_increment : int
If greater than 0 and the input data are multi-integration, a message
will be written to the log every `log_increment` integrations.
Expand Down Expand Up @@ -1543,7 +1588,7 @@ def create_extraction(input_model, slit, output_model,
extract_params = get_extract_parameters(
extract_ref_dict, data_model, slitname, sp_order, input_model.meta,
smoothing_length, bkg_fit,bkg_order, use_source_posn,
subtract_background
subtract_background, extraction_type, specwcs_ref_name, psf_ref_name
)

if extract_params['match'] == NO_MATCH:
Expand Down Expand Up @@ -1784,9 +1829,10 @@ def create_extraction(input_model, slit, output_model,
return profile_model


def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_length,
bkg_fit, bkg_order, log_increment, subtract_background,
use_source_posn, save_profile):
def run_extract1d(input_model, extract_ref_name, apcorr_ref_name,
specwcs_ref_name, psf_ref_name, extraction_type,
smoothing_length, bkg_fit, bkg_order, log_increment,
subtract_background, use_source_posn, save_profile):
"""Extract all 1-D spectra from an input model.
Parameters
Expand All @@ -1798,7 +1844,18 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
The name of the extract1d reference file, or "N/A".
apcorr_ref_name : str
Name of the APCORR reference file. Default is None
Name of the APCORR reference file. Default is None.
specwcs_ref_name : str
The name of the specwcs reference file, or "N/A".
psf_ref_name : str
The name of the PSF reference file, or "N/A".
extraction_type : str
Extraction type ('box' or 'optimal'). Optimal extraction is
only available if `specwcs_ref_name` and `psf_ref_name` are
not 'N/A'.
smoothing_length : int or None
Width of a boxcar function for smoothing the background regions.
Expand Down Expand Up @@ -1863,6 +1920,14 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
if apcorr_ref_name is not None and apcorr_ref_name != 'N/A':
apcorr_ref_model = read_apcorr_ref(apcorr_ref_name, exp_type)

# Check for non-null specwcs and PSF reference files
if (exp_type not in OPTIMAL_EXPTYPES
or specwcs_ref_name == 'N/A' or psf_ref_name == 'N/A'):
if extraction_type != 'box':
log.warning(f'Optimal extraction is not available for EXP_TYPE {exp_type}')
log.warning('Defaulting to box extraction.')
extraction_type = 'box'

# Set up the output model
output_model = datamodels.MultiSpecModel()
if hasattr(meta_source, "int_times"):
Expand Down Expand Up @@ -1944,8 +2009,9 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
input_model, slit, output_model,
extract_ref_dict, slitname, sp_order, smoothing_length,
bkg_fit, bkg_order, use_source_posn, exp_type,
subtract_background, apcorr_ref_model, log_increment,
save_profile)
subtract_background, apcorr_ref_model,
extraction_type, specwcs_ref_name, psf_ref_name,
log_increment, save_profile)
except ContinueError:
pass

Expand Down Expand Up @@ -1979,8 +2045,9 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
input_model, slit, output_model,
extract_ref_dict, slitname, sp_order, smoothing_length,
bkg_fit, bkg_order, use_source_posn, exp_type,
subtract_background, apcorr_ref_model, log_increment,
save_profile)
subtract_background, apcorr_ref_model,
extraction_type, specwcs_ref_name, psf_ref_name,
log_increment, save_profile)
except ContinueError:
pass

Expand Down
27 changes: 23 additions & 4 deletions jwst/extract_1d/extract_1d_step.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,12 @@ class Extract1dStep(Step):
Attributes
----------
extraction_type : str
If 'box', a standard extraction is performed, summing over an
aperture box. If 'optimal', a PSF-based extraction is performed.
Currently, optimal extraction is only available for MIRI LRS Fixed
Slit data.
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
Expand Down Expand Up @@ -154,6 +160,7 @@ class Extract1dStep(Step):
class_alias = "extract_1d"

spec = """
extraction_type = option("box", "optimal", default="box") # Perform box or optimal extraction
use_source_posn = boolean(default=None) # use source coords to center extractions?
apply_apcorr = boolean(default=True) # apply aperture corrections?
log_increment = integer(default=50) # increment for multi-integration log messages
Expand Down Expand Up @@ -186,7 +193,8 @@ class Extract1dStep(Step):
soss_modelname = output_file(default = None) # Filename for optional model output of traces and pixel weights
"""

reference_file_types = ['extract1d', 'apcorr', 'pastasoss', 'specprofile', 'speckernel']
reference_file_types = ['extract1d', 'apcorr', 'pastasoss', 'specprofile',
'speckernel', 'specwcs'] #, 'psf']

def _get_extract_reference_files_by_mode(self, model, exp_type):
"""Get extraction reference files with special handling by exposure type."""
Expand All @@ -207,7 +215,15 @@ def _get_extract_reference_files_by_mode(self, model, exp_type):
if apcorr_ref != 'N/A':
self.log.info(f'Using APCORR file {apcorr_ref}')

return extract_ref, apcorr_ref
if exp_type in extract.OPTIMAL_EXPTYPES:
specwcs_ref = self.get_reference_file(model, 'specwcs')
#psf_ref = self.get_reference_file(model, 'psf')
psf_ref = "MIRI_LRS_SLIT_EPSF_20240602_WAVE_1D.fits"
else:
specwcs_ref = 'N/A'
psf_ref = 'N/A'

return extract_ref, apcorr_ref, specwcs_ref, psf_ref

def _extract_soss(self, model):
"""Extract NIRISS SOSS spectra."""
Expand Down Expand Up @@ -381,8 +397,8 @@ def process(self, input):
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)
extract_ref, apcorr_ref, specwcs_ref, psf_ref = \
self._get_extract_reference_files_by_mode(model, exp_type)

profile = None
if isinstance(model, datamodels.IFUCubeModel):
Expand All @@ -394,6 +410,9 @@ def process(self, input):
model,
extract_ref,
apcorr_ref,
specwcs_ref,
psf_ref,
self.extraction_type,
self.smoothing_length,
self.bkg_fit,
self.bkg_order,
Expand Down
Loading

0 comments on commit e31fcb5

Please sign in to comment.