Add var_sky array to mosaic and enable usage in drizzle weight image.

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(

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)

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)

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,