Configuration Loader

Source Code: ConfigLoader

ConfigLoader.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: ConfigLoader.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: 2026-05-01
 24:Description: Loads run-specific configuration from JAXALFA_RUNDIR/Config.py.
 25              JAXALFA_RUNDIR must be set in the environment before running.
 26
 27Usage
 28-----
 29Before launching JAX-ALFA, set the run directory::
 30
 31    export JAXALFA_RUNDIR=/path/to/run_directory
 32    python -m src.Main
 33
 34The run directory must contain a Config.py with all simulation parameters.
 35"""
 36
 37
 38# ============================================================
 39# Imports
 40# ============================================================
 41
 42import os
 43import numpy as np
 44
 45
 46# ============================================================
 47# Load run-directory configuration (mandatory)
 48# ============================================================
 49
 50_rundir = os.environ.get('JAXALFA_RUNDIR')
 51if _rundir is None:
 52    raise EnvironmentError(
 53        "\n\nJAXALFA_RUNDIR is not set.\n"
 54        "Set it to the run directory before launching JAX-ALFA:\n\n"
 55        "    export JAXALFA_RUNDIR=/path/to/run_directory\n"
 56        "    python $JAXALFA_RUNDIR/CreateInputs*.py\n"
 57        "    python -m src.Main\n"
 58    )
 59
 60_config_path = os.path.join(_rundir, 'Config.py')
 61if not os.path.isfile(_config_path):
 62    raise FileNotFoundError(
 63        f"\nConfig.py not found in JAXALFA_RUNDIR.\n"
 64        f"Expected: {_config_path}\n"
 65    )
 66
 67with open(_config_path) as _f:
 68    exec(_f.read())
 69
 70# Backward-compatible defaults for time-varying surface BC
 71if 'optSurfBC' not in dir():
 72    optSurfBC = 0
 73if 'SurfaceBCFile' not in dir():
 74    SurfaceBCFile = 'input/SurfaceBC.npz'
 75
 76# Backward-compatible defaults for time/height-varying geostrophic wind
 77# optGeoWind = 0: constant Ug2, Vg2 from Config (default, all existing cases)
 78# optGeoWind = 1: time + height varying, loaded from GeoWindFile
 79if 'optGeoWind' not in dir():
 80    optGeoWind = 0
 81if 'GeoWindFile' not in dir():
 82    GeoWindFile = 'input/GeoWind.npz'
 83
 84# Backward-compatible default for screen-level temperature reference height.
 85# zTemperature = 0  : use z0T as the thermal reference height (default, GABLS1/SBL)
 86# zTemperature > 0  : use this height (m) as the reference for the heat-flux
 87#                     denominator and psi_h0 — needed when the prescribed surface
 88#                     temperature is observed at a screen height above z0T (e.g.
 89#                     Wangara: temperature measured at 1.2 m above ground).
 90if 'zTemperature' not in dir():
 91    zTemperature = 0.0
 92
 93# Backward-compatible defaults for time/height-varying large-scale advection.
 94# optAdvection = 0: no mesoscale advection forcing (default, all existing cases)
 95# optAdvection = 1: time + height varying, loaded from AdvectionFile
 96if 'optAdvection' not in dir():
 97    optAdvection = 0
 98if 'AdvectionFile' not in dir():
 99    AdvectionFile = 'input/AdvForcing.npz'
100
101# Backward-compatible defaults for moisture.
102# optMoisture = 0: no moisture (default, all existing cases)
103# optMoisture = 1: specific humidity Q is a prognostic variable
104if 'optMoisture' not in dir():
105    optMoisture = 0
106# zMoisture: screen-level reference height for moisture flux denominator (m).
107#   0  : use z0T as the moisture reference height (same as heat)
108#   > 0: use this height (e.g. 0.25 m when surface Q observed at 0.25 m)
109if 'zMoisture' not in dir():
110    zMoisture = 0.0
111# MoistureFlux: constant surface moisture flux (kg/kg m/s), used when optMoistureSurfBC=0
112if 'MoistureFlux' not in dir():
113    MoistureFlux = 0.0
114# optMoistureSurfBC: 0 = constant flux (MoistureFlux)
115#                   1 = time-varying flux (from MoistureSurfaceBCFile)
116#                   2 = time-varying surface Q (from MoistureSurfaceBCFile)
117if 'optMoistureSurfBC' not in dir():
118    optMoistureSurfBC = 0
119if 'MoistureSurfaceBCFile' not in dir():
120    MoistureSurfaceBCFile = 'input/MoistureSurfaceBC.npz'
121# q_inversion: specific humidity lapse rate above domain top (kg/kg/m).
122#   0 : zero gradient (flat Q profile at top, default)
123if 'q_inversion' not in dir():
124    q_inversion = 0.0
125
126# GPU_ID: which GPU to use when optGPU=1 (0-indexed).
127# Set to 1 in Config.py to run on the second GPU.
128if 'GPU_ID' not in dir():
129    GPU_ID = 0
130
131# Backward-compatible defaults for the stability-dependent Smagorinsky model
132# (STAB-SM, optSgs=5).  Override any of these in Config.py if needed.
133# CsMO_SM : Smagorinsky coefficient (0.23 recommended for finite-difference codes)
134# aMO_SM  : heat stability parameter (= 1/Pr_t_neutral)
135# bMO_SM, cMO_SM : unstable stability function coefficients
136# fMO_SM, gMO_SM, hMO_SM : stable stability function coefficients
137# RicMO_SM : critical Richardson number
138# rMO_SM   : stable stability function exponent
139if 'CsMO_SM'  not in dir(): CsMO_SM  = 0.17
140if 'aMO_SM'   not in dir(): aMO_SM   = 1.0 / 0.7
141if 'bMO_SM'   not in dir(): bMO_SM   = 40.0
142if 'cMO_SM'   not in dir(): cMO_SM   = 16.0
143if 'fMO_SM'   not in dir(): fMO_SM   = 1.0 / 0.7
144if 'gMO_SM'   not in dir(): gMO_SM   = 1.2
145if 'hMO_SM'   not in dir(): hMO_SM   = 0.0
146if 'RicMO_SM' not in dir(): RicMO_SM = 0.25
147if 'rMO_SM'   not in dir(): rMO_SM   = 4.0
148
149# Backward-compatible default for float precision
150if 'use_double_precision' not in dir():
151    use_double_precision = False
152
153# Pressure solver selection.
154# optPressureSolver = 0: LU    — dense matrix + jnp.linalg.solve (original)
155# optPressureSolver = 1: Thomas — tridiagonal Thomas algorithm (faster)
156if 'optPressureSolver' not in dir():
157    optPressureSolver = 0
158
159del _rundir, _config_path, _f