Skip to content

Vorticity Convection Nonlinear Function¤

2D: Streamfunction-Vorticity Convection¤

exponax.nonlin_fun.VorticityConvection2d ¤

Bases: BaseNonlinearFun

Source code in exponax/nonlin_fun/_vorticity_convection.py 
 8
 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
class VorticityConvection2d(BaseNonlinearFun):
    convection_scale: float
    derivative_operator: Complex[Array, "D ... (N//2)+1"]
    inv_laplacian: Complex[Array, "1 ... (N//2)+1"]

    def __init__(
        self,
        num_spatial_dims: int,
        num_points: int,
        *,
        convection_scale: float = 1.0,
        derivative_operator: Complex[Array, "D ... (N//2)+1"],
        dealiasing_fraction: float,
    ):
        """
        Performs a pseudo-spectral evaluation of the nonlinear vorticity
        convection, e.g., found in the 2d Navier-Stokes equations in
        streamfunction-vorticity formulation. In state space, it reads

        ```
            𝒩(u) = - b ([1, -1]ᵀ ⊙ ∇(Δ⁻¹u)) ⋅ ∇u
        ```

        with `b` the convection scale, `⊙` the Hadamard product, `∇` the
        derivative operator, `Δ⁻¹` the inverse Laplacian, and `u` the vorticity.

        The minus arises because `Exponax` follows the convention that all
        nonlinear and linear differential operators are on the right-hand side
        of the equation. Typically, the vorticity convection term is on the
        left-hand side. Hence, the minus is required to move the term to the
        right-hand side.

        Since the inverse Laplacian is required, it internally performs a
        Poisson solve which is straightforward in Fourier space.

        **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.
        - `convection_scale`: The scale `b` of the convection term. Defaults
            to `1.0`.
        - `derivative_operator`: A complex array of shape `(d, ..., N//2+1)`
            that represents the derivative operator in Fourier space.
        - `dealiasing_fraction`: The fraction of the highest resolved modes
            that are not aliased. Defaults to `2/3` which corresponds to
            Orszag's 2/3 rule.
        """
        if num_spatial_dims != 2:
            raise ValueError(f"Expected num_spatial_dims = 2, got {num_spatial_dims}.")

        super().__init__(
            num_spatial_dims,
            num_points,
            dealiasing_fraction=dealiasing_fraction,
        )

        self.convection_scale = convection_scale
        self.derivative_operator = derivative_operator

        laplacian = build_laplace_operator(derivative_operator, order=2)

        # Uses the UNCHANGED mean solution to the Poisson equation (hence, the
        # mean of the "right-hand side" will be the mean of the solution).
        # However, this does not matter because we subsequently take the
        # gradient which would annihilate any mean energy anyway.
        self.inv_laplacian = jnp.where(laplacian == 0, 1.0, 1 / laplacian)

    def __call__(
        self, u_hat: Complex[Array, "1 ... (N//2)+1"]
    ) -> Complex[Array, "1 ... (N//2)+1"]:
        vorticity_hat = u_hat
        stream_function_hat = self.inv_laplacian * vorticity_hat

        u_hat = +self.derivative_operator[1:2] * stream_function_hat
        v_hat = -self.derivative_operator[0:1] * stream_function_hat
        del_vorticity_del_x_hat = self.derivative_operator[0:1] * vorticity_hat
        del_vorticity_del_y_hat = self.derivative_operator[1:2] * vorticity_hat

        u = self.ifft(u_hat)
        v = self.ifft(v_hat)
        del_vorticity_del_x = self.ifft(del_vorticity_del_x_hat)
        del_vorticity_del_y = self.ifft(del_vorticity_del_y_hat)

        convection = u * del_vorticity_del_x + v * del_vorticity_del_y

        convection_hat = self.fft(convection)

        # Need minus to bring term to the right-hand side
        return -self.convection_scale * convection_hat
__init__ ¤
__init__(
    num_spatial_dims: int,
    num_points: int,
    *,
    convection_scale: float = 1.0,
    derivative_operator: Complex[Array, "D ... (N//2)+1"],
    dealiasing_fraction: float
)

Performs a pseudo-spectral evaluation of the nonlinear vorticity convection, e.g., found in the 2d Navier-Stokes equations in streamfunction-vorticity formulation. In state space, it reads

    𝒩(u) = - b ([1, -1]ᵀ ⊙ ∇(Δ⁻¹u)) ⋅ ∇u

with b the convection scale, ⊙ the Hadamard product, ∇ the derivative operator, Δ⁻¹ the inverse Laplacian, and u the vorticity.

The minus arises because Exponax follows the convention that all nonlinear and linear differential operators are on the right-hand side of the equation. Typically, the vorticity convection term is on the left-hand side. Hence, the minus is required to move the term to the right-hand side.

