Scalar All Terms

Source Code: SCL-AllTerms

SCL_AllTerms.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_AllTerms.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 all the right hand side terms of scalar equations
 25"""
 26
 27# ============================================================
 28#  Imports
 29# ============================================================
 30
 31import jax
 32
 33# Import configuration from namelist
 34from ..config.ConfigLoader import *
 35
 36# Import derived variables
 37from ..config.DerivedVars import *
 38
 39# Import utility functions
 40from ..utilities.Utilities import PlanarMean
 41
 42
 43# ============================================================
 44# Right hand side terms for scalar equations
 45# ============================================================
 46
 47@jax.jit
 48def RHS_Scalar(TH, THAdvectionSum, divq, RayleighDampCoeff_stag, THadv):
 49    """
 50    Computes the right-hand side terms for the scalar transport equation.
 51
 52    Parameters:
 53    -----------
 54    THAdvectionSum : ndarray of shape (nx, ny, nz)
 55        Sum of advection terms for potential temperature
 56    divq : ndarray of shape (nx, ny, nz)
 57        Divergence of subgrid-scale scalar flux
 58    THadv : ndarray of shape (nx, ny, nz)
 59        Non-dimensional large-scale (mesoscale) advection tendency for potential
 60        temperature.  Pass ZeRo3D when optAdvection == 0.
 61
 62    Returns:
 63    --------
 64    RHS_TH : ndarray of shape (nx, ny, nz)
 65        Right-hand side terms for potential temperature equation
 66    """
 67
 68    RHS_TH = - THAdvectionSum - divq
 69
 70    if optAdvection >= 1:
 71        RHS_TH = RHS_TH + THadv
 72
 73    if optDamping == 1:
 74        # Compute mean potential temperature at each vertical level
 75        TH_bar = PlanarMean(TH)
 76
 77        # Compute temperature fluctuations from mean
 78        TH_fluc = TH - TH_bar.reshape(1, 1, -1)
 79
 80        RHS_TH = RHS_TH - RayleighDampCoeff_stag * TH_fluc
 81
 82    return RHS_TH
 83
 84
 85@jax.jit
 86def RHS_Moisture(Q, QAdvectionSum, divqm, RayleighDampCoeff_stag, Qadv):
 87    """
 88    Right-hand side for the specific humidity (Q) transport equation.
 89
 90    Parameters:
 91    -----------
 92    Q : ndarray (nx, ny, nz) — specific humidity (kg/kg)
 93    QAdvectionSum : ndarray (nx, ny, nz) — advection sum for Q
 94    divqm : ndarray (nx, ny, nz) — divergence of SGS moisture flux
 95    RayleighDampCoeff_stag : ndarray (nx, ny, nz) — Rayleigh damping coefficients
 96    Qadv : ndarray (nx, ny, nz) — large-scale moisture advection tendency
 97                                   (ZeRo3D when no explicit forcing is prescribed)
 98
 99    Returns:
100    --------
101    RHS_Q : ndarray (nx, ny, nz)
102    """
103    RHS_Q = - QAdvectionSum - divqm + Qadv
104
105    if optDamping == 1:
106        Q_bar  = PlanarMean(Q)
107        Q_fluc = Q - Q_bar.reshape(1, 1, -1)
108        RHS_Q  = RHS_Q - RayleighDampCoeff_stag * Q_fluc
109
110    return RHS_Q