Scalar SGS Flux Computations: Wong-Lilly
Computes SGS scalar (potential temperature) fluxes using the Wong-Lilly (WL) base formulation:
\[q_i = -\frac{C_{WL}}{Pr_t} \Delta^{4/3} \frac{\partial \bar{\theta}}{\partial x_i}\]
Note the absence of the strain rate magnitude factor compared to the
Smagorinsky formulation. Used for optSgs = 2 (LASDD-WL) and
optSgs = 4 (LAD-WL).
Source Code: ScalarSGSFluxes_WL
ScalarSGSFluxes_WL.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: ScalarSGSFluxes_WL.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: 2026-5-9
24:Description: computes SGS scalar fluxes using the Wong-Lilly (1994)
25 SGS base model. Flux formula:
26 q_i = -(C_WL/Pr_t) * Delta^(4/3) * (dTH/dx_i)
27 (no strain-rate magnitude factor; filter width exponent
28 4/3 instead of 2; no leading factor of 2).
29"""
30
31# ============================================================
32# Imports
33# ============================================================
34
35import jax
36import jax.numpy as jnp
37
38# Import derived variables
39from ..config.DerivedVars import *
40
41# Import FFT modules
42from ..operations.FFT import FFT, FFT_pad
43
44# Import helper functions
45from ..utilities.Utilities import StagGridAvg
46
47# Import dealiasing functions
48from ..operations.Dealiasing import Dealias1, Dealias2
49
50
51# ======================================================
52# Compute SGS scalar fluxes on UVP nodes with dealiasing
53# ======================================================
54
55@jax.jit
56def ScalarFluxesUVPnodes_Dealias(
57 dTHdx_pad, dTHdy_pad,
58 CwlPrRatio_3D_pad,
59 ZeRo3D_fft):
60 """
61 Parameters:
62 -----------
63 dTHdx_pad, dTHdy_pad : ndarray
64 Dealiased potential temperature gradients
65 CwlPrRatio_3D_pad : ndarray
66 Dealiased C_WL/Pr_t field
67 ZeRo3D_fft : ndarray
68 Pre-allocated zero array for dealiasing
69
70 Returns:
71 --------
72 qx, qy : ndarray
73 SGS scalar flux components in x and y directions
74 """
75
76 preCompute = -(L ** (4 / 3)) * CwlPrRatio_3D_pad
77 qx_pad = preCompute * dTHdx_pad
78 qy_pad = preCompute * dTHdy_pad
79
80 qx_pad = qx_pad.at[:, :, nz - 1].set(0)
81 qy_pad = qy_pad.at[:, :, nz - 1].set(0)
82
83 qx = Dealias2(FFT_pad(qx_pad), ZeRo3D_fft)
84 qy = Dealias2(FFT_pad(qy_pad), ZeRo3D_fft)
85
86 return qx, qy
87
88
89# =========================================================
90# Compute SGS scalar fluxes on UVP nodes without dealiasing
91# =========================================================
92
93@jax.jit
94def ScalarFluxesUVPnodes_NoDealias(
95 dTHdx, dTHdy,
96 CwlPrRatio_3D):
97 """
98 Parameters:
99 -----------
100 dTHdx, dTHdy : ndarray
101 Potential temperature gradients at UVP nodes
102 CwlPrRatio_3D : ndarray
103 C_WL/Pr_t field
104
105 Returns:
106 --------
107 qx, qy : ndarray
108 SGS scalar flux components in x and y directions
109 """
110
111 preCompute = -(L ** (4 / 3)) * CwlPrRatio_3D
112 qx = preCompute * dTHdx
113 qy = preCompute * dTHdy
114
115 qx = qx.at[:, :, nz - 1].set(0)
116 qy = qy.at[:, :, nz - 1].set(0)
117
118 return qx, qy
119
120
121# ====================================================
122# Compute SGS scalar fluxes on W nodes with dealiasing
123# ====================================================
124
125@jax.jit
126def ScalarFluxesWnodes_Dealias(
127 dTHdz_pad,
128 CwlPrRatio_3D_pad,
129 qz_sfc,
130 ZeRo3D_fft):
131 """
132 Parameters:
133 -----------
134 dTHdz_pad : ndarray
135 Dealiased potential temperature gradient in z-direction
136 CwlPrRatio_3D_pad : ndarray
137 Dealiased C_WL/Pr_t field
138 qz_sfc : ndarray
139 Surface heat flux
140 ZeRo3D_fft : ndarray
141 Pre-allocated zero array for dealiasing
142
143 Returns:
144 --------
145 qz : ndarray
146 SGS scalar flux component in z-direction
147 """
148
149 qz_pad = jnp.zeros_like(dTHdz_pad)
150
151 qz_pad = qz_pad.at[:, :, 1:nz - 1].set(
152 -(L ** (4 / 3)) *
153 StagGridAvg(CwlPrRatio_3D_pad[:, :, :nz - 1]) *
154 dTHdz_pad[:, :, 1:nz - 1]
155 )
156
157 qz_pad = qz_pad.at[:, :, nz - 1].set(0)
158
159 qz = Dealias2(FFT_pad(qz_pad), ZeRo3D_fft)
160 qz = qz.at[:, :, 0].set(qz_sfc)
161
162 return qz
163
164
165# =======================================================
166# Compute SGS scalar fluxes on W nodes without dealiasing
167# =======================================================
168
169@jax.jit
170def ScalarFluxesWnodes_NoDealias(
171 dTHdz,
172 CwlPrRatio_3D,
173 qz_sfc):
174 """
175 Parameters:
176 -----------
177 dTHdz : ndarray
178 Potential temperature gradient in z-direction
179 CwlPrRatio_3D : ndarray
180 C_WL/Pr_t field
181 qz_sfc : ndarray
182 Surface heat flux
183
184 Returns:
185 --------
186 qz : ndarray
187 SGS scalar flux component in z-direction
188 """
189
190 qz = jnp.zeros_like(dTHdz)
191
192 qz = qz.at[:, :, 1:nz - 1].set(
193 -(L ** (4 / 3)) *
194 StagGridAvg(CwlPrRatio_3D[:, :, :nz - 1]) *
195 dTHdz[:, :, 1:nz - 1]
196 )
197
198 qz = qz.at[:, :, nz - 1].set(0)
199 qz = qz.at[:, :, 0].set(qz_sfc)
200
201 return qz