JP-3782: Fix improper setting of AREA extension in coron3 pipeline (s…

…pacetelescope#8907)

changes/8907.pipeline.rst

@@ -0,0 +1 @@
+Fixed a bug leading to incorrect area extensions, and sometimes crashes, in the coron3 pipeline

jwst/pipeline/calwebb_coron3.py

@@ -30,7 +30,14 @@ def to_container(model):
                 'var_poisson', 'var_rnoise', 'var_flat'
         ]:
             try:
-                setattr(image, attribute, model.getarray_noinit(attribute)[plane])
+                arr = model.getarray_noinit(attribute)
+                if arr.ndim == 3:
+                    setattr(image, attribute, arr[plane])
+                elif arr.ndim == 2:
+                    setattr(image, attribute, arr)
+                else:
+                    msg = f"Unexpected array shape for extension {attribute}: {arr.shape}"
+                    raise ValueError(msg)
             except AttributeError:
                 pass
         image.update(model)

jwst/pipeline/tests/test_calwebb_coron3.py

@@ -3,6 +3,7 @@
 from stdatamodels.jwst.datamodels import CubeModel
 
 from jwst.pipeline.calwebb_coron3 import to_container
+import numpy as np
 
 
 # Generate data
@@ -22,6 +23,26 @@ def make_container():
 ]
 
 
+def test_to_container():
+    """Cover bug where IndexError would be raised when area extension of CubeModel
+    has shape (x,y) instead of (nints,x,y). In this case area extension should
+    be copied to each ImageModel in the ModelContainer.
+    """
+    shp = (10,5,5)
+    cube = CubeModel(shp)
+    extensions_3d = ['data', 'dq', 'err',]
+    for extension in extensions_3d:
+        setattr(cube, extension, np.random.rand(*shp))
+    cube.area = np.random.rand(5,5)
+
+    container = to_container(cube)
+    for i, model in enumerate(container):
+        for extension in extensions_3d:
+            assert np.all(getattr(model, extension) == getattr(cube, extension)[i])
+        assert hasattr(model, 'area')
+        assert np.all(model.area == cube.area)
+
+
 @pytest.mark.parametrize('cube, container', [(cube, container)])
 def test_container_shape(cube, container):
     """Test container shape"""