ENH: signal.welch: Add Array API support#25177
Conversation
| welch_extra_note = \ | ||
| """Support for CuPy and JAX is provided by delegation to | ||
| ``cupyx.scipy.signal.welch`` and ``jax.scipy.signal.welch``. | ||
|
|
||
| """ |
There was a problem hiding this comment.
Are we keeping this note?
| assert_allclose(np.sqrt(trapezoid(p_dens, freq)), A*np.sqrt(2)/2, | ||
| rtol=1e-3) | ||
| rms_sq = xp.sum((p_dens[1:] + p_dens[:-1]) | ||
| * (freq[1:] - freq[:-1]) / 2) | ||
| xp_assert_close(xp.asarray(xp.sqrt(rms_sq)), | ||
| xp.asarray(A * np.sqrt(2) / 2), rtol=1e-3) |
There was a problem hiding this comment.
This reimplements trapezoid. I do no think there is support as per data-apis/array-api-extra#490.
j-bowhay
left a comment
There was a problem hiding this comment.
A few comments:
Blocking: The dtype of the expected value shouldn't be, for example, res.dtype.
Non-blocking: please try to minimise unrelated formatting changes in these PRs as they make the diff much larger. I know it's hard to resist the temptation...
| r"nperseg\s*=\s*256 is greater than (?:input length|signal).*", | ||
| UserWarning | ||
| ) | ||
| f, p = welch(x, window='hann') # default nperseg |
There was a problem hiding this comment.
what is the reason for removing the user-specified nperseg
There was a problem hiding this comment.
I removed it because that user-specified nperseg=256 case was failing in the GPU backend CI after the xp conversion at the beginning. Should I re-add it with a backend-specific skip/xfail?
| For single-precision input (float32 or complex64), JAX returns the sample | ||
| frequencies in float32, whereas SciPy and CuPy always return them in float64. |
There was a problem hiding this comment.
| For single-precision input (float32 or complex64), JAX returns the sample | |
| frequencies in float32, whereas SciPy and CuPy always return them in float64. | |
| For single-precision input (``float32`` or ``complex64``), JAX returns the sample | |
| frequencies in ``float32``, whereas SciPy and CuPy always return them in ``float64``. |
Reference issue
@j-bowhay's comment
What does this implement/fix?
signal.welchAdditional information
There are a lot of changes here so
signal.welchmight not have been the best candidate to start with. This is my first PR adding array-api support for a function. Feel free to let me know if this is the correct approach. I asked @ev-br for some feedback while I was working on this.AI Generation Disclosure
I've started rewriting the tests myself and Codex helped me fixing tests locally.