def wiener_deconvolution(image, psf, snr=30, add_pad=0):
""" A GPU accelerated implementation of a linear Wiener filter. Some effort is made
to allow processing even relatively large images, but some kind of block-based processing
(as in the RL implementation) may be required in some cases."""
assert isinstance(image, Image)
assert isinstance(psf, Image)
image_s = Image(image.copy(), image.spacing)
orig_shape = image.shape
if image.ndim != psf.ndim:
raise ValueError("Image and psf dimensions do not match")
if psf.spacing != image.spacing:
psf = imops.zoom_to_spacing(psf, image.spacing)
if add_pad != 0:
new_shape = list(i + 2 * add_pad for i in image_s.shape)
image_s = imops.zero_pad_to_shape(image_s, new_shape)
if psf.shape != image_s.shape:
psf = imops.zero_pad_to_shape(psf, image_s.shape)
psf /= psf.max()
psf = fftshift(psf)
psf_dev = cp.asarray(psf.astype(np.complex64))
with get_fft_plan(psf_dev):
psf_dev = fftn(psf_dev, overwrite_x=True)
below = cp.asnumpy(psf_dev)
psf_abs = cp.abs(psf_dev)**2
psf_abs /= (psf_abs + snr)
above = cp.asnumpy(psf_abs)
psf_abs = None
psf_dev = None
image_dev = cp.asarray(image_s.astype(np.complex64))
with get_fft_plan(image_dev):
image_dev = fftn(image_dev, overwrite_x=True)
wiener_dev = cp.asarray(arrayops.safe_divide(above, below))
image_dev *= wiener_dev
result = cp.asnumpy(cp.abs(ifftn(image_dev, overwrite_x=True)).real)
result = Image(result, image.spacing)
return imops.remove_zero_padding(result, orig_shape)
def rfft2(self,z):
m = z.shape[0]
if self._plan_fft2 is None:
self._plan_fft2 = cpxFFT.get_fft_plan(z + 1j*cp.zeros_like(z), shape=z.shape)
temp = cpxFFT.fft2(z.astype(cp.complex64),shape=z.shape,plan=self._plan_fft2,overwrite_x=True)
temp = cp.fft.fftshift(temp,axes=0)
return temp[m//2 -(self.basis_number-1):m//2 +self.basis_number,:self.basis_number ]/(2*self.extended_basis_number-1)
def __enter__(self):
farplane = cp.empty(
(self.nwaves, self.detector_shape, self.detector_shape),
dtype='complex64')
self.plan = get_fft_plan(farplane, axes=(-2, -1))
del farplane
return self
def ifft2(x):
global plan2D
if plan2D == None and Plan:
plan2D = fftpack.get_fft_plan(x, axes=(-2, -1))
# return
return fftpack.ifft2(x, overwrite_x=owrite, plan=plan2D)
def test_fftn_error_on_wrong_plan(self, dtype, enable_nd):
# This test ensures the context manager plan is picked up
from cupyx.scipy.fftpack import get_fft_plan
from cupy.fft import fftn
assert config.enable_nd_planning == enable_nd
# can't get a plan, so skip
if self.axes is not None:
if self.s is not None:
if len(self.s) != len(self.axes):
return
elif len(self.shape) != len(self.axes):
return
a = testing.shaped_random(self.shape, cupy, dtype)
bad_in_shape = tuple(2 * i for i in self.shape)
if self.s is None:
bad_out_shape = bad_in_shape
else:
bad_out_shape = tuple(2 * i for i in self.s)
b = testing.shaped_random(bad_in_shape, cupy, dtype)
plan_wrong = get_fft_plan(b, bad_out_shape, self.axes)
with pytest.raises(ValueError) as ex, plan_wrong:
fftn(a, s=self.s, axes=self.axes, norm=self.norm)
# targeting a particular error
assert 'The CUFFT plan and a.shape do not match' in str(ex.value)
def fft(x):
global plan1D
#print("Plan1D",plan1D)
if plan1D == None and Plan:
plan1D = fftpack.get_fft_plan(x, axes=(-1))
return fftpack.fft(x, overwrite_x=owrite, plan=plan1D)
def run(self):
max_shape = self._find_max_shape()
# compute FT, assuming they are the same size
fft1 = cp.asarray(self.image1, dtype=cp.complex64)
fft2 = cp.asarray(self.image2, dtype=cp.complex64)
plan = get_fft_plan(fft1, value_type="C2C")
fft1 = fftn(fft1, overwrite_x=True, plan=plan)
fft2 = fftn(fft2, overwrite_x=True, plan=plan)
print(f"shape: {fft1.shape}, dtype: {fft1.dtype}")
@cp.fuse
def normalize(fft_image):
re, im = cp.real(fft_image), cp.imag(fft_image)
length = cp.sqrt(re * re + im * im)
return fft_image / length
fft1 = normalize(fft1)
fft2 = cp.conj(normalize(fft2))
# phase correlation spectrum
pcm = fft1 * fft2
pcm = ifftn(pcm, overwrite_x=True, plan=plan)
pcm = cp.real(pcm)
from skimage.morphology import disk
from skimage.filters import median
pcm = cp.asnumpy(pcm)
pcm = median(pcm, disk(3))
pcm = cp.asarray(pcm)
peak_list = self._extract_correlation_peaks(pcm)
def _get_fft_plan(self, a, axes, **kwargs):
"""Cache multiple FFT plans at the same time."""
key = (*a.shape, *axes)
if key in self.plan_cache:
plan = self.plan_cache[key]
else:
plan = get_fft_plan(a, axes=axes)
self.plan_cache[key] = plan
return plan
def _get_fft_plan(self, a, axes=None, **kwargs):
"""Cache multiple FFT plans at the same time."""
axes = tuple(range(a.ndim)) if axes is None else axes
key = (*a.shape, *axes)
if key in self.plan_cache:
plan = self.plan_cache[key]
else:
plan = get_fft_plan(a, axes=axes)
self.plan_cache[key] = plan
return plan
def cupy_prop(self, signal):
if not signal.is_pol:
raise Exception("only dp signal supported at this time")
step_number = self.length / self.step_length
step_number = int(np.floor(step_number))
temp = np.zeros_like(signal[:])
freq = fftfreq(signal.shape[1], 1 / signal.fs_in_fiber)
freq_gpu = cp.asarray(freq)
omeg = 2 * np.pi * freq_gpu
D = -1j / 2 * self.beta2(signal.center_wavelength) * omeg**2
N = 8 / 9 * 1j * self.gamma
atten = -self.alpha_lin / 2
time_x = cp.asarray(signal[0, :])
time_y = cp.asarray(signal[1, :])
plan = get_fft_plan(time_x)
for i in range(step_number):
time_x, time_y = self.linear_prop_cupy(D, time_x, time_y,
self.step_length / 2, plan)
time_x, time_y = self.nonlinear_prop_cupy(N, time_x, time_y)
time_x = time_x * math.exp(atten * self.step_length)
time_y = time_y * math.exp(atten * self.step_length)
time_x, time_y = self.linear_prop_cupy(D, time_x, time_y,
self.step_length / 2, plan)
last_step = self.length - self.step_length * step_number
last_step_eff = (1 -
np.exp(-self.alpha_lin * last_step)) / self.alpha_lin
if last_step == 0:
time_x = cp.asnumpy(time_x)
time_y = cp.asnumpy(time_y)
temp[0, :] = time_x
temp[1, :] = time_y
return temp
else:
time_x, time_y = self.linear_prop_cupy(D, time_x, time_y,
last_step / 2, plan)
time_x, time_y = self.nonlinear_prop_cupy(N, time_x, time_y,
last_step_eff)
time_x = time_x * math.exp(atten * last_step)
time_y = time_y * math.exp(atten * last_step)
time_x, time_y = self.linear_prop_cupy(D, time_x, time_y,
last_step / 2, plan)
temp[0, :] = cp.asnumpy(time_x)
temp[1, :] = cp.asnumpy(time_y)
return temp
def test_fftn_multiple_plan_error(self, dtype):
import cupy
import cupyx.scipy.fftpack as fftpack
x = testing.shaped_random(self.shape, cupy, dtype)
# hack: avoid testing the cases when getting a cuFFT plan is impossible
if _default_fft_func(x, s=self.s, axes=self.axes) is not _fftn:
return
plan = fftpack.get_fft_plan(x, shape=self.s, axes=self.axes)
with pytest.raises(RuntimeError) as ex, plan:
fftpack.fftn(x, shape=self.s, axes=self.axes, plan=plan)
assert 'Use the cuFFT plan either as' in str(ex.value)
def __init__(self, image, psf, writer, options):
"""
:param image: the image as a Image object
:param psf: the psf as an Image object
:param options: command line options that control the behavior
of the fusion algorithm
"""
deconvolve.DeconvolutionRL.__init__(self, image, psf, writer, options)
self._fft_plan = fftpack.get_fft_plan(cp.zeros(self.block_size, dtype=cp.complex64))
self.__get_fourier_psfs()
def test_fft_multiple_plan_error(self, dtype):
# hack: avoid testing the cases when the output array is of size 0
# because cuFFT and numpy raise different kinds of exceptions
if self.n == 0:
return
import cupy
import cupyx.scipy.fftpack as fftpack
x = testing.shaped_random(self.shape, cupy, dtype)
plan = fftpack.get_fft_plan(x, shape=self.n, axes=self.axis)
with pytest.raises(RuntimeError) as ex, plan:
fftpack.fft(x, n=self.n, axis=self.axis, plan=plan)
assert 'Use the cuFFT plan either as' in str(ex.value)
def irfft2(self,uHalf2D):
"""
Fourier transform of one dimensional signal
ut = 1D signal
num = Ut length - 1
dt = timestep
(now using cp.fft.fft) in the implementation
"""
u2D = util.extend2D(util.symmetrize_2D(uHalf2D),self.extended_basis_number)
if self._plan_ifft2 is None:
self._plan_ifft2 = cpxFFT.get_fft_plan(u2D, shape=u2D.shape)
u2D = cp.fft.ifftshift(u2D)
temp = cpxFFT.ifft2(u2D,shape=u2D.shape,plan=self._plan_ifft2,overwrite_x=True)
return temp.real*(2*self.extended_basis_number-1)
def test_ifftn(self, xp, dtype, enable_nd):
assert config.enable_nd_planning == enable_nd
a = testing.shaped_random(self.shape, xp, dtype)
if xp == cupy:
from cupyx.scipy.fftpack import get_fft_plan
plan = get_fft_plan(a, self.s, self.axes)
with plan:
out = xp.fft.ifftn(a, s=self.s, axes=self.axes, norm=self.norm)
else:
out = xp.fft.ifftn(a, s=self.s, axes=self.axes, norm=self.norm)
if xp == np and dtype is np.complex64:
out = out.astype(np.complex64)
return out
def test_rfftn(self, xp, dtype, enable_nd):
assert config.enable_nd_planning == enable_nd
a = testing.shaped_random(self.shape, xp, dtype)
if xp is cupy:
from cupyx.scipy.fftpack import get_fft_plan
plan = get_fft_plan(a, self.s, self.axes, value_type='R2C')
with plan:
out = xp.fft.rfftn(a, s=self.s, axes=self.axes, norm=self.norm)
else:
out = xp.fft.rfftn(a, s=self.s, axes=self.axes, norm=self.norm)
if xp is np and dtype in [np.float16, np.float32, np.complex64]:
out = out.astype(np.complex64)
return out
def __init__(self, data, writer, options):
"""
:param data: a ImageData object
:param options: command line options that control the behavior
of the fusion algorithm
:param writer: a writer object that can save intermediate results
"""
fusion.MultiViewFusionRL.__init__(self, data, writer, options)
padded_block_size = self.block_size + 2 * self.options.block_pad
self._fft_plan = fftpack.get_fft_plan(
cp.zeros(padded_block_size, dtype=cp.complex64))
self.__get_fourier_psfs()
def test_fft_error_on_wrong_plan(self, dtype):
# This test ensures the context manager plan is picked up
from cupyx.scipy.fftpack import get_fft_plan
from cupy.fft import fft
a = testing.shaped_random(self.shape, cupy, dtype)
bad_shape = tuple(5 * i for i in self.shape)
b = testing.shaped_random(bad_shape, cupy, dtype)
plan_wrong = get_fft_plan(b)
assert isinstance(plan_wrong, cupy.cuda.cufft.Plan1d)
with pytest.raises(ValueError) as ex, plan_wrong:
fft(a, n=self.n, norm=self.norm)
# targeting a particular error
assert 'Target array size does not match the plan.' in str(ex.value)
def test_ifft(self, xp, dtype):
a = testing.shaped_random(self.shape, xp, dtype)
if xp == cupy:
from cupyx.scipy.fftpack import get_fft_plan
shape = (self.n, ) if self.n is not None else None
plan = get_fft_plan(a, shape=shape)
assert isinstance(plan, cupy.cuda.cufft.Plan1d)
with plan:
out = xp.fft.ifft(a, n=self.n, norm=self.norm)
else:
out = xp.fft.ifft(a, n=self.n, norm=self.norm)
if xp == np and dtype is np.complex64:
out = out.astype(np.complex64)
return out
def test_irfft(self, xp, dtype):
a = testing.shaped_random(self.shape, xp, dtype)
if xp is cupy:
from cupyx.scipy.fftpack import get_fft_plan
shape = (self.n,) if self.n is not None else None
plan = get_fft_plan(a, shape=shape, value_type='C2R')
assert isinstance(plan, cupy.cuda.cufft.Plan1d)
with plan:
out = xp.fft.irfft(a, n=self.n, norm=self.norm)
else:
out = xp.fft.irfft(a, n=self.n, norm=self.norm)
if xp is np and dtype in [np.float16, np.float32, np.complex64]:
out = out.astype(np.float32)
return out
def __init__(self, context, data, axes):
self.context = context
self.axes = axes
assert len(data.shape) > max(axes)
from cupyx.scipy import fftpack as cufftp
if data.flags.f_contiguous:
self._ax = [data.ndim - 1 - aa for aa in axes]
_dat = data.T
self.f_contiguous = True
else:
self._ax = axes
_dat = data
self.f_contiguous = False
self._fftplan = cufftp.get_fft_plan(_dat,
axes=self._ax,
value_type="C2C")
def ifft(x):
global plan1D
if plan1D == None and Plan:
plan1D = fftpack.get_fft_plan(x, axes=(-1))
return fftpack.ifft(x, overwrite_x=owrite, plan=plan1D)
def channelize_poly(x, h, n_chans):
"""
Polyphase channelize signal into n channels
Parameters
----------
x : array_like
The input data to be channelized
h : array_like
The 1-D input filter; will be split into n
channels of int number of taps
n_chans : int
Number of channels for channelizer
Returns
----------
yy : channelized output matrix
Notes
----------
Currently only supports simple channelizer where channel
spacing is equivalent to the number of channels used (zero overlap).
Number of filter taps (len of filter / n_chans) must be <=32.
"""
dtype = cp.promote_types(x.dtype, h.dtype)
x = asarray(x, dtype=dtype)
h = asarray(h, dtype=dtype)
# number of taps in each h_n filter
n_taps = int(len(h) / n_chans)
if n_taps > 32:
raise NotImplementedError(
"The number of calculated taps ({}) in \
each filter is currently capped at 32. Please reduce filter \
length or number of channels".format(
n_taps
)
)
if n_taps > 32:
raise NotImplementedError(
"Number of taps ({}) must be less than (32).".format(n_taps)
)
# number of outputs
n_pts = int(len(x) / n_chans)
if x.dtype == cp.float32 or x.dtype == cp.complex64:
y = cp.empty((n_pts, n_chans), dtype=cp.complex64)
elif x.dtype == cp.float64 or x.dtype == cp.complex128:
y = cp.empty((n_pts, n_chans), dtype=cp.complex128)
_channelizer(x, h, y, n_chans, n_taps, n_pts)
# Remove with CuPy v8
if (x.dtype, n_pts, n_taps, n_chans) in _cupy_fft_cache:
plan = _cupy_fft_cache[(x.dtype, n_pts, n_taps, n_chans)]
else:
plan = _cupy_fft_cache[
(x.dtype, n_pts, n_taps, n_chans)
] = fftpack.get_fft_plan(y, axes=-1)
return cp.conj(fftpack.fft(y, overwrite_x=True, plan=plan)).T
def fftconvolve(in1, in2, mode="full", axes=None):
"""Convolve two N-dimensional arrays using FFT.
Convolve `in1` and `in2` using the fast Fourier transform method, with
the output size determined by the `mode` argument.
This is generally much faster than `convolve` for large arrays (n > ~500),
but can be slower when only a few output values are needed, and can only
output float arrays (int or object array inputs will be cast to float).
As of v0.19, `convolve` automatically chooses this method or the direct
method based on an estimation of which is faster.
Parameters
----------
in1 : array_like
First input.
in2 : array_like
Second input. Should have the same number of dimensions as `in1`.
mode : str {'full', 'valid', 'same'}, optional
A string indicating the size of the output:
``full``
The output is the full discrete linear convolution
of the inputs. (Default)
``valid``
The output consists only of those elements that do not
rely on the zero-padding. In 'valid' mode, either `in1` or `in2`
must be at least as large as the other in every dimension.
``same``
The output is the same size as `in1`, centered
with respect to the 'full' output.
axis : tuple, optional
axes : int or array_like of ints or None, optional
Axes over which to compute the convolution.
The default is over all axes.
Returns
-------
out : array
An N-dimensional array containing a subset of the discrete linear
convolution of `in1` with `in2`.
Examples
--------
Autocorrelation of white noise is an impulse.
>>> import cusignal
>>> import cupy as cp
>>> import numpy as np
>>> sig = cp.random.randn(1000)
>>> autocorr = cusignal.fftconvolve(sig, sig[::-1], mode='full')
>>> import matplotlib.pyplot as plt
>>> fig, (ax_orig, ax_mag) = plt.subplots(2, 1)
>>> ax_orig.plot(cp.asnumpy(sig))
>>> ax_orig.set_title('White noise')
>>> ax_mag.plot(np.arange(-len(sig)+1,len(sig)), autocorr)
>>> ax_mag.set_title('Autocorrelation')
>>> fig.tight_layout()
>>> fig.show()
Gaussian blur implemented using FFT convolution. Notice the dark borders
around the image, due to the zero-padding beyond its boundaries.
The `convolve2d` function allows for other types of image boundaries,
but is far slower.
>>> from scipy import misc
>>> face = misc.face(gray=True)
>>> kernel = cp.outer(cusignal.gaussian(70, 8), cusignal.gaussian(70, 8))
>>> blurred = cusignal.fftconvolve(face, kernel, mode='same')
>>> fig, (ax_orig, ax_kernel, ax_blurred) = plt.subplots(3, 1,
... figsize=(6, 15))
>>> ax_orig.imshow(face, cmap='gray')
>>> ax_orig.set_title('Original')
>>> ax_orig.set_axis_off()
>>> ax_kernel.imshow(cp.asnumpy(kernel), cmap='gray')
>>> ax_kernel.set_title('Gaussian kernel')
>>> ax_kernel.set_axis_off()
>>> ax_blurred.imshow(cp.asnumpy(blurred), cmap='gray')
>>> ax_blurred.set_title('Blurred')
>>> ax_blurred.set_axis_off()
>>> fig.show()
"""
in1 = cp.ascontiguousarray(in1)
in2 = cp.ascontiguousarray(in2)
noaxes = axes is None
if in1.ndim == in2.ndim == 0: # scalar inputs
return in1 * in2
elif in1.ndim != in2.ndim:
raise ValueError("in1 and in2 should have the same dimensionality")
elif in1.size == 0 or in2.size == 0: # empty arrays
return cp.array([])
_, axes = _init_nd_shape_and_axes_sorted(in1, shape=None, axes=axes)
# axes needs to be numpy type for proper execution in FFT
axes = cp.asnumpy(axes)
if not noaxes and not axes.size:
raise ValueError("when provided, axes cannot be empty")
if noaxes:
other_axes = np.array([], dtype=cp.intc)
else:
other_axes = np.setdiff1d(np.arange(in1.ndim), axes)
s1 = np.array(in1.shape)
s2 = np.array(in2.shape)
if not np.all((s1[other_axes] == s2[other_axes])
| (s1[other_axes] == 1)
| (s2[other_axes] == 1)):
raise ValueError("incompatible shapes for in1 and in2:"
" {0} and {1}".format(in1.shape, in2.shape))
complex_result = np.issubdtype(in1.dtype,
np.complexfloating) or np.issubdtype(
in2.dtype, cp.complexfloating)
shape = np.maximum(s1, s2)
shape[axes] = s1[axes] + s2[axes] - 1
# Check that input sizes are compatible with 'valid' mode
if _inputs_swap_needed(mode, s1, s2):
# Convolution is commutative; order doesn't have any effect on output
in1, s1, in2, s2 = in2, s2, in1, s1
# Speed up FFT by padding to optimal size for FFTPACK
fshape = [next_fast_len(d) for d in shape[axes]]
fslice = tuple([slice(sz) for sz in shape])
if not complex_result:
if (str(fshape), str(axes), "R2C") in _cupy_fft_cache:
rplan = _cupy_fft_cache[(str(fshape), str(axes), "R2C")]
else:
rplan = _cupy_fft_cache[(str(fshape), str(axes),
"R2C")] = fftpack.get_fft_plan(
in1,
fshape,
axes=axes,
value_type="R2C")
try:
with rplan:
sp1 = cp.fft.rfftn(in1, fshape, axes=axes)
sp2 = cp.fft.rfftn(in2, fshape, axes=axes)
except Exception:
sp1 = cp.fft.rfftn(in1, fshape, axes=axes)
sp2 = cp.fft.rfftn(in2, fshape, axes=axes)
ret = cp.fft.irfftn(sp1 * sp2, fshape, axes=axes)[fslice].copy()
else:
# Need to move to cupyx.scipy.fft with CuPy v8
if (str(fshape), str(axes)) in _cupy_fft_cache:
plan = _cupy_fft_cache[(str(fshape), str(axes))]
else:
plan = _cupy_fft_cache[(str(fshape),
str(axes))] = fftpack.get_fft_plan(
in1, fshape, axes=axes)
try:
with plan:
sp1 = fftpack.fftn(in1, fshape, axes=axes)
sp2 = fftpack.fftn(in2, fshape, axes=axes)
ret = fftpack.ifftn(sp1 * sp2, axes=axes)[fslice].copy()
except Exception:
sp1 = fftpack.fftn(in1, fshape, axes=axes)
sp2 = fftpack.fftn(in2, fshape, axes=axes)
ret = fftpack.ifftn(sp1 * sp2, axes=axes)[fslice].copy()
if mode == "full":
return ret
elif mode == "same":
return _centered(ret, s1)
elif mode == "valid":
shape_valid = shape.copy()
shape_valid[axes] = s1[axes] - s2[axes] + 1
return _centered(ret, shape_valid)
else:
raise ValueError("acceptable mode flags are \
'valid',"
" 'same', or 'full'")
def fiber_cupy(param: FiberParam, samples: np.ndarray, fs: float,
center_wavelength: float, device: str):
import cupy as np
from cupyx.scipy.fftpack import fft as improved_fft
from cupyx.scipy.fftpack import ifft as improved_ifft
from cupyx.scipy.fftpack import get_fft_plan
samples = np.array(samples)
plan = get_fft_plan(samples[0])
def step(samples, step_length, step_eff):
xpol = samples[0]
ypol = samples[1]
xpol, ypol = linear_prop(xpol, ypol, step_length / 2)
xpol, ypol = nonlinear_prop(xpol, ypol, step_eff)
xpol, ypol = atteunation(xpol, ypol, step_length)
xpol, ypol = linear_prop(xpol, ypol, step_length / 2)
return np.vstack((xpol, ypol))
def atteunation(xpol, ypol, step_length):
xpol = xpol * np.exp(atten * step_length)
ypol = ypol * np.exp(atten * step_length)
return xpol, ypol
def nonlinear_prop(xpol, ypol, step_length):
amplitude_xpol = (xpol[:, 0]**2 + xpol[:, 1]**2)
amplitude_ypol = (ypol[:, 0]**2 + ypol[:, 1]**2)
phase_rotation_xpol = 8 / 9 * (amplitude_xpol +
amplitude_ypol) * gamma * step_length
phase_rotation_ypol = phase_rotation_xpol
xpol = xpol * np.exp(1j * 2 * np.pi * phase_rotation_xpol)
ypol = ypol * np.exp(1j * 2 * np.pi * phase_rotation_ypol)
return xpol, ypol
def linear_prop(xpol, ypol, length):
frequency_domain_xpol = improved_fft(xpol, overwrite_x=True, plan=plan)
frequency_domain_ypol = improved_fft(ypol, overwrite_x=True, plan=plan)
frequency_domain_xpol = frequency_domain_xpol * np.exp(1j * D * length)
frequency_domain_ypol = frequency_domain_ypol * np.exp(1j * D * length)
xpol_time_domain = improved_ifft(frequency_domain_xpol,
overwrite_x=True,
plan=plan)
ypol_time_domain = improved_ifft(frequency_domain_ypol,
overwrite_x=True,
plan=plan)
return xpol_time_domain, ypol_time_domain
gamma = param.gamma
length = param.length
step_length = param.step_length
xpol = samples[0]
freq = np.fft.fftfreq(len(xpol), 1 / fs)
omeg = 2 * np.pi * freq
D = -1 / 2 * param.beta2(center_wavelength) * omeg**2
atten = -param.alphalin / 2
step_number = length / step_length
step_number = int(np.floor(step_number))
last_length = length - step_number * step_length
step_eff = 1 - np.exp(-param.alphalin * step_length)
step_eff = step_eff / param.alphalin
last_length_eff = 1 - np.exp(-param.alphalin * last_length)
last_length_eff = last_length_eff / param.alphalin
for step_index in range(step_number):
samples = step(samples, step_length, step_eff)
if last_length_eff:
samples = step(samples, last_length, last_length_eff)
return samples
