From 5fccc9f106d901d932cbc170a557ba5833997ab6 Mon Sep 17 00:00:00 2001 From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com> Date: Wed, 9 Oct 2024 07:33:46 +0200 Subject: [PATCH] Reduced the number of redistributions in spectral helper (#494) --- pySDC/helpers/spectral_helper.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/pySDC/helpers/spectral_helper.py b/pySDC/helpers/spectral_helper.py index 5c45b75871..0fc885a357 100644 --- a/pySDC/helpers/spectral_helper.py +++ b/pySDC/helpers/spectral_helper.py @@ -1560,14 +1560,13 @@ def transform_single_component(self, u, axes=None, padding=None): if self.comm is not None: _out *= np.prod([self.axes[i].N for i in _axes]) - axes_next_base = axes_collapsed[(trf + 1) % len(axes_collapsed)] + axes_next_base = (axes_collapsed + [(-1,)])[trf + 1] alignment = alignment if len(axes_next_base) == 0 else self.ndim + axes_next_base[-1] result = self.get_aligned( _out, axis_in=self.ndim + _axes[0], axis_out=alignment, fft=fft, forward=True, shape=shape ) - fft = self.get_fft(axes=axes, padding=padding) - return self.get_aligned(result, axis_in=alignment, axis_out=self.ndim - 1, fft=fft, forward=True, shape=shape) + return result def transform(self, u, axes=None, padding=None): """ @@ -1720,14 +1719,13 @@ def itransform_single_component(self, u, axes=None, padding=None): else: shape[_ax] = _out.shape[_ax] - axes_next_base = axes_collapsed[(trf + 1) % len(axes_collapsed)] + axes_next_base = (axes_collapsed + [(-1,)])[trf + 1] alignment = alignment if len(axes_next_base) == 0 else self.ndim + axes_next_base[0] result = self.get_aligned( _out, axis_in=self.ndim + _axes[-1], axis_out=alignment, fft=fft, forward=False, shape=shape ) - fft = self.get_fft(axes=axes, padding=padding) - return self.get_aligned(result, axis_in=alignment, axis_out=self.ndim - 1, fft=fft, shape=shape) + return result def get_aligned(self, u, axis_in, axis_out, fft=None, forward=False, **kwargs): """