Advection Terms
Source Code: NSE_AdvectionTerms
NSE_AdvectionTerms.py
1# Copyright (C) 2025 Sukanta Basu
2#
3# This program is free software: you can redistribute it and/or modify
4# it under the terms of the GNU General Public License as published by
5# the Free Software Foundation, either version 3 of the License, or
6# (at your option) any later version.
7#
8# This program is distributed in the hope that it will be useful,
9# but WITHOUT ANY WARRANTY; without even the implied warranty of
10# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11# GNU General Public License for more details.
12#
13# You should have received a copy of the GNU General Public License
14# along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16"""
17File: NSE_AdvectionTerms.py
18==================
19
20:Author: Sukanta Basu
21:AI Assistance: Claude Code (Anthropic) and Codex (OpenAI) are used for documentation,
22 code restructuring, and performance optimization
23:Date: 2025-4-3
24:Description: computes the advective terms as::
25
26.. math::
27
28 C_x = v\\left(\\frac{\\partial u}{\\partial y} - \\frac{\\partial v}{\\partial x}\\right) + w\\left(\\frac{\\partial u}{\\partial z} - \\frac{\\partial w}{\\partial x}\\right)
29
30 C_y = u\\left(\\frac{\\partial v}{\\partial x} - \\frac{\\partial u}{\\partial y}\\right) + w\\left(\\frac{\\partial v}{\\partial z} - \\frac{\\partial w}{\\partial y}\\right)
31
32 C_z = u\\left(\\frac{\\partial w}{\\partial x} - \\frac{\\partial u}{\\partial z}\\right) + v\\left(\\frac{\\partial w}{\\partial y} - \\frac{\\partial v}{\\partial z}\\right)
33"""
34
35
36# ============================================
37# Imports
38# ============================================
39
40import jax
41from ..config.ConfigLoader import *
42from ..config import ConfigLoader as Config
43
44# Import derived variables
45from ..config.DerivedVars import *
46
47from ..utilities.Utilities import StagGridAvg
48
49# Import FFT modules
50from ..operations.FFT import FFT, FFT_pad
51
52# Import dealias functions
53from ..operations.Dealiasing import Dealias1, Dealias2
54
55# Import constants
56from ..initialization.Preprocess import Constant
57mx, my, nx_rfft, ny_rfft, mx_rfft, my_rfft = Constant()
58
59
60# ------------------------------------------------------------------------------
61# Use dealiasing in computing the advection terms
62# ------------------------------------------------------------------------------
63
64@jax.jit
65def Advection_Dealias(u, v, w,
66 dudy, dudz,
67 dvdx, dvdz,
68 dwdx, dwdy,
69 ZeRo3D_fft, ZeRo3D_pad, ZeRo3D_pad_fft):
70 """
71 Parameters:
72 -----------
73 u, v, w : ndarray
74 Velocity components in x, y, and z directions respectively.
75 Shape: (nx, ny, nz)
76 dudy, dudz : ndarray
77 Derivatives of u-velocity with respect to y and z.
78 dvdx, dvdz : ndarray
79 Derivatives of v-velocity with respect to x and z.
80 dwdx, dwdy : ndarray
81 Derivatives of w-velocity with respect to x and y.
82
83 Returns:
84 --------
85 tuple[ndarray, ndarray, ndarray]
86 Cx : Convective term in x-direction
87 Cy : Convective term in y-direction
88 Cz : Convective term in z-direction
89 Each array has shape (nx, ny, nz)
90 """
91
92 # Dealias
93 u_pad = Dealias1(FFT(u), ZeRo3D_pad_fft)
94 v_pad = Dealias1(FFT(v), ZeRo3D_pad_fft)
95 w_pad = Dealias1(FFT(w), ZeRo3D_pad_fft)
96 dudy_pad = Dealias1(FFT(dudy), ZeRo3D_pad_fft)
97 dudz_pad = Dealias1(FFT(dudz), ZeRo3D_pad_fft)
98 dvdx_pad = Dealias1(FFT(dvdx), ZeRo3D_pad_fft)
99 dvdz_pad = Dealias1(FFT(dvdz), ZeRo3D_pad_fft)
100 dwdx_pad = Dealias1(FFT(dwdx), ZeRo3D_pad_fft)
101 dwdy_pad = Dealias1(FFT(dwdy), ZeRo3D_pad_fft)
102
103 # Compute Cx: advection term in x-direction
104 # First term: v*(∂u/∂y - ∂v/∂x)
105 arg1 = v_pad * (dudy_pad - dvdx_pad)
106 # Second term: w*(∂u/∂z - ∂w/∂x)
107 arg2 = w_pad * (dudz_pad - dwdx_pad)
108
109 # Apply staggered grid averaging to second term
110 arg2G = StagGridAvg(arg2)
111
112 # Combine terms with boundary condition handling
113 cc = ZeRo3D_pad.copy()
114 cc = cc.at[:, :, 1:nz-1].set(arg1[:, :, 1:nz-1] + arg2G[:, :, 1:nz-1]) # Interior points
115 cc = cc.at[:, :, 0].set(arg1[:, :, 0] + 0.5 * arg2[:, :, 1]) # Bottom boundary
116 cc = cc.at[:, :, nz-1].set(arg1[:, :, nz-1] + arg2[:, :, nz-1]) # Top boundary
117
118 Cx = Dealias2(FFT_pad(cc), ZeRo3D_fft)
119
120 # Compute Cy: advection term in y-direction
121 # First term: u*(∂v/∂x - ∂u/∂y)
122 arg1 = u_pad * (dvdx_pad - dudy_pad)
123 # Second term: w*(∂v/∂z - ∂w/∂y)
124 arg2 = w_pad * (dvdz_pad - dwdy_pad)
125
126 # Apply staggered grid averaging to second term
127 arg2G = StagGridAvg(arg2)
128
129 cc = ZeRo3D_pad.copy()
130 cc = cc.at[:, :, 1:nz-1].set(arg1[:, :, 1:nz-1] + arg2G[:, :, 1:nz-1]) # Interior points
131 cc = cc.at[:, :, 0].set(arg1[:, :, 0] + 0.5 * arg2[:, :, 1]) # Bottom boundary
132 cc = cc.at[:, :, nz-1].set(arg1[:, :, nz-1] + arg2[:, :, nz-1]) # Top boundary
133
134 Cy = Dealias2(FFT_pad(cc), ZeRo3D_fft)
135
136 # Compute Cz: advection term in z-direction
137 arg1 = ZeRo3D_pad.copy()
138 arg2 = ZeRo3D_pad.copy()
139
140 # Average u and v components to staggered grid points
141 u_pad_G = StagGridAvg(u_pad)
142 v_pad_G = StagGridAvg(v_pad)
143
144 # Compute interior points for both terms
145 arg1 = arg1.at[:, :, 1:nz-1].set(u_pad_G[:, :, 0:nz-2] * (dwdx_pad[:, :, 1:nz-1] - dudz_pad[:, :, 1:nz-1]))
146 arg2 = arg2.at[:, :, 1:nz-1].set(v_pad_G[:, :, 0:nz-2] * (dwdy_pad[:, :, 1:nz-1] - dvdz_pad[:, :, 1:nz-1]))
147
148 # Combine terms and set boundary conditions
149 cc = arg1 + arg2
150 cc = cc.at[:, :, 0].set(0) # Bottom boundary
151 cc = cc.at[:, :, nz-1].set(0) # Top boundary
152
153 Cz = Dealias2(FFT_pad(cc), ZeRo3D_fft)
154
155 return Cx, Cy, Cz
156
157
158# ------------------------------------------------------------------------------
159# Compute the advection terms without using any dealiasing
160# ------------------------------------------------------------------------------
161
162@jax.jit
163def Advection_NoDealias(u, v, w,
164 dudy, dudz,
165 dvdx, dvdz,
166 dwdx, dwdy,
167 ZeRo3D):
168 """
169 Parameters:
170 -----------
171 u, v, w : ndarray, shape: (nx, ny, nz)
172 Velocity components in x, y, and z directions respectively.
173 dudy, dudz : ndarray
174 Derivatives of u-velocity with respect to y and z.
175 dvdx, dvdz : ndarray
176 Derivatives of v-velocity with respect to x and z.
177 dwdx, dwdy : ndarray
178 Derivatives of w-velocity with respect to x and y.
179
180 Returns:
181 --------
182 tuple[ndarray, ndarray, ndarray]
183 Cx : advection term in x-direction
184 Cy : advection term in y-direction
185 Cz : advection term in z-direction
186 Each array has shape (nx, ny, nz)
187
188 Notes:
189 ------
190 - Function uses staggered grid averaging (StagGridAvg) for proper
191 interpolation of velocities and derivatives
192 """
193 # Compute Cx: advection term in x-direction
194 # First term: v*(∂u/∂y - ∂v/∂x)
195 arg1 = v * (dudy - dvdx)
196 # Second term: w*(∂u/∂z - ∂w/∂x)
197 arg2 = w * (dudz - dwdx)
198
199 # Apply staggered grid averaging to second term
200 arg2G = StagGridAvg(arg2)
201
202 # Combine terms with boundary condition handling
203 Cx = ZeRo3D.copy()
204 Cx = Cx.at[:, :, 1:nz-1].set(arg1[:, :, 1:nz-1] + arg2G[:, :, 1:nz-1]) # Interior points
205 Cx = Cx.at[:, :, 0].set(arg1[:, :, 0] + 0.5 * arg2[:, :, 1]) # Bottom boundary
206 Cx = Cx.at[:, :, nz-1].set(arg1[:, :, nz-1] + arg2[:, :, nz-1]) # Top boundary
207
208 # Compute Cy: advection term in y-direction
209 # First term: u*(∂v/∂x - ∂u/∂y)
210 arg1 = u * (dvdx - dudy)
211 # Second term: w*(∂v/∂z - ∂w/∂y)
212 arg2 = w * (dvdz - dwdy)
213 arg2G = StagGridAvg(arg2)
214
215 Cy = ZeRo3D.copy()
216 Cy = Cy.at[:, :, 1:nz-1].set(arg1[:, :, 1:nz-1] + arg2G[:, :, 1:nz-1]) # Interior points
217 Cy = Cy.at[:, :, 0].set(arg1[:, :, 0] + 0.5 * arg2[:, :, 1]) # Bottom boundary
218 Cy = Cy.at[:, :, nz-1].set(arg1[:, :, nz-1] + arg2[:, :, nz-1]) # Top boundary
219
220 # Compute Cz: advection term in z-direction
221 # Special handling required due to staggered grid arrangement
222 arg1 = ZeRo3D.copy()
223 arg2 = ZeRo3D.copy()
224
225 # Average u and v components to staggered grid points
226 u_G = StagGridAvg(u)
227 v_G = StagGridAvg(v)
228
229 # Compute interior points for both terms
230 arg1 = arg1.at[:, :, 1:nz-1].set(u_G[:, :, 0:nz-2] * (dwdx[:, :, 1:nz-1] - dudz[:, :, 1:nz-1]))
231 arg2 = arg2.at[:, :, 1:nz-1].set(v_G[:, :, 0:nz-2] * (dwdy[:, :, 1:nz-1] - dvdz[:, :, 1:nz-1]))
232
233 # Combine terms and set boundary conditions
234 Cz = arg1 + arg2
235 Cz = Cz.at[:, :, 0].set(0) # Bottom boundary
236 Cz = Cz.at[:, :, nz-1].set(0) # Top boundary
237
238 return Cx, Cy, Cz
239
240
241# ------------------------------------------------------------------------------
242# Select one of the above functions based on dealias flag
243# ------------------------------------------------------------------------------
244
245@jax.jit
246def Advection(
247 u, v, w,
248 dudy, dudz,
249 dvdx, dvdz,
250 dwdx, dwdy,
251 ZeRo3D, ZeRo3D_fft,
252 ZeRo3D_pad, ZeRo3D_pad_fft):
253
254 if optDealias == 1:
255
256 Cx, Cy, Cz = (
257 Advection_Dealias(
258 u, v, w,
259 dudy, dudz,
260 dvdx, dvdz,
261 dwdx, dwdy,
262 ZeRo3D_fft,
263 ZeRo3D_pad, ZeRo3D_pad_fft))
264
265 else:
266
267 Cx, Cy, Cz = (
268 Advection_NoDealias(
269 u, v, w,
270 dudy, dudz,
271 dvdx, dvdz,
272 dwdx, dwdy,
273 ZeRo3D))
274
275 return Cx, Cy, Cz