Fix gain_correct and flat_correct on a non-default device

CHANGES.rst

@@ -94,6 +94,10 @@ Bug Fixes
 - Keep the mask of the result of ``Combiner.average_combine``,
   ``median_combine``, ``sum_combine`` and ``combine`` in the array namespace
   and on the device of the data instead of converting it to NumPy. [#992]
+- Fix ``gain_correct`` and ``flat_correct`` for images on a non-default
+  device: put the gain and flat normalization on the device of the data, cast
+  an integer gain to float, and keep the uncertainty propagation for
+  multiplication and division out of NumPy. [#993]
 
 2.5.1 (2025-07-05)
 ------------------

ccdproc/_ccddata_wrapper_for_array_api.py

@@ -305,7 +305,7 @@ def _propagate_subtract(self, other_uncert, result_data, correlation):
     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)
-        return super()._propagate_multiply_divide(
+        return self._propagate_multiply_divide(
             other_uncert,
             result_data,
             correlation,
@@ -317,7 +317,7 @@ def _propagate_multiply(self, other_uncert, result_data, correlation):
     def _propagate_divide(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_multiply_divide(
+        return self._propagate_multiply_divide(
             other_uncert,
             result_data,
             correlation,
@@ -326,6 +326,93 @@ def _propagate_divide(self, other_uncert, result_data, correlation):
             from_variance=from_variance,
         )
 
+    def _propagate_multiply_divide(
+        self,
+        other_uncert,
+        result_data,
+        correlation,
+        divide=False,
+        to_variance=lambda x: x,
+        from_variance=lambda x: x,
+    ):
+        """
+        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`
+            The uncertainty of the other operand. Its ``array`` and
+            ``parent_nddata.data`` 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.
+        correlation : float or array-like
+            Correlation coefficient between the two operands, ``0`` for
+            uncorrelated.
+        divide : bool, optional
+            ``True`` for division, ``False`` (default) for multiplication.
+        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``).
+        """
+        del result_data  # accepted only for compatibility with astropy
+        xp = array_api_compat.array_namespace(self.array, other_uncert.array)
+
+        correlation_sign = -1 if divide else 1
+
+        if other_uncert.array is not None:
+            d_b = to_variance(other_uncert.array)
+            # Formula: sigma**2 = |A|**2 * d_b
+            right = xp.abs(self.parent_nddata.data**2 * d_b)
+        else:
+            right = 0
+
+        if self.array is not None:
+            # Just the reversed case
+            d_a = to_variance(self.array)
+            # Formula: sigma**2 = |B|**2 * d_a
+            left = xp.abs(other_uncert.parent_nddata.data**2 * d_a)
+        else:
+            left = 0
+
+        if isinstance(correlation, np.ndarray) or correlation != 0:
+            corr = (
+                2
+                * correlation
+                * xp.sqrt(d_a * d_b)
+                * self.parent_nddata.data
+                * other_uncert.parent_nddata.data
+            )
+        else:
+            corr = 0
+
+        if divide:
+            return from_variance(
+                (left + right + correlation_sign * corr)
+                / other_uncert.parent_nddata.data**4
+            )
+        else:
+            return from_variance(left + right + correlation_sign * corr)
+
 
 class _StdDevUncertaintyWrapper(
     _CupyOperationNamesMixin, _ArrayAPIPropagationMixin, StdDevUncertainty

ccdproc/core.py

@@ -965,7 +965,15 @@ def gain_correct(ccd, gain, gain_unit=None, xp=None):
     gain_value = (
         gain_value.decompose().value if isinstance(gain_value, Quantity) else gain_value
     )
-    _result = _ccd.multiply(xp.asarray(gain_value), xp=xp, handle_mask=xp.logical_or)
+    if isinstance(gain_value, numbers.Real):
+        # The array API standard does not promote integer arrays with
+        # floating-point arrays, so make sure a plain-number gain is a float.
+        gain_value = float(gain_value)
+    _result = _ccd.multiply(
+        xp.asarray(gain_value, device=array_api_compat.device(_ccd.data)),
+        xp=xp,
+        handle_mask=xp.logical_or,
+    )
     if gain_unit:
         # Set unit of the data
         _result.unit = _ccd.unit * gain_unit
@@ -1054,7 +1062,10 @@ def flat_correct(ccd, flat, min_value=None, norm_value=None, xp=None):
     # Make sure flat_mean is a plain python float so that we
     # can use it with the array namespace. -- actually, we need to cast
     # flat_mean to the array namespace.
-    _flat_normed = _use_flat.divide(xp.asarray(flat_mean), xp=xp)
+    _flat_normed = _use_flat.divide(
+        xp.asarray(flat_mean, device=array_api_compat.device(_use_flat.data)),
+        xp=xp,
+    )
 
     # We need to fix up the unit now since we stripped the unit from
     # the flat_mean above.

ccdproc/tests/test_ccddata_wrapper_for_array_api.py

@@ -50,10 +50,6 @@ def test_trim_image_returns_plain_ccddata():
     assert result.wcs.wcs.compare(ccd.wcs.wcs)
 
 
-@pytest.mark.backend_xfail(
-    "array-api-strict",
-    reason="Astropy uncertainty propagation mixes NumPy and strict arrays",
-)
 def test_flat_correct_returns_public_uncertainty():
     ccd = CCDData(
         xp.ones((2, 2)),
@@ -141,13 +137,6 @@ def test_unwrap_rejects_non_ccddata():
         _unwrap_ccddata_for_array_api(object())
 
 
-_STRICT_STDDEV_MULDIV_XFAIL = pytest.mark.backend_xfail(
-    "array-api-strict",
-    reason="astropy's _propagate_multiply_divide applies np.sqrt/np.abs to the "
-    "std-dev result, which fails on a non-default strict device (see #940)",
-)
-
-
 def test_propagation_mixin_requires_variance_hooks():
     """The mixin is abstract: a subclass that forgets ``_variance_hooks`` fails
     loudly rather than silently propagating with the wrong conversions."""
@@ -156,21 +145,9 @@ def test_propagation_mixin_requires_variance_hooks():
 
 
 @pytest.mark.parametrize(
-    ("uncertainty_type", "operation"),
-    [
-        pytest.param(
-            unc,
-            op,
-            marks=(
-                [_STRICT_STDDEV_MULDIV_XFAIL]
-                if unc is StdDevUncertainty and op in ("multiply", "divide")
-                else []
-            ),
-        )
-        for unc in (StdDevUncertainty, VarianceUncertainty, InverseVariance)
-        for op in ("add", "subtract", "multiply", "divide")
-    ],
+    "uncertainty_type", [StdDevUncertainty, VarianceUncertainty, InverseVariance]
 )
+@pytest.mark.parametrize("operation", ["add", "subtract", "multiply", "divide"])
 def test_wrapped_arithmetic_keeps_uncertainty_in_namespace(uncertainty_type, operation):
     data1 = [[1.0, 2.0], [3.0, 4.0]]
     data2 = [[2.0, 2.0], [4.0, 8.0]]
@@ -196,3 +173,67 @@ def make(data, unc, asarray):
     assert xp.all(
         xpx.isclose(result.uncertainty.array, xp.asarray(expected.uncertainty.array))
     )
+
+
+@pytest.mark.parametrize(
+    "uncertainty_type", [StdDevUncertainty, VarianceUncertainty, InverseVariance]
+)
+@pytest.mark.parametrize("operation", ["add", "subtract", "multiply", "divide"])
+def test_wrapped_arithmetic_uncertainty_only_on_operand(uncertainty_type, operation):
+    # When only the operand has an uncertainty astropy propagates from an
+    # empty uncertainty on the first operand, which is the case in, e.g.,
+    # flat_correct of an image without an uncertainty by a flat with one.
+    data1 = [[1.0, 2.0], [3.0, 4.0]]
+    data2 = [[2.0, 2.0], [4.0, 8.0]]
+    unc2 = [[0.2, 0.1], [0.4, 0.3]]
+
+    def make(asarray):
+        ccd1 = CCDData(asarray(data1), unit=u.adu)
+        ccd2 = CCDData(
+            asarray(data2), unit=u.adu, uncertainty=uncertainty_type(asarray(unc2))
+        )
+        return ccd1, ccd2
+
+    ccd1, ccd2 = (_wrap_ccddata_for_array_api(ccd) for ccd in make(xp.asarray))
+    result = getattr(ccd1, operation)(ccd2)
+
+    # Reference: astropy's own propagation on plain numpy CCDData.
+    ref1, ref2 = make(np.asarray)
+    expected = getattr(ref1, operation)(ref2)
+
+    assert array_api_compat.array_namespace(result.uncertainty.array) is xp
+    assert isinstance(result.uncertainty, uncertainty_type)
+    assert xp.all(
+        xpx.isclose(result.uncertainty.array, xp.asarray(expected.uncertainty.array))
+    )
+
+
+@pytest.mark.parametrize(
+    "uncertainty_type", [StdDevUncertainty, VarianceUncertainty, InverseVariance]
+)
+@pytest.mark.parametrize("operation", ["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]]
+    unc1 = [[0.1, 0.2], [0.3, 0.4]]
+    unc2 = [[0.2, 0.1], [0.4, 0.3]]
+
+    def make(data, unc, asarray):
+        return CCDData(
+            asarray(data), unit=u.adu, uncertainty=uncertainty_type(asarray(unc))
+        )
+
+    ccd1 = _wrap_ccddata_for_array_api(make(data1, unc1, xp.asarray))
+    ccd2 = _wrap_ccddata_for_array_api(make(data2, unc2, xp.asarray))
+    result = getattr(ccd1, operation)(ccd2, uncertainty_correlation=0.5)
+
+    # Reference: astropy's own propagation on plain numpy CCDData.
+    ref1 = make(data1, unc1, np.asarray)
+    ref2 = make(data2, unc2, np.asarray)
+    expected = getattr(ref1, operation)(ref2, uncertainty_correlation=0.5)
+
+    assert array_api_compat.array_namespace(result.uncertainty.array) is xp
+    assert isinstance(result.uncertainty, uncertainty_type)
+    assert xp.all(
+        xpx.isclose(result.uncertainty.array, xp.asarray(expected.uncertainty.array))
+    )

ccdproc/tests/test_ccdproc.py

@@ -611,7 +611,11 @@ def test_flat_correct():
     size = ccd_data.shape[0]
     # create the flat, with some scatter
     data = 2 * RNG().normal(loc=1.0, scale=0.05, size=(size, size))
-    flat = CCDData(xp.asarray(data), meta=fits.header.Header(), unit=ccd_data.unit)
+    flat = CCDData(
+        xp.asarray(data, device=xp_device),
+        meta=fits.header.Header(),
+        unit=ccd_data.unit,
+    )
     flat_data = flat_correct(ccd_data, flat, add_keyword=None)
 
     # Check that the flat was normalized
@@ -719,7 +723,9 @@ def test_flat_correct_norm_value():
     # the mean of the flat data.
     flat_mean = 5.0
     data = RNG().normal(loc=1.0, scale=0.05, size=ccd_data.shape)
-    flat = CCDData(xp.asarray(data), meta=fits.Header(), unit=ccd_data.unit)
+    flat = CCDData(
+        xp.asarray(data, device=xp_device), meta=fits.Header(), unit=ccd_data.unit
+    )
     flat_data = flat_correct(ccd_data, flat, add_keyword=None, norm_value=flat_mean)
 
     # Check that the flat was normalized
@@ -757,7 +763,7 @@ def test_flat_correct_deviation():
     ccd_data.unit = u.electron
     ccd_data = create_deviation(ccd_data, readnoise=5 * u.electron)
     # Create the flat
-    data = 2 * xp.ones((size, size))
+    data = 2 * xp.ones((size, size), device=xp_device)
     flat = CCDData(data, meta=fits.header.Header(), unit=ccd_data.unit)
     flat = create_deviation(flat, readnoise=0.5 * u.electron)
     ccd_data = flat_correct(ccd_data, flat)
@@ -766,16 +772,16 @@ def test_flat_correct_deviation():
 # Test the uncertainty on the data after flat correction
 def test_flat_correct_data_uncertainty():
     # Regression test for #345
-    # TODO: remove when fix that NDUncertainty explicitly checks
-    # whether the value is a numpy array.
     dat = CCDData(
-        xp.ones([100, 100]), unit="adu", uncertainty=np_array(xp.ones([100, 100]))
+        xp.ones([100, 100]),
+        unit="adu",
+        uncertainty=StdDevUncertainty(xp.ones([100, 100])),
     )
     # Note flat is set to 10, error, if present, is set to one.
     flat = CCDData(10 * xp.ones([100, 100]), unit="adu")
     res = flat_correct(dat, flat)
-    assert (res.data == dat.data).all()
-    assert (res.uncertainty.array == dat.uncertainty.array).all()
+    assert xp.all(res.data == dat.data)
+    assert xp.all(res.uncertainty.array == dat.uncertainty.array)
 
 
 # Tests for gain correction

ccdproc/tests/test_gain.py

@@ -32,12 +32,16 @@ def test_linear_gain_correct(gain):
     if isinstance(gain, Keyword):
         gain = gain.value  # convert to Quantity...
     try:
-        gain_value = gain.value
+        # Make this a python float so that it can be multiplied by an array
+        # from any array namespace.
+        gain_value = float(gain.value)
     except AttributeError:
         gain_value = gain
 
-    xp.all(xpx.isclose(ccd.data, gain_value * orig_data))
-    xp.all(xpx.isclose(ccd.uncertainty.array, gain_value * ccd_data.uncertainty.array))
+    assert xp.all(xpx.isclose(ccd.data, gain_value * orig_data))
+    assert xp.all(
+        xpx.isclose(ccd.uncertainty.array, gain_value * ccd_data.uncertainty.array)
+    )
 
     if isinstance(gain, u.Quantity):
         assert ccd.unit == ccd_data.unit * gain.unit
@@ -57,6 +61,6 @@ def test_linear_gain_unit_keyword():
     gain = 3.0
     gain_unit = u.electron / u.adu
     ccd = gain_correct(ccd_data, gain, gain_unit=gain_unit)
-    xp.all(xpx.isclose(ccd.data, gain * orig_data))
-    xp.all(xpx.isclose(ccd.uncertainty.array, gain * ccd_data.uncertainty.array))
+    assert xp.all(xpx.isclose(ccd.data, gain * orig_data))
+    assert xp.all(xpx.isclose(ccd.uncertainty.array, gain * ccd_data.uncertainty.array))
     assert ccd.unit == ccd_data.unit * gain_unit