Implement MPIFFT2D and MPIFFTND by rohanbabbar04 · Pull Request #195 · PyLops/pylops-mpi

rohanbabbar04 · 2026-05-08T19:02:20Z

Implement FFT using mpi4py-fft
Introduce MPIFFTND, MPIFFT2D and _MPIBaseFFTND
Basic Test to compare with pylops operators.
Working example in examples\
Update GA to include fftw libraries.
Add mpi implementations of fftshift and ifftshift.

Note: mpi4py-fft works only with numpy arrays, and PFFT works with multi-dimensional arrays.

rohanbabbar04 · 2026-05-13T05:40:02Z

Worked on implementing fft using mpi4py-fft. Overall, the implementation is straightforward — getting the data into PFFT and retrieving it back into our Distributed Array. There are some key considerations to keep in mind.

Since we have axes as a parameter in the class, we need to ensure that x (the data) is redistributed across the axes[0] axis before pushing it into the DistArray (mpi4py-fft). This is a key step which acts as the starting point.
I tried implementing fftshift and ifftshift. I think the best approach is: if the distribution axis differs from the required axis, we can compute it directly; otherwise, we redistribute to a new axis first and then compute.
PFFT doesn't support 1-D arrays, which makes sense since there is no free axis available for computation. We can explore a workaround for this.
The nffts parameter from the PyLops version is currently missing. PFFT does have a padding parameter, but we will need to investigate how it can be mapped to nffts.
mpi4py-fft works with NumPy arrays, since DistArray is built on top of np.ndarray. For CuPy/NCCL support, we should look into using nvidia.distributed.fft.

Once we agree on the implementation, I will go ahead and update the documentation including the mpi4py-fft.

mrava87

@rohanbabbar04 this is a great addition to pylops-mpi!

I have gone through the PR mostly trying to understand the rationale of your code and the decision made in the implementation - left some comments and questions?

Once we agree on the right approach, I'll do a more focused review on the actual code 😄

mrava87 · 2026-05-13T22:05:47Z

