add Array API support to denoise_tv_chambolle

[grlee77]

Description

Here is a demo of converting an existing non-trivial function denoise_tv_chambolle to support any array object implementing the Array API as implemented for the upcoming NumPy 1.22.

Demo code

The following is demo code that can be run with NumPy 1.22 (currently available via pip install -U --pre numpy):

import matplotlib.pyplot as plt
import numpy as np
from numpy import array_api

from skimage import data
from skimage.restoration import denoise_tv_chambolle

# generate a noisy cameraman image
x = data.camera()
rng = np.random.default_rng()
x = x + 20 * rng.standard_normal(x.shape)
x /= 255

weight = 0.05

# first make sure it works for standard NumPy arrays
denoised = denoise_tv_chambolle(x, weight=weight)

# now try with a numpy.array_api array
x_api = array_api.asarray(x)
denoised_array_api = denoise_tv_chambolle(x_api, weight=weight)

# confirm that the results for these two cases are equivalent
np.testing.assert_allclose(x, x_api._array)

# view the result
fig, axes = plt.subplots(1, 2)
axes[0].imshow(x, cmap=plt.cm.gray)
axes[0].set_title('original')
axes[1].imshow(denoised, cmap=plt.cm.gray)
axes[1].set_title('denoised')
for ax in axes:
    ax.set_axis_off()
plt.show()

Some specific points of this conversion:

• np.diff has no equivalent in the array API so had to use explicit difference via slicing instead

• Had to use functions (e.g. xp.sum(image)) instead of calling methods on the array object (e.g. image.sum())

• Had to use expand_dims instead of np.newaxis

• Had to use an ugly if/else case to get an image.astype equivalent working across both np and np.array_api arrays. This is because array_api defines a function numpy.array_api.astype while in numpy it is an astype method on the array.

• Currently only supports floating point images for the Array API case as img_as_float hasn't been converted to support Array API inputs. That could be resolved, though.

Other Limitations

This function is nearly the only one under skimage.restoration that is compatible with acceleration via the Array API. The various deconvolution functions require FFTs and complex dtypes which are not currently in the Array API. The other denoising functions, inpainting and phase unwrapping functions all rely on compiled Cython code and so can't be converted in this manner.

Probably the higher level calibrate_denoiser is the only other functions under skimage.restoration that we could currently convert to use the Array API and even then it would be restricted to use a denoise_function that has Array API support (only denoise_tv_chambolle among the ones we provide)

Outside of skimage.restoration there are some other functions that could be implemented purely using the Array API, but I think this are in the minority (probably below 10% of all functions in the library).

Checklist

• Docstrings for all functions
• Gallery example in ./doc/examples (new features only)
• Benchmark in ./benchmarks, if your changes aren't covered by an
  existing benchmark
• Unit tests
• Clean style in the spirit of PEP8
• Descriptive commit messages (see below)

For reviewers

• Check that the PR title is short, concise, and will make sense 1 year
  later.
• Check that new functions are imported in corresponding __init__.py.
• Check that new features, API changes, and deprecations are mentioned in
  doc/release/release_dev.rst.

[pep8speaks]

Hello @grlee77! Thanks for opening this PR. We checked the lines you've touched for PEP 8 issues, and found:

• In the file skimage/restoration/_denoise.py:

Line 384:7: E221 multiple spaces before operator

[thomasjpfan]

Seeing array_api used is very informative. Thank you!

[thomasjpfan]

At a glance, this is similar to what np.diff does so there should not be a performance issue.

To actually make these workarounds maintainable, I see introducing a utils.array_api module which has utils.array_api.diff, utils.array_api.astype and all other workarounds.

[thomasjpfan]

I think this is reasonable alternative. Although, there are 3 nested xp calls, so I think it is a little less readable.

[github-actions]

Hey, there hasn't been any activity on this for more than 180 days, so we have marked it as "dormant" for now until there is some new activity. You are welcome to reach out to people by mentioning them here or on our forum if you need more feedback! Otherwise, we would be thankful for a short update, for example if you would like to continue or if you are okay with someone else doing so. In any case, thank you for your contributions so far!