Pass axis argument when transforming back from Fourier space

Ensures transposed arrays (which don't necessarily have the 'zeta' dimension last) are handled correctly.
boutproject · Oct 16, 2021 · d01d28c · d01d28c
1 parent fa8b2d4
commit d01d28c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py
@@ -118,7 +118,7 @@ def _shift_z(self, zShift):
 
         data_shifted_fft = data_fft * np.exp(phase.data)
 
-        data_shifted = fft.irfft(data_shifted_fft, n=nz)
+        data_shifted = fft.irfft(data_shifted_fft, n=nz, axis=axis)
 
         # Return a DataArray with the same attributes as self, but values from
         # data_shifted