All Terms

Source Code: NSE_AllTerms

NSE_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: NSE_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-3
 24:Description: computes all the right hand side terms of NSE
 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 momentum equations
 45# ============================================================
 46
 47@jax.jit
 48def RHS_Momentum(u, v, w,
 49                 Ug, Vg,
 50                 Cx, Cy, Cz,
 51                 buoyancy,
 52                 divtx, divty, divtz,
 53                 RayleighDampCoeff, RayleighDampCoeff_stag,
 54                 Uadv, Vadv):
 55    """
 56    Parameters:
 57    -----------
 58    u, v, w : ndarray of shape (nx, ny, nz)
 59        Velocity components in x, y, and z directions
 60    Ug, Vg : ndarray of shape (nx, ny, nz)
 61        Geostrophic winds in x and y directions
 62    Cx, Cy, Cz : ndarray of shape (nx, ny, nz)
 63        Advection terms
 64    buoyancy : ndarray of shape (nx, ny, nz)
 65        Buoyancy term
 66    divtx, divty, divtz : ndarray of shape (nx, ny, nz)
 67        SGS stress divergence terms
 68    Uadv, Vadv : ndarray of shape (nx, ny, nz)
 69        Non-dimensional large-scale (mesoscale) advection tendencies for u and v.
 70        Pass ZeRo3D when optAdvection == 0.
 71
 72    Returns:
 73    --------
 74    RHS_u : ndarray of shape (nx, ny, nz)
 75        RHS for u component
 76    RHS_v : ndarray of shape (nx, ny, nz)
 77        RHS for v component
 78    RHS_w : ndarray of shape (nx, ny, nz)
 79        RHS for w component
 80    """
 81
 82    RHS_u = - Cx - divtx - f_coriolis_nondim * (Vg - v)
 83    RHS_v = - Cy - divty + f_coriolis_nondim * (Ug - Ugal - u)
 84    RHS_w = - Cz - divtz + buoyancy
 85
 86    if optAdvection >= 1:
 87        RHS_u = RHS_u + Uadv
 88        RHS_v = RHS_v + Vadv
 89
 90    if optDamping == 1:
 91
 92        # Compute mean velocity components at each vertical level
 93        u_bar = PlanarMean(u)
 94        v_bar = PlanarMean(v)
 95
 96        # Compute velocity fluctuations from mean
 97        u_fluc = u - u_bar.reshape(1, 1, -1)
 98        v_fluc = v - v_bar.reshape(1, 1, -1)
 99
100        RHS_u = RHS_u - RayleighDampCoeff_stag * u_fluc
101        RHS_v = RHS_v - RayleighDampCoeff_stag * v_fluc
102        RHS_w = RHS_w - RayleighDampCoeff * w  # assume w_bar = 0
103
104    return RHS_u, RHS_v, RHS_w