Adopt the array api #885

mwcraig · 2025-03-19T15:17:48Z

This PR makes the necessary changes to adopt the array API. There are still some things to sort out:

Run tests mostly against numpy, but do one run against jax (or some other non-numpy array library)
Check for remaining invocations of numpy in the code.
Check whether changing .data_arr.mask to .data_arr_mask constitutes an API change. If it does (and I think it does) then patch things up so that .data_arr.mask points to .data_arr_mask
Clean up any remaining handling of immutable arrays (array_api_extra makes handling those easy)

Once those are done the next step will be to find an outside reviewer for this...

codecov · 2025-03-20T15:04:20Z

Codecov Report

Attention: Patch coverage is 92.43499% with 32 lines in your changes missing coverage. Please review.

Please upload report for BASE (main@5aa9219). Learn more about missing BASE report.
Report is 13 commits behind head on main.

Files with missing lines	Patch %	Lines
ccdproc/_ccddata_wrapper_for_array_api.py	86.72%	15 Missing ⚠️
ccdproc/combiner.py	91.91%	11 Missing ⚠️
ccdproc/core.py	97.26%	4 Missing ⚠️
ccdproc/image_collection.py	92.30%	2 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #885   +/-   ##
=======================================
  Coverage        ?   96.11%           
=======================================
  Files           ?        8           
  Lines           ?     1520           
  Branches        ?        0           
=======================================
  Hits            ?     1461           
  Misses          ?       59           
  Partials        ?        0

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

eteq · 2025-06-24T21:51:22Z

As a quick FYI I ran the test suite in this PR on a machine with a NVidia GPU with CCDPROC_ARRAY_LIBRARY=cupy and got a whole pile of errors (sent the log directly to @mwcraig but can also share snippets here if someone else wants to see them)

pyproject.toml

ccdproc/combiner.py

nstarman · 2025-06-25T18:23:58Z

pyproject.toml

@@ -8,17 +8,19 @@ dynamic = ["version"]
 description = "Astropy affiliated package"
 readme = "README.rst"
 license = { text = "BSD-3-Clause" }
-requires-python = ">=3.8"
+requires-python = ">=3.11"


Adopt Spec 0? Has that been adopted?

Good catch!

Wait, misread this -- what is Spectrum 0?

https://scientific-python.org/specs/spec-0000/

So the move here would be to bump it up to 3.12?

3.11 for a few more months. Then 3.12. I was also thinking about adding the badge saying "Spec 0", e.g. you can see in the README in https://github.com/GalacticDynamics/unxt

pyproject.toml

ccdproc/combiner.py

nstarman · 2025-06-26T22:49:40Z

pyproject.toml

@@ -8,17 +8,19 @@ dynamic = ["version"]
 description = "Astropy affiliated package"
 readme = "README.rst"
 license = { text = "BSD-3-Clause" }
-requires-python = ">=3.8"
+requires-python = ">=3.11"


https://scientific-python.org/specs/spec-0000/

ccdproc/core.py

pyproject.toml

These are necessary because upstream dependencies (astropy) explicitly check for numpy arrays rather than using the array API.

Copilot

Pull Request Overview

This PR updates the codebase to use the Array API through array_api_compat and array_api_extra, adds support for alternative array libraries (e.g., JAX, Dask), and bumps minimum Python and NumPy versions.

Introduce xp parameter and array-namespace logic in core routines to support multiple backends
Replace direct NumPy usages with xp calls and use array_api_extra for immutable updates
Update testing infrastructure (tox.ini, pyproject.toml, conftest.py) to run tests against NumPy, JAX, and Dask

Reviewed Changes

