diff --git a/CHANGES.rst b/CHANGES.rst index a8f57b6c..51377695 100644 --- a/CHANGES.rst +++ b/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 diff --git a/ccdproc/__init__.py b/ccdproc/__init__.py index d95ccefa..81104753 100644 --- a/ccdproc/__init__.py +++ b/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: diff --git a/ccdproc/combiner.py b/ccdproc/combiner.py index 5020f92f..1a7439b2 100644 --- a/ccdproc/combiner.py +++ b/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 diff --git a/ccdproc/tests/test_combiner.py b/ccdproc/tests/test_combiner.py index f7149883..05c8dc1b 100644 --- a/ccdproc/tests/test_combiner.py +++ b/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=)`` 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]