KuoEliassen 0.2.4

KuoEliassen: A High-Performance Solver for the Kuo-Eliassen Equation

A High-Performance Solver for the Kuo-Eliassen Equation

A high-performance, production-ready solver for the Kuo-Eliassen equation describing meridional atmospheric circulation forced by diabatic heating and eddy fluxes. Combines optimized Fortran 90 backend with intuitive Python interface.

The Kuo-Eliassen Equation

The Kuo-Eliassen equation governs the zonal-mean meridional mass streamfunction response to forcing. This is an elliptic partial differential equation relating the streamfunction to diabatic heating and eddy momentum/heat flux convergences.

Governing Equation

The compact form of the Kuo-Eliassen equation:

$$\frac{f^2 g}{2\pi a \cos\phi} \frac{\partial^2\psi}{\partial p^2} + \frac{S^2 g}{2\pi a} \frac{\partial}{\partial\phi}\left(\frac{1}{a\cos\phi}\frac{\partial\psi}{\partial\phi}\right) = D$$

Complete expanded form (RHS fully decomposed):

$$f^2 \frac{g}{2\pi a\cos\phi} \frac{\partial^2\psi}{\partial p^2} + S^2 \frac{g}{2\pi a} \frac{\partial}{\partial\phi}\left(\frac{1}{a\cos\phi}\frac{\partial\psi}{\partial\phi}\right) = \frac{R}{p}\left(\frac{1}{a}\frac{\partial\overline{Q}}{\partial\phi} - \frac{1}{a}\frac{\partial}{\partial\phi}\left(\frac{1}{a\cos\phi}\frac{\partial(\overline{v'T'}\cos\phi)}{\partial\phi}\right)\right) + f\left(\frac{1}{a\cos^2\phi}\frac{\partial^2(\overline{u'v'}\cos^2\phi)}{\partial p \partial\phi} - \frac{\partial\overline{X}}{\partial p}\right)$$

Operator form (alternative notation):

$$\mathcal{L}(\psi) = D$$

where the elliptic operator $\mathcal{L}$ is defined as:

$$\mathcal{L}(\psi) = \frac{f^2}{2\pi a \cos\phi} g\frac{\partial^2}{\partial p^2} + \frac{S^2 g}{2\pi a \cos\phi} \frac{\partial}{\partial \phi}\left(\frac{1}{a\cos\phi}\frac{\partial}{\partial \phi}\right)$$

Component breakdown of RHS:

$$D = \underbrace{\frac{R}{p a}\frac{\partial\overline{Q}}{\partial\phi}}{\text{Diabatic Heating}} - \underbrace{\frac{R}{pa}\frac{\partial}{\partial\phi}\left(\frac{1}{a\cos\phi}\frac{\partial(\overline{v'T'}\cos\phi)}{\partial\phi}\right)}{\text{Eddy Heat Flux}} + \underbrace{\frac{f}{a\cos^2\phi}\frac{\partial^2(\overline{u'v'}\cos^2\phi)}{\partial p \partial\phi}}{\text{Eddy Momentum}} - \underbrace{f\frac{\partial\overline{X}}{\partial p}}{\text{Friction}}$$

Meridional velocity diagnostic:

$$\bar{v} = -\frac{1}{a}\frac{\partial\psi}{\partial p}$$

Vertical velocity diagnostic (from continuity):

$$\bar{\omega} = -\frac{1}{a\cos\phi}\frac{\partial(\psi\cos\phi)}{\partial\phi}$$

where:

• ψ — meridional mass streamfunction [kg/s]
• f = 2Ω sin(φ) — Coriolis parameter [s⁻¹]
• Ω — Earth's rotation rate = 7.29 × 10⁻⁵ [rad/s]
• a — Earth's radius ≈ 6.371 × 10⁶ [m]
• φ — latitude [rad]
• p — pressure [Pa]
• g — gravitational acceleration ≈ 9.81 [m/s²]
• S² — static stability [s⁻²]
• R — specific gas constant ≈ 287 [J/(kg·K)]
• $\overline{Q}$ — zonal-mean diabatic heating rate [K/s]
• $\overline{v'T'}$ — meridional eddy heat flux [K·m/s]
• $\overline{u'v'}$ — eddy momentum flux [m²/s²]
• $\overline{X}$ — friction/dissipation function [K/s]
• Overbars represent zonal and monthly mean, and primes represent deviations from zonal and montly mean

