Scalar Advection Terms

Source Code: SCL-AdvectionTerms

SCL_AdvectionTerms.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_AdvectionTerms.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 advective terms of scalar transport
 25"""
 26
 27
 28# ============================================
 29# Imports
 30# ============================================
 31
 32import jax
 33from ..config.ConfigLoader import *
 34from ..config import ConfigLoader as Config
 35
 36# Import derived variables
 37from ..config.DerivedVars import *
 38
 39from ..utilities.Utilities import StagGridAvg
 40
 41# Import FFT modules
 42from ..operations.FFT import FFT, FFT_pad
 43
 44# Import dealias functions
 45from ..operations.Dealiasing import Dealias1, Dealias2
 46
 47# Import constants
 48from ..initialization.Preprocess import Constant
 49mx, my, nx_rfft, ny_rfft, mx_rfft, my_rfft = Constant()
 50
 51
 52# -----------------------------------------------------------------------------
 53# Use dealiasing in computing the advection terms of scalar transport
 54# -----------------------------------------------------------------------------
 55
 56@jax.jit
 57def ScalarAdvection_Dealias(u, v, w,
 58                            dsdx, dsdy, dsdz,
 59                            ZeRo3D_fft, ZeRo3D_pad, ZeRo3D_pad_fft):
 60    """
 61    Parameters:
 62    -----------
 63    u, v, w : ndarray of shape (nx, ny, nz)
 64        Velocity components in x, y, and z directions respectively
 65    dsdx, dsdy, dsdz : ndarray of shape (nx, ny, nz)
 66        Derivatives of a scalar with respect to x, y and z
 67    ZeRo3D_fft : ndarray
 68        Pre-allocated zero array for Fourier operations
 69    ZeRo3D_pad : ndarray
 70        Pre-allocated zero-padded array
 71    ZeRo3D_pad_fft : ndarray
 72        Pre-allocated zero-padded array for Fourier operations
 73
 74    Returns:
 75    --------
 76    scalarAdvectionSum : ndarray of shape (nx, ny, nz)
 77        Sum of convective terms with dealiasing applied
 78    """
 79
 80    # Dealias - forward operation
 81    u_pad = Dealias1(FFT(u), ZeRo3D_pad_fft)
 82    v_pad = Dealias1(FFT(v), ZeRo3D_pad_fft)
 83    w_pad = Dealias1(FFT(w), ZeRo3D_pad_fft)
 84    dsdx_pad = Dealias1(FFT(dsdx), ZeRo3D_pad_fft)
 85    dsdy_pad = Dealias1(FFT(dsdy), ZeRo3D_pad_fft)
 86    dsdz_pad = Dealias1(FFT(dsdz), ZeRo3D_pad_fft)
 87
 88    # Compute advection terms
 89    cc1 = u_pad * dsdx_pad + v_pad * dsdy_pad
 90    cc2 = w_pad * dsdz_pad
 91
 92    # Initialize array for combined convective terms
 93    cc = ZeRo3D_pad.copy()
 94
 95    # Compute convective terms with grid averaging for interior points
 96    cc = cc.at[:, :, 0:nz - 1].set(cc1[:, :, 0:nz - 1] +
 97                                   StagGridAvg(cc2))
 98
 99    # Handle top boundary
100    cc = cc.at[:, :, nz - 1].set(cc1[:, :, nz - 1] +
101                                 cc2[:, :, nz - 1])
102
103    # Dealias the convective terms - inverse operation
104    scalarAdvectionSum = Dealias2(FFT_pad(cc), ZeRo3D_fft)
105
106    return scalarAdvectionSum
107
108
109# ------------------------------------------------------------------------------
110# Compute the advection terms of scalar transport without any dealiasing
111# ------------------------------------------------------------------------------
112
113@jax.jit
114def ScalarAdvection_NoDealias(u, v, w,
115                              dsdx, dsdy, dsdz,
116                              ZeRo3D):
117    """
118    Parameters:
119    -----------
120    u, v, w : ndarray of shape (nx, ny, nz)
121        Velocity components in x, y, and z directions respectively
122    dsdx, dsdy, dsdz : ndarray of shape (nx, ny, nz)
123        Derivatives of a scalar with respect to x, y and z
124    ZeRo3D : ndarray of shape (nx, ny, nz)
125        Pre-allocated zero array for calculations
126
127    Returns:
128    --------
129    scalarAdvectionSum : ndarray of shape (nx, ny, nz)
130        Sum of convective terms without dealiasing
131    """
132
133    # Compute advection terms
134    cc1 = u * dsdx + v * dsdy
135    cc2 = w * dsdz
136
137    # Initialize array for combined convective terms
138    scalarAdvectionSum = ZeRo3D.copy()
139
140    # Compute convective terms with grid averaging for interior points
141    scalarAdvectionSum = (
142        scalarAdvectionSum.at[:, :, 0:nz - 1].set(cc1[:, :, 0:nz - 1]
143                                                  + StagGridAvg(cc2)))
144
145    # Handle top boundary
146    scalarAdvectionSum = (
147        scalarAdvectionSum.at[:, :, nz - 1].set(cc1[:, :, nz - 1]
148                                                + cc2[:, :, nz - 1]))
149
150    return scalarAdvectionSum
151
152
153# ------------------------------------------------------------------------------
154# Select one of the above functions based on dealias flag
155# ------------------------------------------------------------------------------
156
157@jax.jit
158def ScalarAdvection(
159        u, v, w,
160        dsdx, dsdy, dsdz,
161        ZeRo3D, ZeRo3D_fft, ZeRo3D_pad,
162        ZeRo3D_pad_fft):
163    """
164    Parameters:
165    -----------
166    u, v, w : ndarray of shape (nx, ny, nz)
167        Velocity components in x, y, and z directions respectively
168    dsdx, dsdy, dsdz : ndarray of shape (nx, ny, nz)
169        Derivatives of a scalar with respect to x, y and z
170    ZeRo3D : ndarray of shape (nx, ny, nz)
171        Pre-allocated zero array for calculations
172    ZeRo3D_fft : ndarray
173        Pre-allocated zero array for Fourier operations
174    ZeRo3D_pad : ndarray
175        Pre-allocated zero-padded array
176    ZeRo3D_pad_fft : ndarray
177        Pre-allocated zero-padded array for Fourier operations
178
179    Returns:
180    --------
181    scalarAdvectionSum : ndarray of shape (nx, ny, nz)
182        Sum of advection terms from either dealiased or non-dealiased method
183    """
184
185    if optDealias == 1:
186
187        scalarAdvectionSum = (
188            ScalarAdvection_Dealias(
189                u, v, w,
190                dsdx, dsdy, dsdz,
191                ZeRo3D_fft, ZeRo3D_pad, ZeRo3D_pad_fft))
192
193    else:
194
195        scalarAdvectionSum = (
196            ScalarAdvection_NoDealias(
197                u, v, w,
198                dsdx, dsdy, dsdz,
199                ZeRo3D))
200
201    return scalarAdvectionSum