Reference: Fourier Space#
Fast Fourier transforms#
- pystella.DFT(decomp, context, queue, grid_shape, dtype, backend=None, **kwargs)[source]#
A wrapper to the creation of various FFT class options which determines whether to use
pystella.fourier.pyclDFT
(for single-device, OpenCL-based FFTs viaVkFFT
orclfft
) orpystella.fourier.pDFT
(for distributed, CPU FFTs viampi4py_fft.mpifft.PFFT
), based on the processor shapeproc_shape
attribute ofdecomp
and a flagbackend
.- Parameters:
decomp – A
DomainDecomposition
.context – A
pyopencl.Context
.queue – A
pyopencl.CommandQueue
.grid_shape – A 3-
tuple
specifying the shape of position-space arrays to be transformed.dtype – The datatype of position-space arrays to be transformed. The complex datatype for momentum-space arrays is chosen to have the same precision.
The following keyword-only arguments are recognized:
- Parameters:
backend – A
str
specifying the FFT backend to use. Valid options are"vkfft"
,"fftw"
, or"clfft"
. Defaults to None, in which case"vkfft"
is selected for single-rankDomainDecomposition
s and"fftw"
otherwise.
Any remaining keyword arguments are passed to
pystella.fourier.pDFT
, should this function return such an object.Changed in version 2020.1: Keyword argument
use_fftw
deprecated in favor ofbackend
.
- class pystella.fourier.dft.BaseDFT[source]#
Base class for all FFT options.
- dft(fx=None, fk=None, **kwargs)[source]#
Computes the forward Fourier transform.
- Parameters:
fx – The array to be transformed. Can be a
pyopencl.array.Array
with or without halo padding (which will be removed bypystella.DomainDecomposition.remove_halos()
if needed) or anumpy.ndarray
without halo padding. Arrays are copied as necessary. Defaults to None, in which casefx
(attached to the transform) is used.fk – The array in which to output the result of the transform. Can be a
pyopencl.array.Array
or anumpy.ndarray
. Arrays are copied as necessary. Defaults to None, in which casefk
(attached to the transform) is used.
- Returns:
The forward Fourier transform of
fx
. Eitherfk
if supplied orfk
.
Any remaining keyword arguments are passed to
forward_transform()
.Note
If you need the result of multiple Fourier transforms at once, you must either supply an
fk
array or copy the output. Namely, without passingfk
the same memory (attached to the transform object) will be used for output, overwriting any prior results.
- idft(fk=None, fx=None, **kwargs)[source]#
Computes the backward Fourier transform.
- Parameters:
fk – The array to be transformed. Can be a
pyopencl.array.Array
or anumpy.ndarray
. Arrays are copied as necessary. Defaults to None, in which casefk
(attached to the transform) is used.fx – The array in which to output the result of the transform. Can be a
pyopencl.array.Array
with or without halo padding (which will be restored bypystella.DomainDecomposition.restore_halos()
if needed) or anumpy.ndarray
without halo padding. Arrays are copied as necessary. Defaults to None, in which casefx
(attached to the transform) is used.
- Returns:
The forward Fourier transform of
fx
. Eitherfk
if supplied orfk
.
Any remaining keyword arguments are passed to
backward_transform()
.Note
If you need the result of multiple Fourier transforms at once, you must either supply an
fx
array or copy the output. Namely, without passingfx
the same memory (attached to the transform object) will be used for output, overwriting any prior results.
- zero_corner_modes(array, only_imag=False)[source]#
Zeros the “corner” modes (modes where each component of its integral wavenumber is either zero or the Nyquist along that axis) of
array
(or just the imaginary part).- Parameters:
array – The array to operate on. May be a
pyopencl.array.Array
or anumpy.ndarray
.only_imag – A
bool
determining whether to only set the imaginary part of the array to zero. Defaults to False, i.e., setting the mode to0+0j
.
- fx#
The (default) position-space array for transforms.
- fk#
The (default) momentum-space array for transforms.
- class pystella.fourier.pDFT(decomp, queue, grid_shape, dtype, **kwargs)[source]#
A wrapper to
mpi4py_fft.mpifft.PFFT
to compute distributed Fast Fourier transforms.See
pystella.fourier.dft.BaseDFT
.- Parameters:
decomp – A
pystella.DomainDecomposition
. The shape of the MPI processor grid is determined by theproc_shape
attribute of this object.queue – A
pyopencl.CommandQueue
.grid_shape – A 3-
tuple
specifying the shape of position-space arrays to be transformed.dtype – The datatype of position-space arrays to be transformed. The complex datatype for momentum-space arrays is chosen to have the same precision.
Any keyword arguments are passed to
mpi4py_fft.mpifft.PFFT
.Changed in version 2020.1: Support for complex-to-complex transforms.
- class pystella.fourier.pyclDFT(decomp, context, queue, grid_shape, dtype, backend)[source]#
A wrapper to
pycl_fft
to compute Fast Fourier transforms withVkFFT
orclFFT
.See
pystella.fourier.dft.BaseDFT
.- Parameters:
decomp – A
pystella.DomainDecomposition
.context – A
pyopencl.Context
.queue – A
pyopencl.CommandQueue
.grid_shape – A 3-
tuple
specifying the shape of position-space arrays to be transformed.dtype – The datatype of position-space arrays to be transformed. The complex datatype for momentum-space arrays is chosen to have the same precision.
backend – Which
pycl_fft
backend to use. One of"vkfft"
or"clfft"
.
New in version 2021.1: Supereseded the now-deprecated
gDFT
.Changed in version 2020.1: Support for complex-to-complex transforms.
Field power spectra#
- class pystella.PowerSpectra(decomp, fft, dk, volume, **kwargs)[source]#
A class for computing power spectra of fields.
- Parameters:
decomp – A
DomainDecomposition
.fft – An FFT object as returned by
DFT()
. The datatype of position-space arrays will match that of the passed FFT object.dk – A 3-
tuple
of the momentum-space grid spacing of each axis (i.e., the infrared cutoff of the grid in each direction).volume – The physical volume of the grid.
The following keyword-only arguments are also recognized:
- Parameters:
bin_width – A
float
specifying the bin width to use. Defaults tomin(dk)
.
- __call__(fx, queue=None, k_power=3, allocator=None)[source]#
Computes the power spectrum of the position-space field
fx
,\[\Delta_f^2(k) = \frac{1}{2 \pi^2 V} \int \mathrm{d} \Omega \, \left\vert \mathbf{k} \right\vert^n \left\vert f(\mathbf{k}) \right\vert^2\]by first Fourier transforming
fx
and then callingbin_power()
.- Parameters:
fx – The array containing the position-space field whose power spectrum is to be computed. If
fx
has more than three axes, all the outer axes are looped over. As an example, iff
has shape(2, 3, 130, 130, 130)
, this method loops over the outermost two axes with shape(2, 3)
, and the resulting output data would have the shape(2, 3, num_bins)
.
The following keyword arguments are recognized:
- Parameters:
queue – A
pyopencl.CommandQueue
. Defaults tofx.queue
.k_power – The exponent \(n\) to use for the weight \(\vert \mathbf{k} \vert^n\). Defaults to 3 (to compute the “dimensionless” power spectrum).
allocator – A
pyopencl
allocator used to allocate temporary arrays, i.e., most usefully apyopencl.tools.MemoryPool
. See the note in the documentation ofSpectralCollocator()
.
- Returns:
A
numpy.ndarray
containing the binned momentum-space power spectrum offx
, with shapefx.shape[:-3]+(num_bins,)
.
- bin_power(fk, queue=None, k_power=3, allocator=None)[source]#
Computes the binned power spectrum of a momentum-space field, weighted by \(k^n\) where
k_power
specifies the value of \(n\).- Parameters:
fk – The array containing the complex-valued, momentum-space field whose power spectrum is to be computed.
The following keyword arguments are recognized:
- Parameters:
queue – A
pyopencl.CommandQueue
. Defaults tofk.queue
.k_power – The exponent \(n\) to use for the weight \(\vert \mathbf{k} \vert^n\). Defaults to 3 (to compute the “dimensionless” power spectrum).
allocator – A
pyopencl
allocator used to allocate temporary arrays, i.e., most usefully apyopencl.tools.MemoryPool
. See the note in the documentation ofSpectralCollocator()
.
- Returns:
A
numpy.ndarray
containing the unnormalized, binned power spectrum offk
.
- polarization(vector, projector, queue=None, k_power=3, allocator=None)[source]#
Computes the power spectra of the plus and minus polarizations of a vector field.
- Parameters:
vector – The array containing the position-space vector field whose power spectrum is to be computed. If
vector
has more than four axes, all the outer axes are looped over. As an example, ifvector
has shape(2, 3, 3, 130, 130, 130)
(where the fourth-to-last axis is the vector-component axis), this method loops over the outermost two axes with shape(2, 3)
, and the resulting output data would have the shape(2, 3, 2, num_bins)
(where the second-to-last axis is the polarization axis).projector – A
Projector
.
The remaining arguments are the same as those to
__call__()
.- Returns:
A
numpy.ndarray
containing the polarization spectra with shapevector.shape[:-4]+(2, num_bins)
.
- gw(hij, projector, hubble, queue=None, k_power=3, allocator=None)[source]#
Computes the spectral abundance of (transverse-traceless) gravitational waves,
\[\Delta_h^2(k) = \frac{1}{24 \pi^{2} \mathcal{H}^{2}} \frac{1}{V} \sum_{i, j} \int \mathrm{d} \Omega \, \left\vert \mathbf{k} \right\vert^3 \left\vert h_{i j}^{\prime}(k) \right \vert^{2}\]- Parameters:
hij – The array containing the position-space tensor field whose power spectrum is to be computed. Must be 4-dimensional, with the first axis being length-6.
projector – A
Projector
.hubble – The current value of the conformal Hubble parameter.
The remaining arguments are the same as those to
__call__()
.- Returns:
A
numpy.ndarray
containing \(\Delta_{h}^2(k)\).
- gw_polarization(hij, projector, hubble, queue=None, k_power=3, allocator=None)[source]#
Computes the spectral abundance of gravitational waves decomposed onto a circular polarization basis.
\[\Delta_{h_\lambda}^2(k) = \frac{1}{24 \pi^{2} \mathcal{H}^{2}} \frac{1}{V} \int \mathrm{d} \Omega \, \left\vert \mathbf{k} \right\vert^3 \left\vert h_\lambda^{\prime}(k) \right \vert^{2}\]- Parameters:
hij – The array containing the position-space tensor field whose power spectrum is to be computed. Must be 4-dimensional, with the first axis being length-6.
projector – A
Projector
.hubble – The current value of the conformal Hubble parameter.
The remaining arguments are the same as those to
__call__()
.- Returns:
A
numpy.ndarray
containing \(\Delta_{h_\lambda}^2(k)\) with shape(2, num_bins)
.
New in version 2020.1.
Changed in version 2020.1: Support for complex fields.
Generating Gaussian-random fields#
- class pystella.RayleighGenerator(context, fft, dk, volume, seed=13298)[source]#
Constructs kernels to generate Gaussian-random fields with a chosen power spectrum in Fourier space by drawing modes according to the corresponding Rayleigh distribution.
- Parameters:
context – A
pyopencl.Context
.fft – An FFT object as returned by
DFT()
. The datatype of position-space arrays will match that of the passed FFT object.dk – A 3-
tuple
of the momentum-space grid spacing of each axis (i.e., the infrared cutoff of the grid in each direction).volume – The physical volume of the grid.
The following keyword-only arguments are recognized:
- Parameters:
seed – The seed to the random number generator. Defaults to
13298
.
- generate(queue, random=True, field_ps=<function RayleighGenerator.<lambda>>, norm=1, window=<function RayleighGenerator.<lambda>>)[source]#
Generate a 3-D array of Fourier modes with a given power spectrum and random phases.
- Parameters:
queue – A
pyopencl.CommandQueue
.random – Whether to randomly sample the Rayleigh distribution of mode amplitudes. Defaults to True.
field_ps – A
Callable
returning the desired power spectrum of the field as a function of momentumkmag
. Defaults to the Bunch-Davies vacuum,lambda kmag: 1/2/kmag
.norm – A constant normalization factor by which to multiply all power spectra. Defaults to
1
.window – A
Callable
window function filtering initial mode amplitudes. Defaults tolambda kmag: 1
, i.e., no filter.
- Returns:
An
numpy.ndarray
containing the generated Fourier modes of the field.
- init_field(fx, queue=None, **kwargs)[source]#
A wrapper which calls
generate()
to initialize a field in Fourier space and returns its inverse Fourier transform.- Parameters:
fx – The array in which the field will be stored.
The following keyword arguments are recognized:
- Parameters:
queue – A
pyopencl.CommandQueue
. Defaults tofx.queue
.
Any additional keyword arguments are passed to
generate()
.
- init_transverse_vector(projector, vector, queue=None, **kwargs)[source]#
A wrapper which calls
generate()
to initialize a transverse three-vector field in Fourier space and returns its inverse Fourier transform. Each component will have the same power spectrum.- Parameters:
projector – A
Projector
used to project out longitudinal components of the vector field.vector – The array in which the vector field will be stored.
The following keyword arguments are recognized:
- Parameters:
queue – A
pyopencl.CommandQueue
. Defaults tovector.queue
.
Any additional keyword arguments are passed to
generate()
.
- init_vector_from_pol(projector, vector, plus_ps, minus_ps, queue=None, **kwargs)[source]#
A wrapper which calls
generate()
to initialize a transverse three-vector field in Fourier space and returns its inverse Fourier transform. In contrast toinit_transverse_vector()
, modes are generated for the plus and minus polarizations of the vector field, from which the vector field itself is constructed.- Parameters:
projector – A
Projector
used to project out longitudinal components of the vector field.vector – The array in which the vector field will be stored.
plus_ps – A
Callable
returning the power spectrum of the plus polarization as a function of momentumkmag
.minus_ps – A
Callable
returning the power spectrum of the minus polarization as a function of momentumkmag
.
The following keyword arguments are recognized:
- Parameters:
queue – A
pyopencl.CommandQueue
. Defaults tovector.queue
.
Any additional keyword arguments are passed to
generate()
.
In addition, the following methods apply the WKB approximation to initialize a field and its (conformal-) time derivative in FLRW spacetime.
- generate_WKB(queue, random=True, field_ps=<function RayleighGenerator.<lambda>>, norm=1, omega_k=<function RayleighGenerator.<lambda>>, hubble=0.0, window=<function RayleighGenerator.<lambda>>)[source]#
Generate a 3-D array of Fourier modes with a given power spectrum and random phases, along with that of its time derivative according to the WKB approximation (for Klein-Gordon fields in conformal FLRW spacetime).
Arguments match those of
generate()
, with the following exceptions/additions:- Parameters:
field_ps – A
Callable
returning the desired power spectrum of the field as a function of \(\omega(k)`\). Defaults to the Bunch-Davies vacuum,lambda wk: 1/2/wk
, wherewk=omega_k(kmag)
.omega_k – A
Callable
defining the dispersion relation of the field. Defaults tolambda kmag: kmag
.hubble – The value of the (conformal) Hubble parameter to use in generating modes for the field’s time derivative. Only used when
WKB=True
. Defaults to0
.
- Returns:
A tuple
(fk, dfk)
containing the generated Fourier modes of the field and its time derivative.
- init_WKB_fields(fx, dfx, queue=None, **kwargs)[source]#
A wrapper which calls
generate_WKB()
to initialize a field and its time derivative in Fourier space and inverse Fourier transform the results.- Parameters:
fx – The array in which the field will be stored.
dfx – The array in which the field’s time derivative will be stored.
queue – A
pyopencl.CommandQueue
. Defaults tofx.queue
.
Any additional keyword arguments are passed to
generate_WKB()
.
Changed in version 2020.1: Support for generating complex fields.
Vector and tensor projections#
- class pystella.Projector(fft, effective_k, dk, dx)[source]#
Constructs kernels to projector vectors and tensors to and from their polarization basis, to project out the longitudinal component of a vector, and to project a tensor field to its transverse and traceless component.
- Parameters:
fft – An FFT object as returned by
DFT()
.grid_shape
anddtype
are determined byfft
’s attributes.effective_k – A
Callable
with signature(k, dx)
returning the effective momentum of the corresponding stencil \(\Delta\), i.e., \(k_\mathrm{eff}\) such that \(\Delta e^{i k x} = i k_\mathrm{eff} e^{i k x}\). That is, projections are implemented relative to the stencil whose eigenvalues (divided by \(i\)) are returned by this function.dk – A 3-
tuple
of the momentum-space grid spacing of each axis (i.e., the infrared cutoff of the grid in each direction).dx – A 3-
tuple
specifying the grid spacing of each axis.
Changed in version 2020.2: Added new required arguments
dk
anddx
.- transversify(queue, vector, vector_T=None)[source]#
Projects out the longitudinal component of a vector field.
- Parameters:
queue – A
pyopencl.CommandQueue
.vector – The array containing the momentum-space vector field to be projected. Must have shape
(3,)+k_shape
, wherek_shape
is the shape of a single momentum-space field array.vector_T – The array in which the resulting projected vector field will be stored. Must have the same shape as
vector
. Defaults to None, in which case the projection is performed in-place.
- Returns:
The
pyopencl.Event
associated with the kernel invocation.
- pol_to_vec(queue, plus, minus, vector)[source]#
Projects the plus and minus polarizations of a vector field onto its vector components.
- Parameters:
queue – A
pyopencl.CommandQueue
.plus – The array containing the momentum-space field of the plus polarization.
minus – The array containing the momentum-space field of the minus polarization.
vector – The array into which the vector field will be stored. Must have shape
(3,)+k_shape
, wherek_shape
is the shape of a single momentum-space field array.
- Returns:
The
pyopencl.Event
associated with the kernel invocation.
- vec_to_pol(queue, plus, minus, vector)[source]#
Projects the components of a vector field onto the basis of plus and minus polarizations.
- Parameters:
queue – A
pyopencl.CommandQueue
.plus – The array into which will be stored the momentum-space field of the plus polarization.
minus – The array into which will be stored the momentum-space field of the minus polarization.
vector – The array whose polarization components will be computed. Must have shape
(3,)+k_shape
, wherek_shape
is the shape of a single momentum-space field array.
- Returns:
The
pyopencl.Event
associated with the kernel invocation.
- transverse_traceless(queue, hij, hij_TT=None)[source]#
Projects a tensor field to be transverse and traceless.
- Parameters:
queue – A
pyopencl.CommandQueue
.hij – The array containing the momentum-space tensor field to be projected. Must have shape
(6,)+k_shape
, wherek_shape
is the shape of a single momentum-space field array.hij_TT – The array in wihch the resulting projected tensor field will be stored. Must have the same shape as
hij
. Defaults to None, in which case the projection is performed in-place.
- Returns:
The
pyopencl.Event
associated with the kernel invocation.
- tensor_to_pol(queue, plus, minus, hij)[source]#
Projects the components of a rank-2 tensor field onto the basis of plus and minus polarizations.
- Parameters:
queue – A
pyopencl.CommandQueue
.plus – The array into which will be stored the momentum-space field of the plus polarization.
minus – The array into which will be stored the momentum-space field of the minus polarization.
hij – The array containing the momentum-space tensor field to be projected. Must have shape
(6,)+k_shape
, wherek_shape
is the shape of a single momentum-space field array.
- Returns:
The
pyopencl.Event
associated with the kernel invocation.
New in version 2020.1.
- pol_to_tensor(queue, plus, minus, hij)[source]#
Projects the plus and minus polarizations of a rank-2 tensor field onto its tensor components.
- Parameters:
queue – A
pyopencl.CommandQueue
.plus – The array into which will be stored the momentum-space field of the plus polarization.
minus – The array into which will be stored the momentum-space field of the minus polarization.
hij – The array containing the momentum-space tensor field to be projected. Must have shape
(6,)+k_shape
, wherek_shape
is the shape of a single momentum-space field array.
- Returns:
The
pyopencl.Event
associated with the kernel invocation.
New in version 2020.1.
Spectral solvers#
- class pystella.SpectralCollocator(fft, dk)[source]#
Interface (analagous to
FiniteDifferencer
) for computing spatial derivatives via spectral collocation.The following arguments are required:
- Parameters:
- __call__(queue, fx, *, lap=None, pdx=None, pdy=None, pdz=None, grd=None, allocator=None)[source]#
Computes requested derivatives of the input
fx
. Provides the same interface aspystella.FiniteDifferencer.__call__()
, while additionally accepting the following arguments:- Parameters:
allocator – A
pyopencl
allocator used to allocate temporary arrays, i.e., most usefully apyopencl.tools.MemoryPool
.
Note
This method allocates extra temporary arrays (when computing more than one derivative), since in-place DFTs are not yet supported. Passing a
pyopencl.tools.MemoryPool
is highly recommended to amortize the cost of memory allocation at each invocation, e.g.:>>> derivs = SpectralCollocator(fft, dk) >>> import pyopencl.tools as clt >>> pool = clt.MemoryPool(clt.ImmediateAllocator(queue)) >>> derivs(queue, fx, lap=lap, allocator=pool)
- divergence(queue, vec, div, allocator=None)[source]#
Computes the divergence of the input
vec
. Provides the same interface aspystella.FiniteDifferencer.divergence()
, while additionally accepting the following arguments:- Parameters:
allocator – A
pyopencl
allocator used to allocate temporary arrays, i.e., most usefully apyopencl.tools.MemoryPool
.
- class pystella.SpectralPoissonSolver(fft, dk, dx, effective_k)[source]#
A Fourier-space solver for the Poisson equation,
\[\nabla^2 f - m^2 f = \rho,\]allowing a term linear in \(f\) with coefficient \(m^2\).
The following arguments are required:
- Parameters:
fft – An FFT object as returned by
DFT()
.grid_shape
anddtype
are determined byfft
’s attributes.dk – A 3-
tuple
of the momentum-space grid spacing of each axis (i.e., the infrared cutoff of the grid in each direction).dx – A 3-
tuple
specifying the grid spacing of each axis.effective_k – A
Callable
with signature(k, dx)
returning the eigenvalue of the second-difference stencil corresponding to \(\nabla^2\). That is, the solver is implemented relative to the stencil whose eigenvalues are returned by this function.
- __call__(queue, fx, rho, m_squared=0, allocator=None)[source]#
Executes the Poisson solver.
- Parameters:
queue – A
pyopencl.CommandQueue
.fx – The array in which the solution \(f\) is stored.
rho – The array containing the right-hand–side, \(\rho\).
m_squared – The value of the coefficient \(m^2\) of the linear term in the Poisson equation to be solved. Defaults to
0
.
Changed in version 2020.1: Added m_squared to support solving with a linear term \(m^2 f\).