From 2acf397fdad34787ded919fccab5b2e3b42cf731 Mon Sep 17 00:00:00 2001 From: "M. Teodoro" Date: Mon, 9 Dec 2024 15:29:32 -0500 Subject: [PATCH] Add var_sky array to mosaic and enable usage in drizzle weight image. --- romancal/regtest/test_resample.py | 6 +++-- romancal/resample/resample.py | 11 ++++----- romancal/resample/resample_utils.py | 31 ++++++++++++++++++++++++ romancal/resample/tests/test_resample.py | 2 +- 4 files changed, 41 insertions(+), 9 deletions(-) diff --git a/romancal/regtest/test_resample.py b/romancal/regtest/test_resample.py index d669eecc6..fad1deda8 100644 --- a/romancal/regtest/test_resample.py +++ b/romancal/regtest/test_resample.py @@ -60,6 +60,7 @@ def test_resample_single_file(rtdata, ignore_asdf_paths): "var_poisson", "var_rnoise", "var_flat", + "var_sky", ] ) }""" @@ -76,6 +77,7 @@ def test_resample_single_file(rtdata, ignore_asdf_paths): np.sum(~np.isnan(getattr(resample_out, x))) for x in [ "var_poisson", "var_rnoise", + "var_sky", ] ) }""" # noqa: E501 @@ -94,14 +96,14 @@ def test_resample_single_file(rtdata, ignore_asdf_paths): np.isnan(getattr(resample_out, x)), np.equal(getattr(resample_out, x), 0) ) - ) > 0 for x in ["var_poisson", "var_rnoise", "var_flat"] + ) > 0 for x in ["var_poisson", "var_rnoise", "var_flat", "var_sky"] ) }""" ) assert all( np.sum(np.isnan(getattr(resample_out, x))) - for x in ["var_poisson", "var_rnoise", "var_flat"] + for x in ["var_poisson", "var_rnoise", "var_flat", "var_sky"] ) step.log.info( diff --git a/romancal/resample/resample.py b/romancal/resample/resample.py index 9e05d0055..7944b1f9c 100644 --- a/romancal/resample/resample.py +++ b/romancal/resample/resample.py @@ -88,6 +88,8 @@ def __init__( ) self.input_models = input_models + # add sky variance array + resample_utils.add_var_sky_array(self.input_models) self.output_filename = output self.pscale_ratio = pscale_ratio self.single = single @@ -158,9 +160,6 @@ def __init__( # f'Model cannot be instantiated.' # ) - # NOTE: wait for William to fix bug in datamodels' init and then - # use datamodels.ImageModel(shape=(nx, ny)) instead of mk_datamodel() - # n_images sets the number of context image planes. # This should be 1 to start (not the default of 2). self.blank_output = maker_utils.mk_datamodel( @@ -198,6 +197,8 @@ def __init__( self.blank_output["individual_image_cal_logs"] = [ model.meta.cal_logs for model in models ] + # add sky variance array + self.blank_output["var_sky"] = np.zeros_like(self.blank_output.var_flat) for i, m in enumerate(models): self.input_models.shelve(m, i, modify=False) @@ -397,11 +398,11 @@ def resample_many_to_one(self): self.resample_variance_array("var_rnoise", output_model) self.resample_variance_array("var_poisson", output_model) self.resample_variance_array("var_flat", output_model) + self.resample_variance_array("var_sky", output_model) # Make exposure time image exptime_tot = self.resample_exposure_time(output_model) - # TODO: fix unit here output_model.err = np.sqrt( np.nansum( [ @@ -415,7 +416,6 @@ def resample_many_to_one(self): self.update_exposure_times(output_model, exptime_tot) - # TODO: fix RAD to expect a context image datatype of int32 output_model.context = output_model.context.astype(np.uint32) return ModelLibrary([output_model]) @@ -494,7 +494,6 @@ def resample_variance_array(self, name, output_model): # We now have a sum of the inverse resampled variances. We need the # inverse of that to get back to units of variance. - # TODO: fix unit here output_variance = np.reciprocal(inverse_variance_sum) setattr(output_model, name, output_variance) diff --git a/romancal/resample/resample_utils.py b/romancal/resample/resample_utils.py index 0b399e68b..858ae6061 100644 --- a/romancal/resample/resample_utils.py +++ b/romancal/resample/resample_utils.py @@ -11,6 +11,7 @@ from stcal.alignment.util import wcs_from_footprints from romancal.assign_wcs.utils import wcs_bbox_from_shape +from romancal.datamodels.library import ModelLibrary log = logging.getLogger(__name__) log.setLevel(logging.DEBUG) @@ -163,6 +164,20 @@ def build_driz_weight( elif weight_type == "exptime": exptime = model.meta.exposure.exposure_time result = exptime * dqmask + elif weight_type == "ivsky": + if ( + hasattr(model, "var_sky") + and model.var_sky is not None + and model.var_sky.shape == model.data.shape + ): + with np.errstate(divide="ignore", invalid="ignore"): + inv_sky_variance = model.var_sky**-1 + inv_sky_variance[~np.isfinite(inv_sky_variance)] = 1 + else: + warnings.warn( + "var_rnoise and var_poisson arrays are not available. Setting drizzle weight map to 1" + ) + result = inv_sky_variance * dqmask elif weight_type is None: result = np.ones(model.data.shape, dtype=model.data.dtype) * dqmask else: @@ -401,3 +416,19 @@ def resample_range(data_shape, bbox=None): ymax = min(data_shape[0] - 1, int(y2 + 0.5)) return xmin, xmax, ymin, ymax + + +def add_var_sky_array(input_models: ModelLibrary): + input_models = ( + input_models + if isinstance(input_models, ModelLibrary) + else ModelLibrary([input_models]) + ) + with input_models: + ref_img = input_models.borrow(index=0) + var_sky = np.zeros_like(ref_img.data) + input_models.shelve(model=ref_img, index=0) + for i, img in enumerate(input_models): + var_sky += img.var_rnoise + img.var_poisson / img.data * np.median(img.data) + img["var_sky"] = var_sky + input_models.shelve(img, i, modify=True) diff --git a/romancal/resample/tests/test_resample.py b/romancal/resample/tests/test_resample.py index 0606162fd..4119d79a2 100644 --- a/romancal/resample/tests/test_resample.py +++ b/romancal/resample/tests/test_resample.py @@ -679,7 +679,7 @@ def get_footprint(model, index): np.testing.assert_allclose(output_max_value, expected_max_value) -@pytest.mark.parametrize("weight_type", ["ivm", "exptime"]) +@pytest.mark.parametrize("weight_type", ["ivm", "exptime", "ivsky"]) def test_resampledata_do_drizzle_default_single_exposure_weight_array( exposure_1, weight_type,