SGS Stress Computations: Smagorinsky

Computes SGS momentum stresses using the Smagorinsky (SM) base formulation:

\[\tau_{ij} = -2 C_s^2 \Delta^2 |\bar{S}| \bar{S}_{ij}\]

Used for optSgs = 1 (LASDD-SM) and optSgs = 3 (LAD-SM), both on dynamic call steps and on cached (non-call) steps when dynamicSGS_call_time > 1.

Source Code: SGSStresses_SM

SGSStresses_SM.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: SGSStresses.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: computes SGS stresses
 25"""
 26
 27# ============================================================
 28#  Imports
 29# ============================================================
 30
 31import jax
 32import jax.numpy as jnp
 33
 34# Import configuration from namelist
 35from ..config.ConfigLoader import *
 36
 37# Import derived variables
 38from ..config.DerivedVars import *
 39
 40# Import FFT modules
 41from ..operations.FFT import FFT, FFT_pad
 42
 43# Import dealiasing functions
 44from ..operations.Dealiasing import Dealias1, Dealias2
 45
 46# Import strain rates functions
 47from .StrainRates import StrainsUVPnodes_Dealias, StrainsWnodes_Dealias
 48from .StrainRates import StrainsUVPnodes_NoDealias, StrainsWnodes_NoDealias
 49
 50# Import LASDD
 51from .DynamicSGS_LASDD_SM import LASDD
 52from .DynamicSGS_ScalarLASDD_SM import ScalarLASDD
 53
 54# Import helper functions
 55from ..utilities.Utilities import StagGridAvg
 56
 57
 58# =========================================
 59# Wall function to compute surface stresses
 60# =========================================
 61
 62@jax.jit
 63def Wall(u, v, M_sfc_loc, psi2D_m, psi2D_m0):
 64    """
 65    Parameters:
 66    -----------
 67    u, v : ndarray
 68        Velocity components near the surface (z = dz/2)
 69    M_sfc_loc : ndarray
 70        Near-surface (z = dz/2) wind speed
 71    psi2D_m, psi2D_m0 : ndarray
 72        Stability correction functions at first level and surface
 73
 74    Returns:
 75    --------
 76    txz, tyz : ndarray
 77        Surface stresses in x and y directions
 78    """
 79
 80    # Compute denominator for friction velocity
 81    denom = jnp.log(0.5 * dz * z_scale / z0m) + psi2D_m - psi2D_m0
 82
 83    # Compute friction velocity
 84    ustar = vonk * M_sfc_loc / denom
 85
 86    # Compute wall stresses
 87    txz = -(ustar ** 2) * (u[:, :, 0] + Ugal) / M_sfc_loc
 88    tyz = -(ustar ** 2) * v[:, :, 0] / M_sfc_loc
 89
 90    return txz, tyz
 91
 92
 93# =================================================
 94# Compute SGS stresses on UVP nodes with dealiasing
 95# =================================================
 96
 97@jax.jit
 98def StressesUVPnodes_Dealias(
 99        S11_pad, S22_pad, S33_pad, S12_pad,
100        S_uvp_pad,
101        Cs2_3D_pad,
102        ZeRo3D_fft):
103    """
104    Parameters:
105    -----------
106    S11_pad, S22_pad, S33_pad, S12_pad : ndarray
107        Dealiased strain rate components
108    S_uvp_pad : ndarray
109        Dealiased strain rate magnitude
110    Cs2_3D_pad : ndarray
111        Dealiased Smagorinsky coefficient field
112    ZeRo3D_fft : ndarray
113        Pre-allocated zero array for dealiasing
114
115    Returns:
116    --------
117    txx, tyy, tzz, txy : ndarray
118        SGS stress components at UVP nodes
119    """
120
121    # Compute SGS stresses on uvp nodes
122    preCompute = -2 * (L ** 2) * Cs2_3D_pad * S_uvp_pad
123    txx_pad = preCompute * S11_pad
124    tyy_pad = preCompute * S22_pad
125    tzz_pad = preCompute * S33_pad
126    txy_pad = preCompute * S12_pad
127
128    # Set top boundary conditions
129    txx_pad = txx_pad.at[:, :, nz - 1].set(0)
130    tyy_pad = tyy_pad.at[:, :, nz - 1].set(0)
131    tzz_pad = tzz_pad.at[:, :, nz - 1].set(0)
132    txy_pad = txy_pad.at[:, :, nz - 1].set(0)
133
134    # Apply dealiasing
135    txx = Dealias2(FFT_pad(txx_pad), ZeRo3D_fft)
136    tyy = Dealias2(FFT_pad(tyy_pad), ZeRo3D_fft)
137    tzz = Dealias2(FFT_pad(tzz_pad), ZeRo3D_fft)
138    txy = Dealias2(FFT_pad(txy_pad), ZeRo3D_fft)
139
140    return txx, tyy, tzz, txy
141
142
143# ====================================================
144# Compute SGS stresses on UVP nodes without dealiasing
145# ====================================================
146
147
148@jax.jit
149def StressesUVPnodes_NoDealias(
150        S11, S22, S33, S12,
151        S_uvp,
152        Cs2_3D):
153    """
154    Parameters:
155    -----------
156    S11, S22, S33, S12 : ndarray
157        Strain rate components
158    S_uvp : ndarray
159        Strain rate magnitude
160    Cs2_3D : ndarray
161        Smagorinsky coefficient field
162
163    Returns:
164    --------
165    txx, tyy, tzz, txy : ndarray
166        SGS stress components at UVP nodes
167    """
168
169    # Compute SGS stresses on uvp nodes
170    preCompute = -2 * (L ** 2) * Cs2_3D * S_uvp
171    txx = preCompute * S11
172    tyy = preCompute * S22
173    tzz = preCompute * S33
174    txy = preCompute * S12
175
176    # Set top boundary conditions
177    txx = txx.at[:, :, nz - 1].set(0)
178    tyy = tyy.at[:, :, nz - 1].set(0)
179    tzz = tzz.at[:, :, nz - 1].set(0)
180    txy = txy.at[:, :, nz - 1].set(0)
181
182    return txx, tyy, tzz, txy
183
184
185# ===============================================
186# Compute SGS stresses on W nodes with dealiasing
187# ===============================================
188
189@jax.jit
190def StressesWnodes_Dealias(
191        S13_pad, S23_pad,
192        S_w_pad,
193        Cs2_3D_pad,
194        u, v, M_sfc_loc, psi2D_m, psi2D_m0,
195        ZeRo3D_fft):
196    """
197    Parameters:
198    -----------
199    S13_pad, S23_pad : ndarray
200        Dealiased strain rate components
201    S_w_pad : ndarray
202        Dealiased strain rate magnitude at W nodes
203    Cs2_3D_pad : ndarray
204        Dealiased Smagorinsky coefficient field
205    u, v : ndarray
206        Velocity components for wall model
207    M_sfc_loc, psi2D_m, psi2D_m0 : ndarray
208        Surface parameters for wall model
209    ZeRo3D_fft : ndarray
210        Pre-allocated zero array for dealiasing
211
212    Returns:
213    --------
214    txz, tyz : ndarray
215        SGS stress components at W nodes
216    """
217
218    # Initialize arrays
219    txz_pad = jnp.zeros_like(S_w_pad)
220    tyz_pad = jnp.zeros_like(S_w_pad)
221
222    # Interior points
223    preCompute = (-2 * (L ** 2) * StagGridAvg(Cs2_3D_pad[:, :, :nz - 1]) *
224                  S_w_pad[:, :, 1:nz - 1])
225    txz_pad = txz_pad.at[:, :, 1:nz - 1].set(preCompute *
226                                             S13_pad[:, :, 1:nz - 1])
227    tyz_pad = tyz_pad.at[:, :, 1:nz - 1].set(preCompute *
228                                             S23_pad[:, :, 1:nz - 1])
229
230    # Top boundary conditions
231    txz_pad = txz_pad.at[:, :, nz - 1].set(0)
232    tyz_pad = tyz_pad.at[:, :, nz - 1].set(0)
233
234    # Apply dealiasing - inverse operations
235    txz = Dealias2(FFT_pad(txz_pad), ZeRo3D_fft)
236    tyz = Dealias2(FFT_pad(tyz_pad), ZeRo3D_fft)
237
238    # Apply wall model for bottom boundary
239    txz_wall, tyz_wall = Wall(u, v, M_sfc_loc, psi2D_m, psi2D_m0)
240    txz = txz.at[:, :, 0].set(txz_wall)
241    tyz = tyz.at[:, :, 0].set(tyz_wall)
242
243    return txz, tyz
244
245
246# ============================================================
247# Compute SGS stresses on W nodes without dealiasing
248# ============================================================
249
250@jax.jit
251def StressesWnodes_NoDealias(
252        S13, S23,
253        S_w,
254        Cs2_3D,
255        u, v, M_sfc_loc, psi2D_m, psi2D_m0):
256    """
257    Computes SGS stresses at W nodes without dealiasing.
258
259    Parameters:
260    -----------
261    S13, S23 : ndarray
262        Strain rate components
263    S_w : ndarray
264        Strain rate magnitude at W nodes
265    Cs2_3D : ndarray
266        Smagorinsky coefficient field
267    u, v : ndarray
268        Velocity components for wall model
269    M_sfc_loc, psi2D_m, psi2D_m0 : ndarray
270        Surface parameters for wall model
271
272    Returns:
273    --------
274    txz, tyz : ndarray
275        SGS stress components at W nodes
276    """
277
278    # Initialize arrays
279    txz = jnp.zeros_like(S_w)
280    tyz = jnp.zeros_like(S_w)
281
282    # Interior points
283    preCompute = (-2 * (L ** 2) * StagGridAvg(Cs2_3D[:, :, :nz - 1]) *
284                  S_w[:, :, 1:nz - 1])
285    txz = txz.at[:, :, 1:nz - 1].set(preCompute * S13[:, :, 1:nz - 1])
286    tyz = tyz.at[:, :, 1:nz - 1].set(preCompute * S23[:, :, 1:nz - 1])
287
288    # Top boundary conditions
289    txz = txz.at[:, :, nz - 1].set(0)
290    tyz = tyz.at[:, :, nz - 1].set(0)
291
292    # Apply wall model for bottom boundary
293    txz_wall, tyz_wall = Wall(u, v, M_sfc_loc, psi2D_m, psi2D_m0)
294    txz = txz.at[:, :, 0].set(txz_wall)
295    tyz = tyz.at[:, :, 0].set(tyz_wall)
296
297    return txz, tyz