+
+        # Axis along which PFFT decomposes the output array across MPI processes
+        dist_axis = [i for i, s in enumerate(u_hat.subcomm) if s.Get_size() > 1]
+        y = DistributedArray(global_shape=self.dimsd, dtype=self.dtype, axis=dist_axis[0] if dist_axis else 0,


Does this mean that the output could have a different distribution of the input?

Yes, it could be, as it depends on the axes parameter.

@rohanbabbar04 ok, this is not ideal as it is not really predictable.

A user having a global 2d array chunked over axis=0 and flattened and then put into a 1D distributed array may get back a new 1D distributed array that contains local arrays that if reshaped to 2d have chunks over axis=1 and entire axis 0?

If so, how do you see a user knowing this, is there at least a parameter in the FFT operator they can consult?

mrava87 · 2026-06-01T22:26:48Z

@rohanbabbar04 thanks for the updates!

I left a few additional comments (and unresolved some conversations to make sure I fully undestand some of your design choices... should be clear now but please confirm).

Once you have handled these last few points, I think this is good to go!

Next, we can look into enabling N-D Distributed arrays into MPI Linear Operators as in some cases like FFT we could benefit from not having to always distribute the first axis... I had a little play and it seems doable with minor code changes, I may create a draft PR to document what I did so far

rohanbabbar04 · 2026-06-03T03:37:36Z

@rohanbabbar04 thanks for the updates!

I left a few additional comments (and unresolved some conversations to make sure I fully undestand some of your design choices... should be clear now but please confirm).

Once you have handled these last few points, I think this is good to go!

Next, we can look into enabling N-D Distributed arrays into MPI Linear Operators as in some cases like FFT we could benefit from not having to always distribute the first axis... I had a little play and it seems doable with minor code changes, I may create a draft PR to document what I did so far

Thanks @mrava87 for the review, I have added the comments to the respective questions. Do let me know if you need any more information 🙂 ?

mrava87 · 2026-06-03T20:20:51Z

@rohanbabbar04 thanks for the updates!
I left a few additional comments (and unresolved some conversations to make sure I fully undestand some of your design choices... should be clear now but please confirm).
Once you have handled these last few points, I think this is good to go!
Next, we can look into enabling N-D Distributed arrays into MPI Linear Operators as in some cases like FFT we could benefit from not having to always distribute the first axis... I had a little play and it seems doable with minor code changes, I may create a draft PR to document what I did so far

Thanks @mrava87 for the review, I have added the comments to the respective questions. Do let me know if you need any more information 🙂 ?

Quite a few comments unanswered (you may have missed completely my comments from the second review I did... I pinged you in each of them, so you can easily find them 😄

mrava87 · 2026-06-03T20:19:55Z

+
+        # Axis along which PFFT decomposes the output array across MPI processes
+        dist_axis = [i for i, s in enumerate(u_hat.subcomm) if s.Get_size() > 1]
+        y = DistributedArray(global_shape=self.dimsd, dtype=self.dtype, axis=dist_axis[0] if dist_axis else 0,


@rohanbabbar04 ok, this is not ideal as it is not really predictable.

A user having a global 2d array chunked over axis=0 and flattened and then put into a 1D distributed array may get back a new 1D distributed array that contains local arrays that if reshaped to 2d have chunks over axis=1 and entire axis 0?

If so, how do you see a user knowing this, is there at least a parameter in the FFT operator they can consult?

mrava87

Thanks @rohanbabbar04, this looks great now.

Happy with nfft being addressed later in a second PR, and perhaps we will also revisit things a bit when we will be able to pass ND-distributed arrays to solvers... right now I think what you do with internal redistributions and final reflattening with distribution over the axis=0 is the only possible way 😄

I will merge once the CI is finished.

rohanbabbar04 marked this pull request as ready for review May 13, 2026 05:40

rohanbabbar04 requested a review from mrava87 May 13, 2026 05:40

mrava87 requested changes May 13, 2026

View reviewed changes

mrava87 reviewed May 13, 2026

View reviewed changes

Comment thread pylops_mpi/signalprocessing/FFTND.py Outdated

rohanbabbar04 added 12 commits May 25, 2026 00:07

Implement MPIFFT2D and MPIFFTND

1ebf46d

Update GA to include fftw libraries

074489c

Fallback to dims if dimsd does not exist

a7743ca

Add __truediv__

13d35a8

Add example and fftshifts

de364da

Add FFTND and FFT2D to docs

bcdd74f

Use assert_array_allmost_equal and minor changes

c3253de

Add optional dependencies and update code for self.real

7b7e892

Add scale_real_fft and set subcomm_axis=0

f87b5b7

Update documentation

4d76502

Minor change in plot_ffts.py

d78311b

Update notes section

887dad6

rohanbabbar04 force-pushed the fft branch from e133ab6 to 887dad6 Compare May 24, 2026 18:38

mrava87 requested changes Jun 4, 2026

View reviewed changes

rohanbabbar04 added 2 commits June 6, 2026 23:33

Comments update

f37ec59

Remove nffts as a param in BaseFFT, and update notes section

bdfd301

rohanbabbar04 requested a review from mrava87 June 8, 2026 16:57

minor: removed plt.show in plot_ffts

78d3719

mrava87 approved these changes Jun 8, 2026

View reviewed changes

Comment thread examples/plot_ffts.py Outdated

mrava87 merged commit ecf5cf2 into PyLops:main Jun 9, 2026
65 checks passed

rohanbabbar04 deleted the fft branch June 9, 2026 20:25

mrava87 mentioned this pull request Jun 18, 2026

FFTs #110

Closed

Conversation

rohanbabbar04 commented May 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rohanbabbar04 commented May 13, 2026

Uh oh!

mrava87 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mrava87 May 13, 2026

Choose a reason for hiding this comment

Uh oh!

rohanbabbar04 Jun 3, 2026 • 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!

mrava87 Jun 3, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mrava87 commented Jun 1, 2026

Uh oh!

rohanbabbar04 commented Jun 3, 2026

Uh oh!

mrava87 commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mrava87 Jun 3, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mrava87 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

rohanbabbar04 commented May 8, 2026 •

edited

Loading

rohanbabbar04 Jun 3, 2026 •

edited

Loading

mrava87 commented Jun 3, 2026 •

edited

Loading