Static Stability

The static stability parameter $S^2$ characterizes atmospheric resistance to vertical motion and is defined as:

$$S^2 = -\frac{1}{\rho\theta}\frac{\partial\theta}{\partial p}$$

Alternatively, in terms of absolute temperature:

$$S^2 = -\left( \frac{R_d}{p} \right) \left[ \frac{\partial T}{\partial p} - \left( \frac{R_d}{c_p} \right) \frac{T}{p} \right]$$

where:

• ρ — air density [kg/m³]
• θ — potential temperature [K]
• T — absolute temperature [K]
• c_p — specific heat at constant pressure ≈ 1005 [J/(kg·K)]
• g — gravitational acceleration ≈ 9.81 [m/s²]
• $\frac{\partial\theta}{\partial p}$ — vertical potential temperature gradient [K/Pa]
• $\frac{\partial T}{\partial p}$ — vertical absolute temperature gradient [K/Pa]

Example Solution

The solver produces meridional streamfunction fields that reveal Hadley and Ferrel cell circulations:

Meridional mass streamfunction and Kuo-Eliassen equation solution

Features

• Optimized Fortran Backend: Core solver implemented in production-grade Fortran 90 with advanced numerical methods
• Component Decomposition: Separate diagnostic contributions from latent heating, radiative heating, eddy heat flux, and eddy momentum flux
• Flexible Interface: NumPy and xarray-compatible APIs for seamless integration with scientific workflows
• Cross-Platform Support: Pre-built wheels for Windows, macOS (Apple Silicon), and Linux
• Extensively Tested: >90% code coverage with comprehensive test suite
• High Performance: Dual solver architecture with LU decomposition (exact) and SOR (memory-efficient iterative)
• Numerical Robustness: Robust handling of geometric singularities near poles (requires grid excluding exact $\pm 90^\circ$)
• Solver Flexibility: Choose between direct sparse LU solver or optimized SOR with configurable relaxation parameters

Installation

From PyPI (Recommended)

Pre-built binary wheels are available for Python 3.9-3.13 on all major platforms:

pip install kuoeliassen

No compiler required! Wheels are provided for:

• Linux: x86_64 (manylinux_2_28)
• macOS: arm64 (Apple Silicon M1/M2/M3)
• Windows: x86_64

From Source (Development)

If you want to contribute or modify the code:

1. Prerequisites: Install a Fortran compiler

   • Windows: conda install -c conda-forge gcc-gfortran -y
   • macOS: brew install gcc
   • Linux: sudo apt-get install gfortran (Ubuntu/Debian)

2. Clone and install:

   git clone https://github.com/QianyeSu/KuoEliassen.git
   cd KuoEliassen
   pip install -e .
   

[!IMPORTANT] Grid Selection: The input latitude grid must not include the exact poles ($\pm 90^\circ$). The Kuo-Eliassen equation contains $1/\cos\phi$ terms that are singular at the poles. Always slice your data (e.g., lat = slice(-89.9, 89.9)) before solving.

Solver Methods

KuoEliassen provides two high-performance numerical solvers:

1. LU Decomposition (Default)

• Method: Direct sparse matrix solver using SuperLU
• Pros: Exact solution (within machine precision), guaranteed convergence
• Cons: Higher memory usage for large grids
• Best for: Small to medium grids (< 100×100), when precision is critical

result = solve_ke(v, temperature, vt_eddy, vu_eddy, pressure, latitude,
                  solver='lu')  # Default

2. SOR (Successive Over-Relaxation)

