Source code for pde.tools.spectral

"""
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, Optional

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 """ if rng is None: rng = np.random.default_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 arr = np_irfftn(arr, shape) return arr return noise_colored