Support immutable arrays in flat correction and masking

CHANGES.rst

@@ -16,6 +16,8 @@ Bug Fixes
 
 - Exclude masked and clipped pixels, and their weights, when computing weighted
   average combinations. [#952]
+- Make ``flat_correct`` and ``ccdmask`` use functional array updates so they
+  support immutable array-API backends. [#956]
 - Fix dtype conversion in ``Combiner._weighted_sum`` to use the array-API
   namespace form ``xp.astype(weights, xp.float64)`` instead of the deprecated
   string-based ``.astype("float64")``. This resolves the ``DeprecationWarning``

ccdproc/core.py

@@ -971,7 +971,9 @@ def flat_correct(ccd, flat, min_value=None, norm_value=None, xp=None):
     _use_flat = _flat
     if min_value is not None:
         _flat_min = _flat.copy()
-        xpx.at(_flat_min.data)[_flat_min.data < min_value].set(min_value)
+        _flat_min.data = xpx.at(_flat_min.data)[_flat_min.data < min_value].set(
+            min_value
+        )
         _use_flat = _flat_min
 
     # If a norm_value was input and is positive, use it to scale the flat
@@ -1007,7 +1009,7 @@ def flat_correct(ccd, flat, min_value=None, norm_value=None, xp=None):
     # value is set to unity to avoid runtime divide-by-zero errors that are due
     # to a masked value being set to 0.
     if _flat_normed.mask is not None and _flat_normed.mask.any():
-        _flat_normed.data[_flat_normed.mask] = 1.0
+        _flat_normed.data = xpx.at(_flat_normed.data)[_flat_normed.mask].set(1.0)
 
     # divide through the flat
     _flat_corrected = _ccd.divide(_flat_normed, xp=xp, handle_mask=xp.logical_or)
@@ -2250,12 +2252,8 @@ def _sigma_mask(baseline, one_sigma_value, lower_sigma, upper_sigma):
                 # function that first tries percentile in case a particular
                 # array package has it but otherwise falls back to a sort.
                 # This is the case at least as of the 2023.12 API.
-                high = _percentile_fallback(
-                    xp.reshape(block, (xp.prod(block.shape),)), 69.1
-                )
-                low = _percentile_fallback(
-                    xp.reshape(block, (xp.prod(block.shape),)), 30.9
-                )
+                high = _percentile_fallback(xp.reshape(block, (-1,)), 69.1)
+                low = _percentile_fallback(xp.reshape(block, (-1,)), 30.9)
                 block_sigma = (high - low) / 2.0
                 block_mask = _sigma_mask(block, block_sigma, lsigma, hsigma)
                 # mblock = np.ma.MaskedArray(block, mask=block_mask, copy=False)
@@ -2271,17 +2269,18 @@ def _sigma_mask(baseline, one_sigma_value, lower_sigma, upper_sigma):
                         subset = block[:, k]
                         csum.append(xp.sum(subset[~block_mask[:, k]]))
                         all_masked.append(xp.all(block_mask[:, k]))
-                    csum = xp.asarray(csum)
-                    csum[csum <= 0] = 0
+                    csum = xp.stack(csum)
+                    all_masked = xp.stack(all_masked)
+                    csum = xpx.at(csum)[csum <= 0].set(0)
                     csum_sigma = xp.asarray(xp.sqrt(c2 - c1 - csum))
                     # The prior code filled the csum array with the value 1, which
                     # only affects those cases where all of the input values to
                     # the csum were masked, so we fill those with 1.
-                    csum[all_masked] = 1
+                    csum = xpx.at(csum)[all_masked].set(1)
                     colmask = _sigma_mask(csum, csum_sigma, lsigma, hsigma)
-                    block_mask[:, :] |= colmask[xp.newaxis, :]
+                    block_mask = xp.logical_or(block_mask, colmask[xp.newaxis, :])
 
-                mask[l1:l2, c1:c2] = block_mask
+                mask = xpx.at(mask)[l1:l2, c1:c2].set(block_mask)
     else:
         high = ndimage.percentile_filter(medsub, 69.1, size=(nlsig, ncsig))
         low = ndimage.percentile_filter(medsub, 30.9, size=(nlsig, ncsig))
@@ -2298,7 +2297,7 @@ def _sigma_mask(baseline, one_sigma_value, lower_sigma, upper_sigma):
                     for i in range(2, ngood + 2):
                         lend = line + i
                         if mask[lend, col] and not xp.all(mask[line : lend + 1, col]):
-                            mask[line:lend, col] = True
+                            mask = xpx.at(mask)[line:lend, col].set(True)
     return mask
 
 

ccdproc/tests/test_ccdmask.py

@@ -5,9 +5,40 @@
 from astropy.nddata import CCDData
 from numpy.testing import assert_array_equal
 
+from ccdproc.conftest import testing_array_device as xp_device
+from ccdproc.conftest import testing_array_library as xp
 from ccdproc.core import ccdmask
 
 
+# Construct the CCD in the selected test namespace so every backend exercises
+# the functional-update paths, and verify that ccdmask leaves its input intact.
+def _ccdmask_in_active_namespace(data, **kwargs):
+    ratio = CCDData(xp.asarray(data, device=xp_device), unit="adu")
+    data_before = xp.asarray(data, device=xp_device, copy=True)
+
+    result = ccdmask(ratio, xp=xp, **kwargs)
+
+    assert xp.all(ratio.data == data_before)
+    return result
+
+
+def _ones_like_filter(data, size):
+    assert size == (3, 3)
+    return xp.ones_like(data)
+
+
+def _zeros_like_percentile_filter(data, percentile, size):
+    assert percentile in (30.9, 69.1)
+    assert size == (3, 3)
+    return xp.zeros_like(data)
+
+
+def _fixed_percentile(array, percentile):
+    assert percentile in (30.9, 69.1)
+    value = 2 if percentile == 69.1 else -2
+    return array[0] * 0 + value
+
+
 def test_ccdmask_no_ccddata():
     # Fails when a simple list is given.
     with pytest.raises(ValueError):
@@ -215,3 +246,75 @@ def test_ccdmask_pixels():
     mask = ccdmask(ratio, ncsig=11, nlsig=15, findbadcolumns=True)
     target_mask[:, 2] = True
     assert_array_equal(mask, target_mask)
+
+
+@pytest.mark.parametrize("findbadcolumns", [False, True])
+def test_ccdmask_byblocks_with_immutable_array(monkeypatch, findbadcolumns):
+    monkeypatch.setattr(
+        "ccdproc.core.ndimage.median_filter",
+        _ones_like_filter,
+    )
+    monkeypatch.setattr(
+        "ccdproc.core._percentile_fallback",
+        _fixed_percentile,
+    )
+    data = np.ones((8, 8))
+    # In the first 4x4 block, column 0 has an unmasked residual sum of 3.5,
+    # which is large enough for column masking but not per-pixel masking.
+    data[:4, 0] = 1.875
+    # Column 1 has an unmasked residual sum of -20. It must be clamped to zero
+    # before the column test or the whole column would be masked.
+    data[:4, 1] = -4
+    # Column 2 is individually masked, exercising—but not separately
+    # asserting—the all-masked fill in this non-degenerate block.
+    data[:4, 2] = 1000
+
+    mask = _ccdmask_in_active_namespace(
+        data,
+        byblocks=True,
+        findbadcolumns=findbadcolumns,
+        ncsig=4,
+        nlsig=4,
+        ncmed=3,
+        nlmed=3,
+        lsigma=3,
+        hsigma=3,
+        ngood=3,
+    )
+
+    expected = np.zeros(data.shape, dtype=bool)
+    expected[:4, 2] = True
+    if findbadcolumns:
+        expected[:4, 0] = True
+    assert xp.all(mask == xp.asarray(expected, device=xp_device))
+
+
+def test_ccdmask_column_gap_with_immutable_array(monkeypatch):
+    monkeypatch.setattr(
+        "ccdproc.core.ndimage.median_filter",
+        _ones_like_filter,
+    )
+    monkeypatch.setattr(
+        "ccdproc.core.ndimage.percentile_filter",
+        _zeros_like_percentile_filter,
+    )
+    data = np.ones((8, 8))
+    data[1, 3] = 1000
+    data[3, 3] = 1000
+
+    mask = _ccdmask_in_active_namespace(
+        data,
+        byblocks=False,
+        findbadcolumns=True,
+        ncsig=3,
+        nlsig=3,
+        ncmed=3,
+        nlmed=3,
+        lsigma=3,
+        hsigma=3,
+        ngood=4,
+    )
+
+    expected = np.zeros(data.shape, dtype=bool)
+    expected[1:4, 3] = True
+    assert xp.all(mask == xp.asarray(expected, device=xp_device))

ccdproc/tests/test_ccdproc.py

@@ -22,6 +22,7 @@
 from numpy import mgrid as np_mgrid
 from numpy import random as np_random
 
+from ccdproc.conftest import testing_array_device as xp_device
 from ccdproc.conftest import testing_array_library as xp
 from ccdproc.core import (
     Keyword,
@@ -614,19 +615,52 @@ def test_flat_correct():
     assert flat_data.header == ccd_data.header
 
 
+@pytest.mark.backend_xfail(
+    "array-api-strict",
+    reason="the CCDData array-API wrapper passes the Python bool type to "
+    "array-api-strict when copying a mask",
+)
+def test_flat_correct_masked_flat_with_immutable_array():
+    ccd_values = [[8.0, 8.0], [8.0, 8.0]]
+    flat_values = [[2.0, 0.0], [4.0, 8.0]]
+    mask_values = [[False, True], [False, False]]
+    ccd_data = CCDData(xp.asarray(ccd_values, device=xp_device), unit="adu")
+    flat = CCDData(xp.asarray(flat_values, device=xp_device), unit="adu")
+    flat._mask = xp.asarray(mask_values, device=xp_device)
+
+    ccd_before = xp.asarray(ccd_values, device=xp_device)
+    flat_before = xp.asarray(flat_values, device=xp_device)
+    mask_before = xp.asarray(mask_values, device=xp_device)
+
+    result = flat_correct(ccd_data, flat, norm_value=1, add_keyword=None)
+
+    expected_data = xp.asarray([[4.0, 8.0], [2.0, 1.0]], device=xp_device)
+    assert xp.all(xpx.isclose(result.data, expected_data))
+    assert xp.all(result.mask == mask_before)
+    assert xp.all(ccd_data.data == ccd_before)
+    assert xp.all(flat.data == flat_before)
+    assert xp.all(flat.mask == mask_before)
+
+
 # Test for flat correction with min_value
 def test_flat_correct_min_value():
     ccd_data = ccd_data_func()
     size = ccd_data.shape[0]
 
     # Create the flat
     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_orig_data = flat.data.copy()
+    flat = CCDData(
+        xp.asarray(data, device=xp_device),
+        meta=fits.header.Header(),
+        unit=ccd_data.unit,
+    )
+    flat_orig_data = xp.asarray(data, device=xp_device, copy=True)
     min_value = 2.1  # Should replace some, but not all, values
     flat_corrected_data = flat_correct(ccd_data, flat, min_value=min_value)
     flat_with_min = flat.copy()
-    xpx.at(flat_with_min.data)[flat_with_min.data < min_value].set(min_value)
+    flat_with_min.data = xpx.at(flat_with_min.data)[flat_with_min.data < min_value].set(
+        min_value
+    )
 
     # Check that the flat was normalized. The asserts below, which look a
     # little odd, are correctly testing that
@@ -649,7 +683,7 @@ def test_flat_correct_min_value():
     )
 
     # Test that flat is not modified.
-    assert (flat_orig_data == flat.data).all()
+    assert xp.all(flat_orig_data == flat.data)
     assert flat_orig_data is not flat.data