"""Functions making use of spectral decompositions.
.. autosummary::
:nosignatures:
make_colored_noise
.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
"""
from __future__ import annotations
from typing import Callable
import numpy as np
try:
from pyfftw.interfaces.numpy_fft import irfftn as np_irfftn
from pyfftw.interfaces.numpy_fft import rfftn as np_rfftn
except ImportError:
from numpy.fft import irfftn as np_irfftn
from numpy.fft import rfftn as np_rfftn
[docs]
def make_colored_noise(
shape: tuple[int, ...],
dx=1.0,
exponent: float = 0,
scale: float = 1,
rng: np.random.Generator | None = None,
) -> Callable[[], np.ndarray]:
r"""Return a function creating an array of random values that obey.
.. math::
\langle c(\boldsymbol k) c(\boldsymbol k’) \rangle =
\Gamma^2 |\boldsymbol k|^\nu \delta(\boldsymbol k-\boldsymbol k’)
in spectral space on a Cartesian grid. The special case :math:`\nu = 0` corresponds
to white noise. For simplicity, the correlations respect periodic boundary
conditions.
Args:
shape (tuple of ints):
Number of supports points in each spatial dimension. The number of the list
defines the spatial dimension.
dx (float or list of floats):
Discretization along each dimension. A uniform discretization in each
direction can be indicated by a single number.
exponent:
Exponent :math:`\nu` of the power spectrum
scale:
Scaling factor :math:`\Gamma` determining noise strength
rng (:class:`~numpy.random.Generator`):
Random number generator (default: :func:`~numpy.random.default_rng()`)
Returns:
callable: a function returning a random realization
"""
rng = np.random.default_rng(rng)
# extract some information about the grid
dim = len(shape)
dx = np.broadcast_to(dx, (dim,))
if exponent == 0:
# fast case of white noise
def noise_normal():
"""Return array of colored noise."""
return scale * rng.normal(size=shape)
return noise_normal
# deal with colored noise in the following
# prepare wave vectors
k2s = np.array(0)
for i in range(dim):
if i == dim - 1:
k = np.fft.rfftfreq(shape[i], dx[i])
else:
k = np.fft.fftfreq(shape[i], dx[i])
k2s = np.add.outer(k2s, k**2)
# scaling of all modes with k != 0
k2s.flat[0] = 1
scaling = 2 * np.pi * scale * k2s ** (exponent / 4)
scaling.flat[0] = 0
def noise_colored() -> np.ndarray:
"""Return array of colored noise."""
# random field
arr: np.ndarray = rng.normal(size=shape)
# forward transform
arr = np_rfftn(arr)
# scale according to frequency
arr *= scaling
# backwards transform
return np_irfftn(arr, shape) # type: ignore
return noise_colored