Array Quantities

CHANGES.rst

@@ -1,6 +1,8 @@
 0.8.2 (unreleased)
 ==================
 
+- Added support for Quantities for data arrays. [#602]
+
 - New Roman's RTD page layout [#596]
 
 - pin ``numpy`` to ``>=1.20`` [#592]

romancal/dark_current/dark_current_step.py

@@ -5,6 +5,8 @@
 from roman_datamodels import datamodels as rdd
 from stcal.dark_current import dark_sub
 from roman_datamodels.testing import utils as testutil
+from astropy import units as u
+from roman_datamodels import units as ru
 
 
 __all__ = ["DarkCurrentStep"]
@@ -43,7 +45,7 @@ def process(self, input):
                 dark_model.meta.exposure['groupgap'] = input_model.meta.exposure.groupgap
 
             # Reshaping data variables for stcal compatibility
-            input_model.data = input_model.data.astype(np.float32)[np.newaxis, :]
+            input_model.data = input_model.data[np.newaxis, :]
             input_model.groupdq = input_model.groupdq[np.newaxis, :]
             input_model.err = input_model.err[np.newaxis, :]
 
@@ -133,10 +135,10 @@ def dark_output_data_as_ramp_model(out_data, input_model):
     # Removing integration dimension from variables (added for stcal
     # compatibility)
     # Roman 3D
-    out_model.data = out_data.data[0]
+    out_model.data = u.Quantity(out_data.data[0], ru.DN, dtype=out_data.data.dtype)
     out_model.groupdq = out_data.groupdq[0]
     # Roman 2D
     out_model.pixeldq = out_data.pixeldq
-    out_model.err = out_data.err[0]
+    out_model.err = u.Quantity(out_data.err[0], ru.DN, dtype=out_data.err.dtype)
 
     return out_model

romancal/dark_current/tests/test_dark.py

@@ -5,12 +5,15 @@
 import pytest
 import numpy as np
 import os
+from astropy import units as u
 
 from romancal.dark_current import DarkCurrentStep
 
 from roman_datamodels.datamodels import RampModel, DarkRefModel
 import roman_datamodels as rdm
 from roman_datamodels.testing import utils as testutil
+from roman_datamodels import units as ru
+
 
 
 @pytest.mark.parametrize(
@@ -71,17 +74,17 @@ def test_dark_step_subtraction(instrument, exptype):
 
     # populate data array of science cube
     for i in range(0, 20):
-        ramp_model.data[0, 0, i] = i
+        ramp_model.data.value[0, 0, i] = i
         darkref_model.data[0, 0, i] = i * 0.1
 
     # Perform Dark Current subtraction step
     result = DarkCurrentStep.call(ramp_model, override_dark=darkref_model)
 
     # check that the dark file is subtracted frame by frame from the science data
-    diff = ramp_model.data - darkref_model.data
+    diff = ramp_model.data.value - darkref_model.data
 
     # test that the output data file is equal to the difference found when subtracting reffile from sci file
-    np.testing.assert_array_equal(result.data, diff, err_msg='dark file should be subtracted from sci file ')
+    np.testing.assert_array_equal(result.data.value, diff, err_msg='dark file should be subtracted from sci file ')
 
 
 @pytest.mark.parametrize(
@@ -127,7 +130,7 @@ def create_ramp_and_dark(shape, instrument, exptype):
     ramp.meta.instrument.detector = 'WFI01'
     ramp.meta.instrument.optical_element = 'F158'
     ramp.meta.exposure.type = exptype
-    ramp.data = np.ones(shape, dtype=np.float32)
+    ramp.data = u.Quantity(np.ones(shape, dtype=np.float32) , ru.DN, dtype=np.float32)
     ramp_model = RampModel(ramp)
 
     # Create dark model

romancal/dq_init/tests/test_dq_init.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pytest
 import warnings
+from astropy import units as u
 
 from romancal.lib import dqflags
 from romancal.dq_init import DQInitStep
@@ -11,6 +12,7 @@
 from roman_datamodels import stnode
 from roman_datamodels.datamodels import MaskRefModel, ScienceRawModel
 from roman_datamodels.testing import utils as testutil
+from roman_datamodels import units as ru
 
 # Set parameters for multiple runs of data
 args = "xstart, ystart, xsize, ysize, ngroups, instrument, exp_type"
@@ -192,7 +194,7 @@ def test_dqinit_step_interface(instrument, exptype):
     wfi_sci_raw.meta.instrument.detector = 'WFI01'
     wfi_sci_raw.meta.instrument.optical_element = 'F158'
     wfi_sci_raw.meta.exposure.type = exptype
-    wfi_sci_raw.data = np.ones(shape, dtype=np.uint16)
+    wfi_sci_raw.data = u.Quantity(np.ones(shape, dtype=np.uint16), ru.DN, dtype=np.uint16)
     wfi_sci_raw_model = ScienceRawModel(wfi_sci_raw)
 
     # Create mask model
@@ -242,7 +244,7 @@ def test_dqinit_refpix(instrument, exptype):
     wfi_sci_raw.meta.instrument.detector = 'WFI01'
     wfi_sci_raw.meta.instrument.optical_element = 'F158'
     wfi_sci_raw.meta.exposure.type = exptype
-    wfi_sci_raw.data = np.ones(shape, dtype=np.uint16)
+    wfi_sci_raw.data = u.Quantity(np.ones(shape, dtype=np.uint16), ru.DN, dtype=np.uint16)
     wfi_sci_raw_model = ScienceRawModel(wfi_sci_raw)
 
     # Create mask model
@@ -264,7 +266,7 @@ def test_dqinit_refpix(instrument, exptype):
 
     # check if reference pixels are correct
     assert result.data.shape == (2, 20, 20)  # no pixels should be trimmed
-    assert result.amp33.shape == (2, 4096 ,128)
+    assert result.amp33.value.shape == (2, 4096 ,128)
     assert result.border_ref_pix_right.shape == (2, 20, 4)
     assert result.border_ref_pix_left.shape == (2, 20, 4)
     assert result.border_ref_pix_top.shape == (2, 4, 20)

romancal/flatfield/flat_field.py

@@ -4,6 +4,8 @@
 
 import logging
 import numpy as np
+from astropy import units as u
+from roman_datamodels import units as ru
 
 from romancal.lib import dqflags
 
@@ -109,21 +111,21 @@ def apply_flat_field(science, flat):
     flat_data[np.where(flat_bad)] = 1.0
     # Now let's apply the correction to science data and error arrays.  Rely
     # on array broadcasting to handle the cubes
-    science.data /= flat_data
+    science.data = u.Quantity((science.data.value / flat_data), ru.electron / u.s, dtype=science.data.dtype)
 
     # Update the variances using BASELINE algorithm.  For guider data, it has
     # not gone through ramp fitting so there is no Poisson noise or readnoise
     flat_data_squared = flat_data**2
-    science.var_poisson /= flat_data_squared
-    science.var_rnoise /= flat_data_squared
-    try:
-        science.var_flat = science.data**2 / flat_data_squared * flat_err**2
-    except AttributeError:
-        science['var_flat'] = np.zeros(shape=science.data.shape,
-                                       dtype=np.float32)
-        science.var_flat = science.data**2 / flat_data_squared * flat_err**2
-    science.err = np.sqrt(science.var_poisson +
-                          science.var_rnoise + science.var_flat)
+    science.var_poisson = u.Quantity((science.var_poisson.value / flat_data_squared), ru.electron**2 / u.s**2,
+                                     dtype = science.var_poisson.dtype)
+    science.var_rnoise = u.Quantity((science.var_rnoise / flat_data_squared), ru.electron**2 / u.s**2,
+                                    dtype = science.var_rnoise.dtype)
+
+    science['var_flat'] = u.Quantity((science.data.value ** 2 / flat_data_squared * flat_err ** 2),
+                                  ru.electron ** 2 / u.s ** 2, dtype=science.data.value.dtype)
+
+    err_sqrt = np.sqrt(science.var_poisson.value + science.var_rnoise.value + science.var_flat)
+    science.err = u.Quantity(err_sqrt, ru.electron / u.s, dtype=err_sqrt.dtype)
 
     # Combine the science and flat DQ arrays
     science.dq = np.bitwise_or(science.dq, flat_dq)

romancal/flatfield/tests/test_flatfield.py

@@ -6,9 +6,10 @@
 from roman_datamodels import stnode
 from roman_datamodels.datamodels import ImageModel, FlatRefModel
 from roman_datamodels.testing import utils as testutil
+from roman_datamodels import units as ru
 from romancal.flatfield import FlatFieldStep
 from astropy.time import Time
-
+from astropy import units as u
 
 @pytest.mark.parametrize(
     "instrument, exptype",
@@ -31,12 +32,17 @@ def test_flatfield_step_interface(instrument, exptype):
     wfi_image.meta.instrument.detector = 'WFI01'
     wfi_image.meta.instrument.optical_element = 'F158'
     wfi_image.meta.exposure.type = exptype
-    wfi_image.data = np.ones(shape, dtype=np.float32)
+    wfi_image.data = u.Quantity(np.ones(shape, dtype=np.float32),
+                                ru.electron / u.s, dtype=np.float32)
     wfi_image.dq = np.zeros(shape, dtype=np.uint32)
-    wfi_image.err = np.zeros(shape, dtype=np.float32)
-    wfi_image.var_poisson = np.zeros(shape, dtype=np.float32)
-    wfi_image.var_rnoise = np.zeros(shape, dtype=np.float32)
-    wfi_image.var_flat = np.zeros(shape, dtype=np.float32)
+    wfi_image.err = u.Quantity(np.zeros(shape, dtype=np.float32),
+                               ru.electron / u.s, dtype=np.float32)
+    wfi_image.var_poisson = u.Quantity(np.zeros(shape, dtype=np.float32),
+                                       ru.electron**2 / u.s**2, dtype=np.float32)
+    wfi_image.var_rnoise = u.Quantity(np.zeros(shape, dtype=np.float32),
+                                      ru.electron**2 / u.s**2, dtype=np.float32)
+    wfi_image.var_flat = u.Quantity(np.zeros(shape, dtype=np.float32),
+                                    ru.electron**2 / u.s**2, dtype=np.float32)
     wfi_image_model = ImageModel(wfi_image)
     flatref = stnode.FlatRef()
     meta = {}

romancal/jump/jump_step.py

@@ -58,10 +58,10 @@ def process(self, input):
             frames_per_group = meta.exposure.nframes
 
             # Modify the arrays for input into the 'common' jump (4D)
-            data = np.copy(r_data[np.newaxis, :])
+            data = np.copy(r_data[np.newaxis, :].value)
             gdq = r_gdq[np.newaxis, :]
             pdq = r_pdq[np.newaxis, :]
-            err = np.copy(r_err[np.newaxis, :])
+            err = np.copy(r_err[np.newaxis, :].value)
 
             tstart = time.time()
 

romancal/jump/tests/test_jump_step.py

@@ -7,12 +7,14 @@
 import pytest
 import numpy as np
 from astropy.time import Time
+from astropy import units as u
 
 import roman_datamodels.stnode as rds
 from roman_datamodels import datamodels as rdm
 from roman_datamodels.datamodels import GainRefModel
 from roman_datamodels.datamodels import ReadnoiseRefModel
 from roman_datamodels.testing import utils as testutil
+from roman_datamodels import units as ru
 
 from romancal.jump import JumpStep
 
@@ -49,7 +51,7 @@ def generate_wfi_reffiles(tmpdir_factory):
     gain_ref['meta'] = meta
     gain_ref['data'] = np.ones(shape, dtype=np.float32) * ingain
     gain_ref['dq'] = np.zeros(shape, dtype=np.uint16)
-    gain_ref['err'] = (np.random.random(shape) * 0.05).astype(np.float64)
+    gain_ref['err'] = (np.random.random(shape) * 0.05).astype(np.float32)
 
     gain_ref_model = GainRefModel(gain_ref)
     gain_ref_model.save(gainfile)
@@ -72,7 +74,7 @@ def generate_wfi_reffiles(tmpdir_factory):
     rn_ref['meta'] = meta
     rn_ref['data'] = np.ones(shape, dtype=np.float32)
     rn_ref['dq'] = np.zeros(shape, dtype=np.uint16)
-    rn_ref['err'] = (np.random.random(shape) * 0.05).astype(np.float64)
+    rn_ref['err'] = (np.random.random(shape) * 0.05).astype(np.float32)
 
     rn_ref_model = ReadnoiseRefModel(rn_ref)
     rn_ref_model.save(readnoisefile)
@@ -101,10 +103,10 @@ def _setup(ngroups=10, readnoise=10, nrows=20, ncols=20,
         dm_ramp.meta.instrument.name = 'WFI'
         dm_ramp.meta.instrument.optical_element = 'F158'
 
-        dm_ramp.data = data + 6.
+        dm_ramp.data = u.Quantity(data + 6., ru.DN, dtype=data.dtype)
         dm_ramp.pixeldq = pixdq
         dm_ramp.groupdq = gdq
-        dm_ramp.err = err
+        dm_ramp.err = u.Quantity(err, ru.DN, dtype=err.dtype)
 
         dm_ramp.meta.exposure.type = 'WFI_IMAGE'
         dm_ramp.meta.exposure.group_time = deltatime
@@ -142,7 +144,7 @@ def test_one_CR(generate_wfi_reffiles, max_cores, setup_inputs):
                           deltatime=grouptime)
 
     for i in range(ngroups):
-        model1.data[i, :, :] = deltaDN * i
+        model1.data.value[i, :, :] = deltaDN * i
     first_CR_group_locs = [x for x in range(1, 7) if x % 5 == 0]
 
     CR_locs = [x for x in range(xsize * ysize) if x % CR_fraction == 0]
@@ -153,8 +155,8 @@ def test_one_CR(generate_wfi_reffiles, max_cores, setup_inputs):
     # Add CRs to the SCI data
     for i in range(len(CR_x_locs)):
         CR_group = next(CR_pool)
-        model1.data[CR_group:, CR_y_locs[i], CR_x_locs[i]] = \
-            model1.data[CR_group:, CR_y_locs[i], CR_x_locs[i]] + 500.
+        model1.data.value[CR_group:, CR_y_locs[i], CR_x_locs[i]] = \
+            model1.data.value[CR_group:, CR_y_locs[i], CR_x_locs[i]] + 500.
 
     out_model = JumpStep.call(model1, override_gain=override_gain,
                               override_readnoise=override_readnoise,
@@ -190,7 +192,7 @@ def test_two_CRs(generate_wfi_reffiles, max_cores, setup_inputs):
                           deltatime=grouptime)
 
     for i in range(ngroups):
-        model1.data[i, :, :] = deltaDN * i
+        model1.data.value[i, :, :] = deltaDN * i
 
     first_CR_group_locs = [x for x in range(1, 7) if x % 5 == 0]
     CR_locs = [x for x in range(xsize * ysize) if x % CR_fraction == 0]
@@ -201,10 +203,10 @@ def test_two_CRs(generate_wfi_reffiles, max_cores, setup_inputs):
     for i in range(len(CR_x_locs)):
         CR_group = next(CR_pool)
 
-        model1.data[CR_group:, CR_y_locs[i], CR_x_locs[i]] = \
-            model1.data[CR_group:, CR_y_locs[i], CR_x_locs[i]] + 500
-        model1.data[CR_group + 8:, CR_y_locs[i], CR_x_locs[i]] = \
-            model1.data[CR_group + 8:, CR_y_locs[i], CR_x_locs[i]] + 700
+        model1.data.value[CR_group:, CR_y_locs[i], CR_x_locs[i]] = \
+            model1.data.value[CR_group:, CR_y_locs[i], CR_x_locs[i]] + 500
+        model1.data.value[CR_group + 8:, CR_y_locs[i], CR_x_locs[i]] = \
+            model1.data.value[CR_group + 8:, CR_y_locs[i], CR_x_locs[i]] + 700
 
     out_model = JumpStep.call(model1, override_gain=override_gain,
                               override_readnoise=override_readnoise,

romancal/linearity/linearity_step.py

@@ -9,6 +9,7 @@
 from romancal.stpipe import RomanStep
 from romancal.lib import dqflags
 from stcal.linearity.linearity import linearity_correction
+from astropy import units as u
 
 __all__ = ["LinearityStep"]
 
@@ -56,11 +57,11 @@ def process(self, input):
             # Call linearity correction function in stcal
             # The third return value is the procesed zero frame which
             # Roman does not use.
-            new_data, new_pdq, _ = linearity_correction(output_model.data,
+            new_data, new_pdq, _ = linearity_correction(output_model.data.value,
                                                      gdq, pdq, lin_coeffs,
                                                      lin_dq, dqflags.pixel)
 
-            output_model.data = new_data[0, :, :, :]
+            output_model.data = u.Quantity(new_data[0, :, :, :], u.DN, dtype=new_data.dtype)
             output_model.pixeldq = new_pdq
 
             # Close the reference file and update the step status

romancal/ramp_fitting/ramp_fit_step.py

@@ -8,6 +8,8 @@
 from roman_datamodels import datamodels as rdd
 from roman_datamodels import stnode as rds
 from roman_datamodels.testing import utils as testutil
+from roman_datamodels import units as ru
+from astropy import units as u
 
 from stcal.ramp_fitting import ramp_fit
 
@@ -42,15 +44,15 @@ def create_optional_results_model(input_model, opt_info):
     crmag.dtype = np.float32
 
     inst = {'meta': meta,
-            'slope': np.squeeze(slope),
-            'sigslope': np.squeeze(sigslope),
-            'var_poisson': np.squeeze(var_poisson),
-            'var_rnoise': np.squeeze(var_rnoise),
-            'yint': np.squeeze(yint),
-            'sigyint': np.squeeze(sigyint),
-            'pedestal': np.squeeze(pedestal),
+            'slope': u.Quantity(np.squeeze(slope), ru.electron / u.s, dtype=slope.dtype),
+            'sigslope': u.Quantity(np.squeeze(sigslope), ru.electron / u.s, dtype=sigslope.dtype),
+            'var_poisson': u.Quantity(np.squeeze(var_poisson), ru.electron**2 / u.s**2, dtype=var_poisson.dtype),
+            'var_rnoise': u.Quantity(np.squeeze(var_rnoise), ru.electron**2 / u.s**2, dtype=var_rnoise.dtype),
+            'yint': u.Quantity(np.squeeze(yint), ru.electron, dtype=yint.dtype),
+            'sigyint': u.Quantity(np.squeeze(sigyint), ru.electron, dtype=sigyint.dtype),
+            'pedestal': u.Quantity(np.squeeze(pedestal), ru.electron, dtype=pedestal.dtype),
             'weights': np.squeeze(weights),
-            'crmag': crmag
+            'crmag': u.Quantity(crmag, ru.electron, dtype=pedestal.dtype),
             }
 
     out_node = rds.RampFitOutput(inst)
@@ -80,6 +82,11 @@ def create_image_model(input_model, image_info):
     """
     data, dq, var_poisson, var_rnoise, err = image_info
 
+    data = u.Quantity(data, ru.electron / u.s, dtype=data.dtype)
+    var_poisson = u.Quantity(var_poisson, ru.electron**2 / u.s**2, dtype=var_poisson.dtype)
+    var_rnoise = u.Quantity(var_rnoise, ru.electron**2 / u.s**2, dtype=var_rnoise.dtype)
+    err = u.Quantity(err, ru.electron / u.s, dtype=err.dtype)
+
     # Create output datamodel
     # ... and add all keys from input
     meta = {}
@@ -142,7 +149,6 @@ def process(self, input):
             readnoise_filename = self.get_reference_file(input_model, 'readnoise')
             gain_filename = self.get_reference_file(input_model, 'gain')
             input_model.data = input_model.data[np.newaxis, :]
-            input_model.data.dtype=np.float32
             input_model.groupdq = input_model.groupdq[np.newaxis, :]
             input_model.err = input_model.err[np.newaxis, :]