Skip to content

Gaussian Random Field¤

exponax.ic.GaussianRandomField ¤

Bases: BaseRandomICGenerator

Source code in exponax/ic/_gaussian_random_field.py
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
class GaussianRandomField(BaseRandomICGenerator):
    num_spatial_dims: int
    domain_extent: float
    powerlaw_exponent: float
    zero_mean: bool
    std_one: bool
    max_one: bool
    white_noise: WhiteNoise

    def __init__(
        self,
        num_spatial_dims: int,
        *,
        domain_extent: float = 1.0,
        powerlaw_exponent: float = 3.0,
        zero_mean: bool = True,
        std_one: bool = False,
        max_one: bool = False,
    ):
        """
        Random generator for initial states following a power-law spectrum in
        Fourier space, i.e., it decays polynomially with the wavenumber.

        **Arguments:**

        - `num_spatial_dims`: The number of spatial dimensions.
        - `domain_extent`: The extent of the domain in each spatial direction.
        - `powerlaw_exponent`: The exponent of the power-law spectrum.
        - `zero_mean`: Whether the field should have zero mean.
        - `std_one`: Whether to normalize the state to have a standard
            deviation of one. Defaults to `False`. Only works if the offset is
            zero.
        - `max_one`: Whether to normalize the state to have the maximum
            absolute value of one. Defaults to `False`. Only one of `std_one`
            and `max_one` can be `True`.
        """
        validate_normalization_options(
            zero_mean=zero_mean, std_one=std_one, max_one=max_one
        )
        self.num_spatial_dims = num_spatial_dims
        self.domain_extent = domain_extent
        self.powerlaw_exponent = powerlaw_exponent
        self.zero_mean = zero_mean
        self.std_one = std_one
        self.max_one = max_one
        self.white_noise = WhiteNoise(num_spatial_dims)

    def __call__(
        self, num_points: int, *, key: PRNGKeyArray
    ) -> Float[Array, "1 ... N"]:
        noise = self.white_noise(num_points, key=key)

        noise_hat = fft(noise, num_spatial_dims=self.num_spatial_dims)

        wavenumber_grid = build_scaled_wavenumbers(
            self.num_spatial_dims, self.domain_extent, num_points
        )
        wavenumber_norm_grid = jnp.linalg.norm(wavenumber_grid, axis=0, keepdims=True)
        # Further division by 2.0 in the exponent is because we want to have the
        # **power-spectrum** follow a **power-law**. See
        # https://github.com/Ceyron/exponax/issues/9 for more details.
        amplitude = jnp.power(wavenumber_norm_grid, -self.powerlaw_exponent / 2.0)
        amplitude = (
            amplitude.flatten().at[0].set(1.0).reshape(wavenumber_norm_grid.shape)
        )

        noise_hat = noise_hat * amplitude

        ic = ifft(
            noise_hat, num_spatial_dims=self.num_spatial_dims, num_points=num_points
        )

        ic = normalize_ic(
            ic, zero_mean=self.zero_mean, std_one=self.std_one, max_one=self.max_one
        )

        return ic
__init__ ¤
__init__(
    num_spatial_dims: int,
    *,
    domain_extent: float = 1.0,
    powerlaw_exponent: float = 3.0,
    zero_mean: bool = True,
    std_one: bool = False,
    max_one: bool = False
)

Random generator for initial states following a power-law spectrum in Fourier space, i.e., it decays polynomially with the wavenumber.

Arguments:

  • num_spatial_dims: The number of spatial dimensions.
  • domain_extent: The extent of the domain in each spatial direction.
  • powerlaw_exponent: The exponent of the power-law spectrum.
  • zero_mean: Whether the field should have zero mean.
  • std_one: Whether to normalize the state to have a standard deviation of one. Defaults to False. Only works if the offset is zero.
  • max_one: Whether to normalize the state to have the maximum absolute value of one. Defaults to False. Only one of std_one and max_one can be True.
Source code in exponax/ic/_gaussian_random_field.py
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def __init__(
    self,
    num_spatial_dims: int,
    *,
    domain_extent: float = 1.0,
    powerlaw_exponent: float = 3.0,
    zero_mean: bool = True,
    std_one: bool = False,
    max_one: bool = False,
):
    """
    Random generator for initial states following a power-law spectrum in
    Fourier space, i.e., it decays polynomially with the wavenumber.

    **Arguments:**

    - `num_spatial_dims`: The number of spatial dimensions.
    - `domain_extent`: The extent of the domain in each spatial direction.
    - `powerlaw_exponent`: The exponent of the power-law spectrum.
    - `zero_mean`: Whether the field should have zero mean.
    - `std_one`: Whether to normalize the state to have a standard
        deviation of one. Defaults to `False`. Only works if the offset is
        zero.
    - `max_one`: Whether to normalize the state to have the maximum
        absolute value of one. Defaults to `False`. Only one of `std_one`
        and `max_one` can be `True`.
    """
    validate_normalization_options(
        zero_mean=zero_mean, std_one=std_one, max_one=max_one
    )
    self.num_spatial_dims = num_spatial_dims
    self.domain_extent = domain_extent
    self.powerlaw_exponent = powerlaw_exponent
    self.zero_mean = zero_mean
    self.std_one = std_one
    self.max_one = max_one
    self.white_noise = WhiteNoise(num_spatial_dims)
__call__ ¤
__call__(
    num_points: int, *, key: PRNGKeyArray
) -> Float[Array, "1 ... N"]
Source code in exponax/ic/_gaussian_random_field.py
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def __call__(
    self, num_points: int, *, key: PRNGKeyArray
) -> Float[Array, "1 ... N"]:
    noise = self.white_noise(num_points, key=key)

    noise_hat = fft(noise, num_spatial_dims=self.num_spatial_dims)

    wavenumber_grid = build_scaled_wavenumbers(
        self.num_spatial_dims, self.domain_extent, num_points
    )
    wavenumber_norm_grid = jnp.linalg.norm(wavenumber_grid, axis=0, keepdims=True)
    # Further division by 2.0 in the exponent is because we want to have the
    # **power-spectrum** follow a **power-law**. See
    # https://github.com/Ceyron/exponax/issues/9 for more details.
    amplitude = jnp.power(wavenumber_norm_grid, -self.powerlaw_exponent / 2.0)
    amplitude = (
        amplitude.flatten().at[0].set(1.0).reshape(wavenumber_norm_grid.shape)
    )

    noise_hat = noise_hat * amplitude

    ic = ifft(
        noise_hat, num_spatial_dims=self.num_spatial_dims, num_points=num_points
    )

    ic = normalize_ic(
        ic, zero_mean=self.zero_mean, std_one=self.std_one, max_one=self.max_one
    )

    return ic