Provide an 'out' parameter for numpy.fft.fft #25399

serge-sans-paille · 2023-12-15T07:21:00Z

As the first parameter is always copied to the output, it doesn't have much impact performance wise.

It is useful, however, for those who need fine-grain control over memory allocation and cannot afford the cost of a temporary allocation.

serge-sans-paille · 2023-12-15T07:21:51Z

Note: this is just to test the waters. In case of positive feedback, I'll provide the same parameter for other numpy.fft.* functions.

serge-sans-paille · 2023-12-15T14:20:11Z

cc @stefanv , but really that could be anyone :-)

mhvk

I think it would be very useful to have an out argument! However, one needs to take care that it has the right dtype and shape - see in-line comments.

p.s. Looking at the actual code, I'm somewhat surprised it is does not use the iterator, since then that kind of stuff could be dealt with by it (as well as possibly the axis). Indeed, the fft routines would seem easily implemented as a gufunc. Though that may be better done as follow-up!

mhvk · 2023-12-15T15:40:44Z

numpy/fft/_pocketfft.py

    else:
        a = swapaxes(a, axis, -1)
-        r = pfi.execute(a, is_real, is_forward, fct)
+        r = pfi.execute(a, is_real, is_forward, fct, out)


This may be risky - here, the axes of a have been altered but those of out have not. For complex-to-complex, it is possible to copy beforehand, though generally I think it is better to swap the axis for out as well - for real-to-complex, this will be required.

I think you have to do exactly the same operation on out as you do on a. .resize definitely is not right as that can allocate new memory.

The only thing that would seem safe is the following:

if out is not None: out_swapped = swapaxes(out, axis, -1) pfi.execute(a, is_real, is_forward, fct, out_swapped) return out

numpy/fft/_pocketfft.c

mhvk · 2023-12-15T15:51:15Z

numpy/fft/_pocketfft.c

+      if (!data) return NULL;
+    }
+    else {
+      data = (PyArrayObject*)PyArray_EnsureArray(out);


In this branch, one needs to be sure the dtype and shape are both correct. Does PyArray_CopyObject take care of that?

That's my understanding, yes.

serge-sans-paille · 2023-12-17T09:10:50Z

++ extra test cases

serge-sans-paille · 2023-12-17T09:43:34Z

(the macos issue seems unrelated)

serge-sans-paille · 2023-12-19T06:58:16Z

@mhvk gentle ping :-)

mhvk

I fear the code you have is wrong: CopyObject converts data types, so, e.g., someone could have passed in a float32 array as out, to which the input gets copied correctly, but the data is interpreted incorrectly in the calculation below, where it is assumed to be float64.

It also looks like below the data array is assumed to be in C order, which does not have to be the case for an arbitrary out array.

My own sense is that one should write this as a gufunc so that iteration and input/output is done correctly automatically.

mhvk · 2023-12-19T23:42:30Z

I also had a quick look at scipy, which provides PocketFFT as well (but as c++). It has an overwrite_x argument instead of an out one. I like out better, but deviating further from scipy is perhaps not ideal, not sure.

serge-sans-paille · 2023-12-21T09:07:20Z

@mhvk : I've kept the 'out' name to prepare the (future) move to ufunc. I also added the proper checks before copying.

numpy/fft/_pocketfft.c

serge-sans-paille · 2023-12-22T23:33:31Z

Extra checks and tests added, thanks @mhvk for the hints

numpy/fft/_pocketfft.py

As the first parameter is always copied to the output, it doesn't have much impact performance wise. It is useful, however, for those who need fine-grain control over memory allocation and cannot afford the cost of a temporary allocation.

serge-sans-paille · 2024-01-03T10:21:04Z

@mhvk looks good now?

mhvk

As you'll see from the in-line comments, I still think you have a problem here... The main problem is that in every other case in numpy, out will be used to store the result only, with no change to shape or strides, but here for all but axis=-1, you need to swap axes which means that if you pass in something C-contiguous, it will not be C-contiguous afterwards. Although it still would work if one is re-using a previous result, since that is swapped just the same way already. Since this is arguably one of the more important use cases, you could still just go for that (see in-line comment).

Otherwise, I think you are sort-of stuck actually rewriting the current simple loop using the iterator. Though at that point, writing it as a gufunc is almost certainly less work and it would avoid all problems... The one tricky thing there would be to precalculate the fft plan and pass that to the inner loop (via *data).

mhvk · 2024-01-03T15:15:48Z

numpy/fft/_pocketfft.py

    else:
        a = swapaxes(a, axis, -1)
-        r = pfi.execute(a, is_real, is_forward, fct)
+        r = pfi.execute(a, is_real, is_forward, fct, out)


I think you have to do exactly the same operation on out as you do on a. .resize definitely is not right as that can allocate new memory.

The only thing that would seem safe is the following:

if out is not None: out_swapped = swapaxes(out, axis, -1) pfi.execute(a, is_real, is_forward, fct, out_swapped) return out

mhvk · 2024-01-03T15:18:55Z

numpy/fft/tests/test_pocketfft.py

+        # tests below only test the out parameter
+        y = random((30, 20)) + 1j*random((30, 20))
+
+        out = np.zeros_like(x, dtype=complex)


Why pass in dtype=complex here?

mhvk · 2024-01-03T15:19:17Z

numpy/fft/tests/test_pocketfft.py