Copilot reviewed 18 out of 18 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tox.ini	Added JAX and Dask environments, updated oldestdeps
pyproject.toml	Bumped `requires-python` to ≥3.11; updated dependencies
docs/image_combination.rst	Updated example to use `combiner.mask`
docs/conf.py	Removed Python-version conditional import for `tomllib`
ccdproc/tests/test_rebin.py	Wrapped `rebin` calls in try/except for unsupported libs
ccdproc/tests/test_memory_use.py	Added `USING_NUMPY_ARRAY_LIBRARY` flag and skip logic
ccdproc/tests/test_image_collection.py	Enforced `zip(..., strict=True)` for more robust tests
ccdproc/tests/test_cosmicray.py	Switched to `xp` namespace and NumPy fallbacks
ccdproc/tests/test_combiner.py	Replaced `data_arr` with `_data_arr` and `xp` namespace
ccdproc/tests/test_ccdproc.py	Added `xp` namespace in core tests
ccdproc/tests/run_for_memory_profile.py	Removed unused `combine_uncertainty_function` param
ccdproc/tests/pytest_fixtures.py	Use `array_library.asarray`
ccdproc/log_meta.py	Tightened `zip(..., strict=True)` and updated type check
ccdproc/image_collection.py	Minor doc comment added about NumPy usage
ccdproc/core.py	Introduced `_is_array`, `xp` parameters in key funcs
ccdproc/conftest.py	Environment-based selection of `testing_array_library`
ccdproc/combiner.py	Refactored to use `xp`, `xpx`, and added `xp` params
.github/workflows/ci_tests.yml	Updated matrix to Python 3.11–3.13 and array backends

Comments suppressed due to low confidence (3)

ccdproc/image_collection.py:22

Typo in comment: wantrs should be wants.

# is fine to implement its internal tables however it wantrs, regardless

tox.ini:39

[nitpick] The comment reads "currently oldest support numpy version" but is missing "supported" and is a bit unclear. Consider updating to "# currently oldest supported NumPy version".

    numpy126: numpy==1.26.*  # currently oldest support numpy version

ccdproc/tests/test_memory_use.py:44

The skipif condition references platform but platform is not imported in this file, which will cause a NameError. Please add import platform at the top.

    not platform.startswith("linux"),

nstarman · 2025-07-05T20:12:57Z

ccdproc/combiner.py


 from .core import sigma_func

 __all__ = ["Combiner", "combine"]


-def _default_median():  # pragma: no cover
+def _default_median(xp=None):  # pragma: no cover
    if HAS_BOTTLENECK:


Does Bottleneck work with the array api? If not, would it be better to change the order of this logic and only use bottleneck if the namespace is numpy / None?

Yeah, good call. I can check later what libraries bottlneck works with. Mostly I don't want to change the behavior for current users

nstarman · 2025-07-05T20:13:54Z

ccdproc/combiner.py



-def _default_average():  # pragma: no cover
+def _default_average(xp=None):  # pragma: no cover
    if HAS_BOTTLENECK:


nstarman · 2025-07-05T20:14:04Z

ccdproc/combiner.py


-def _default_sum():  # pragma: no cover
+
+def _default_sum(xp=None):  # pragma: no cover
    if HAS_BOTTLENECK:


nstarman · 2025-07-05T20:14:15Z

ccdproc/combiner.py



-def _default_std():  # pragma: no cover
+def _default_std(xp=None):  # pragma: no cover
    if HAS_BOTTLENECK:


ccdproc/combiner.py

nstarman · 2025-07-05T20:22:47Z

ccdproc/combiner.py

