SGS Terms
Source Code: NSE_SGSTerms
NSE_SGSTerms.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_SGSTerms.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-6
24:Description: computes the SGS stress divergence terms
25"""
26
27# ============================================
28# Imports
29# ============================================
30
31import jax
32
33# Import configuration from namelist
34from ..config import ConfigLoader as Config
35
36# Import derived variables
37from ..config.DerivedVars import *
38
39# Import FFT modules
40from ..operations.FFT import FFT
41
42# Import derivative functions
43from ..operations.Derivatives import Derivxy, Derivz_Generic_uvp, Derivz_Generic_w
44
45# Import SGS models
46from ..subgridscale.DynamicSGS_Main import DynamicSGS
47from ..subgridscale.StaticSGS_Main import StaticSGS
48
49
50# =================================================
51# Compute Divergence of Stress Tensor
52# =================================================
53
54@jax.jit
55def DivStress(
56 txx, tyy, tzz,
57 txy, txz, tyz,
58 ZeRo3D,
59 kx2, ky2):
60 """
61 Parameters:
62 -----------
63 txx, tyy, tzz : ndarray of shape (nx, ny, nz)
64 Normal stress components
65 txy, txz, tyz : ndarray of shape (nx, ny, nz)
66 Shear stress components
67 kx2, ky2 : ndarray
68 Wavenumber arrays for spectral derivatives
69 ZeRo3D : ndarray of shape (nx, ny, nz)
70 Pre-allocated zero array for intermediate calculations
71
72 Returns:
73 --------
74 tuple containing:
75 divtx : ndarray of shape (nx, ny, nz)
76 x-component of stress divergence
77 divty : ndarray of shape (nx, ny, nz)
78 y-component of stress divergence
79 divtz : ndarray of shape (nx, ny, nz)
80 z-component of stress divergence
81 """
82 # Compute derivatives in x-direction
83 dxtxx = Derivxy(FFT(txx), kx2)
84 dxtxy = Derivxy(FFT(txy), kx2)
85 dxtxz = Derivxy(FFT(txz), kx2)
86
87 # Compute derivatives in y-direction
88 dytyy = Derivxy(FFT(tyy), ky2)
89 dytxy = Derivxy(FFT(txy), ky2)
90 dytyz = Derivxy(FFT(tyz), ky2)
91
92 # Compute derivatives in z-direction
93 # For txz and tyz, use w-grid derivative
94 dztxz = Derivz_Generic_w(txz, ZeRo3D)
95 dztyz = Derivz_Generic_w(tyz, ZeRo3D)
96
97 # For tzz, use uvp-grid derivative
98 dztzz = Derivz_Generic_uvp(tzz, ZeRo3D)
99
100 # Compute divergence components
101 divtx = dxtxx + dytxy + dztxz
102 divty = dxtxy + dytyy + dztyz
103 divtz = dxtxz + dytyz + dztzz
104
105 return divtx, divty, divtz
106
107
108# =================================================
109# Compute stress divergence using dynamic SGS models
110# =================================================
111
112@jax.jit
113def DivStressDynamicSGS(
114 dudx, dvdx, dwdx,
115 dudy, dvdy, dwdy,
116 dudz, dvdz, dwdz,
117 u, v, w, M_sfc_loc, MOSTfunctions,
118 ZeRo3D, ZeRo3D_fft, ZeRo3D_pad_fft,
119 kx2, ky2):
120 """
121 Compute stress divergence terms using dynamic SGS models.
122
123 Parameters:
124 -----------
125 dudx, dvdx, dwdx : ndarray of shape (nx, ny, nz)
126 Derivatives of velocity components in x-direction
127 dudy, dvdy, dwdy : ndarray of shape (nx, ny, nz)
128 Derivatives of velocity components in y-direction
129 dudz, dvdz, dwdz : ndarray of shape (nx, ny, nz)
130 Derivatives of velocity components in z-direction
131 u, v, w : ndarray of shape (nx, ny, nz)
132 Velocity components
133 M_sfc_loc : ndarray of shape (nx, ny)
134 Near-surface wind speed
135 psi2D_m, psi2D_m0 : ndarray of shape (nx, ny)
136 Stability correction functions
137 ZeRo3D, ZeRo3D_fft, ZeRo3D_pad_fft : ndarray
138 Pre-allocated arrays for calculations
139 kx2, ky2 : ndarray
140 Wavenumber arrays for spectral derivatives
141
142 Returns:
143 --------
144 divtx, divty, divtz : ndarray of shape (nx, ny, nz)
145 Components of stress divergence
146 Cs2_1D_avg1, Cs2_1D_avg2 : ndarray of shape (nz)
147 Profiles of dynamic Smagorinsky coefficient
148 beta1_1D : ndarray of shape (nz)
149 Profile of filter width ratio
150 dynamicSGSmomentum : tuple
151 Complete set of results from dynamic SGS calculation
152 """
153
154 # Unpack MOSTfunctions
155 (psi2D_m, psi2D_m0, _, _, _, _) = MOSTfunctions
156
157 # Call DynamicSGS to get stresses and other variables
158 dynamicSGSmomentum = DynamicSGS(
159 dudx, dvdx, dwdx,
160 dudy, dvdy, dwdy,
161 dudz, dvdz, dwdz,
162 u, v, w, M_sfc_loc, psi2D_m, psi2D_m0,
163 ZeRo3D, ZeRo3D_fft, ZeRo3D_pad_fft)
164
165 # Unpack only the variables we need
166 txx, tyy, tzz, txy, txz, tyz = dynamicSGSmomentum[0:6]
167 Cs2_1D_avg1, Cs2_1D_avg2, beta1_1D = dynamicSGSmomentum[6:9]
168 Cs2_3D = dynamicSGSmomentum[9]
169
170 # Compute divergence of stress
171 divtx, divty, divtz = DivStress(txx, tyy, tzz,
172 txy, txz, tyz,
173 ZeRo3D,
174 kx2, ky2)
175
176 # Return both the divergence terms and the complete momentum results,
177 # so they can be reused for scalar calculations
178 return divtx, divty, divtz, Cs2_1D_avg1, Cs2_1D_avg2, beta1_1D, Cs2_3D, dynamicSGSmomentum
179
180
181# =================================================
182# Compute stress divergence using static SGS models
183# =================================================
184
185@jax.jit
186def DivStressStaticSGS(
187 dudx, dvdx, dwdx,
188 dudy, dvdy, dwdy,
189 dudz, dvdz, dwdz,
190 Cs2_3D,
191 u, v, M_sfc_loc, MOSTfunctions,
192 ZeRo3D, ZeRo3D_fft, ZeRo3D_pad_fft,
193 kx2, ky2):
194 """
195 Compute stress divergence terms using static SGS models.
196
197 Parameters:
198 -----------
199 dudx, dvdx, dwdx : ndarray of shape (nx, ny, nz)
200 Derivatives of velocity components in x-direction
201 dudy, dvdy, dwdy : ndarray of shape (nx, ny, nz)
202 Derivatives of velocity components in y-direction
203 dudz, dvdz, dwdz : ndarray of shape (nx, ny, nz)
204 Derivatives of velocity components in z-direction
205 Cs2_3D : ndarray of shape (nx, ny, nz)
206 Static Smagorinsky coefficient
207 u, v : ndarray of shape (nx, ny, nz)
208 Velocity components
209 M_sfc_loc : ndarray of shape (nx, ny)
210 Near-surface wind speed
211 psi2D_m, psi2D_m0 : ndarray of shape (nx, ny)
212 Stability correction functions
213 ZeRo3D, ZeRo3D_fft, ZeRo3D_pad_fft : ndarray
214 Pre-allocated arrays for calculations
215 kx2, ky2 : ndarray
216 Wavenumber arrays for spectral derivatives
217
218 Returns:
219 --------
220 divtx, divty, divtz : ndarray of shape (nx, ny, nz)
221 Components of stress divergence
222 staticSGSmomentum : tuple
223 Complete set of results from static SGS calculation
224 """
225
226 # Unpack MOSTfunctions
227 (psi2D_m, psi2D_m0, _, _, _, _) = MOSTfunctions
228
229 # Static SGS model
230 staticSGSmomentum = StaticSGS(
231 dudx, dvdx, dwdx,
232 dudy, dvdy, dwdy,
233 dudz, dvdz, dwdz,
234 Cs2_3D,
235 u, v, M_sfc_loc, psi2D_m, psi2D_m0,
236 ZeRo3D, ZeRo3D_fft, ZeRo3D_pad_fft)
237
238 # Unpack only the stress components we need
239 txx, tyy, tzz, txy, txz, tyz = staticSGSmomentum[0:6]
240
241 # Compute divergence of stress
242 divtx, divty, divtz = DivStress(txx, tyy, tzz,
243 txy, txz, tyz,
244 ZeRo3D,
245 kx2, ky2)
246
247 # Return both the divergence terms and the complete momentum results,
248 # so they can be reused for scalar calculations
249 return divtx, divty, divtz, staticSGSmomentum