Reference: Fourier Space¶
Fast Fourier transforms¶
-
pystella.
DFT
(decomp, context, queue, grid_shape, dtype, **kwargs)[source]¶ A wrapper to the creation of various FFT class options which determines whether to use
pystella.fourier.gDFT
(for single-device, OpenCL-based FFTs viagpyfft
) orpystella.fourier.pDFT
(for distributed, CPU FFTs viampi4py_fft.PFFT
), based on the processor shapeproc_shape
attribute ofdecomp
and a flaguse_fftw
.- 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
use_fftw – A
bool
dictating whether to usepystella.fourier.pDFT
. Defaults to False, i.e., this flag must be set to True to override the default choice to usepystella.fourier.gDFT
on a single rank.
Any remaining keyword arguments are passed to
pystella.fourier.pDFT
, should this function return such an object.
-
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
.
-
-
class
pystella.fourier.
pDFT
(decomp, queue, grid_shape, dtype, **kwargs)[source]¶ A wrapper to
mpi4py_fft.PFFT
to compute distributed Fast Fourier transforms.See
pystella.fourier.dft.BaseDFT
.-
__init__
(decomp, queue, grid_shape, dtype, **kwargs)[source]¶ - 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.PFFT.__init__()
.
Changed in version 2020.1: Support for complex-to-complex transforms.
-
-
class
pystella.fourier.
gDFT
(decomp, context, queue, grid_shape, dtype)[source]¶ A wrapper to
gpyfft
to compute Fast Fourier transforms withclfft
.See
pystella.fourier.dft.BaseDFT
.-
__init__
(decomp, context, queue, grid_shape, dtype)[source]¶ - 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.
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.
-
__init__
(decomp, fft, dk, volume, **kwargs)[source]¶ - 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_with – 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 present, transverse-traceless gravitational wave power spectrum.
\[\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 polarization components of the present gravitational wave power spectrum.
\[\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, **kwargs)[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.
-
__init__
(context, fft, dk, volume, **kwargs)[source]¶ - 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)[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.
-
__init__
(fft, effective_k)[source]¶ - 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.
-
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.-
__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\).
-
__init__
(fft, dk, dx, effective_k)[source]¶ 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\).
-