@@ -99,15 +142,12 @@ class Combiner:
                 [ 0.66666667,  0.66666667,  0.66666667,  0.66666667]]...)
    """

-    def __init__(self, ccd_iter, dtype=None):
+    def __init__(self, ccd_iter, dtype=None, xp=None):


Not necessarily for this PR, but the validation logic in the later for-loop can be improved, e.g.

ccd_list = list(ccd_iter) if not all(isinstance(x, CCDData) for x in ccd_list): raise TypeError... default_shape = ccd_list[0].shape if not all(x.shape == default_shape): raise TypeError... default_unit = ...

nstarman · 2025-07-05T20:24:34Z

ccdproc/combiner.py

+        xp = xp or array_api_compat.array_namespace(ccd_list[0].data)
+        self._xp = xp
+        if dtype is None:
+            dtype = xp.float64
        self.ccd_list = ccd_list


Having both ccd_list and _data_arr can double the memory footprint. Are both needed?

ccdproc/log_meta.py

Co-authored-by: Nathaniel Starkman <[email protected]>

nstarman · 2025-07-06T14:53:43Z

ccdproc/conftest.py

+    case _:
+        raise ValueError(
+            f"Unsupported array library: {array_library}. "
+            "Supported libraries are 'numpy' and 'jax'."


incomplete enumeration.

nstarman · 2025-07-06T14:56:41Z

ccdproc/log_meta.py

@@ -134,7 +134,7 @@ def _replace_array_with_placeholder(value):
    return_type_not_value = False
    if isinstance(value, u.Quantity):
        return_type_not_value = not value.isscalar
-    elif isinstance(value, (NDData, np.ndarray)):
+    elif isinstance(value, (NDData, np.ndarray)):  # noqa: UP038


Should this be a expanded to be

Suggested change

elif isinstance(value, (NDData, np.ndarray)): # noqa: UP038

elif isinstance(value, NDData) or hasattr(value, "__array_namespace__"): # noqa: UP038

nstarman · 2025-07-06T16:20:31Z

ccdproc/tests/test_cosmicray.py

-    crflux = 10 * scale * rng.random(NCRAYS) + (threshold + 15) * scale
+    crflux = xp.asarray(10 * scale * rng.random(ncrays) + (threshold + 15) * scale)
+
+    # Some array libraries (Dask) do not support setting individual elements,


Does array_api_extras address this? I haven't looked into the details

nstarman · 2025-07-06T16:49:20Z

ccdproc/tests/test_cosmicray.py

+    data_as_np = np_array(data.data)
    for i in range(ncrays):
        y, x = crrays[i]
-        data.data[y, x] = crflux[i]
+        data_as_np[y, x] = crflux[i]
+    data.data = xp.asarray(data_as_np)


Does this work for vectorized assignment?

data_as_np = np_array(data.data) data_as_np[crrays.T] = crflux data.data = xp.asarray(data_as_np)

nstarman · 2025-07-06T16:50:44Z

ccdproc/tests/test_cosmicray.py

-    noise = DATA_SCALE * np.ones_like(ccd_data.data)
-    ccd_data.uncertainty = noise
+    noise = DATA_SCALE * xp.ones_like(ccd_data.data)
+    # Workaround for the fact that upstream checks for numpy array


upstream is Astropy? So we can fix this to be array-namespace.

One more fix

This is necessary because CCDData does not implement the array API currently.

One more fix FIx the fix

Turns out there were several places where the functions in the array's namespace were not actually being used. This fixes that. Clean up uncertainty wrapper

nstarman

Hi @mwcraig. LMK how I can be of help with reviews. This PR looks great, and it seems to be mostly down to small line-by-line suggestions.

nstarman · 2025-07-17T16:12:42Z

pyproject.toml

@@ -34,6 +36,7 @@ test = [
    "pre-commit",
    "pytest-astropy>=0.10.0",
    "ruff",
+    "jax",


Suggested change

"jax",

"jax>=0.6.3",

or another version.

nstarman · 2025-07-17T16:15:06Z

ccdproc/tests/test_cosmicray.py

-    noise = DATA_SCALE * np.ones_like(ccd_data.data)
-    ccd_data.uncertainty = noise
+    noise = DATA_SCALE * xp.ones_like(ccd_data.data)
+    ccd_data.uncertainty = np_array(noise)


Mark this as a workaround?
It'll take a careful pass to find all these forced numpy arrays when object is eventually made array-api compatible.

mwcraig · 2025-07-17T21:34:55Z

Hi @mwcraig. LMK how I can be of help with reviews. This PR looks great, and it seems to be mostly down to small line-by-line suggestions.

Thanks, @nstarman! I'll want one last look once I've tied down a few loose ends. One thing I learned since the coordination meeting is that using np.testing.assert_allclose is a bad idea because it converts any array-like arguments into numpy arrays for the comparison, hiding any mis-match in array name space between the two things being compared. As a result, there are still a few places that I thought were working that were actually converting the data or uncertainty or mask to numpy arrays.

I've switch to using xp.all(xpx.isclose(a, b)) for numerical array comparisons, which raises an exception if a and b are in different namespaces.

I'll ping you when I'm confident that I have all of those cases fixed up.

mwcraig added this to the 2.5.0 milestone Mar 19, 2025

mwcraig added the enhancement label Mar 19, 2025

mwcraig force-pushed the explore-array-api branch from 85625f0 to 04293b7 Compare March 27, 2025 18:05

mwcraig marked this pull request as ready for review June 24, 2025 14:55

mwcraig changed the title ~~WIP: Adopt the array api~~ Adopt the array api Jun 24, 2025

mwcraig requested a review from Copilot June 24, 2025 16:42

This comment was marked as outdated.

Sign in to view

nstarman reviewed Jun 25, 2025

View reviewed changes

nstarman reviewed Jun 26, 2025

View reviewed changes

ccdproc/core.py Outdated Show resolved Hide resolved

pyproject.toml Show resolved Hide resolved

mwcraig force-pushed the explore-array-api branch 3 times, most recently from 2fa76b7 to 3b2b9fc Compare July 3, 2025 18:28

mwcraig added 15 commits July 5, 2025 10:05

Initial attempt to adopt array API

6dd9166

A couple more changes

d3b6203

Add function for check if something is an array

f0bc6b8

Remove numpy from core

424ebdc

WIP rewrite of a boolean index

7f5050d

WIP testing updates

6a6bfa2

Change more almost_equal to allclose

1d2d171

Use the random number generator from numpy

22bc04d

Remove more in-place array operations

245d641

Set a reasonable tolerance value in float comparisons

55b90be

Continue to use numpy arrays in a few places

23de261

These are necessary because upstream dependencies (astropy) explicitly check for numpy arrays rather than using the array API.

Refactor warning handling for numerical warnings

dffadce

Avoid explicit use of numpy masked arrays

12ff8b0

Rewrite test to not modify array in-place

63f81ba

Use assert_allclose instead of older alternatives

4ae3fc6

mwcraig added 2 commits July 5, 2025 10:05

cast number to float to avoid multiple namespaces

d0bcdad

Change internal data and mask to private properties

baae84f

mwcraig force-pushed the explore-array-api branch from 2e0a544 to baae84f Compare July 5, 2025 15:06

Add properties for accessing the data and mask to be used in combination

655e844

mwcraig requested a review from Copilot July 5, 2025 16:10

Copilot AI reviewed Jul 5, 2025

View reviewed changes

nstarman reviewed Jul 5, 2025

View reviewed changes

mwcraig and others added 2 commits July 5, 2025 16:09

Apply suggestions from code review

17c0dca

Co-authored-by: Nathaniel Starkman <[email protected]>

Choose performance over style

a4e6d4a

nstarman reviewed Jul 6, 2025

View reviewed changes

mwcraig added 3 commits July 12, 2025 13:25

Changes to testing for cupy

5843955

One more fix

Make sure to use CCDData.data instead of CCDData in comparisons

bbc1249

This is necessary because CCDData does not implement the array API currently.

Add more robust handling of open file test

6af1e1e

One more fix FIx the fix

mwcraig force-pushed the explore-array-api branch from 8137626 to 6af1e1e Compare July 12, 2025 20:25

mwcraig added 7 commits July 12, 2025 16:27

Add several workarounds for non-compliance of CCData with Array API

5ec5759

Minor changes

41e17e6

Add and use wrapper classes to ensure array API use

7664bc0

Turns out there were several places where the functions in the array's namespace were not actually being used. This fixes that. Clean up uncertainty wrapper

Fix for immutable array types

fb6e7d6

Ensure array namespace is used throughout core tests

c4c1277

Include more uncertainty types in tests

048f079

Use array API correctly in combiner tests

359126a

nstarman reviewed Jul 17, 2025

View reviewed changes

mwcraig added 7 commits July 17, 2025 17:06

Add array namespace input to combine function

bb42d5b

Add an array namespace conversion in one more place

ba83726

Add argument for desired array namespace

698c6fa

Specify namespace in a couple of tests where data is read from disk

7c7ce31

Fix broken links in docstrings

9a5ccd8

Switch more tests away from np.testing or explain why not switching

4503ddf

WIP array API docs [ci skip]

747d4a4

	elif isinstance(value, (NDData, np.ndarray)): # noqa: UP038
	elif isinstance(value, NDData) or hasattr(value, "__array_namespace__"): # noqa: UP038

Uh oh!

Adopt the array api #885

Are you sure you want to change the base?

Adopt the array api #885

Uh oh!

Conversation

mwcraig commented Mar 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Mar 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

This comment was marked as outdated.

Uh oh!

eteq commented Jun 24, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nstarman Jul 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nstarman left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mwcraig commented Mar 19, 2025 •

edited

Loading

codecov bot commented Mar 20, 2025 •

edited

Loading

nstarman Jul 6, 2025 •

edited

Loading