Skip to content

Base Nonlinear Function¤

exponax.nonlin_fun.BaseNonlinearFun ¤

Bases: Module, ABC

Source code in exponax/nonlin_fun/_base.py 
  9
 10
 11
 12
 13
 14
 15
 16
 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
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
class BaseNonlinearFun(eqx.Module, ABC):
    num_spatial_dims: int
    num_points: int
    dealiasing_mask: Bool[Array, "1 ... (N//2)+1"] | None

    def __init__(
        self,
        num_spatial_dims: int,
        num_points: int,
        *,
        dealiasing_fraction: float | None = None,
    ):
        """
        Base class for all nonlinear functions. This class provides the basic
        functionality to dealias the nonlinear terms and perform forward and
        inverse Fourier transforms.

        **Arguments:**

        - `num_spatial_dims`: The number of spatial dimensions `D`.
        - `num_points`: The number of points `N` used to discretize the domain.
            This **includes** the left boundary point and **excludes** the right
            boundary point. In higher dimensions; the number of points in each
            dimension is the same.
        - `dealiasing_fraction`: The fraction of the highest resolved mode to
            keep for dealiasing. For example, `2/3` corresponds to Orszag's 2/3
            rule typically used for quadratic nonlinearities. If `None`, no
            dealiasing is performed.

        !!! info
            Dealiasing is applied both before and after the nonlinear
            evaluation: the `ifft` method zeros out high modes before
            transforming to physical space (pre-dealiasing), and the `fft`
            method zeros out high modes after transforming back to Fourier
            space (post-dealiasing). Pre-dealiasing ensures that the
            physical-space representation is free of aliased modes before
            computing nonlinear products. Post-dealiasing removes any aliases
            created by those products, so that the returned nonlinear term is
            spectrally clean.

        !!! info
            Some dealiasing strategies (like Orszag's 2/3 rule) are designed to
            not fully remove aliasing (which would require 1/2 in the case of
            quadratic nonlinearities), rather to only have aliases being created
            in those modes that will be zeroed out anyway. See also [Orszag
            (1971)](https://doi.org/10.1175/1520-0469(1971)028%3C1074:OTEOAI%3E2.0.CO;2)
        """
        self.num_spatial_dims = num_spatial_dims
        self.num_points = num_points

        if dealiasing_fraction is None:
            self.dealiasing_mask = None
        else:
            # Can be done because num_points is identical in all spatial dimensions
            nyquist_mode = (num_points // 2) + 1
            highest_resolved_mode = nyquist_mode - 1
            start_of_aliased_modes = dealiasing_fraction * highest_resolved_mode

            self.dealiasing_mask = low_pass_filter_mask(
                num_spatial_dims,
                num_points,
                cutoff=start_of_aliased_modes - 1,
            )

    def dealias(
        self, u_hat: Complex[Array, "C ... (N//2)+1"]
    ) -> Complex[Array, "C ... (N//2)+1"]:
        """
        Dealias the Fourier representation of a state `u_hat` by zeroing out all
        the coefficients associated with modes beyond `dealiasing_fraction` set
        in the constructor.

        !!! note
            In most cases you do not need to call this method directly. The
            `fft` and `ifft` methods apply the dealiasing mask automatically
            (post-dealiasing and pre-dealiasing, respectively).

        **Arguments:**

        - `u_hat`: The Fourier representation of the state `u`.

        **Returns:**

        - `u_hat_dealiased`: The dealiased Fourier representation of the state
            `u`.
        """
        if self.dealiasing_mask is None:
            raise ValueError("Nonlinear function was set up without dealiasing")
        return self.dealiasing_mask * u_hat

    def fft(self, u: Float[Array, "C ... N"]) -> Complex[Array, "C ... (N//2)+1"]:
        """
        Correctly wrapped **real-valued** Fourier transform for the shape of the
        state vector associated with this nonlinear function. If a dealiasing
        mask is set, the output is dealiased (post-dealiasing).

        **Arguments:**

        - `u`: The state vector in real space.

        **Returns:**

        - `u_hat`: The (real-valued) Fourier transform of the state vector.
        """
        u_hat = fft(u, num_spatial_dims=self.num_spatial_dims)
        if self.dealiasing_mask is not None:
            u_hat = self.dealiasing_mask * u_hat
        return u_hat

    def ifft(self, u_hat: Complex[Array, "C ... (N//2)+1"]) -> Float[Array, "C ... N"]:
        """
        Correctly wrapped **real-valued** inverse Fourier transform for the shape
        of the state vector associated with this nonlinear function. If a
        dealiasing mask is set, the input is dealiased before transforming
        (pre-dealiasing).

        **Arguments:**

        - `u_hat`: The (real-valued) Fourier transform of the state vector.

        **Returns:**

        - `u`: The state vector in real space.
        """
        if self.dealiasing_mask is not None:
            u_hat = self.dealiasing_mask * u_hat
        return ifft(
            u_hat, num_spatial_dims=self.num_spatial_dims, num_points=self.num_points
        )

    @abstractmethod
    def __call__(
        self,
        u_hat: Complex[Array, "C ... (N//2)+1"],
    ) -> Complex[Array, "C ... (N//2)+1"]:
        """
        Evaluates the nonlinear function with a pseudo-spectral treatment and
        accounts for dealiasing.

        Use this in combination with `exponax.etdrk` routines to solve
        semi-linear PDEs in Fourier space.

        **Arguments:**

        - `u_hat`: The Fourier representation of the state `u`.

        **Returns:**

        - `𝒩(u_hat)`: The Fourier representation of the nonlinear term.
        """
        pass
__init__ ¤
__init__(
    num_spatial_dims: int,
    num_points: int,
    *,
    dealiasing_fraction: float | None = None
)

Base class for all nonlinear functions. This class provides the basic functionality to dealias the nonlinear terms and perform forward and inverse Fourier transforms.

Arguments:

  • num_spatial_dims: The number of spatial dimensions D.
  • num_points: The number of points N used to discretize the domain. This includes the left boundary point and excludes the right boundary point. In higher dimensions; the number of points in each dimension is the same.
  • dealiasing_fraction: The fraction of the highest resolved mode to keep for dealiasing. For example, 2/3 corresponds to Orszag's 2/3 rule typically used for quadratic nonlinearities. If None, no dealiasing is performed.

Info

Dealiasing is applied both before and after the nonlinear evaluation: the ifft method zeros out high modes before transforming to physical space (pre-dealiasing), and the fft method zeros out high modes after transforming back to Fourier space (post-dealiasing). Pre-dealiasing ensures that the physical-space representation is free of aliased modes before computing nonlinear products. Post-dealiasing removes any aliases created by those products, so that the returned nonlinear term is spectrally clean.

Info

Some dealiasing strategies (like Orszag's 2/3 rule) are designed to not fully remove aliasing (which would require 1/2 in the case of quadratic nonlinearities), rather to only have aliases being created in those modes that will be zeroed out anyway. See also Orszag (1971)

Source code in exponax/nonlin_fun/_base.py 
14
15
16
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
def __init__(
    self,
    num_spatial_dims: int,
    num_points: int,
    *,
    dealiasing_fraction: float | None = None,
):
    """
    Base class for all nonlinear functions. This class provides the basic
    functionality to dealias the nonlinear terms and perform forward and
    inverse Fourier transforms.

    **Arguments:**

    - `num_spatial_dims`: The number of spatial dimensions `D`.
    - `num_points`: The number of points `N` used to discretize the domain.
        This **includes** the left boundary point and **excludes** the right
        boundary point. In higher dimensions; the number of points in each
        dimension is the same.
    - `dealiasing_fraction`: The fraction of the highest resolved mode to
        keep for dealiasing. For example, `2/3` corresponds to Orszag's 2/3
        rule typically used for quadratic nonlinearities. If `None`, no
        dealiasing is performed.

    !!! info
        Dealiasing is applied both before and after the nonlinear
        evaluation: the `ifft` method zeros out high modes before
        transforming to physical space (pre-dealiasing), and the `fft`
        method zeros out high modes after transforming back to Fourier
        space (post-dealiasing). Pre-dealiasing ensures that the
        physical-space representation is free of aliased modes before
        computing nonlinear products. Post-dealiasing removes any aliases
        created by those products, so that the returned nonlinear term is
        spectrally clean.

    !!! info
        Some dealiasing strategies (like Orszag's 2/3 rule) are designed to
        not fully remove aliasing (which would require 1/2 in the case of
        quadratic nonlinearities), rather to only have aliases being created
        in those modes that will be zeroed out anyway. See also [Orszag
        (1971)](https://doi.org/10.1175/1520-0469(1971)028%3C1074:OTEOAI%3E2.0.CO;2)
    """
    self.num_spatial_dims = num_spatial_dims
    self.num_points = num_points

    if dealiasing_fraction is None:
        self.dealiasing_mask = None
    else:
        # Can be done because num_points is identical in all spatial dimensions
        nyquist_mode = (num_points // 2) + 1
        highest_resolved_mode = nyquist_mode - 1
        start_of_aliased_modes = dealiasing_fraction * highest_resolved_mode

        self.dealiasing_mask = low_pass_filter_mask(
            num_spatial_dims,
            num_points,
            cutoff=start_of_aliased_modes - 1,
        )
__call__ abstractmethod ¤
__call__(
    u_hat: Complex[Array, "C ... (N//2)+1"],
) -> Complex[Array, "C ... (N//2)+1"]

Evaluates the nonlinear function with a pseudo-spectral treatment and accounts for dealiasing.

Use this in combination with exponax.etdrk routines to solve semi-linear PDEs in Fourier space.

Arguments:

  • u_hat: The Fourier representation of the state u.

Returns:

  • 𝒩(u_hat): The Fourier representation of the nonlinear term.
Source code in exponax/nonlin_fun/_base.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
@abstractmethod
def __call__(
    self,
    u_hat: Complex[Array, "C ... (N//2)+1"],
) -> Complex[Array, "C ... (N//2)+1"]:
    """
    Evaluates the nonlinear function with a pseudo-spectral treatment and
    accounts for dealiasing.

    Use this in combination with `exponax.etdrk` routines to solve
    semi-linear PDEs in Fourier space.

    **Arguments:**

    - `u_hat`: The Fourier representation of the state `u`.

    **Returns:**

    - `𝒩(u_hat)`: The Fourier representation of the nonlinear term.
    """
    pass