Skip to content

Commit b403ba1

Browse files
stscirijtaylorbell57
stscirij
authored and
taylorbell57
committed
JP-3861: Code style changes for extract_2d step (spacetelescope#9305)
1 parent fc0ca43 commit b403ba1

File tree

7 files changed

+305
-227
lines changed

7 files changed

+305
-227
lines changed

‎.pre-commit-config.yaml

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,6 @@ repos:
4343
jwst/cube_build/.* |
4444
jwst/cube_skymatch/.* |
4545
jwst/dark_current/.* |
46-
jwst/extract_2d/.* |
4746
jwst/flatfield/.* |
4847
jwst/fringe/.* |
4948
jwst/ipc/.* |

‎.ruff.toml

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@ exclude = [
2525
"jwst/cube_build/**.py",
2626
"jwst/cube_skymatch/**.py",
2727
"jwst/dark_current/**.py",
28-
"jwst/extract_2d/**.py",
2928
"jwst/flatfield/**.py",
3029
"jwst/fringe/**.py",
3130
"jwst/ipc/**.py",
@@ -115,7 +114,6 @@ ignore-fully-untyped = true # Turn off annotation checking for fully untyped co
115114
"jwst/cube_build/**.py" = ["D", "N", "A", "ARG", "B", "C4", "ICN", "INP", "ISC", "LOG", "NPY", "PGH", "PTH", "S", "SLF", "SLOT", "T20", "TRY", "UP", "YTT", "E501"]
116115
"jwst/cube_skymatch/**.py" = ["D", "N", "A", "ARG", "B", "C4", "ICN", "INP", "ISC", "LOG", "NPY", "PGH", "PTH", "S", "SLF", "SLOT", "T20", "TRY", "UP", "YTT", "E501"]
117116
"jwst/dark_current/**.py" = ["D", "N", "A", "ARG", "B", "C4", "ICN", "INP", "ISC", "LOG", "NPY", "PGH", "PTH", "S", "SLF", "SLOT", "T20", "TRY", "UP", "YTT", "E501"]
118-
"jwst/extract_2d/**.py" = ["D", "N", "A", "ARG", "B", "C4", "ICN", "INP", "ISC", "LOG", "NPY", "PGH", "PTH", "S", "SLF", "SLOT", "T20", "TRY", "UP", "YTT", "E501"]
119117
"jwst/flatfield/**.py" = ["D", "N", "A", "ARG", "B", "C4", "ICN", "INP", "ISC", "LOG", "NPY", "PGH", "PTH", "S", "SLF", "SLOT", "T20", "TRY", "UP", "YTT", "E501"]
120118
"jwst/fringe/**.py" = ["D", "N", "A", "ARG", "B", "C4", "ICN", "INP", "ISC", "LOG", "NPY", "PGH", "PTH", "S", "SLF", "SLOT", "T20", "TRY", "UP", "YTT", "E501"]
121119
"jwst/ipc/**.py" = ["D", "N", "A", "ARG", "B", "C4", "ICN", "INP", "ISC", "LOG", "NPY", "PGH", "PTH", "S", "SLF", "SLOT", "T20", "TRY", "UP", "YTT", "E501"]

‎jwst/extract_2d/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
"""Perform extract_2d calibration step."""
2+
13
from .extract_2d_step import Extract2dStep
24

35
__all__ = ["Extract2dStep"]

‎jwst/extract_2d/extract_2d.py

Lines changed: 48 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -11,18 +11,20 @@
1111
log.setLevel(logging.DEBUG)
1212

1313

14-
def extract2d(input_model,
15-
slit_names=None,
16-
source_ids=None,
17-
reference_files={},
18-
grism_objects=None,
19-
tsgrism_extract_height=None,
20-
wfss_extract_half_height=None,
21-
extract_orders=None,
22-
mmag_extract=None,
23-
nbright=None):
14+
def extract2d(
15+
input_model,
16+
slit_names=None,
17+
source_ids=None,
18+
reference_files=None,
19+
grism_objects=None,
20+
tsgrism_extract_height=None,
21+
wfss_extract_half_height=None,
22+
extract_orders=None,
23+
mmag_extract=None,
24+
nbright=None,
25+
):
2426
"""
25-
The main extract_2d function
27+
Extract rectangular cutouts around each spectrum from a spectral dataset.
2628
2729
Parameters
2830
----------
@@ -36,7 +38,7 @@ def extract2d(input_model,
3638
Reference files.
3739
grism_objects : list
3840
A list of grism objects.
39-
tsgrism_extract_height: int
41+
tsgrism_extract_height : int
4042
Cross-dispersion extraction height to use for time series grisms.
4143
This will override the default which for NRC_TSGRISM is a set
4244
size of 64 pixels.
@@ -54,45 +56,57 @@ def extract2d(input_model,
5456
-------
5557
output_model : `~jwst.datamodels.ImageModel` or `~jwst.datamodelsCubeModel`
5658
A copy of the input_model that has been processed.
57-
5859
"""
59-
nrs_modes = ['NRS_FIXEDSLIT', 'NRS_MSASPEC', 'NRS_BRIGHTOBJ',
60-
'NRS_LAMP', 'NRS_AUTOFLAT', 'NRS_AUTOWAVE']
61-
slitless_modes = ['NIS_WFSS', 'NRC_WFSS', 'NRC_TSGRISM']
60+
nrs_modes = [
61+
"NRS_FIXEDSLIT",
62+
"NRS_MSASPEC",
63+
"NRS_BRIGHTOBJ",
64+
"NRS_LAMP",
65+
"NRS_AUTOFLAT",
66+
"NRS_AUTOWAVE",
67+
]
68+
slitless_modes = ["NIS_WFSS", "NRC_WFSS", "NRC_TSGRISM"]
6269

6370
exp_type = input_model.meta.exposure.type.upper()
64-
log.info(f'EXP_TYPE is {exp_type}')
71+
log.info(f"EXP_TYPE is {exp_type}")
72+
73+
if reference_files is None:
74+
reference_files = {}
6575

6676
if exp_type in nrs_modes:
6777
if input_model.meta.instrument.grating.lower() == "mirror":
6878
# Catch the case of EXP_TYPE=NRS_LAMP and grating=MIRROR
69-
log.info(f'EXP_TYPE {exp_type} with grating=MIRROR not supported for extract 2D')
70-
input_model.meta.cal_step.extract_2d = 'SKIPPED'
79+
log.info(f"EXP_TYPE {exp_type} with grating=MIRROR not supported for extract 2D")
80+
input_model.meta.cal_step.extract_2d = "SKIPPED"
7181
return input_model
7282
output_model = nrs_extract2d(input_model, slit_names=slit_names, source_ids=source_ids)
7383
elif exp_type in slitless_modes:
74-
if exp_type == 'NRC_TSGRISM':
84+
if exp_type == "NRC_TSGRISM":
7585
if tsgrism_extract_height is None:
7686
tsgrism_extract_height = 64
77-
output_model = extract_tso_object(input_model,
78-
reference_files=reference_files,
79-
tsgrism_extract_height=tsgrism_extract_height,
80-
extract_orders=extract_orders)
87+
output_model = extract_tso_object(
88+
input_model,
89+
reference_files=reference_files,
90+
tsgrism_extract_height=tsgrism_extract_height,
91+
extract_orders=extract_orders,
92+
)
8193
else:
82-
output_model = extract_grism_objects(input_model,
83-
grism_objects=grism_objects,
84-
reference_files=reference_files,
85-
extract_orders=extract_orders,
86-
mmag_extract=mmag_extract,
87-
wfss_extract_half_height=wfss_extract_half_height,
88-
nbright=nbright)
94+
output_model = extract_grism_objects(
95+
input_model,
96+
grism_objects=grism_objects,
97+
reference_files=reference_files,
98+
extract_orders=extract_orders,
99+
mmag_extract=mmag_extract,
100+
wfss_extract_half_height=wfss_extract_half_height,
101+
nbright=nbright,
102+
)
89103

90104
else:
91-
log.info(f'EXP_TYPE {exp_type} not supported for extract 2D')
92-
input_model.meta.cal_step.extract_2d = 'SKIPPED'
105+
log.info(f"EXP_TYPE {exp_type} not supported for extract 2D")
106+
input_model.meta.cal_step.extract_2d = "SKIPPED"
93107
return input_model
94108

95109
# Set the step status to COMPLETE
96-
output_model.meta.cal_step.extract_2d = 'COMPLETE'
110+
output_model.meta.cal_step.extract_2d = "COMPLETE"
97111
del input_model
98112
return output_model

‎jwst/extract_2d/extract_2d_step.py

Lines changed: 29 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,7 @@
99

1010

1111
class Extract2dStep(Step):
12-
"""
13-
This Step performs a 2D extraction of spectra.
14-
"""
12+
"""Class that provides method to perform a 2D extraction of spectra."""
1513

1614
class_alias = "extract_2d"
1715

@@ -24,24 +22,40 @@ class Extract2dStep(Step):
2422
wfss_extract_half_height = integer(default=5) # extraction half height in pixels, WFSS mode
2523
wfss_mmag_extract = float(default=None) # minimum abmag to extract, WFSS mode
2624
wfss_nbright = integer(default=1000) # number of brightest objects to extract, WFSS mode
27-
""" # noqa: E501
25+
""" # noqa: E501
2826

29-
reference_file_types = ['wavelengthrange']
27+
reference_file_types = ["wavelengthrange"]
3028

31-
def process(self, input_model, *args, **kwargs):
29+
def process(self, input_model):
30+
"""
31+
Perform the extract_2d calibration step.
32+
33+
Parameters
34+
----------
35+
input_model : DataModel
36+
DataModel to be processed
37+
38+
Returns
39+
-------
40+
output_model : DataModel
41+
The resulting DataModel of the extract_2d step
42+
"""
3243
reference_file_names = {}
3344
for reftype in self.reference_file_types:
3445
reffile = self.get_reference_file(input_model, reftype)
3546
reference_file_names[reftype] = reffile if reffile else ""
3647
with datamodels.open(input_model) as dm:
37-
output_model = extract_2d.extract2d(dm, self.slit_names,
38-
self.source_ids,
39-
reference_files=reference_file_names,
40-
grism_objects=self.grism_objects,
41-
tsgrism_extract_height=self.tsgrism_extract_height,
42-
wfss_extract_half_height=self.wfss_extract_half_height,
43-
extract_orders=self.extract_orders,
44-
mmag_extract=self.wfss_mmag_extract,
45-
nbright=self.wfss_nbright)
48+
output_model = extract_2d.extract2d(
49+
dm,
50+
self.slit_names,
51+
self.source_ids,
52+
reference_files=reference_file_names,
53+
grism_objects=self.grism_objects,
54+
tsgrism_extract_height=self.tsgrism_extract_height,
55+
wfss_extract_half_height=self.wfss_extract_half_height,
56+
extract_orders=self.extract_orders,
57+
mmag_extract=self.wfss_mmag_extract,
58+
nbright=self.wfss_nbright,
59+
)
4660

4761
return output_model

0 commit comments

Comments
 (0)