SGS Stress Computations: Wong-Lilly

Computes SGS momentum stresses using the Wong-Lilly (WL) base formulation:

\[\tau_{ij} = -2 C_{WL} \Delta^{4/3} \bar{S}_{ij}\]

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), both on dynamic call steps and on cached (non-call) steps when dynamicSGS_call_time > 1.

Source Code: SGSStresses_WL

SGSStresses_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: SGSStresses_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 stresses using the Wong-Lilly (1994) SGS
 25              base model.  Stress formula:
 26                tau_ij = -2 * C_WL * Delta^(4/3) * S_ij
 27              (no strain-rate magnitude factor; filter width exponent
 28              4/3 instead of 2).
 29"""
 30
 31# ============================================================
 32#  Imports
 33# ============================================================
 34
 35import jax
 36import jax.numpy as jnp
 37
 38# Import configuration from namelist
 39from ..config.ConfigLoader import *
 40
 41# Import derived variables
 42from ..config.DerivedVars import *
 43
 44# Import FFT modules
 45from ..operations.FFT import FFT, FFT_pad
 46
 47# Import dealiasing functions
 48from ..operations.Dealiasing import Dealias1, Dealias2
 49
 50# Import strain rates functions
 51from .StrainRates import StrainsUVPnodes_Dealias, StrainsWnodes_Dealias
 52from .StrainRates import StrainsUVPnodes_NoDealias, StrainsWnodes_NoDealias
 53
 54# Import LASDD-WL
 55from .DynamicSGS_LASDD_WL import LASDD
 56from .DynamicSGS_ScalarLASDD_WL import ScalarLASDD
 57
 58# Import helper functions
 59from ..utilities.Utilities import StagGridAvg
 60
 61# Import wall model from SM file (model-agnostic surface BC)
 62from .SGSStresses_SM import Wall
 63
 64
 65# =================================================
 66# Compute SGS stresses on UVP nodes with dealiasing
 67# =================================================
 68
 69@jax.jit
 70def StressesUVPnodes_Dealias(
 71        S11_pad, S22_pad, S33_pad, S12_pad,
 72        Cwl_3D_pad,
 73        ZeRo3D_fft):
 74    """
 75    Parameters:
 76    -----------
 77    S11_pad, S22_pad, S33_pad, S12_pad : ndarray
 78        Dealiased strain rate components
 79    Cwl_3D_pad : ndarray
 80        Dealiased Wong-Lilly coefficient C_WL field
 81    ZeRo3D_fft : ndarray
 82        Pre-allocated zero array for dealiasing
 83
 84    Returns:
 85    --------
 86    txx, tyy, tzz, txy : ndarray
 87        SGS stress components at UVP nodes
 88    """
 89
 90    preCompute = -2 * (L ** (4 / 3)) * Cwl_3D_pad
 91    txx_pad = preCompute * S11_pad
 92    tyy_pad = preCompute * S22_pad
 93    tzz_pad = preCompute * S33_pad
 94    txy_pad = preCompute * S12_pad
 95
 96    txx_pad = txx_pad.at[:, :, nz - 1].set(0)
 97    tyy_pad = tyy_pad.at[:, :, nz - 1].set(0)
 98    tzz_pad = tzz_pad.at[:, :, nz - 1].set(0)
 99    txy_pad = txy_pad.at[:, :, nz - 1].set(0)
100
101    txx = Dealias2(FFT_pad(txx_pad), ZeRo3D_fft)
102    tyy = Dealias2(FFT_pad(tyy_pad), ZeRo3D_fft)
103    tzz = Dealias2(FFT_pad(tzz_pad), ZeRo3D_fft)
104    txy = Dealias2(FFT_pad(txy_pad), ZeRo3D_fft)
105
106    return txx, tyy, tzz, txy
107
108
109# ====================================================
110# Compute SGS stresses on UVP nodes without dealiasing
111# ====================================================
112
113@jax.jit
114def StressesUVPnodes_NoDealias(
115        S11, S22, S33, S12,
116        Cwl_3D):
117    """
118    Parameters:
119    -----------
120    S11, S22, S33, S12 : ndarray
121        Strain rate components
122    Cwl_3D : ndarray
123        Wong-Lilly coefficient C_WL field
124
125    Returns:
126    --------
127    txx, tyy, tzz, txy : ndarray
128        SGS stress components at UVP nodes
129    """
130
131    preCompute = -2 * (L ** (4 / 3)) * Cwl_3D
132    txx = preCompute * S11
133    tyy = preCompute * S22
134    tzz = preCompute * S33
135    txy = preCompute * S12
136
137    txx = txx.at[:, :, nz - 1].set(0)
138    tyy = tyy.at[:, :, nz - 1].set(0)
139    tzz = tzz.at[:, :, nz - 1].set(0)
140    txy = txy.at[:, :, nz - 1].set(0)
141
142    return txx, tyy, tzz, txy
143
144
145# ===============================================
146# Compute SGS stresses on W nodes with dealiasing
147# ===============================================
148
149@jax.jit
150def StressesWnodes_Dealias(
151        S13_pad, S23_pad,
152        Cwl_3D_pad,
153        u, v, M_sfc_loc, psi2D_m, psi2D_m0,
154        ZeRo3D_fft):
155    """
156    Parameters:
157    -----------
158    S13_pad, S23_pad : ndarray
159        Dealiased strain rate components
160    Cwl_3D_pad : ndarray
161        Dealiased Wong-Lilly coefficient C_WL field
162    u, v : ndarray
163        Velocity components for wall model
164    M_sfc_loc, psi2D_m, psi2D_m0 : ndarray
165        Surface parameters for wall model
166    ZeRo3D_fft : ndarray
167        Pre-allocated zero array for dealiasing
168
169    Returns:
170    --------
171    txz, tyz : ndarray
172        SGS stress components at W nodes
173    """
174
175    txz_pad = jnp.zeros_like(S13_pad)
176    tyz_pad = jnp.zeros_like(S13_pad)
177
178    preCompute = (-2 * (L ** (4 / 3)) *
179                  StagGridAvg(Cwl_3D_pad[:, :, :nz - 1]))
180    txz_pad = txz_pad.at[:, :, 1:nz - 1].set(preCompute *
181                                              S13_pad[:, :, 1:nz - 1])
182    tyz_pad = tyz_pad.at[:, :, 1:nz - 1].set(preCompute *
183                                              S23_pad[:, :, 1:nz - 1])
184
185    txz_pad = txz_pad.at[:, :, nz - 1].set(0)
186    tyz_pad = tyz_pad.at[:, :, nz - 1].set(0)
187
188    txz = Dealias2(FFT_pad(txz_pad), ZeRo3D_fft)
189    tyz = Dealias2(FFT_pad(tyz_pad), ZeRo3D_fft)
190
191    txz_wall, tyz_wall = Wall(u, v, M_sfc_loc, psi2D_m, psi2D_m0)
192    txz = txz.at[:, :, 0].set(txz_wall)
193    tyz = tyz.at[:, :, 0].set(tyz_wall)
194
195    return txz, tyz
196
197
198# ============================================================
199# Compute SGS stresses on W nodes without dealiasing
200# ============================================================
201
202@jax.jit
203def StressesWnodes_NoDealias(
204        S13, S23,
205        Cwl_3D,
206        u, v, M_sfc_loc, psi2D_m, psi2D_m0):
207    """
208    Parameters:
209    -----------
210    S13, S23 : ndarray
211        Strain rate components
212    Cwl_3D : ndarray
213        Wong-Lilly coefficient C_WL field
214    u, v : ndarray
215        Velocity components for wall model
216    M_sfc_loc, psi2D_m, psi2D_m0 : ndarray
217        Surface parameters for wall model
218
219    Returns:
220    --------
221    txz, tyz : ndarray
222        SGS stress components at W nodes
223    """
224
225    txz = jnp.zeros_like(S13)
226    tyz = jnp.zeros_like(S13)
227
228    preCompute = (-2 * (L ** (4 / 3)) *
229                  StagGridAvg(Cwl_3D[:, :, :nz - 1]))
230    txz = txz.at[:, :, 1:nz - 1].set(preCompute * S13[:, :, 1:nz - 1])
231    tyz = tyz.at[:, :, 1:nz - 1].set(preCompute * S23[:, :, 1:nz - 1])
232
233    txz = txz.at[:, :, nz - 1].set(0)
234    tyz = tyz.at[:, :, nz - 1].set(0)
235
236    txz_wall, tyz_wall = Wall(u, v, M_sfc_loc, psi2D_m, psi2D_m0)
237    txz = txz.at[:, :, 0].set(txz_wall)
238    tyz = tyz.at[:, :, 0].set(tyz_wall)
239
240    return txz, tyz