Static SGS Modeling: Main Code
Source Code: StaticSGS_Main
StaticSGS_Main.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: StaticSGS_Main.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-29
24:Description: static SGS modeling - main code
25"""
26
27# ============================================================
28# Imports
29# ============================================================
30
31import jax
32
33# Import derived variables
34from ..config.DerivedVars import *
35
36# Import FFT modules
37from ..operations.FFT import FFT
38
39# Import dealiasing functions
40from ..operations.Dealiasing import Dealias1
41
42# Import strain rates functions
43from .StrainRates import StrainsUVPnodes_Dealias, StrainsWnodes_Dealias
44from .StrainRates import StrainsUVPnodes_NoDealias, StrainsWnodes_NoDealias
45
46# Import stress functions (SM and WL)
47from .SGSStresses_SM import (
48 StressesUVPnodes_Dealias as StressesUVPnodes_Dealias_SM,
49 StressesUVPnodes_NoDealias as StressesUVPnodes_NoDealias_SM,
50 StressesWnodes_Dealias as StressesWnodes_Dealias_SM,
51 StressesWnodes_NoDealias as StressesWnodes_NoDealias_SM)
52from .SGSStresses_WL import (
53 StressesUVPnodes_Dealias as StressesUVPnodes_Dealias_WL,
54 StressesUVPnodes_NoDealias as StressesUVPnodes_NoDealias_WL,
55 StressesWnodes_Dealias as StressesWnodes_Dealias_WL,
56 StressesWnodes_NoDealias as StressesWnodes_NoDealias_WL)
57
58# Import scalar flux functions (SM and WL)
59from .ScalarSGSFluxes_SM import (
60 ScalarFluxesUVPnodes_Dealias as ScalarFluxesUVPnodes_Dealias_SM,
61 ScalarFluxesUVPnodes_NoDealias as ScalarFluxesUVPnodes_NoDealias_SM,
62 ScalarFluxesWnodes_Dealias as ScalarFluxesWnodes_Dealias_SM,
63 ScalarFluxesWnodes_NoDealias as ScalarFluxesWnodes_NoDealias_SM)
64from .ScalarSGSFluxes_WL import (
65 ScalarFluxesUVPnodes_Dealias as ScalarFluxesUVPnodes_Dealias_WL,
66 ScalarFluxesUVPnodes_NoDealias as ScalarFluxesUVPnodes_NoDealias_WL,
67 ScalarFluxesWnodes_Dealias as ScalarFluxesWnodes_Dealias_WL,
68 ScalarFluxesWnodes_NoDealias as ScalarFluxesWnodes_NoDealias_WL)
69
70
71# ============================================================
72# Static SGS: compute all the SGS stresses on proper nodes
73# ============================================================
74
75@jax.jit
76def StaticSGS(
77 dudx, dvdx, dwdx,
78 dudy, dvdy, dwdy,
79 dudz, dvdz, dwdz,
80 Cs2_3D,
81 u, v, M_sfc_loc, psi2D_m, psi2D_m0,
82 ZeRo3D, ZeRo3D_fft, ZeRo3D_pad_fft):
83 """
84 Parameters:
85 -----------
86 dudx, dvdx, dwdx, dudy, dvdy, dwdy, dudz, dvdz, dwdz : ndarray
87 Velocity gradients
88 Cs2_3D : ndarray
89 Smagorinsky coefficient field
90 u, v : ndarray
91 Velocity components for wall model
92 M_sfc_loc, psi2D_m, psi2D_m0 : ndarray
93 Surface parameters for wall model
94 ZeRo3D, ZeRo3D_fft, ZeRo3D_pad_fft : ndarray
95 Pre-allocated zero arrays
96
97 Returns:
98 --------
99 txx, tyy, tzz, txy, txz, tyz : ndarray
100 SGS stress components
101 S_uvp, S_uvp_pad, S_w, S_w_pad : ndarray
102 Strain rate fields
103 """
104
105 # ----------------------------------------
106 # Compute txx, tyy, tzz and txy components
107 # ----------------------------------------
108 if optDealias == 1:
109
110 # --------------------------------------
111 # Compute strain rates
112 # --------------------------------------
113 (S11, S22, S33,
114 S12, S13, S23,
115 S_uvp,
116 S11_pad, S22_pad, S33_pad,
117 S12_pad, S13_pad, S23_pad,
118 S_uvp_pad) = (
119 StrainsUVPnodes_Dealias(
120 dudx, dvdx, dwdx,
121 dudy, dvdy, dwdy,
122 dudz, dvdz, dwdz,
123 ZeRo3D, ZeRo3D_pad_fft))
124
125 # --------------------------------------
126 # Dealias Cs2 field
127 # --------------------------------------
128 Cs2_3D_pad = Dealias1(FFT(Cs2_3D), ZeRo3D_pad_fft)
129
130 # --------------------------------------
131 # Compute SGS stresses
132 # --------------------------------------
133 if optSgs in [2, 4]: # WL variants
134 (txx, tyy, tzz, txy) = (
135 StressesUVPnodes_Dealias_WL(
136 S11_pad, S22_pad, S33_pad, S12_pad,
137 Cs2_3D_pad,
138 ZeRo3D_fft))
139 else: # SM variants (optSgs in [0, 1, 3])
140 (txx, tyy, tzz, txy) = (
141 StressesUVPnodes_Dealias_SM(
142 S11_pad, S22_pad, S33_pad, S12_pad,
143 S_uvp_pad,
144 Cs2_3D_pad,
145 ZeRo3D_fft))
146
147 else:
148
149 # --------------------------------------
150 # Compute strain rates
151 # --------------------------------------
152 (S11, S22, S33,
153 S12, S13, S23,
154 S_uvp) = (
155 StrainsUVPnodes_NoDealias(
156 dudx, dvdx, dwdx,
157 dudy, dvdy, dwdy,
158 dudz, dvdz, dwdz,
159 ZeRo3D))
160
161 # create a dummy variable for passing
162 S_uvp_pad = S_uvp
163
164 # --------------------------------------
165 # Compute SGS stresses
166 # --------------------------------------
167 if optSgs in [2, 4]: # WL variants
168 (txx, tyy, tzz, txy) = (
169 StressesUVPnodes_NoDealias_WL(
170 S11, S22, S33, S12,
171 Cs2_3D))
172 else: # SM variants (optSgs in [0, 1, 3])
173 (txx, tyy, tzz, txy) = (
174 StressesUVPnodes_NoDealias_SM(
175 S11, S22, S33, S12,
176 S_uvp,
177 Cs2_3D))
178
179 # ------------------------------------------------------------
180 # Compute txz, tyz components
181 # ------------------------------------------------------------
182 if optDealias == 1:
183
184 # --------------------------------------
185 # Compute strain rates
186 # --------------------------------------
187 (S13_pad, S23_pad,
188 S_w_pad) = (
189 StrainsWnodes_Dealias(
190 dudx, dvdx, dwdx,
191 dudy, dvdy, dwdy,
192 dudz, dvdz, dwdz,
193 ZeRo3D, ZeRo3D_pad_fft))
194
195 # create a dummy variable for passing
196 S_w = S_w_pad
197
198 # --------------------------------------
199 # Compute SGS stresses
200 # --------------------------------------
201 if optSgs in [2, 4]: # WL variants
202 (txz, tyz) = (
203 StressesWnodes_Dealias_WL(
204 S13_pad, S23_pad,
205 Cs2_3D_pad,
206 u, v, M_sfc_loc, psi2D_m, psi2D_m0,
207 ZeRo3D_fft))
208 else: # SM variants (optSgs in [0, 1, 3])
209 (txz, tyz) = (
210 StressesWnodes_Dealias_SM(
211 S13_pad, S23_pad,
212 S_w_pad,
213 Cs2_3D_pad,
214 u, v, M_sfc_loc, psi2D_m, psi2D_m0,
215 ZeRo3D_fft))
216
217 else:
218
219 # --------------------------------------
220 # Compute strain rates
221 # --------------------------------------
222 (S13, S23,
223 S_w) = (
224 StrainsWnodes_NoDealias(
225 dudx, dvdx, dwdx,
226 dudy, dvdy, dwdy,
227 dudz, dvdz, dwdz,
228 ZeRo3D))
229
230 # create a dummy variable for passing
231 S_w_pad = S_w
232
233 # --------------------------------------
234 # Compute SGS stresses
235 # --------------------------------------
236 if optSgs in [2, 4]: # WL variants
237 (txz, tyz) = (
238 StressesWnodes_NoDealias_WL(
239 S13, S23,
240 Cs2_3D,
241 u, v, M_sfc_loc, psi2D_m, psi2D_m0))
242 else: # SM variants (optSgs in [0, 1, 3])
243 (txz, tyz) = (
244 StressesWnodes_NoDealias_SM(
245 S13, S23,
246 S_w,
247 Cs2_3D,
248 u, v, M_sfc_loc, psi2D_m, psi2D_m0))
249
250 # Return stresses along with strain rates for scalar calculations
251 return (txx, tyy, tzz, txy, txz, tyz,
252 S_uvp, S_uvp_pad,
253 S_w, S_w_pad)
254
255
256# =====================================================
257# Static SGS: compute scalar SGS fluxes on proper nodes
258# =====================================================
259
260@jax.jit
261def StaticSGSscalar(
262 S_uvp, S_uvp_pad,
263 S_w, S_w_pad,
264 Cs2PrRatio_3D,
265 dTHdx, dTHdy, dTHdz,
266 qz_sfc,
267 ZeRo3D_fft, ZeRo3D_pad_fft):
268 """
269 Parameters:
270 -----------
271 S_uvp, S_uvp_pad : ndarray
272 Strain rate magnitudes at UVP nodes
273 S_w, S_w_pad : ndarray
274 Strain rate magnitudes at W nodes
275 Cs2PrRatio_3D : ndarray
276 Cs^2/Pr_t field
277 dTHdx, dTHdy, dTHdz : ndarray
278 Potential temperature gradients
279 qz_sfc : ndarray
280 Surface heat flux
281 ZeRo3D_fft, ZeRo3D_pad_fft : ndarray
282 Pre-allocated zero arrays
283
284 Returns:
285 --------
286 qx, qy, qz : ndarray
287 SGS scalar flux components
288 """
289
290 # ------------------------------------------------------------
291 # Compute qx, qy and qz components
292 # ------------------------------------------------------------
293 if optDealias == 1:
294
295 dTHdx_pad = Dealias1(FFT(dTHdx), ZeRo3D_pad_fft)
296 dTHdy_pad = Dealias1(FFT(dTHdy), ZeRo3D_pad_fft)
297 dTHdz_pad = Dealias1(FFT(dTHdz), ZeRo3D_pad_fft)
298
299 Cs2PrRatio_3D_pad = Dealias1(FFT(Cs2PrRatio_3D), ZeRo3D_pad_fft)
300
301 # Compute fluxes at UVP nodes
302 if optSgs in [2, 4]: # WL variants
303 (qx, qy) = (
304 ScalarFluxesUVPnodes_Dealias_WL(
305 dTHdx_pad, dTHdy_pad,
306 Cs2PrRatio_3D_pad,
307 ZeRo3D_fft))
308 else: # SM variants
309 (qx, qy) = (
310 ScalarFluxesUVPnodes_Dealias_SM(
311 dTHdx_pad, dTHdy_pad,
312 S_uvp_pad,
313 Cs2PrRatio_3D_pad,
314 ZeRo3D_fft))
315
316 # Compute flux on W nodes
317 if optSgs in [2, 4]: # WL variants
318 qz = (
319 ScalarFluxesWnodes_Dealias_WL(
320 dTHdz_pad,
321 Cs2PrRatio_3D_pad,
322 qz_sfc,
323 ZeRo3D_fft))
324 else: # SM variants
325 qz = (
326 ScalarFluxesWnodes_Dealias_SM(
327 dTHdz_pad,
328 S_w_pad,
329 Cs2PrRatio_3D_pad,
330 qz_sfc,
331 ZeRo3D_fft))
332
333 else:
334
335 # Compute fluxes at UVP nodes
336 if optSgs in [2, 4]: # WL variants
337 (qx, qy) = (
338 ScalarFluxesUVPnodes_NoDealias_WL(
339 dTHdx, dTHdy,
340 Cs2PrRatio_3D))
341 else: # SM variants
342 (qx, qy) = (
343 ScalarFluxesUVPnodes_NoDealias_SM(
344 dTHdx, dTHdy,
345 S_uvp,
346 Cs2PrRatio_3D))
347
348 # Compute flux on W nodes
349 if optSgs in [2, 4]: # WL variants
350 qz = (
351 ScalarFluxesWnodes_NoDealias_WL(
352 dTHdz,
353 Cs2PrRatio_3D,
354 qz_sfc))
355 else: # SM variants
356 qz = (
357 ScalarFluxesWnodes_NoDealias_SM(
358 dTHdz,
359 S_w,
360 Cs2PrRatio_3D,
361 qz_sfc))
362
363 return qx, qy, qz