Style refactoring.

romancal/tweakreg/tests/test_tweakreg.py

@@ -685,7 +685,13 @@ def test_tweakreg_raises_error_on_invalid_abs_refcat(tmp_path, base_image):
  add_tweakreg_catalog_attribute(tmp_path, img)
 
  with pytest.raises(Exception) as exec_info:
- trs.TweakRegStep.call([img], save_abs_catalog=True, abs_refcat="my_ref_cat")
+ trs.TweakRegStep.call(
+ [img],
+ save_abs_catalog=True,
+ abs_refcat="my_ref_cat",
+ catalog_path=str(tmp_path),
+ output_dir=str(tmp_path),
+ )
 
  assert type(exec_info.value) == TypeError
 
@@ -1052,8 +1058,7 @@ def test_tweakreg_handles_mixed_exposure_types(tmp_path, base_image):
  """Test that TweakReg can handle mixed exposure types
  (non-WFI_IMAGE data will be marked as SKIPPED only and won't be processed)."""
  img1 = base_image(shift_1=1000, shift_2=1000)
- add_tweakreg_catalog_attribute(tmp_path, img1, catalog_filename="img1")
- img1.meta.exposure.type = "WFI_IMAGE"
+ img1.meta.exposure.type = "WFI_GRISM"
 
  img2 = base_image(shift_1=1000, shift_2=1000)
  add_tweakreg_catalog_attribute(tmp_path, img2, catalog_filename="img2")
@@ -1062,9 +1067,18 @@ def test_tweakreg_handles_mixed_exposure_types(tmp_path, base_image):
  img3 = base_image(shift_1=1000, shift_2=1000)
  img3.meta.exposure.type = "WFI_GRISM"
 
- res = trs.TweakRegStep.call([img1, img2, img3])
+ img4 = base_image(shift_1=1000, shift_2=1000)
+ add_tweakreg_catalog_attribute(tmp_path, img4, catalog_filename="img4")
+ img4.meta.exposure.type = "WFI_IMAGE"
+
+ img5 = base_image(shift_1=1000, shift_2=1000)
+ img5.meta.exposure.type = "WFI_GRISM"
+
+ res = trs.TweakRegStep.call([img1, img2, img3, img4, img5])
 
- assert len(res) == 3
- assert img1.meta.cal_step.tweakreg == "COMPLETE"
+ assert len(res) == 5
+ assert img1.meta.cal_step.tweakreg == "SKIPPED"
  assert img2.meta.cal_step.tweakreg == "COMPLETE"
  assert img3.meta.cal_step.tweakreg == "SKIPPED"
+ assert img4.meta.cal_step.tweakreg == "COMPLETE"
+ assert img5.meta.cal_step.tweakreg == "SKIPPED"

romancal/tweakreg/tweakreg_step.py

@@ -9,13 +9,12 @@
 from astropy.table import Table
 from roman_datamodels import datamodels as rdm
 from stcal.tweakreg import tweakreg
+from stcal.tweakreg.tweakreg import _SINGLE_GROUP_REFCAT_STR, SINGLE_GROUP_REFCAT
 
 # LOCAL
 from ..datamodels import ModelLibrary
 from ..stpipe import RomanStep
 
-SINGLE_GROUP_REFCAT = tweakreg.SINGLE_GROUP_REFCAT
-_SINGLE_GROUP_REFCAT_STR = tweakreg._SINGLE_GROUP_REFCAT_STR
 DEFAULT_ABS_REFCAT = SINGLE_GROUP_REFCAT[0]
 
 __all__ = ["TweakRegStep"]
@@ -204,34 +203,29 @@ def process(self, input):
  catalog_table = Table(image_model.meta.tweakreg_catalog)
  catalog_table.meta["name"] = catalog_name
 
- imcats.append(
- {
- "model_index": i,
- "imcat": tweakreg.construct_wcs_corrector(
- wcs=image_model.meta.wcs,
- refang=image_model.meta.wcsinfo,
- catalog=catalog_table,
- group_id=image_model.meta.group_id,
- ),
- }
+ imcat = tweakreg.construct_wcs_corrector(
+ wcs=image_model.meta.wcs,
+ refang=image_model.meta.wcsinfo,
+ catalog=catalog_table,
+ group_id=image_model.meta.group_id,
  )
+ imcat.meta["model_index"] = i
+ imcats.append(imcat)
  images.shelve(image_model, i)
 
  # run alignment only if it was possible to build image catalogs
  if len(imcats):
  # extract WCS correctors to use for image alignment
- correctors = [x["imcat"] for x in imcats]
  if len(images.group_indices) > 1:
- self.do_relative_alignment(correctors)
+ self.do_relative_alignment(imcats)
 
  if self.abs_refcat in SINGLE_GROUP_REFCAT:
- self.do_absolute_alignment(ref_image, correctors)
+ self.do_absolute_alignment(ref_image, imcats)
 
  # finalize step
  with images:
- for item in imcats:
- imcat = item["imcat"]
- image_model = images.borrow(item["model_index"])
+ for imcat in imcats:
+ image_model = images.borrow(imcat.meta["model_index"])
  image_model.meta.cal_step["tweakreg"] = "COMPLETE"
  # remove source catalog
  del image_model.meta["tweakreg_catalog"]
@@ -276,7 +270,7 @@ def process(self, input):
  del image_model.meta["wcs_fit_results"][k]
 
  image_model.meta.wcs = imcat.wcs
- images.shelve(image_model, item["model_index"])
+ images.shelve(image_model, imcat.meta["model_index"])
 
  return images