Since the inverse Laplacian is required, it internally performs a Poisson solve which is straightforward in Fourier space.

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.
  • convection_scale: The scale b of the convection term. Defaults to 1.0.
  • derivative_operator: A complex array of shape (d, ..., N//2+1) that represents the derivative operator in Fourier space.
  • dealiasing_fraction: The fraction of the highest resolved modes that are not aliased. Defaults to 2/3 which corresponds to Orszag's 2/3 rule.
Source code in exponax/nonlin_fun/_vorticity_convection.py 
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
def __init__(
    self,
    num_spatial_dims: int,
    num_points: int,
    *,
    convection_scale: float = 1.0,
    derivative_operator: Complex[Array, "D ... (N//2)+1"],
    dealiasing_fraction: float,
):
    """
    Performs a pseudo-spectral evaluation of the nonlinear vorticity
    convection, e.g., found in the 2d Navier-Stokes equations in
    streamfunction-vorticity formulation. In state space, it reads

    ```
        𝒩(u) = - b ([1, -1]ᵀ ⊙ ∇(Δ⁻¹u)) ⋅ ∇u
    ```

    with `b` the convection scale, `⊙` the Hadamard product, `∇` the
    derivative operator, `Δ⁻¹` the inverse Laplacian, and `u` the vorticity.

    The minus arises because `Exponax` follows the convention that all
    nonlinear and linear differential operators are on the right-hand side
    of the equation. Typically, the vorticity convection term is on the
    left-hand side. Hence, the minus is required to move the term to the
    right-hand side.

    Since the inverse Laplacian is required, it internally performs a
    Poisson solve which is straightforward in Fourier space.

    **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.
    - `convection_scale`: The scale `b` of the convection term. Defaults
        to `1.0`.
    - `derivative_operator`: A complex array of shape `(d, ..., N//2+1)`
        that represents the derivative operator in Fourier space.
    - `dealiasing_fraction`: The fraction of the highest resolved modes
        that are not aliased. Defaults to `2/3` which corresponds to
        Orszag's 2/3 rule.
    """
    if num_spatial_dims != 2:
        raise ValueError(f"Expected num_spatial_dims = 2, got {num_spatial_dims}.")

    super().__init__(
        num_spatial_dims,
        num_points,
        dealiasing_fraction=dealiasing_fraction,
    )

    self.convection_scale = convection_scale
    self.derivative_operator = derivative_operator

    laplacian = build_laplace_operator(derivative_operator, order=2)

    # Uses the UNCHANGED mean solution to the Poisson equation (hence, the
    # mean of the "right-hand side" will be the mean of the solution).
    # However, this does not matter because we subsequently take the
    # gradient which would annihilate any mean energy anyway.
    self.inv_laplacian = jnp.where(laplacian == 0, 1.0, 1 / laplacian)
__call__ ¤
__call__(
    u_hat: Complex[Array, "1 ... (N//2)+1"],
) -> Complex[Array, "1 ... (N//2)+1"]
Source code in exponax/nonlin_fun/_vorticity_convection.py 
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def __call__(
    self, u_hat: Complex[Array, "1 ... (N//2)+1"]
) -> Complex[Array, "1 ... (N//2)+1"]:
    vorticity_hat = u_hat
    stream_function_hat = self.inv_laplacian * vorticity_hat

    u_hat = +self.derivative_operator[1:2] * stream_function_hat
    v_hat = -self.derivative_operator[0:1] * stream_function_hat
    del_vorticity_del_x_hat = self.derivative_operator[0:1] * vorticity_hat
    del_vorticity_del_y_hat = self.derivative_operator[1:2] * vorticity_hat

    u = self.ifft(u_hat)
    v = self.ifft(v_hat)
    del_vorticity_del_x = self.ifft(del_vorticity_del_x_hat)
    del_vorticity_del_y = self.ifft(del_vorticity_del_y_hat)

    convection = u * del_vorticity_del_x + v * del_vorticity_del_y

    convection_hat = self.fft(convection)

    # Need minus to bring term to the right-hand side
    return -self.convection_scale * convection_hat

exponax.nonlin_fun.VorticityConvection2dKolmogorov ¤

Bases: VorticityConvection2d

Source code in exponax/nonlin_fun/_vorticity_convection.py 
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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
class VorticityConvection2dKolmogorov(VorticityConvection2d):
    injection: Complex[Array, "1 ... (N//2)+1"]

    def __init__(
        self,
        num_spatial_dims: int,
        num_points: int,
        *,
        convection_scale: float = 1.0,
        injection_mode: int = 4,
        injection_scale: float = 1.0,
        derivative_operator: Complex[Array, "D ... (N//2)+1"],
        dealiasing_fraction: float,
    ):
        """
        Performs a pseudo-spectral evaluation of the nonlinear vorticity
        convection together with a Kolmogorov-like injection term, e.g., found
        in the 2d Navier-Stokes equations in streamfunction-vorticity
        formulation used for simulating a Kolmogorov flow.

        In state space, it reads

        ```
            𝒩(u) = - b ([1, -1]ᵀ ⊙ ∇(Δ⁻¹u)) ⋅ ∇u - f
        ```

        For details on the vorticity convective term, see
        `exponax.nonlin_fun.VorticityConvection2d`. The forcing term has the
        form

        ```
            f = -k (2π/L) γ cos(k (2π/L) x₁)
        ```

        i.e., energy of intensity `γ` is injected at wavenumber `k`. Note that
        the forcing is on the **vorticity**. As such, we get the prefactor `k
        (2π/L)` and the `sin(...)` turns into a `-cos(...)` (minus sign because
        the vorticity is derived via the curl).

        **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.
        - `convection_scale`: The scale `b` of the convection term. Defaults
            to `1.0`.
        - `injection_mode`: The wavenumber `k` at which energy is injected.
        - `injection_scale`: The intensity `γ` of the injection term.
        - `derivative_operator`: A complex array of shape `(d, ..., N//2+1)`
            that represents the derivative operator in Fourier space.
        - `dealiasing_fraction`: The fraction of the highest resolved modes
            that are not aliased. Defaults to `2/3` which corresponds to
            Orszag's 2/3 rule.
        """
        super().__init__(
            num_spatial_dims,
            num_points,
            convection_scale=convection_scale,
            derivative_operator=derivative_operator,
            dealiasing_fraction=dealiasing_fraction,
        )

        wavenumbers = build_wavenumbers(num_spatial_dims, num_points)
        injection_mask = (wavenumbers[0] == 0) & (wavenumbers[1] == injection_mode)
        self.injection = jnp.where(
            injection_mask,
            # Need to additional scale the `injection_scale` with the
            # `injection_mode`, because we apply the forcing on the vorticity.
            -injection_mode
            * injection_scale
            * build_scaling_array(num_spatial_dims, num_points, mode="coef_extraction"),
            0.0,
        )

    def __call__(
        self, u_hat: Complex[Array, "1 ... (N//2)+1"]
    ) -> Complex[Array, "1 ... (N//2)+1"]:
        neg_convection_hat = super().__call__(u_hat)
        return neg_convection_hat + self.injection
__init__ ¤
__init__(
    num_spatial_dims: int,
    num_points: int,
    *,
    convection_scale: float = 1.0,
    injection_mode: int = 4,
    injection_scale: float = 1.0,
    derivative_operator: Complex[Array, "D ... (N//2)+1"],
    dealiasing_fraction: float
)

Performs a pseudo-spectral evaluation of the nonlinear vorticity convection together with a Kolmogorov-like injection term, e.g., found in the 2d Navier-Stokes equations in streamfunction-vorticity formulation used for simulating a Kolmogorov flow.

In state space, it reads

    𝒩(u) = - b ([1, -1]ᵀ ⊙ ∇(Δ⁻¹u)) ⋅ ∇u - f

For details on the vorticity convective term, see exponax.nonlin_fun.VorticityConvection2d. The forcing term has the form

    f = -k (2π/L) γ cos(k (2π/L) x₁)

i.e., energy of intensity γ is injected at wavenumber k. Note that the forcing is on the vorticity. As such, we get the prefactor k (2π/L) and the sin(...) turns into a -cos(...) (minus sign because the vorticity is derived via the curl).

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.
  • convection_scale: The scale b of the convection term. Defaults to 1.0.
  • injection_mode: The wavenumber k at which energy is injected.
  • injection_scale: The intensity γ of the injection term.
  • derivative_operator: A complex array of shape (d, ..., N//2+1) that represents the derivative operator in Fourier space.
  • dealiasing_fraction: The fraction of the highest resolved modes that are not aliased. Defaults to 2/3 which corresponds to Orszag's 2/3 rule.
Source code in exponax/nonlin_fun/_vorticity_convection.py 
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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def __init__(
    self,
    num_spatial_dims: int,
    num_points: int,
    *,
    convection_scale: float = 1.0,
    injection_mode: int = 4,
    injection_scale: float = 1.0,
    derivative_operator: Complex[Array, "D ... (N//2)+1"],
    dealiasing_fraction: float,
):
    """
    Performs a pseudo-spectral evaluation of the nonlinear vorticity
    convection together with a Kolmogorov-like injection term, e.g., found
    in the 2d Navier-Stokes equations in streamfunction-vorticity
    formulation used for simulating a Kolmogorov flow.

    In state space, it reads

    ```
        𝒩(u) = - b ([1, -1]ᵀ ⊙ ∇(Δ⁻¹u)) ⋅ ∇u - f
    ```

    For details on the vorticity convective term, see
    `exponax.nonlin_fun.VorticityConvection2d`. The forcing term has the
    form

    ```
        f = -k (2π/L) γ cos(k (2π/L) x₁)
    ```

    i.e., energy of intensity `γ` is injected at wavenumber `k`. Note that
    the forcing is on the **vorticity**. As such, we get the prefactor `k
    (2π/L)` and the `sin(...)` turns into a `-cos(...)` (minus sign because
    the vorticity is derived via the curl).

    **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.
    - `convection_scale`: The scale `b` of the convection term. Defaults
        to `1.0`.
    - `injection_mode`: The wavenumber `k` at which energy is injected.
    - `injection_scale`: The intensity `γ` of the injection term.
    - `derivative_operator`: A complex array of shape `(d, ..., N//2+1)`
        that represents the derivative operator in Fourier space.
    - `dealiasing_fraction`: The fraction of the highest resolved modes
        that are not aliased. Defaults to `2/3` which corresponds to
        Orszag's 2/3 rule.
    """
    super().__init__(
        num_spatial_dims,
        num_points,
        convection_scale=convection_scale,
        derivative_operator=derivative_operator,
        dealiasing_fraction=dealiasing_fraction,
    )

    wavenumbers = build_wavenumbers(num_spatial_dims, num_points)
    injection_mask = (wavenumbers[0] == 0) & (wavenumbers[1] == injection_mode)
    self.injection = jnp.where(
        injection_mask,
        # Need to additional scale the `injection_scale` with the
        # `injection_mode`, because we apply the forcing on the vorticity.
        -injection_mode
        * injection_scale
        * build_scaling_array(num_spatial_dims, num_points, mode="coef_extraction"),
        0.0,
    )
__call__ ¤
__call__(
    u_hat: Complex[Array, "1 ... (N//2)+1"],
) -> Complex[Array, "1 ... (N//2)+1"]
Source code in exponax/nonlin_fun/_vorticity_convection.py 
178
179
180
181
182
def __call__(
    self, u_hat: Complex[Array, "1 ... (N//2)+1"]
) -> Complex[Array, "1 ... (N//2)+1"]:
    neg_convection_hat = super().__call__(u_hat)
    return neg_convection_hat + self.injection

3D: Projected Convection (Rotational Form)¤

exponax.nonlin_fun.ProjectedConvection3d ¤

Bases: BaseNonlinearFun

Source code in exponax/nonlin_fun/_projected_convection.py 
 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
class ProjectedConvection3d(BaseNonlinearFun):
    derivative_operator: Complex[Array, " 3 N N (N//2+1) "]
    leray_projection: Leray

    def __init__(
        self,
        num_spatial_dims: int,
        num_points: int,
        *,
        derivative_operator: Complex[Array, " 3 N N (N//2+1) "],
        dealiasing_fraction: float = 2 / 3,
    ):
        """
        Performs a pseudo-spectral evaluation of the nonlinear convection term
        for the 3d incompressible Navier-Stokes equations in velocity
        formulation using the rotational form. In state space, it reads

        ```
            𝒩(u) = 𝒫(u × ω)
        ```

        with `u` the velocity, `ω = ∇ × u` the vorticity, and `𝒫` the Leray
        projection onto divergence-free fields.

        The incompressible Navier-Stokes equations read

        ```
            uₜ = -(u ⋅ ∇)u - ∇p + ν Δu
        ```

        The [Lamb vector identity](https://en.wikipedia.org/wiki/Lamb_vector)
        allows rewriting the convection term as

        ```
            (u ⋅ ∇)u = ∇(|u|²/2) + ω × u
        ```

        Substituting and rearranging gives

        ```
            uₜ = u × ω - ∇(|u|²/2 + p) + ν Δu
        ```

        where `u × ω = -(ω × u)`. The Leray projection `𝒫` projects onto the
        space of divergence-free fields by removing any gradient component
        (i.e., `𝒫(∇φ) = 0` for any scalar `φ`). Since both the pressure
        gradient `∇p` and the kinetic energy gradient `∇(|u|²/2)` are
        gradients of scalar fields, the projection annihilates them:

        ```
            𝒫(uₜ) = 𝒫(u × ω) + ν Δu
        ```

        Hence, the nonlinear term reduces to `𝒩(u) = 𝒫(u × ω)`, and the
        pressure never needs to be computed explicitly.

        The curl is computed in Fourier space, the cross product `u × ω` in
        physical space (pseudo-spectral), and the result is projected back
        in Fourier space.

        Based on https://arxiv.org/pdf/1602.03638

        **Arguments:**

        - `num_spatial_dims`: The number of spatial dimensions `D`. Must be
            `3`.
        - `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.
        - `derivative_operator`: A complex array of shape `(3, ..., N//2+1)`
            that represents the derivative operator in Fourier space.
        - `dealiasing_fraction`: The fraction of the highest resolved modes
            that are not aliased. Defaults to `2/3` which corresponds to
            Orszag's 2/3 rule.
        """
        if num_spatial_dims != 3:
            raise ValueError(
                "ProjectedConvection3d only supports 3 spatial dimensions."
            )

        super().__init__(
            num_spatial_dims=num_spatial_dims,
            num_points=num_points,
            dealiasing_fraction=dealiasing_fraction,
        )

        self.derivative_operator = derivative_operator

        self.leray_projection = Leray(
            num_spatial_dims=num_spatial_dims,
            num_points=num_points,
            derivative_operator=derivative_operator,
        )

    def __call__(
        self,
        u_hat: Complex[Array, " 3 N N (N//2+1) "],
    ) -> Complex[Array, " 3 N N (N//2+1) "]:
        velocity_hat = u_hat
        curl_hat = _cross_product_3d(
            self.derivative_operator,
            velocity_hat,
        )

        curl = self.ifft(curl_hat)
        velocity = self.ifft(velocity_hat)

        convection = _cross_product_3d(
            velocity,
            curl,
        )

        convection_hat = self.fft(convection)

        convection_projected_hat = self.leray_projection(convection_hat)

        return convection_projected_hat
__init__ ¤
__init__(
    num_spatial_dims: int,
    num_points: int,
    *,
    derivative_operator: Complex[Array, " 3 N N (N//2+1) "],
    dealiasing_fraction: float = 2 / 3
)

Performs a pseudo-spectral evaluation of the nonlinear convection term for the 3d incompressible Navier-Stokes equations in velocity formulation using the rotational form. In state space, it reads

    𝒩(u) = 𝒫(u × ω)

with u the velocity, ω = ∇ × u the vorticity, and 𝒫 the Leray projection onto divergence-free fields.

The incompressible Navier-Stokes equations read

    uₜ = -(u ⋅ ∇)u - ∇p + ν Δu

The Lamb vector identity allows rewriting the convection term as

    (u ⋅ ∇)u = ∇(|u|²/2) + ω × u

Substituting and rearranging gives

    uₜ = u × ω - ∇(|u|²/2 + p) + ν Δu

where u × ω = -(ω × u). The Leray projection 𝒫 projects onto the space of divergence-free fields by removing any gradient component (i.e., 𝒫(∇φ) = 0 for any scalar φ). Since both the pressure gradient ∇p and the kinetic energy gradient ∇(|u|²/2) are gradients of scalar fields, the projection annihilates them:

    𝒫(uₜ) = 𝒫(u × ω) + ν Δu

Hence, the nonlinear term reduces to 𝒩(u) = 𝒫(u × ω), and the pressure never needs to be computed explicitly.

The curl is computed in Fourier space, the cross product u × ω in physical space (pseudo-spectral), and the result is projected back in Fourier space.

Based on https://arxiv.org/pdf/1602.03638

Arguments:

  • num_spatial_dims: The number of spatial dimensions D. Must be 3.
  • 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.
  • derivative_operator: A complex array of shape (3, ..., N//2+1) that represents the derivative operator in Fourier space.
  • dealiasing_fraction: The fraction of the highest resolved modes that are not aliased. Defaults to 2/3 which corresponds to Orszag's 2/3 rule.
Source code in exponax/nonlin_fun/_projected_convection.py 
 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
def __init__(
    self,
    num_spatial_dims: int,
    num_points: int,
    *,
    derivative_operator: Complex[Array, " 3 N N (N//2+1) "],
    dealiasing_fraction: float = 2 / 3,
):
    """
    Performs a pseudo-spectral evaluation of the nonlinear convection term
    for the 3d incompressible Navier-Stokes equations in velocity
    formulation using the rotational form. In state space, it reads

    ```
        𝒩(u) = 𝒫(u × ω)
    ```

    with `u` the velocity, `ω = ∇ × u` the vorticity, and `𝒫` the Leray
    projection onto divergence-free fields.

    The incompressible Navier-Stokes equations read

    ```
        uₜ = -(u ⋅ ∇)u - ∇p + ν Δu
    ```

    The [Lamb vector identity](https://en.wikipedia.org/wiki/Lamb_vector)
    allows rewriting the convection term as

    ```
        (u ⋅ ∇)u = ∇(|u|²/2) + ω × u
    ```

    Substituting and rearranging gives

    ```
        uₜ = u × ω - ∇(|u|²/2 + p) + ν Δu
    ```

    where `u × ω = -(ω × u)`. The Leray projection `𝒫` projects onto the
    space of divergence-free fields by removing any gradient component
    (i.e., `𝒫(∇φ) = 0` for any scalar `φ`). Since both the pressure
    gradient `∇p` and the kinetic energy gradient `∇(|u|²/2)` are
    gradients of scalar fields, the projection annihilates them:

    ```
        𝒫(uₜ) = 𝒫(u × ω) + ν Δu
    ```

    Hence, the nonlinear term reduces to `𝒩(u) = 𝒫(u × ω)`, and the
    pressure never needs to be computed explicitly.

    The curl is computed in Fourier space, the cross product `u × ω` in
    physical space (pseudo-spectral), and the result is projected back
    in Fourier space.

    Based on https://arxiv.org/pdf/1602.03638

    **Arguments:**

    - `num_spatial_dims`: The number of spatial dimensions `D`. Must be
        `3`.
    - `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.
    - `derivative_operator`: A complex array of shape `(3, ..., N//2+1)`
        that represents the derivative operator in Fourier space.
    - `dealiasing_fraction`: The fraction of the highest resolved modes
        that are not aliased. Defaults to `2/3` which corresponds to
        Orszag's 2/3 rule.
    """
    if num_spatial_dims != 3:
        raise ValueError(
            "ProjectedConvection3d only supports 3 spatial dimensions."
        )

    super().__init__(
        num_spatial_dims=num_spatial_dims,
        num_points=num_points,
        dealiasing_fraction=dealiasing_fraction,
    )

    self.derivative_operator = derivative_operator

    self.leray_projection = Leray(
        num_spatial_dims=num_spatial_dims,
        num_points=num_points,
        derivative_operator=derivative_operator,
    )
__call__ ¤
__call__(
    u_hat: Complex[Array, " 3 N N (N//2+1) "],
) -> Complex[Array, " 3 N N (N//2+1) "]
Source code in exponax/nonlin_fun/_projected_convection.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
def __call__(
    self,
    u_hat: Complex[Array, " 3 N N (N//2+1) "],
) -> Complex[Array, " 3 N N (N//2+1) "]:
    velocity_hat = u_hat
    curl_hat = _cross_product_3d(
        self.derivative_operator,
        velocity_hat,
    )

    curl = self.ifft(curl_hat)
    velocity = self.ifft(velocity_hat)

    convection = _cross_product_3d(
        velocity,
        curl,
    )

    convection_hat = self.fft(convection)

    convection_projected_hat = self.leray_projection(convection_hat)

    return convection_projected_hat

exponax.nonlin_fun.ProjectedConvection3dKolmogorov ¤

Bases: ProjectedConvection3d

Source code in exponax/nonlin_fun/_projected_convection.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
class ProjectedConvection3dKolmogorov(ProjectedConvection3d):
    injection: Complex[Array, "3 ... (N//2)+1"]

    def __init__(
        self,
        num_spatial_dims: int,
        num_points: int,
        *,
        injection_mode: int = 4,
        injection_scale: float = 1.0,
        derivative_operator: Complex[Array, " 3 ... (N//2)+1"],
        dealiasing_fraction: float,
    ):
        """
        Performs a pseudo-spectral evaluation of the nonlinear convection term
        together with a Kolmogorov-like injection term for the 3d
        incompressible Navier-Stokes equations in velocity formulation. In
        state space, it reads

        ```
            𝒩(u) = 𝒫(u × ω) + f
        ```

        For details on the projected convection term, see
        `exponax.nonlin_fun.ProjectedConvection3d`. The forcing term has the
        form

        ```
            f₀ = γ sin(k (2π/L) x₁)

            f₁ = 0

            f₂ = 0
        ```

        i.e., energy of intensity `γ` is injected at wavenumber `k` in the
        first velocity channel varying over the second spatial axis. This
        forcing is naturally divergence-free because the forced component does
        not vary along its own direction.

        **Arguments:**

        - `num_spatial_dims`: The number of spatial dimensions `D`. Must be
            `3`.
        - `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.
        - `injection_mode`: The wavenumber `k` at which energy is injected.
        - `injection_scale`: The intensity `γ` of the injection term.
        - `derivative_operator`: A complex array of shape `(3, ..., N//2+1)`
            that represents the derivative operator in Fourier space.
        - `dealiasing_fraction`: The fraction of the highest resolved modes
            that are not aliased. Defaults to `2/3` which corresponds to
            Orszag's 2/3 rule.
        """
        super().__init__(
            num_spatial_dims,
            num_points,
            derivative_operator=derivative_operator,
            dealiasing_fraction=dealiasing_fraction,
        )

        wavenumbers = build_wavenumbers(num_spatial_dims, num_points)
        injection_mask = (
            (wavenumbers[0] == 0)
            & (wavenumbers[1] == injection_mode)
            & (wavenumbers[2] == 0)
        )
        # In 3D, we work with velocity directly (not vorticity), so the
        # forcing is f = γ sin(k x₁) ê₀. Only the first velocity channel is
        # forced, and no extra -k factor is needed (unlike the 2D vorticity
        # formulation).
        injection_single = jnp.where(
            injection_mask,
            injection_scale
            * build_scaling_array(num_spatial_dims, num_points, mode="coef_extraction"),
            0.0,
        )
        # Shape (3, N, N, N//2+1): forcing only in the first velocity channel
        zeros = jnp.zeros_like(injection_single)
        self.injection = jnp.concatenate([injection_single, zeros, zeros], axis=0)

    def __call__(
        self, u_hat: Complex[Array, "3 ... (N//2)+1"]
    ) -> Complex[Array, "3 ... (N//2)+1"]:
        neg_convection_hat = super().__call__(u_hat)
        return neg_convection_hat + self.injection
__init__ ¤
__init__(
    num_spatial_dims: int,
    num_points: int,
    *,
    injection_mode: int = 4,
    injection_scale: float = 1.0,
    derivative_operator: Complex[Array, " 3 ... (N//2)+1"],
    dealiasing_fraction: float
)

Performs a pseudo-spectral evaluation of the nonlinear convection term together with a Kolmogorov-like injection term for the 3d incompressible Navier-Stokes equations in velocity formulation. In state space, it reads

    𝒩(u) = 𝒫(u × ω) + f

For details on the projected convection term, see exponax.nonlin_fun.ProjectedConvection3d. The forcing term has the form

    f₀ = γ sin(k (2π/L) x₁)

    f₁ = 0

    f₂ = 0

i.e., energy of intensity γ is injected at wavenumber k in the first velocity channel varying over the second spatial axis. This forcing is naturally divergence-free because the forced component does not vary along its own direction.

Arguments:

  • num_spatial_dims: The number of spatial dimensions D. Must be 3.
  • 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.
  • injection_mode: The wavenumber k at which energy is injected.
  • injection_scale: The intensity γ of the injection term.
  • derivative_operator: A complex array of shape (3, ..., N//2+1) that represents the derivative operator in Fourier space.
  • dealiasing_fraction: The fraction of the highest resolved modes that are not aliased. Defaults to 2/3 which corresponds to Orszag's 2/3 rule.
Source code in exponax/nonlin_fun/_projected_convection.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
def __init__(
    self,
    num_spatial_dims: int,
    num_points: int,
    *,
    injection_mode: int = 4,
    injection_scale: float = 1.0,
    derivative_operator: Complex[Array, " 3 ... (N//2)+1"],
    dealiasing_fraction: float,
):
    """
    Performs a pseudo-spectral evaluation of the nonlinear convection term
    together with a Kolmogorov-like injection term for the 3d
    incompressible Navier-Stokes equations in velocity formulation. In
    state space, it reads

    ```
        𝒩(u) = 𝒫(u × ω) + f
    ```

    For details on the projected convection term, see
    `exponax.nonlin_fun.ProjectedConvection3d`. The forcing term has the
    form

    ```
        f₀ = γ sin(k (2π/L) x₁)

        f₁ = 0

        f₂ = 0
    ```

    i.e., energy of intensity `γ` is injected at wavenumber `k` in the
    first velocity channel varying over the second spatial axis. This
    forcing is naturally divergence-free because the forced component does
    not vary along its own direction.

    **Arguments:**

    - `num_spatial_dims`: The number of spatial dimensions `D`. Must be
        `3`.
    - `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.
    - `injection_mode`: The wavenumber `k` at which energy is injected.
    - `injection_scale`: The intensity `γ` of the injection term.
    - `derivative_operator`: A complex array of shape `(3, ..., N//2+1)`
        that represents the derivative operator in Fourier space.
    - `dealiasing_fraction`: The fraction of the highest resolved modes
        that are not aliased. Defaults to `2/3` which corresponds to
        Orszag's 2/3 rule.
    """
    super().__init__(
        num_spatial_dims,
        num_points,
        derivative_operator=derivative_operator,
        dealiasing_fraction=dealiasing_fraction,
    )

    wavenumbers = build_wavenumbers(num_spatial_dims, num_points)
    injection_mask = (
        (wavenumbers[0] == 0)
        & (wavenumbers[1] == injection_mode)
        & (wavenumbers[2] == 0)
    )
    # In 3D, we work with velocity directly (not vorticity), so the
    # forcing is f = γ sin(k x₁) ê₀. Only the first velocity channel is
    # forced, and no extra -k factor is needed (unlike the 2D vorticity
    # formulation).
    injection_single = jnp.where(
        injection_mask,
        injection_scale
        * build_scaling_array(num_spatial_dims, num_points, mode="coef_extraction"),
        0.0,
    )
    # Shape (3, N, N, N//2+1): forcing only in the first velocity channel
    zeros = jnp.zeros_like(injection_single)
    self.injection = jnp.concatenate([injection_single, zeros, zeros], axis=0)
__call__ ¤
__call__(
    u_hat: Complex[Array, "3 ... (N//2)+1"],
) -> Complex[Array, "3 ... (N//2)+1"]
Source code in exponax/nonlin_fun/_projected_convection.py 
222
223
224
225
226
def __call__(
    self, u_hat: Complex[Array, "3 ... (N//2)+1"]
) -> Complex[Array, "3 ... (N//2)+1"]:
    neg_convection_hat = super().__call__(u_hat)
    return neg_convection_hat + self.injection

Leray Projection¤

exponax.nonlin_fun.Leray ¤

Bases: BaseNonlinearFun

Source code in exponax/nonlin_fun/_leray.py 
  8
  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
class Leray(BaseNonlinearFun):
    inv_laplacian: Complex[Array, " 1 ... (N//2+1) "]
    derivative_operator: Complex[Array, " D ... (N//2+1) "]

    def __init__(
        self,
        num_spatial_dims: int,
        num_points: int,
        *,
        derivative_operator: Complex[Array, " D ... (N//2+1) "],
        order: int = 2,
    ):
        """
        Leray projection operator that projects a vector field onto the space of
        divergence-free fields. In state space, it reads

        ```
            𝒫(u) = u - ∇(Δ⁻¹ ∇ ⋅ u)
        ```

        This is equivalent to removing the irrotational (curl-free) component
        via the Helmholtz decomposition. The projection is trivial to compute in
        Fourier space because the Laplacian is diagonal.

        **Derivation:**

        Consider an operator splitting approach to the incompressible
        Navier-Stokes equations in which advection and diffusion have already
        been integrated, yielding an intermediate velocity `u*` that is
        generally **not** divergence-free (`∇ ⋅ u* ≠ 0`). The pressure substep
        reads

        ```
            (u** - u*) / Δt = -(1/ρ) ∇p
        ```

        We require `u**` to be incompressible (`∇ ⋅ u** = 0`). Applying the
        divergence to both sides gives

        ```
            (1/Δt)(∇ ⋅ u** - ∇ ⋅ u*) = -(1/ρ) Δp
        ```

        Setting `∇ ⋅ u** = 0` yields the pressure-Poisson equation

        ```
            Δp = (ρ/Δt) ∇ ⋅ u*
        ```

        Solving for `p` and substituting back gives

        ```
            u** = u* - (Δt/ρ) ∇(Δ⁻¹ (ρ/Δt) ∇ ⋅ u*)
        ```

        For constant density (as in the incompressible case), `ρ` and `Δt`
        cancel, and we obtain the Leray projection

        ```
            u** = u* - ∇(Δ⁻¹(∇ ⋅ u*))
        ```

        Hence, making a velocity field incompressible is a three-step process:

        1. Compute the divergence: `d = ∇ ⋅ u*`
        2. Solve a Poisson equation for a pseudo-pressure: `p = Δ⁻¹(-d)`
        3. Correct the velocity via the pressure gradient: `u** = u* + ∇p`

        In general, step 2 is the most computationally demanding (e.g.,
        requiring a conjugate gradient solve). However, in Fourier space the
        Laplacian is diagonal, making the Poisson solve a trivial pointwise
        division.

        Note that one can either use a negative sign in the Poisson solve (`p =
        Δ⁻¹(-d)`) or a positive sign (`p = Δ⁻¹(d)`) as long as the opposite sign
        is used in the velocity correction step. The choice of sign is arbitrary
        and does not affect the final result.

        Technically not a nonlinear operator, but it performs channel mixing
        between the velocity channels and hence subclasses `BaseNonlinearFun`.

        **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.
        - `derivative_operator`: A complex array of shape `(D, ..., N//2+1)`
            that represents the derivative operator in Fourier space.
        - `order`: The order of the Laplacian used for the Poisson solve.
            Default is `2`.
        """
        super().__init__(
            num_spatial_dims=num_spatial_dims,
            num_points=num_points,
        )

        laplace_operator = build_laplace_operator(derivative_operator, order=order)

        self.inv_laplacian = jnp.where(
            laplace_operator != 0, 1.0 / laplace_operator, 0.0
        )

        self.derivative_operator = derivative_operator

    def __call__(
        self,
        u_hat: Complex[Array, " D ... (N//2+1) "],
    ) -> Complex[Array, " D ... (N//2+1) "]:
        """Compute the Leray projection of the velocity field u_hat.

        Args:
            u_hat: Velocity field in Fourier space.

        Returns:
            The Leray projection of u_hat in Fourier space.
        """
        div_u_hat = jnp.sum(
            self.derivative_operator * u_hat,
            axis=0,
            keepdims=True,
        )

        pressure_hat = -self.inv_laplacian * div_u_hat

        grad_pressure_hat = self.derivative_operator * pressure_hat

        return u_hat + grad_pressure_hat
__init__ ¤
__init__(
    num_spatial_dims: int,
    num_points: int,
    *,
    derivative_operator: Complex[Array, " D ... (N//2+1) "],
    order: int = 2
)

Leray projection operator that projects a vector field onto the space of divergence-free fields. In state space, it reads

    𝒫(u) = u - ∇(Δ⁻¹ ∇ ⋅ u)

This is equivalent to removing the irrotational (curl-free) component via the Helmholtz decomposition. The projection is trivial to compute in Fourier space because the Laplacian is diagonal.

Derivation:

Consider an operator splitting approach to the incompressible Navier-Stokes equations in which advection and diffusion have already been integrated, yielding an intermediate velocity u* that is generally not divergence-free (∇ ⋅ u* ≠ 0). The pressure substep reads

    (u** - u*) / Δt = -(1/ρ) ∇p

We require u** to be incompressible (∇ ⋅ u** = 0). Applying the divergence to both sides gives

    (1/Δt)(∇ ⋅ u** - ∇ ⋅ u*) = -(1/ρ) Δp

Setting ∇ ⋅ u** = 0 yields the pressure-Poisson equation

    Δp = (ρ/Δt) ∇ ⋅ u*

Solving for p and substituting back gives

    u** = u* - (Δt/ρ) ∇(Δ⁻¹ (ρ/Δt) ∇ ⋅ u*)

For constant density (as in the incompressible case), ρ and Δt cancel, and we obtain the Leray projection

    u** = u* - ∇(Δ⁻¹(∇ ⋅ u*))

Hence, making a velocity field incompressible is a three-step process:

  1. Compute the divergence: d = ∇ ⋅ u*
  2. Solve a Poisson equation for a pseudo-pressure: p = Δ⁻¹(-d)
  3. Correct the velocity via the pressure gradient: u** = u* + ∇p

In general, step 2 is the most computationally demanding (e.g., requiring a conjugate gradient solve). However, in Fourier space the Laplacian is diagonal, making the Poisson solve a trivial pointwise division.

Note that one can either use a negative sign in the Poisson solve (p = Δ⁻¹(-d)) or a positive sign (p = Δ⁻¹(d)) as long as the opposite sign is used in the velocity correction step. The choice of sign is arbitrary and does not affect the final result.

Technically not a nonlinear operator, but it performs channel mixing between the velocity channels and hence subclasses BaseNonlinearFun.

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.
  • derivative_operator: A complex array of shape (D, ..., N//2+1) that represents the derivative operator in Fourier space.
  • order: The order of the Laplacian used for the Poisson solve. Default is 2.
Source code in exponax/nonlin_fun/_leray.py 
 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
def __init__(
    self,
    num_spatial_dims: int,
    num_points: int,
    *,
    derivative_operator: Complex[Array, " D ... (N//2+1) "],
    order: int = 2,
):
    """
    Leray projection operator that projects a vector field onto the space of
    divergence-free fields. In state space, it reads

    ```
        𝒫(u) = u - ∇(Δ⁻¹ ∇ ⋅ u)
    ```

    This is equivalent to removing the irrotational (curl-free) component
    via the Helmholtz decomposition. The projection is trivial to compute in
    Fourier space because the Laplacian is diagonal.

    **Derivation:**

    Consider an operator splitting approach to the incompressible
    Navier-Stokes equations in which advection and diffusion have already
    been integrated, yielding an intermediate velocity `u*` that is
    generally **not** divergence-free (`∇ ⋅ u* ≠ 0`). The pressure substep
    reads

    ```
        (u** - u*) / Δt = -(1/ρ) ∇p
    ```

    We require `u**` to be incompressible (`∇ ⋅ u** = 0`). Applying the
    divergence to both sides gives

    ```
        (1/Δt)(∇ ⋅ u** - ∇ ⋅ u*) = -(1/ρ) Δp
    ```

    Setting `∇ ⋅ u** = 0` yields the pressure-Poisson equation

    ```
        Δp = (ρ/Δt) ∇ ⋅ u*
    ```

    Solving for `p` and substituting back gives

    ```
        u** = u* - (Δt/ρ) ∇(Δ⁻¹ (ρ/Δt) ∇ ⋅ u*)
    ```

    For constant density (as in the incompressible case), `ρ` and `Δt`
    cancel, and we obtain the Leray projection

    ```
        u** = u* - ∇(Δ⁻¹(∇ ⋅ u*))
    ```

    Hence, making a velocity field incompressible is a three-step process:

    1. Compute the divergence: `d = ∇ ⋅ u*`
    2. Solve a Poisson equation for a pseudo-pressure: `p = Δ⁻¹(-d)`
    3. Correct the velocity via the pressure gradient: `u** = u* + ∇p`

    In general, step 2 is the most computationally demanding (e.g.,
    requiring a conjugate gradient solve). However, in Fourier space the
    Laplacian is diagonal, making the Poisson solve a trivial pointwise
    division.

    Note that one can either use a negative sign in the Poisson solve (`p =
    Δ⁻¹(-d)`) or a positive sign (`p = Δ⁻¹(d)`) as long as the opposite sign
    is used in the velocity correction step. The choice of sign is arbitrary
    and does not affect the final result.

    Technically not a nonlinear operator, but it performs channel mixing
    between the velocity channels and hence subclasses `BaseNonlinearFun`.

    **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.
    - `derivative_operator`: A complex array of shape `(D, ..., N//2+1)`
        that represents the derivative operator in Fourier space.
    - `order`: The order of the Laplacian used for the Poisson solve.
        Default is `2`.
    """
    super().__init__(
        num_spatial_dims=num_spatial_dims,
        num_points=num_points,
    )

    laplace_operator = build_laplace_operator(derivative_operator, order=order)

    self.inv_laplacian = jnp.where(
        laplace_operator != 0, 1.0 / laplace_operator, 0.0
    )

    self.derivative_operator = derivative_operator
__call__ ¤
__call__(
    u_hat: Complex[Array, " D ... (N//2+1) "],
) -> Complex[Array, " D ... (N//2+1) "]

Compute the Leray projection of the velocity field u_hat.

Args: u_hat: Velocity field in Fourier space.

Returns: The Leray projection of u_hat in Fourier space.

Source code in exponax/nonlin_fun/_leray.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
def __call__(
    self,
    u_hat: Complex[Array, " D ... (N//2+1) "],
) -> Complex[Array, " D ... (N//2+1) "]:
    """Compute the Leray projection of the velocity field u_hat.

    Args:
        u_hat: Velocity field in Fourier space.

    Returns:
        The Leray projection of u_hat in Fourier space.
    """
    div_u_hat = jnp.sum(
        self.derivative_operator * u_hat,
        axis=0,
        keepdims=True,
    )

    pressure_hat = -self.inv_laplacian * div_u_hat

    grad_pressure_hat = self.derivative_operator * pressure_hat

    return u_hat + grad_pressure_hat