• Method: Iterative relaxation solver
• Pros: Low memory footprint, excellent for large grids, tunable convergence
• Cons: Requires omega parameter tuning for optimal performance
• Best for: Large grids, production runs, memory-constrained environments

result = solve_ke(v, temperature, vt_eddy, vu_eddy, pressure, latitude,
                  solver='sor',
                  omega=1.8,      # Relaxation factor (1.0-2.0, default=1.8)
                  tol=1e-8,       # Convergence tolerance (default=1e-8)
                  max_iter=50000) # Maximum iterations (default=50000)

Omega Parameter Tuning:

The optimal relaxation factor ω depends on your grid geometry:

• ω = 1.0: Gauss-Seidel (slow but stable)
• ω = 1.5-1.9: Typical optimal range for atmospheric grids
• ω → 2.0: Faster convergence but risk of divergence

Use examples/optimize_sor_omega.py to find the optimal ω for your data:

cd examples
python optimize_sor_omega.py

Example output:

Optimal Omega: 1.85
Minimum iterations: 675
Time cost: 0.0222 s

Quick Start

Basic Usage: Solve for Streamfunction

import numpy as np
import xarray as xr
from kuoeliassen import solve_ke

# Load atmospheric data from NetCDF file
# IMPORTANT: Exclude poles to avoid singularities
data = xr.open_dataset("example_data.nc").sel(lat=slice(-89.9, 89.9))

# Extract variables from dataset
# Assuming dimensions are (time, pressure, latitude)
v = data['v'].values                    # Mean meridional wind [m/s]
temperature = data['temperature'].values  # Temperature field [K]
vt_eddy = data['vt_eddy'].values        # Eddy heat flux v'T' [K⋅m/s]
vu_eddy = data['vu_eddy'].values        # Eddy momentum flux u'v' [m²/s²]
diabatic_heating = data['diabatic_heating'].values  # Total heating [K/s]

# Extract coordinate arrays (must be 1D)
pressure = data['pressure'].values      # Pressure levels [Pa]
latitude = data['latitude'].values      # Latitude [degrees]

# Solve Kuo-Eliassen equation for each time step
result = solve_ke(
    v, temperature, vt_eddy, vu_eddy,
    pressure, latitude,
    heating=diabatic_heating,
    solver='lu'  # or 'sor' for iterative solver
)

# Access results
psi_total = result['PSI']               # Total streamfunction [kg/s]
psi_d = result['D']                     # Total RHS forcing

Component Decomposition: Isolate Individual Forcings

# Separate radiative and latent heating contributions
result = solve_ke(
    v_mean, temperature, vt_eddy, vu_eddy, pressure, latitude,
    rad_heating=rad_heating,        # Radiative heating [K/s]
    latent_heating=latent_heating   # Latent heating from convection [K/s]
)

# Examine component-wise circulation response
psi_total = result['PSI']          # Total circulation
psi_rad = result['PSI_rad']        # Radiative forcing response
psi_latent = result['PSI_latent']  # Latent heating response
psi_vt = result['PSI_vt']          # Eddy heat flux response
psi_vu = result['PSI_vu']          # Eddy momentum flux response

Using xarray for Labeled Data

import xarray as xr
from kuoeliassen import solve_ke_xarray

# Load your atmospheric data (must have pressure and latitude dimensions)
ds = xr.open_dataset('atmospheric_data.nc')

# Solve with automatic coordinate handling
result_ds = solve_ke_xarray(
    ds['v_mean'], ds['temperature'], ds['vt_eddy'], ds['vu_eddy'],
    heating=ds['heating'],
    pressure_dim='pressure', latitude_dim='latitude'  # Specify if needed
)

# Result is an xarray Dataset with proper coordinates and attributes
psi = result_ds['PSI']
print(psi)  # Full metadata preserved
psi.plot()  # Easy visualization

Platform Support

