Skip to content

Commit

Permalink
Add option to save spatial profile
Browse files Browse the repository at this point in the history
  • Loading branch information
@melanieclarke
melanieclarke committed Nov 19, 2024
1 parent 177945c commit 378b149
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 13 deletions.
60 changes: 49 additions & 11 deletions jwst/extract_1d/extract.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1383,7 +1383,8 @@ 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):
subtract_background, apcorr_ref_model, 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 @@ -1465,7 +1466,16 @@ def create_extraction(input_model, slit, output_model,
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.
save_profile : bool
If True, the spatial profile created for the aperture will be returned
as an ImageModel. If False, the return value is None.
Returns
-------
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.
"""

if slit is None:
Expand Down Expand Up @@ -1559,6 +1569,14 @@ def create_extraction(input_model, slit, output_model,
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:
Expand Down Expand Up @@ -1763,10 +1781,12 @@ def create_extraction(input_model, slit, output_model,
else:
log.info(f"All {input_model.data.shape[0]} integrations done")

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):
use_source_posn, save_profile):
"""Extract all 1-D spectra from an input model.
Parameters
Expand Down Expand Up @@ -1810,11 +1830,19 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
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.
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.
"""
# Set "meta_source" to either the first model in a container,
# or the individual input model, for convenience
Expand Down Expand Up @@ -1847,8 +1875,8 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
prism_mode = is_prism(meta_source)

# Handle inputs that contain one or more slit models
profile_model = None
if isinstance(input_model, (ModelContainer, datamodels.MultiSlitModel)):

if isinstance(input_model, ModelContainer):
slits = input_model
else:
Expand All @@ -1858,6 +1886,10 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
# 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()

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)}')
Expand All @@ -1875,14 +1907,18 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
continue

try:
create_extraction(
profile = create_extraction(
meta_source, 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)
subtract_background, apcorr_ref_model, log_increment,
save_profile)
except ContinueError:
continue

if save_profile:
profile_model.append(profile)

else:
# Define source of metadata
slit = None
Expand All @@ -1904,11 +1940,12 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
else:
log.info(f'Processing spectral order {sp_order}')
try:
create_extraction(
profile_model = 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)
subtract_background, apcorr_ref_model, log_increment,
save_profile)
except ContinueError:
pass

Expand Down Expand Up @@ -1938,11 +1975,12 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
log.info(f'Processing spectral order {sp_order}')

try:
create_extraction(
profile_model = 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)
subtract_background, apcorr_ref_model, log_increment,
save_profile)
except ContinueError:
pass

Expand All @@ -1968,4 +2006,4 @@ def run_extract1d(input_model, extract_ref_name, apcorr_ref_name, smoothing_leng
# x1d product just to hold this keyword.
output_model.meta.target.source_type = None

return output_model
return output_model, profile_model
26 changes: 24 additions & 2 deletions jwst/extract_1d/extract_1d_step.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,10 @@ class Extract1dStep(Step):
INFO every `log_increment` integrations. This is intended to
provide progress information when invoking the step interactively.
save_profile : bool
If True, the spatial profile containing the extraction aperture
is saved to disk. Ignored for IFU and NIRISS SOSS extractions.
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.
Expand Down Expand Up @@ -153,6 +157,7 @@ class Extract1dStep(Step):
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
save_profile = boolean(default=False) # save spatial profile to disk
subtract_background = boolean(default=None) # subtract background?
smoothing_length = integer(default=None) # background smoothing size
Expand Down Expand Up @@ -379,12 +384,13 @@ def process(self, input):
extract_ref, apcorr_ref = self._get_extract_reference_files_by_mode(
model, exp_type)

profile = 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 = extract.run_extract1d(
extracted, profile = extract.run_extract1d(
model,
extract_ref,
apcorr_ref,
Expand All @@ -393,14 +399,30 @@ def process(self, input):
self.bkg_order,
self.log_increment,
self.subtract_background,
self.use_source_posn
self.use_source_posn,
self.save_profile
)

# Set the step flag to complete in each model
extracted.meta.cal_step.extract_1d = 'COMPLETE'
result.append(extracted)
del extracted

# Save profile if needed
if self.save_profile and profile is not None:
if isinstance(profile, ModelContainer):
# Save the profile with the slit name + suffix 'profile'
for model in profile:
profile_path = self.make_output_path(suffix=f'{model.name}_profile')
self.log.info(f"Saving profile {profile_path}")
model.save(profile_path)
else:
# Only one profile - just use the suffix 'profile'
profile_path = self.make_output_path(suffix='profile')
self.log.info(f"Saving profile {profile_path}")
profile.save(profile_path)
profile.close()

# If only one result, return the model instead of the container
if len(result) == 1:
result = result[0]
Expand Down

0 comments on commit 378b149

Please sign in to comment.