Derived Variables
Source Code: DerivedVars
DerivedVars.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: DerivedVars.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: Derived variables for JAX-ALFA computed from base configuration
25"""
26
27# ============================================================
28# Imports
29# ============================================================
30
31import jax.numpy as jnp
32from .ConfigLoader import *
33
34
35# ============================================================
36# Constants
37# ============================================================
38
39pi = jnp.pi
40g = 9.81 # gravitational acceleration (m/s2)
41vonk = 0.4 # von Karman constant
42
43
44# ============================================================
45# Derived Variables (computed from Config.py)
46# ============================================================
47
48# Velocity, potential temperature and length scales
49u_scale = 1
50TH_scale = 1
51z_scale = l_x / (2 * pi)
52
53# Non-dimensional g
54g_nondim = g * z_scale / (u_scale ** 2)
55
56# Non-dimensional Coriolis term
57f_coriolis_nondim = f_coriolis * (z_scale / u_scale)
58
59# Non-dimensional reference temperature
60T_0_nondim = T_0 / TH_scale
61
62# Non-dimensional inversion term
63inversion_nondim = inversion * (z_scale / TH_scale)
64
65# Non-dimensional surface sensible heat flux
66qz_sfc = SensibleHeatFlux * jnp.ones((nx, ny)) / (u_scale * TH_scale)
67
68# Moisture scale (= 1 for dimensional Q in kg/kg throughout)
69Q_scale = 1
70
71# Non-dimensional constant surface moisture flux (kg/kg m/s)
72qm_sfc = MoistureFlux * jnp.ones((nx, ny)) / (u_scale * Q_scale)
73
74# Non-dimensional moisture gradient above domain top (kg/kg per non-dim length)
75q_inversion_nondim = q_inversion * (z_scale / Q_scale)
76
77# Non-dimensional grid spacing
78dx = 2 * pi / nx
79dy = 2 * pi / ny
80dz = (l_z / z_scale) / (nz - 1)
81idz = 1.0 / dz
82
83# Filter width
84L = FGR * (dx * dy * dz) ** (1 / 3)
85
86# Non-dimensional time-step
87dt_nondim = dt * u_scale / z_scale
88
89# Total time steps
90nsteps = int(np.ceil(SimTime / dt))
91
92# Test filter ratio (TFR) & dealiasing option
93if FGR == 1:
94 TFR = 2
95 optDealias = 1 # 0: no dealias, 1: dealias
96else:
97 TFR = np.sqrt(2)
98 optDealias = 0 # 0: no dealias, 1: dealias
99
100# Initial 3D fields for dynamic SGS coefficients
101if optSgs in [2, 4]: # WL models: use Cwl / CwlPrRatio
102 Cs2_3D = Cwl * jnp.ones((nx, ny, nz))
103 Cs2PrRatio_3D = CwlPrRatio * jnp.ones((nx, ny, nz))
104else: # SM models: use Cs2 / Cs2PrRatio
105 Cs2_3D = Cs2 * jnp.ones((nx, ny, nz))
106 Cs2PrRatio_3D = Cs2PrRatio * jnp.ones((nx, ny, nz))
107
108# Non-dimensional damping height
109z_damping_nondim = z_damping / z_scale
110
111# Non-dimensional relaxation time
112RelaxTime_nondim = RelaxTime * (u_scale / z_scale)
113
114# Collect samples every SampleInterval (steps)
115SampleInterval = int( SampleInterval_sec / dt)
116# Output averages every OutputInterval (steps)
117OutputInterval = int( round(OutputInterval_sec / dt / 5) * 5)
118# Output 3D fields every Output3DInterval (steps)
119Output3DInterval = int( round(Output3DInterval_sec / dt / 5) * 5)