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 via VkFFT or clfft) or pystella.fourier.pDFT (for distributed, CPU FFTs via mpi4py_fft.mpifft.PFFT), based on the processor shape proc_shape attribute of decomp and a flag backend.

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-rank DomainDecompositions 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 of backend.

class pystella.fourier.dft.BaseDFT[source]#

Base class for all FFT options.

shape(forward_output=True)[source]#

Parameters:: forward_output – A bool specifying whether to output the shape for the result of the forward Fourier transform.
Returns:: A 3-tuple of the (per–MPI-rank) shape of the requested array (as specified by forward_output).

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 by pystella.DomainDecomposition.remove_halos() if needed) or a numpy.ndarray without halo padding. Arrays are copied as necessary. Defaults to None, in which case fx (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 a numpy.ndarray. Arrays are copied as necessary. Defaults to None, in which case fk (attached to the transform) is used.

Returns:

The forward Fourier transform of fx. Either fk if supplied or fk.

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 passing fk 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 a numpy.ndarray. Arrays are copied as necessary. Defaults to None, in which case fk (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 by pystella.DomainDecomposition.restore_halos() if needed) or a numpy.ndarray without halo padding. Arrays are copied as necessary. Defaults to None, in which case fx (attached to the transform) is used.

Returns:

The forward Fourier transform of fx. Either fk if supplied or fk.

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 passing fx 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 a numpy.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 to 0+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 the proc_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 with VkFFT or clFFT.

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 to min(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 calling bin_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, if f 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 to fx.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 a pyopencl.tools.MemoryPool. See the note in the documentation of SpectralCollocator().

Returns:

A numpy.ndarray containing the binned momentum-space power spectrum of fx, with shape fx.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 to fk.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 a pyopencl.tools.MemoryPool. See the note in the documentation of SpectralCollocator().

Returns:

A numpy.ndarray containing the unnormalized, binned power spectrum of fk.

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, if vector 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 shape vector.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 momentum kmag. 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 to lambda 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 to fx.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 to vector.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 to init_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 momentum kmag.
minus_ps – A Callable returning the power spectrum of the minus polarization as a function of momentum kmag.

The following keyword arguments are recognized:

Parameters:: queue – A pyopencl.CommandQueue. Defaults to vector.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, where wk=omega_k(kmag).
omega_k – A Callable defining the dispersion relation of the field. Defaults to lambda 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 to 0.

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 to fx.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 and dtype are determined by fft’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 and dx.

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, where k_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, where k_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, where k_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, where k_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, where k_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, where k_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:

fft – An FFT object as returned by DFT(). grid_shape and dtype are determined by fft’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).

__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 as pystella.FiniteDifferencer.__call__(), while additionally accepting the following arguments:

Parameters:: allocator – A pyopencl allocator used to allocate temporary arrays, i.e., most usefully a pyopencl.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 as pystella.FiniteDifferencer.divergence(), while additionally accepting the following arguments:

Parameters:: allocator – A pyopencl allocator used to allocate temporary arrays, i.e., most usefully a pyopencl.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 and dtype are determined by fft’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\).