Platform	Python Versions	Architecture	Status
Linux	3.9 - 3.13	x86_64	✅ Tested
macOS	3.9 - 3.13	arm64 (M1/M2/M3)	✅ Tested
Windows	3.9 - 3.13	x86_64	✅ Tested

[!NOTE] macOS Intel (x86_64) Support: Binary wheels for macOS x86_64 (Intel processors) are no longer provided as of v0.2.0. If you have an Intel-based Mac, you can install from source:

git clone https://github.com/QianyeSu/KuoEliassen.git
cd KuoEliassen
pip install -e .

Requires: brew install gcc

Requirements

• Python: ≥ 3.9
• NumPy: ≥ 1.24.0
• SciPy: ≥ 1.7.0
• xarray: ≥ 0.19.0 (optional, for labeled array interface)

Development Requirements

For building from source, you'll need:

• A Fortran compiler: gfortran, ifort, or flang
• Meson build system
• NumPy's f2py (included with NumPy)

License

This project is licensed under the BSD-3-Clause License - see the LICENSE file for details.

Citation

If you use KuoEliassen in your research, please cite:

@article{Su2025,
  author = {Su, Q. and Liu, C. and Zhang, Y. and Qiu, J. and Li, J. and Xue, Y. and Cao, N. and Liao, X. and Yang, K. and Zheng, R. and Liang, Z. and Jin, L. and Huang, K. and Jin, K. and Zhou, N.},
  title = {Consistency of Changes in the Ascending and Descending Positions of the Hadley Circulation Using Different Methods},
  journal = {Atmosphere},
  year = {2025},
  volume = {16},
  number = {4},
  pages = {367},
  doi = {10.3390/atmos16040367}
}

@software{kuoeliassen2025,
  author = {Su, Qianye},
  title = {KuoEliassen: A High-Performance Solver for the Kuo-Eliassen Equation},
  year = {2025},
  publisher = {Zenodo},
  doi = {10.5281/zenodo.18060064},
  url = {https://doi.org/10.5281/zenodo.18060064}
}

Contact

• Author: Qianye Su
• Email: [e-mail]
• Repository: https://github.com/QianyeSu/KuoEliassen
• Issue Tracker: https://github.com/QianyeSu/KuoEliassen/issues

For questions, bug reports, or feature requests, please open an issue on GitHub.

Acknowledgments & References

This solver implements the Kuo-Eliassen equation, a fundamental tool in atmospheric dynamics for understanding meridional circulation:

• Su, Q., Liu, C., Zhang, Y., Qiu, J., Li, J., Xue, Y., Cao, N., Liao, X., Yang, K., Zheng, R., Liang, Z., Jin, L., Huang, K., Jin, K., & Zhou, N. (2025). Consistency of Changes in the Ascending and Descending Positions of the Hadley Circulation Using Different Methods. Atmosphere, 16(4), 367. https://doi.org/10.3390/atmos16040367
• Kuo, H.-L. (1956). Forced and free meridional circulations in the atmosphere. J. Atmos. Sci., 13, 561–568.
• Kim, H.-K. & Lee, S. (2001). Hadley cell dynamics in a primitive equation model. Part I: axisymmetric flow. J. Atmos. Sci., 58, 2845–2858.
• Chemke, R. & Polvani, L. M. (2019). Opposite tropical circulation trends in climate models and in reanalyses. Nat. Geosci., 12, 528–532.
• Pikovnik, M., Zaplotnik, Ž, Boljka, L. & Žagar, N. (2022). Metrics of the Hadley circulation strength and associated circulation trends. Weather Clim. Dyn., 3, 625–644.
• Held, I. M. & Zurita-Gotor, P. (2025). Misuse of Kuo–Eliassen Equation in Studies of the Climatological Mean Meridional Circulation. J. Atmos. Sci., 82, 1765–1766.

The numerical implementation employs:

• SOR (Successive Over-Relaxation) iterative solver for the elliptic operator
• Centered finite differences for spatial derivatives with pole-aware boundary handling
• Fortran 90 for high performance and portability