Build Combiner data and mask arrays with xp.stack

CHANGES.rst

@@ -21,6 +21,14 @@ Bug Fixes
   average combinations. [#952]
 - Make ``flat_correct`` and ``ccdmask`` use functional array updates so they
   support immutable array-API backends. [#956]
+- Build the ``Combiner`` data and mask arrays with ``xp.stack`` instead of
+  passing a nested list of arrays to ``xp.asarray``, and keep the mask and
+  scaling arrays in the namespace and on the device of the input data.
+  ``Combiner.scaling`` is now cast to the dtype of the data, accepts an
+  array-API array without ``__len__``, and is reshaped to exactly one
+  broadcast axis per image dimension, so a ``(N, 1)`` scaling array or a
+  callable returning a one-element array no longer adds a spurious trailing
+  dimension to the combined image. [#965]
 - Fix the fallback percentile calculation for array namespaces that do not
   provide ``percentile``. [#957]
 - Make array-API escape logging thread-safe so concurrent worker escapes are

ccdproc/__init__.py

@@ -4,6 +4,7 @@
 processing. These steps will allow reduction of basic CCD data as either a
 stand-alone processing or as part of a pipeline.
 """
+
 try:
     from ._version import version as __version__
 except ImportError:

ccdproc/combiner.py

@@ -105,7 +105,8 @@ class Combiner:
     xp : array namespace, optional
         The array namespace to use for the data. If `None` or not provided, it will
         be inferred from the first `~astropy.nddata.CCDData` object in
-        ``ccd_iter``.
+        ``ccd_iter``. A plain module (e.g. ``numpy``) is accepted and is
+        converted to its array-API-compatible namespace.
         Default is `None`.
 
     Raises
@@ -165,8 +166,14 @@ def __init__(self, ccd_iter, dtype=None, xp=None):
                 if not (default_unit == ccd.unit):
                     raise TypeError("CCDData objects don't have the same unit.")
 
-        # Set array namespace
-        xp = xp or array_api_compat.array_namespace(ccd_list[0].data)
+        # Set array namespace. A raw module such as ``numpy`` or ``dask.array``
+        # may lack array-API features that are used below (``xp.bool``, the
+        # ``device`` keyword), so normalise whatever the caller passed to the
+        # array-api-compat namespace of one of its arrays.
+        if xp is None:
+            xp = array_api_compat.array_namespace(ccd_list[0].data)
+        else:
+            xp = array_api_compat.array_namespace(xp.asarray(0))
         self._xp = xp
         if dtype is None:
             dtype = xp.float64
@@ -177,14 +184,30 @@ def __init__(self, ccd_iter, dtype=None, xp=None):
 
         # set up the data array
         # new_shape = (len(ccd_list),) + default_shape
-        self._data_arr = xp.asarray([ccd.data for ccd in ccd_list], dtype=dtype)
+        # Stack the individual images rather than passing a nested list to
+        # xp.asarray: the array API does not allow nested sequences of arrays.
+        # Keep the stack on the device of the input data, but only when the
+        # data already belong to ``xp``; a device object from a different
+        # namespace (e.g. numpy's 'cpu' for a jax namespace) is meaningless
+        # to ``xp``, so let ``xp`` use its default device instead.
+        data_xp = array_api_compat.array_namespace(ccd_list[0].data)
+        device = array_api_compat.device(ccd_list[0].data) if data_xp is xp else None
+        self._data_arr = xp.stack(
+            [xp.asarray(ccd.data, dtype=dtype, device=device) for ccd in ccd_list]
+        )
 
-        # populate self._data_arr_mask
+        # populate self._data_arr_mask. The mask of a CCDData may be a numpy
+        # array even when its data is not, so coerce each mask into the data
+        # namespace and onto the data device before stacking.
         mask_list = [
-            ccd.mask if ccd.mask is not None else xp.zeros(default_shape)
+            (
+                xp.asarray(ccd.mask, dtype=xp.bool, device=device)
+                if ccd.mask is not None
+                else xp.zeros(default_shape, dtype=xp.bool, device=device)
+            )
             for ccd in ccd_list
         ]
-        self._data_arr_mask = xp.asarray(mask_list, dtype=bool)
+        self._data_arr_mask = xp.stack(mask_list)
 
         # Must be after self.data_arr is defined because it checks the
         # length of the data array.
@@ -264,26 +287,46 @@ def scaling(self, value):
             self._scaling = value
         else:
             n_images = self._data_arr.shape[0]
+            device = array_api_compat.device(self._data_arr)
+            dtype = self._data_arr.dtype
             if callable(value):
-                self._scaling = [value(self._data_arr[i]) for i in range(n_images)]
-                self._scaling = xp.asarray(self._scaling)
+                # The callable may return a Python float or a 0-d array of
+                # the backend; stack per-element conversions rather than
+                # passing a list of arrays to asarray, which array-api-strict
+                # rejects as a nested sequence of arrays.
+                # Cast to the data dtype so that scaling by, e.g., an integer
+                # does not require type promotion, which the array API does
+                # not guarantee between integer and floating dtypes.
+                self._scaling = xp.stack(
+                    [
+                        xp.asarray(
+                            value(self._data_arr[i, ...]), dtype=dtype, device=device
+                        )
+                        for i in range(n_images)
+                    ]
+                )
             else:
+                # Array API arrays need not implement __len__, so use the
+                # shape where there is one and fall back to len() for lists
+                # and tuples.
                 try:
-                    len(value)
-                except TypeError as err:
+                    n_values = getattr(value, "shape", None)
+                    n_values = n_values[0] if n_values else len(value)
+                except (TypeError, IndexError) as err:
                     raise TypeError(
                         "scaling must be a function or an array "
                         "the same length as the number of images.",
                     ) from err
-                if len(value) != n_images:
+                if n_values != n_images:
                     raise ValueError(
                         "scaling must be a function or an array "
                         "the same length as the number of images."
                     )
-                self._scaling = xp.asarray(value)
+                self._scaling = xp.asarray(value, dtype=dtype, device=device)
             # reshape so that broadcasting occurs properly
-            for _ in range(len(self._data_arr.shape) - 1):
-                self._scaling = self.scaling[:, xp.newaxis]
+            self._scaling = xp.reshape(
+                self._scaling, (n_images,) + (1,) * (self._data_arr.ndim - 1)
+            )
 
     # set up IRAF-like minmax clipping
     def clip_extrema(self, nlow=0, nhigh=0):
@@ -1021,7 +1064,7 @@ def combine(
                     ccd.uncertainty.array, dtype=ccd.uncertainty.array.dtype.type
                 )
             if ccd.mask is not None:
-                ccd.mask = xp.asarray(ccd.mask, dtype=bool)
+                ccd.mask = xp.asarray(ccd.mask, dtype=xp.bool)
 
     # Get the array namespace; if array_package was not None and files were read in,
     # then xp the ccd.data will be the same as the array_package.
@@ -1048,7 +1091,7 @@ def combine(
     # If the template doesn't have a mask, add one, because the result may have
     # a mask
     if ccd.mask is None:
-        ccd.mask = xp.zeros_like(ccd.data, dtype=bool)
+        ccd.mask = xp.zeros_like(ccd.data, dtype=xp.bool)
 
     size_of_an_img = _calculate_size_of_image(ccd)
 
@@ -1101,11 +1144,16 @@ def combine(
                                 imgccd.uncertainty.array, dtype=dtype
                             )
                         if imgccd.mask is not None:
-                            imgccd.mask = xp.asarray(imgccd.mask, dtype=bool)
+                            imgccd.mask = xp.asarray(imgccd.mask, dtype=xp.bool)
 
                 scalevalues.append(scale(imgccd.data))
 
-            to_set_in_combiner["scaling"] = xp.asarray(scalevalues)
+            # See Combiner.scaling: stack per-element conversions so that a
+            # callable returning 0-d backend arrays works on array-api-strict.
+            device = array_api_compat.device(ccd.data)
+            to_set_in_combiner["scaling"] = xp.stack(
+                [xp.asarray(value, device=device) for value in scalevalues]
+            )
         else:
             to_set_in_combiner["scaling"] = scale
 
@@ -1144,7 +1192,7 @@ def combine(
                                 imgccd.uncertainty.array, dtype=dtype
                             )
                         if imgccd.mask is not None:
-                            imgccd.mask = xp.asarray(imgccd.mask, dtype=bool)
+                            imgccd.mask = xp.asarray(imgccd.mask, dtype=xp.bool)
 
                 # Trim image and copy
                 # The copy is *essential* to avoid having a bunch

ccdproc/tests/test_combiner.py

@@ -2,12 +2,14 @@
 import array_api_compat
 import array_api_extra as xpx
 import astropy.units as u
+import numpy as np
 import numpy.ma as np_ma
 import pytest
 from astropy.nddata import CCDData
 from astropy.stats import median_absolute_deviation as mad
 from astropy.utils.data import get_pkg_data_filename
 from numpy import median as np_median
+from numpy.testing import assert_allclose
 
 from ccdproc import create_deviation
 from ccdproc.combiner import (
@@ -165,6 +167,125 @@ def test_combiner_mask():
     assert not c._data_arr_mask[0, 5, 5]
 
 
+# Regression test for #965: the Combiner must stack the input images and
+# masks instead of passing a nested list of arrays to xp.asarray, and the
+# stacked arrays must live in the namespace and on the device of the inputs.
+def test_combiner_stacks_arrays_on_input_device():
+    data = xp.zeros((4, 4), dtype=xp.float32, device=xp_device)
+    data = xpx.at(data)[1, 2].set(1)
+    mask = xp.zeros((4, 4), dtype=xp.bool, device=xp_device)
+    mask = xpx.at(mask)[3, 3].set(True)
+    # astropy's CCDData.mask setter always coerces to numpy, which is not
+    # possible for arrays on a non-default device, so assign the mask in the
+    # data's namespace directly to emulate an array-API-aware CCDData.
+    ccd_masked = CCDData(data, unit=u.adu)
+    ccd_masked._mask = mask
+    ccd_unmasked = CCDData(data, unit=u.adu)
+    # Masks on CCDData may be plain numpy arrays even when the data is not.
+    ccd_np_mask = CCDData(data, unit=u.adu, mask=np.ones((4, 4), dtype=bool))
+
+    c = Combiner([ccd_masked, ccd_unmasked, ccd_np_mask], dtype=xp.float32)
+
+    for arr in (c.data, c.mask):
+        assert array_api_compat.array_namespace(
+            arr
+        ) is array_api_compat.array_namespace(data)
+        assert array_api_compat.device(arr) == array_api_compat.device(data)
+        assert arr.shape == (3, 4, 4)
+    assert c.data.dtype == xp.float32
+    assert c.mask.dtype == xp.bool
+    assert c.data[0, 1, 2] == 1
+    assert c.mask[0, 3, 3]
+    assert not xp.any(c.mask[1, ...])
+    assert xp.all(c.mask[2, ...])
+
+    # Scaling values should end up on the same device as the data, and in
+    # the dtype of the data: integer scaling must not rely on int/float type
+    # promotion, which the array API does not guarantee.
+    c.scaling = [1.0, 2.0, 3.0]
+    assert array_api_compat.device(c.scaling) == array_api_compat.device(data)
+    c.scaling = [1, 2, 3]
+    assert c.scaling.dtype == c.data.dtype
+    assert float(c.scaling[1, 0, 0]) == 2.0
+    c.scaling = lambda arr: float(arr.shape[0])
+    assert array_api_compat.device(c.scaling) == array_api_compat.device(data)
+    c.scaling = lambda arr: arr.shape[0]
+    assert c.scaling.dtype == c.data.dtype
+    assert float(c.scaling[0, 0, 0]) == 4.0
+    # A backend array, which need not implement __len__, is accepted as
+    # scaling as long as its length matches the number of images.
+    c.scaling = xp.asarray([1.0, 0.5, 2.0], device=xp_device)
+    assert array_api_compat.device(c.scaling) == array_api_compat.device(data)
+    assert float(c.scaling[2, 0, 0]) == 2.0
+    with pytest.raises(ValueError, match="same length"):
+        c.scaling = xp.asarray([1.0, 0.5], device=xp_device)
+    with pytest.raises(TypeError, match="same length"):
+        c.scaling = 2.0
+    # A scaling callable that returns a 0-d array of the backend (rather than
+    # a Python float) must also work; array-api-strict rejects a list of such
+    # arrays passed to asarray.
+    c.scaling = lambda arr: xp.mean(arr) + 1
+    assert array_api_compat.device(c.scaling) == array_api_compat.device(data)
+    assert c.scaling.shape == (3, 1, 1)
+    # all three images share ``data``, whose mean is 1/16
+    assert float(xp.max(xp.abs(c.scaling - (1.0 + 1 / 16)))) < 1e-6
+
+
+@pytest.mark.backend_xfail(
+    "array-api-strict",
+    reason="combine() sizes the image with .nbytes, which array-api-strict "
+    "does not provide",
+)
+def test_combine_scale_callable_returning_backend_scalar():
+    # Added in #976: ``combine(scale=<callable>)`` must accept a callable that
+    # returns a 0-d array of the backend. Only array-api-strict rejects the
+    # previous ``xp.asarray([0-d array, ...])``, and on that backend
+    # ``combine()`` currently fails earlier on ``ccd.data.nbytes``, so this
+    # test cannot yet fail because of the scaling path on any backend; it
+    # will start guarding it once ``combine()`` stops using ``.nbytes``.
+    ccds = [
+        CCDData(xp.full((4, 4), float(i), dtype=xp.float64), unit=u.adu)
+        for i in (1, 2, 4)
+    ]
+    result = combine(ccds, method="average", scale=lambda arr: 1 / xp.mean(arr))
+    assert_allclose(np.asarray(result.data), 1.0)
+
+
+def test_combiner_explicit_namespace_differs_from_data():
+    # Regression test for the review of #976: when the caller passes an ``xp``
+    # that is not the namespace of the input data, the device of the inputs
+    # must not be forced onto ``xp`` (numpy's 'cpu' means nothing to jax or
+    # array-api-strict). The data is converted into ``xp`` on its default
+    # device instead.
+    #
+    # Only array-api-strict actually rejects a foreign device: on numpy the
+    # data namespace *is* ``xp`` so the device is legitimately reused, and
+    # dask accepts ``device='cpu'`` regardless, so this test can fail only in
+    # the array-api-strict job.
+    np_ccds = [CCDData(np.ones((3, 3)) * i, unit=u.adu) for i in range(1, 3)]
+    np_ccds[0].mask = np.zeros((3, 3), dtype=bool)
+    c = Combiner(np_ccds, xp=xp)
+    assert array_api_compat.array_namespace(c.data) is array_api_compat.array_namespace(
+        xp.zeros(1)
+    )
+    assert c.data.shape == (2, 3, 3)
+    assert c.data.dtype == xp.float64
+    assert c.mask.dtype == xp.bool
+    assert float(xp.sum(c.data)) == 27.0
+
+
+def test_combiner_accepts_raw_module_as_namespace():
+    # A plain module (numpy here) passed as ``xp`` is normalised to its
+    # array-api-compat namespace, so array-API-only features such as
+    # ``xp.bool`` and ``device=`` are available to the Combiner.
+    np_ccds = [CCDData(np.ones((2, 2)) * i, unit=u.adu) for i in range(1, 3)]
+    c = Combiner(np_ccds, xp=np)
+    assert c._xp is array_api_compat.array_namespace(np.zeros(1))
+    assert c.data.shape == (2, 2, 2)
+    c.scaling = [1, 2]
+    assert float(np.sum(c.average_combine().data)) == 10.0
+
+
 def test_weights():
     ccd_data = ccd_data_func()
     ccd_list = [ccd_data, ccd_data, ccd_data]