From cb99fba1cda58b90899189c2fad9aca901cdbf9e Mon Sep 17 00:00:00 2001 From: "M. Teodoro" Date: Wed, 18 Sep 2024 10:34:20 -0400 Subject: [PATCH] Style refactoring. --- romancal/tweakreg/tests/test_tweakreg.py | 26 ++++++++++++++----- romancal/tweakreg/tweakreg_step.py | 32 ++++++++++-------------- 2 files changed, 33 insertions(+), 25 deletions(-) diff --git a/romancal/tweakreg/tests/test_tweakreg.py b/romancal/tweakreg/tests/test_tweakreg.py index 025056d28..ac21e447f 100644 --- a/romancal/tweakreg/tests/test_tweakreg.py +++ b/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" diff --git a/romancal/tweakreg/tweakreg_step.py b/romancal/tweakreg/tweakreg_step.py index c80abd9be..d85267ecd 100644 --- a/romancal/tweakreg/tweakreg_step.py +++ b/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