+        y = random((30, 20)) + 1j*random((30, 20))
+
+        out = np.zeros_like(x, dtype=complex)
+        assert_allclose(fft1(x), np.fft.fft(x, out=out), atol=1e-6)


For all these tests, you need to check too that out is actually returned, i.e., have something like

result = np.fft.fft(x, out=out) assert result is out assert_array_equal(fft1(x), result)

I replaced also with assert_array_equal since, hopefuilly, fft code is reproducible on a given machine!

mhvk · 2024-01-03T15:28:56Z

numpy/fft/_pocketfft.py

+            # This extra copy is unfortunately needed if we want `out`
+            # to retain its original shape while having the correct values.
+            copyto(out, r)
+            r = out


This is contrary to the regular behaviour of out - you must ensure out stays the same object with the same memory layout. Note that in principle, that will be automatic -- if you swapped the axis above, then the swapped case is a view of the original out, so data will just be written in there. I.e., it should be possible to remove this stanza (as in the suggestion I gave above).

mhvk · 2024-01-04T02:54:48Z

@serge-sans-paille - As it seemed hard to get it right without the iterator, I went ahead and tried calling pocketfft from ufuncs. See #25536. I hope to add your tests soon.

mreineck · 2024-01-15T15:08:02Z

Just a small comment: pocketfft could even deal with the situation where the input array and out are the same array (i.e. pointing to the same memory, same shape and strides). This could be pretty useful in some situations, reducing memory consumption and avoiding copying, especially in multi-D transforms.

mreineck · 2024-01-15T15:14:30Z

Sorry, my last comment was confusing, since I was thinking about scipy's variant of pocketfft, not numpy's.

Still, re-using the input array as output should be doable with not too much effort. Not sure whether this is worth it ... it probably depends on the long-term plans for numpy.fft and scipy.fft.

mhvk · 2024-01-15T15:46:43Z

@mreineck - since the C version of pocketfft does the FT in-place, this PR with its out parameter will allow making use of that. I have been wondering whether we should not switch to the C++ version as well, mostly to be able to support float32, but that's for another PR!

mhvk · 2024-01-15T17:19:35Z

@mreineck - sorry, I answered in a PR that makes things less obvious - this one is really superseded by #25536. In fact, let me close this one to avoid further confusion.

serge-sans-paille force-pushed the feature/out-parameter-fft branch from 0b6d5be to 73a1d1c Compare December 15, 2023 07:33

mhvk reviewed Dec 15, 2023

View reviewed changes

serge-sans-paille force-pushed the feature/out-parameter-fft branch 3 times, most recently from 37f057e to b11f2ae Compare December 17, 2023 09:00

mhvk reviewed Dec 19, 2023

View reviewed changes

serge-sans-paille force-pushed the feature/out-parameter-fft branch from b11f2ae to 04f8293 Compare December 21, 2023 08:56

mhvk reviewed Dec 21, 2023

View reviewed changes

numpy/fft/_pocketfft.c Show resolved Hide resolved

serge-sans-paille force-pushed the feature/out-parameter-fft branch from 04f8293 to b67dcc7 Compare December 22, 2023 23:33

mhvk reviewed Dec 23, 2023

View reviewed changes

numpy/fft/_pocketfft.py Outdated Show resolved Hide resolved

serge-sans-paille force-pushed the feature/out-parameter-fft branch from b67dcc7 to 1fd1ce4 Compare December 23, 2023 07:27

Provide an 'out' parameter for numpy.fft.fft

6a17489

As the first parameter is always copied to the output, it doesn't have much impact performance wise. It is useful, however, for those who need fine-grain control over memory allocation and cannot afford the cost of a temporary allocation.

serge-sans-paille force-pushed the feature/out-parameter-fft branch from 1fd1ce4 to 6a17489 Compare December 25, 2023 21:56

mhvk reviewed Jan 3, 2024

View reviewed changes

mhvk mentioned this pull request Jan 4, 2024

MAINT, ENH: Implement calling pocketfft via gufunc and allow out argument #25536

Merged

mhvk closed this Jan 15, 2024

Uh oh!

Provide an 'out' parameter for numpy.fft.fft #25399

Provide an 'out' parameter for numpy.fft.fft #25399

Uh oh!

Conversation

serge-sans-paille commented Dec 15, 2023

Uh oh!

serge-sans-paille commented Dec 15, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

serge-sans-paille commented Dec 15, 2023

Uh oh!

mhvk 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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

serge-sans-paille commented Dec 17, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

serge-sans-paille commented Dec 17, 2023

Uh oh!

serge-sans-paille commented Dec 19, 2023

Uh oh!

mhvk left a comment

Choose a reason for hiding this comment

Uh oh!

mhvk commented Dec 19, 2023

Uh oh!

serge-sans-paille commented Dec 21, 2023

Uh oh!

Uh oh!

serge-sans-paille commented Dec 22, 2023

Uh oh!

Uh oh!

serge-sans-paille commented Jan 3, 2024

Uh oh!

mhvk 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!

Choose a reason for hiding this comment

Uh oh!

mhvk commented Jan 4, 2024

Uh oh!

mreineck commented Jan 15, 2024

Uh oh!

mreineck commented Jan 15, 2024

Uh oh!

mhvk commented Jan 15, 2024

Uh oh!

mhvk commented Jan 15, 2024

Uh oh!

Uh oh!

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

serge-sans-paille commented Dec 15, 2023 •

edited

Loading

serge-sans-paille commented Dec 17, 2023 •

edited

Loading