Skip to content

Commit

Permalink
Reduced the number of redistributions in spectral helper
Browse files Browse the repository at this point in the history
  • Loading branch information
brownbaerchen committed Oct 8, 2024
1 parent 04853bb commit 556d9a9
Showing 1 changed file with 4 additions and 6 deletions.
10 changes: 4 additions & 6 deletions pySDC/helpers/spectral_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -1550,14 +1550,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):
"""
Expand Down Expand Up @@ -1710,14 +1709,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):
"""
Expand Down

0 comments on commit 556d9a9

Please sign in to comment.