Normalise combine()'s raw-module array_package and fix correlated add/subtract uncertainty leak

CHANGES.rst

@@ -20,6 +20,8 @@ Other Changes and Additions
 - Add a ``strict`` tox environment for running the ``array-api-strict`` test
   suite locally: ``tox -e strict`` reproduces the CI ``py313-strict`` job
   without having to name the interpreter. [#986]
+- ``combine`` no longer accepts an array as its ``array_package`` argument;
+  pass the array namespace or module instead, as for ``Combiner``. [#997]
 
 Bug Fixes
 ^^^^^^^^^
@@ -105,6 +107,12 @@ Bug Fixes
 - Compute the ``Combiner.clip_extrema`` mask with a rank comparison instead
   of scattering into the mask through per-pixel integer indices, which the
   array API standard does not support. [#994]
+- Accept a plain module such as ``numpy`` or ``dask.array`` as ``array_package``
+  in ``combine``, normalising it to its array-api-compat namespace the way
+  ``Combiner`` already does. [#997]
+- Keep the uncertainty propagation for correlated addition and subtraction
+  through ``_CCDDataWrapperForArrayAPI`` in the array namespace instead of
+  falling back to NumPy. [#997]
 
 2.5.1 (2025-07-05)
 ------------------

ccdproc/_ccddata_wrapper_for_array_api.py

@@ -281,7 +281,7 @@ def _variance_hooks(xp):
     def _propagate_add(self, other_uncert, result_data, correlation):
         xp = array_api_compat.array_namespace(self.array, other_uncert.array)
         to_variance, from_variance = self._variance_hooks(xp)
-        return super()._propagate_add_sub(
+        return self._propagate_add_sub(
             other_uncert,
             result_data,
             correlation,
@@ -293,7 +293,7 @@ def _propagate_add(self, other_uncert, result_data, correlation):
     def _propagate_subtract(self, other_uncert, result_data, correlation):
         xp = array_api_compat.array_namespace(self.array, other_uncert.array)
         to_variance, from_variance = self._variance_hooks(xp)
-        return super()._propagate_add_sub(
+        return self._propagate_add_sub(
             other_uncert,
             result_data,
             correlation,
@@ -302,6 +302,78 @@ def _propagate_subtract(self, other_uncert, result_data, correlation):
             from_variance=from_variance,
         )
 
+    def _propagate_add_sub(
+        self,
+        other_uncert,
+        result_data,
+        correlation,
+        subtract=False,
+        to_variance=lambda x: x,
+        from_variance=lambda x: x,
+    ):
+        """
+        Propagate uncertainty for addition or subtraction.
+
+        Parameters
+        ----------
+        other_uncert : `~astropy.nddata.NDUncertainty`
+            The uncertainty of the other operand. Its ``array`` must be in
+            the same array namespace as ``self.array``.
+        result_data : array-like
+            Accepted only for signature compatibility with astropy; the
+            formulae do not use it.
+        subtract : bool, optional
+            ``True`` for subtraction, ``False`` (default) for addition.
+        correlation : float or array-like
+            Correlation coefficient between the two operands, ``0`` for
+            uncorrelated.
+        to_variance : callable, optional
+            Converts the stored uncertainty array to a variance. Defaults to
+            the identity, i.e. the uncertainty is already a variance.
+        from_variance : callable, optional
+            Converts a variance back to the stored uncertainty type. Defaults
+            to the identity.
+
+        Returns
+        -------
+        array-like
+            The propagated uncertainty array, in the same array namespace and
+            on the same device as the inputs, in the representation of
+            ``self`` (as determined by ``from_variance``).
+
+        Notes
+        -----
+        This is astropy's ``_VariancePropagationMixin._propagate_add_sub``
+        with the NumPy call replaced by its array-namespace equivalent; see
+        the astropy version for the derivation of the formulae. Unlike
+        astropy's version this does not convert the uncertainties between
+        units, because ``_CCDDataWrapperForArrayAPI._arithmetic_wrapper``
+        removes the units from the uncertainties before doing the
+        arithmetic.
+        """
+        del result_data  # accepted only for compatibility with astropy
+        xp = array_api_compat.array_namespace(self.array, other_uncert.array)
+
+        correlation_sign = -1 if subtract else 1
+
+        other = to_variance(other_uncert.array) if other_uncert.array is not None else 0
+        this = to_variance(self.array) if self.array is not None else 0
+
+        # Formula: sigma**2 = dA + dB +/- 2*cor*sqrt(dA*dB)
+        # Only take the correlation term into account when both operands
+        # have an uncertainty; otherwise ``this`` or ``other`` is a bare
+        # Python ``0`` and ``xp.sqrt(0)`` is rejected on strict. The term is
+        # mathematically zero in that case regardless.
+        if (isinstance(correlation, np.ndarray) or correlation != 0) and (
+            other_uncert.array is not None and self.array is not None
+        ):
+            corr = 2 * correlation * xp.sqrt(this * other)
+            result = this + other + correlation_sign * corr
+        else:
+            result = this + other
+
+        return from_variance(result)
+
     def _propagate_multiply(self, other_uncert, result_data, correlation):
         xp = array_api_compat.array_namespace(self.array, other_uncert.array)
         to_variance, from_variance = self._variance_hooks(xp)
@@ -338,14 +410,6 @@ def _propagate_multiply_divide(
         """
         Propagate uncertainty for multiplication or division.
 
-        This is astropy's
-        ``_VariancePropagationMixin._propagate_multiply_divide`` with the
-        NumPy calls replaced by their array-namespace equivalents; see the
-        astropy version for the derivation of the formulae. Unlike astropy's
-        version this does not convert the uncertainties between units,
-        because ``_CCDDataWrapperForArrayAPI._arithmetic_wrapper`` removes
-        the units from the uncertainties before doing the arithmetic.
-
         Parameters
         ----------
         other_uncert : `~astropy.nddata.NDUncertainty`
@@ -373,6 +437,16 @@ def _propagate_multiply_divide(
             The propagated uncertainty array, in the same array namespace and
             on the same device as the inputs, in the representation of
             ``self`` (as determined by ``from_variance``).
+
+        Notes
+        -----
+        This is astropy's
+        ``_VariancePropagationMixin._propagate_multiply_divide`` with the
+        NumPy calls replaced by their array-namespace equivalents; see the
+        astropy version for the derivation of the formulae. Unlike astropy's
+        version this does not convert the uncertainties between units,
+        because ``_CCDDataWrapperForArrayAPI._arithmetic_wrapper`` removes
+        the units from the uncertainties before doing the arithmetic.
         """
         del result_data  # accepted only for compatibility with astropy
         xp = array_api_compat.array_namespace(self.array, other_uncert.array)

ccdproc/combiner.py

@@ -1006,17 +1006,13 @@ def combine(
         has no effect otherwise.
         Default is ``False``.
 
-    array_package : an array namespace, optional
-        The array package to use for the data if the data needs to be
-        read in from files. This argument is ignored if the input ``ccd_list``
-        is already a list of `~astropy.nddata.CCDData` objects.
-
-        If not specified, the array package used will
-        be numpy. The array package can be specified either by passing in
-        an array namespace (e.g. output from ``array_api_compat.array_namespace``),
-        or an imported array package that follows the array API standard
-        (e.g. ``numpy`` or ``jax.numpy``), or an array whose namespace can be
-        determined (e.g. a `numpy.ndarray` or ``jax.numpy.ndarray``).
+    array_package : array namespace or module, optional
+        The array package to use for data read in from files; ignored if
+        ``ccd_list`` is already a list of `~astropy.nddata.CCDData` objects.
+        Either an array namespace or a plain module that follows the array
+        API standard (e.g. ``numpy`` or ``dask.array``); it is normalised to
+        its array-api-compat namespace the same way `~ccdproc.Combiner`
+        handles ``xp``. Default is NumPy.
 
     ccdkwargs : Other keyword arguments for `astropy.nddata.fits_ccddata_reader`.
 
@@ -1069,10 +1065,12 @@ def combine(
         # The ccd object will always read as numpy, so convert it to the
         # requested namespace if there is one.
         if array_package is not None:
-            try:
-                xp = array_api_compat.array_namespace(array_package)
-            except TypeError:
-                xp = array_package
+            # ``array_package`` may be a raw module such as ``numpy`` or
+            # ``dask.array``; normalise it to the array-api-compat namespace
+            # the same way ``Combiner.__init__`` does, so the conversions
+            # below can rely on array-API features (e.g. the ``device``
+            # keyword) that a raw module may not provide.
+            xp = array_api_compat.array_namespace(array_package.asarray(0))
 
             # ccd.data (and its uncertainty, if any) were just read from a
             # FITS file, so they are NumPy arrays, possibly in big-endian

ccdproc/tests/test_ccddata_wrapper_for_array_api.py

@@ -211,7 +211,7 @@ def make(asarray):
 @pytest.mark.parametrize(
     "uncertainty_type", [StdDevUncertainty, VarianceUncertainty, InverseVariance]
 )
-@pytest.mark.parametrize("operation", ["multiply", "divide"])
+@pytest.mark.parametrize("operation", ["add", "subtract", "multiply", "divide"])
 def test_wrapped_arithmetic_correlated_uncertainty(uncertainty_type, operation):
     data1 = [[1.0, 2.0], [3.0, 4.0]]
     data2 = [[2.0, 2.0], [4.0, 8.0]]

ccdproc/tests/test_combiner.py

@@ -1518,3 +1518,55 @@ def sum_func(_, axis=axis):
     expected_result = xpx.at(expected_result)[5, 5].set(2)
 
     assert xp.all(xpx.isclose(expected_result, actual_result.data))
+
+
+# Regression tests for #982: combine()'s ``array_package`` only normalised
+# an already-instantiated array (via array_api_compat.array_namespace), so a
+# raw array module (e.g. plain ``numpy`` or ``dask.array``, as opposed to
+# ``array_api_compat.numpy``/``array_api_compat.dask.array``) passed through
+# unnormalised into code that relies on array-API features the raw module
+# does not provide.
+def test_combine_array_package_raw_module(tmp_path):
+    """A raw array module passed as ``array_package`` should be normalised
+    to its array-api-compat namespace, the same way ``Combiner`` normalises
+    its ``xp`` argument.
+    """
+    ccd = CCDData(np.arange(9, dtype=float).reshape(3, 3), unit=u.adu)
+    files = []
+    for i in range(3):
+        path = tmp_path / f"raw-module-{i}.fits"
+        ccd.write(path)
+        files.append(str(path))
+
+    result = combine(files, array_package=np, unit="adu")
+    assert array_api_compat.is_numpy_array(result.data)
+
+    # On strict and jax, the suite's own ``xp`` fixture is itself a raw
+    # module (``array_api_strict``/``jax.numpy`` imported directly, not
+    # through array_api_compat), so this also exercises the raw-module
+    # path of #982 on those backends without a separate test case.
+    result = combine(files, array_package=xp, unit="adu")
+    expected_xp = array_api_compat.array_namespace(xp.asarray(0))
+    assert array_api_compat.array_namespace(result.data) is expected_xp
+
+
+def test_combine_array_package_dask_module(tmp_path):
+    """Regression test for #982.
+
+    Passing the raw ``dask.array`` module (rather than its
+    array-api-compat wrapper, ``array_api_compat.dask.array``) as
+    ``array_package`` used to reach ``dask.array.from_array`` with an
+    unsupported ``device=`` keyword, raising ``TypeError: from_array() got
+    an unexpected keyword argument 'device'``.
+    """
+    dask = pytest.importorskip("dask.array")
+
+    ccd = CCDData(np.arange(9, dtype=float).reshape(3, 3), unit=u.adu)
+    files = []
+    for i in range(3):
+        path = tmp_path / f"raw-dask-{i}.fits"
+        ccd.write(path)
+        files.append(str(path))
+
+    result = combine(files, array_package=dask, unit="adu")
+    assert array_api_compat.is_dask_array(result.data)