Scalar SGS Terms

Source Code: SCL-SGSTerms

SCL_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: SCL_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-29
 24:Description: computes the SGS scalar flux 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_w
 44
 45# Import SGS models
 46from ..subgridscale.DynamicSGS_Main import DynamicSGSscalar
 47from ..subgridscale.StaticSGS_Main import StaticSGSscalar
 48
 49
 50# =================================================
 51# Compute Divergence of Scalar Flux
 52# =================================================
 53
 54@jax.jit
 55def DivFlux(
 56        qx, qy, qz,
 57        ZeRo3D,
 58        kx2, ky2):
 59    """
 60    Parameters:
 61    -----------
 62    qx, qy : ndarray of shape (nx, ny, nz)
 63        Horizontal scalar flux components
 64    qz : ndarray of shape (nx, ny, nz)
 65        Vertical scalar flux component
 66    kx2, ky2 : ndarray
 67        Wavenumber arrays for spectral derivatives
 68    ZeRo3D : ndarray of shape (nx, ny, nz)
 69        Pre-allocated zero array for intermediate calculations
 70
 71    Returns:
 72    --------
 73    divq : ndarray of shape (nx, ny, nz)
 74        Divergence of scalar flux
 75    """
 76    # Compute derivatives in x-direction
 77    dxqx = Derivxy(FFT(qx), kx2)
 78
 79    # Compute derivatives in y-direction
 80    dyqy = Derivxy(FFT(qy), ky2)
 81
 82    # Compute derivatives in z-direction
 83    # For qz, use w-grid derivative
 84    dzqz = Derivz_Generic_w(qz, ZeRo3D)
 85
 86    # Compute divergence
 87    divq = dxqx + dyqy + dzqz
 88
 89    return divq
 90
 91
 92# ================================================
 93# Compute flux divergence using dynamic SGS models
 94# ================================================
 95
 96@jax.jit
 97def DivFluxDynamicSGS(
 98        dynamicSGSmomentum,
 99        TH,
100        dTHdx, dTHdy, dTHdz,
101        SHFX,
102        ZeRo3D, ZeRo3D_fft, ZeRo3D_pad_fft,
103        kx2, ky2):
104    """
105    Compute flux divergence using dynamic SGS models.
106
107    Parameters:
108    -----------
109    dynamicSGSmomentum : tuple
110        Results from dynamic SGS momentum calculations,
111        containing filtered velocity components and strain rates
112    TH : ndarray of shape (nx, ny, nz)
113        Potential temperature
114    dTHdx, dTHdy, dTHdz : ndarray of shape (nx, ny, nz)
115        Gradients of potential temperature
116    SHFX : ndarray of shape (nx, ny)
117        Surface sensible heat flux
118    ZeRo3D, ZeRo3D_fft, ZeRo3D_pad_fft : ndarray
119        Pre-allocated zero arrays for calculations
120    kx2, ky2 : ndarray
121        Wavenumber arrays for spectral derivatives
122
123    Returns:
124    --------
125    divq : ndarray of shape (nx, ny, nz)
126        Divergence of scalar flux
127    Cs2PrRatio_1D : ndarray of shape (nz)
128        1D profile of dynamic SGS coefficient (Cs2/Pr)
129    beta2_1D : ndarray of shape (nz)
130        1D profile of beta coefficient for scalar SGS model
131    """
132
133    # Unpack the dynamicSGSmomentum tuple
134    (u_, v_, w_,
135     u_hat, v_hat, w_hat,
136     u_hatd, v_hatd, w_hatd,
137     S_uvp, S_uvp_pad,
138     S_w, S_w_pad,
139     S_uvp_hat, S_uvp_hatd) = dynamicSGSmomentum
140
141    (qx, qy, qz,
142     Cs2PrRatio_3D,
143     Cs2PrRatio_1D, beta2_1D) = (
144        DynamicSGSscalar(
145            u_, v_, w_,
146            u_hat, v_hat, w_hat,
147            u_hatd, v_hatd, w_hatd,
148            S_uvp, S_uvp_pad,
149            S_w, S_w_pad,
150            S_uvp_hat, S_uvp_hatd,
151            TH,
152            dTHdx, dTHdy, dTHdz,
153            SHFX,
154            ZeRo3D, ZeRo3D_fft, ZeRo3D_pad_fft))
155
156    # Compute divergence of flux
157    divq = DivFlux(qx, qy, qz,
158                   ZeRo3D,
159                   kx2, ky2)
160
161    return qz, divq, Cs2PrRatio_3D, Cs2PrRatio_1D, beta2_1D
162
163
164# ======================================================
165# Compute scalar flux divergence using static SGS models
166# ======================================================
167
168@jax.jit
169def DivFluxStaticSGS(
170        staticSGSmomentum,
171        Cs2PrRatio_3D,
172        dTHdx, dTHdy, dTHdz,
173        SHFX,
174        ZeRo3D, ZeRo3D_fft, ZeRo3D_pad_fft,
175        kx2, ky2):
176    """
177    Compute scalar flux divergence using static SGS models.
178
179    Parameters:
180    -----------
181    staticSGSmomentum : tuple
182        Results from static SGS momentum calculations, containing strain rates
183    Cs2PrRatio_3D : ndarray of shape (nx, ny, nz)
184        Static coefficient for scalar SGS model
185    dTHdx, dTHdy, dTHdz : ndarray of shape (nx, ny, nz)
186        Gradients of potential temperature
187    SHFX : ndarray of shape (nx, ny)
188        Surface heat flux
189    ZeRo3D, ZeRo3D_fft, ZeRo3D_pad_fft : ndarray
190        Pre-allocated zero arrays for calculations
191    kx2, ky2 : ndarray
192        Wavenumber arrays for spectral derivatives
193
194    Returns:
195    --------
196    divq : ndarray of shape (nx, ny, nz)
197        Divergence of scalar flux
198    """
199
200    # Unpack the staticSGSmomentum tuple
201    (S_uvp, S_uvp_pad,
202     S_w, S_w_pad) = staticSGSmomentum
203
204    # Static SGS model for scalar
205    qx, qy, qz = (
206        StaticSGSscalar(
207            S_uvp, S_uvp_pad,
208            S_w, S_w_pad,
209            Cs2PrRatio_3D,
210            dTHdx, dTHdy, dTHdz,
211            SHFX,
212            ZeRo3D_fft, ZeRo3D_pad_fft))
213
214    # Compute divergence of flux
215    divq = DivFlux(qx, qy, qz,
216                   ZeRo3D,
217                   kx2, ky2)
218
219    return qz, divq