Buoyancy Terms

Source Code: NSE_BuoyancyTerms

NSE_BuoyancyTerms.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_BuoyancyTerms.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-4
 24:Description: computes the buoyancy term for use in the momentum equation
 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 utility functions
 41from ..utilities.Utilities import PlanarMean
 42
 43
 44# ============================================================
 45#  Compute buoyancy terms using reference temperature (Option 0)
 46# ============================================================
 47
 48@jax.jit
 49def BuoyancyOpt1(TH, H, ZeRo3D):
 50    """
 51    Computes the buoyancy term using reference temperature.
 52    buoyancy: g*(Tv - <Tv>)/T_0
 53
 54    Parameters:
 55    -----------
 56    TH : ndarray of shape (nx, ny, nz)
 57        Potential temperature (normalized)
 58    H : ndarray of shape (nx, ny, nz), optional
 59        Specific humidity (normalized)
 60    ZeRo3D : ndarray of shape (nx, ny, nz)
 61        Pre-allocated zero array
 62
 63    Returns:
 64    --------
 65    buoyancy : ndarray of shape (nx, ny, nz)
 66        Buoyancy term for momentum equation
 67    """
 68
 69    # Initialize arrays with zeros
 70    buoyancy = ZeRo3D.copy()
 71
 72    # TH is stored as anomaly (TH - T_0); use absolute temperature so that
 73    # the moisture term 0.61*Q*T_0 is O(~2 K) not O(~0 K).
 74    THv = (TH + T_0_nondim) * (1 + 0.61 * H)
 75
 76    # Compute planar-averaged virtual potential temperature
 77    THv_bar = PlanarMean(THv)
 78
 79    # Reshape THv_bar for broadcasting
 80    THv_bar_3D = THv_bar.reshape(1, 1, -1)
 81
 82    # Compute normalized deviation
 83    THv_normalized = (THv - THv_bar_3D) / T_0_nondim
 84
 85    # Compute the half-level averages for all levels
 86    above = THv_normalized[:, :, 1:nz]
 87    below = THv_normalized[:, :, 0:nz - 1]
 88    buoyancy_interior = 0.5 * g_nondim * (above + below)
 89
 90    # Update the buoyancy array
 91    buoyancy = buoyancy.at[:, :, 1:nz].set(buoyancy_interior)
 92
 93    # Set bottom boundary condition
 94    buoyancy = buoyancy.at[:, :, 0].set(0.0)
 95
 96    return buoyancy
 97
 98
 99# ============================================================
100#  Compute buoyancy terms using local mean temperature (Option 1)
101# ============================================================
102
103@jax.jit
104def BuoyancyOpt2(TH, H, ZeRo3D):
105    """
106    Computes the buoyancy term using local mean virtual
107    potential temperature.
108    buoyancy: g*(Tv - <Tv>)/<Tv>
109
110    Parameters:
111    -----------
112    TH : ndarray of shape (nx, ny, nz)
113        Potential temperature (normalized)
114    H : ndarray of shape (nx, ny, nz), optional
115        Specific humidity (normalized)
116
117    Returns:
118    --------
119    buoyancy : ndarray of shape (nx, ny, nz)
120        Buoyancy term for momentum equation
121    """
122
123    # Initialize arrays with zeros
124    buoyancy = ZeRo3D.copy()
125
126    # TH is stored as anomaly (TH - T_0); use absolute temperature so that
127    # the moisture term 0.61*Q*T_0 is O(~2 K) not O(~0 K).
128    THv = (TH + T_0_nondim) * (1 + 0.61 * H)
129
130    # Compute plane-averaged virtual potential temperature
131    THv_bar = PlanarMean(THv)
132
133    # Reshape THv_bar for broadcasting
134    THv_bar_3D = THv_bar.reshape(1, 1, -1)
135
136    # THv_bar is already ~T_0_nondim*(1+0.61*<Q>), so use it directly.
137    THv_normalized = (THv - THv_bar_3D) / THv_bar_3D
138
139    # Compute the half-level averages for all levels
140    above = THv_normalized[:, :, 1:nz]
141    below = THv_normalized[:, :, 0:nz - 1]
142    buoyancy_interior = 0.5 * g_nondim * (above + below)
143
144    # Update the buoyancy array
145    buoyancy = buoyancy.at[:, :, 1:nz].set(buoyancy_interior)
146
147    # Set bottom boundary condition
148    buoyancy = buoyancy.at[:, :, 0].set(0.0)
149
150    return buoyancy