def _calc_oa_lens(s1, s2):
# See scipy's documentation in scipy.signal.signaltools
# Set up the arguments for the conventional FFT approach.
fallback = (s1+s2-1, None, s1, s2)
# Use conventional FFT convolve if sizes are same.
if s1 == s2 or s1 == 1 or s2 == 1:
return fallback
# Make s1 the larger size
swapped = s2 > s1
if swapped:
s1, s2 = s2, s1
# There cannot be a useful block size if s2 is more than half of s1.
if s2 >= s1//2:
return fallback
# Compute the optimal block size from the overlap
overlap = s2-1
block_size = fft.next_fast_len(_optimal_oa_block_size(overlap))
# Use conventional FFT convolve if there is only going to be one block.
if block_size >= s1:
return fallback
# Get step size for each of the blocks
in1_step, in2_step = block_size-s2+1, s2
if swapped:
in1_step, in2_step = in2_step, in1_step
return block_size, overlap, in1_step, in2_step
def _freq_domain_conv(in1, in2, axes, shape, calc_fast_len=False):
# See scipy's documentation in scipy.signal.signaltools
real = (in1.dtype.kind != 'c' and in2.dtype.kind != 'c')
fshape = ([fft.next_fast_len(shape[a], real) for a in axes]
if calc_fast_len else shape)
fftn, ifftn = (fft.rfftn, fft.irfftn) if real else (fft.fftn, fft.ifftn)
# Perform the convolution
sp1 = fftn(in1, fshape, axes=axes)
sp2 = fftn(in2, fshape, axes=axes)
out = ifftn(sp1 * sp2, fshape, axes=axes)
return out[tuple(slice(x) for x in shape)] if calc_fast_len else out
def test_next_fast_len(self):
for in_value in range(2000):
out_value = cp_fft.next_fast_len(in_value)
assert self._is_fast_len(out_value)
for i in range(in_value, out_value):
assert not self._is_fast_len(i)
