NEDAS.models.qg.python.model module

QGModel: Python implementation of the multi-layer quasi-geostrophic spectral model.

Translates qg_driver.f90 / qg_run_tools.f90 / qg_init_tools.f90.

Physics

PV equation (spectral):

dq/dt = J(ψ, q) + β·∂ψ/∂x + (mean-flow terms) + (dissipation) + (forcing)

PV-streamfunction relation:

q = -|k|²ψ + S·ψ (S = tridiagonal stratification operator, |k| = wavenumber magnitude)

Time integration: leapfrog with Robert (Asselin) filter, adaptive timestep.

Array conventions

Spectral: shape (nz, nky, nkx) — nky = kmax+1, nkx = 2*kmax+1 Physical: shape (nz, ny, nx) — nx = ny = 2*(kmax+1) Wavenumber grids: shape (nky, nkx)

Usage

m = QGModel(kmax=63, nz=4, F=50.0, beta=1.5, ...)
m.initialize()
for _ in range(1000):
    m.step()

class NEDAS.models.qg.python.model.QGModel(kmax: int = 63, nz: int = 1, beta: float = 1.5, F: float = 0.0, Fe: float = 0.0, uscale: float = 0.0, vscale: float = 0.0, strat_type: str = 'linear', deltc: float = 0.2, surface_bc: str = 'rigid_lid', dt: float = 0.0, dt_max: float = 0.0, adapt_dt: bool = True, dt_tune: float = 1.5, dt_step: int = 10, robert: float = 0.01, filter_type: str = 'hyperviscous', filter_exp: float = 8.0, k_cut: float = 0.0, dealiasing: str = 'isotropic', filt_tune: float = 1.0, bot_drag: float = 0.0, top_drag: float = 0.0, therm_drag: float = 0.0, quad_drag: float = 0.0, qd_angle: float = 0.0, use_forcing: bool = False, norm_forcing: bool = False, forc_coef: float = 0.0, forc_corr: float = 0.0, kf_min: float = 0.0, kf_max: float = 0.0, use_topo: bool = False, use_tracer: bool = False, linear: bool = False, idum: int = -7)

Bases: object

Multi-layer QG spectral model (Python / NumPy / JAX).

Parameters mirror the Fortran input namelist (qg_params.f90).

kmax: int = 63

nz: int = 1

beta: float = 1.5

F: float = 0.0

Fe: float = 0.0

uscale: float = 0.0

vscale: float = 0.0

strat_type: str = 'linear'

deltc: float = 0.2

surface_bc: str = 'rigid_lid'

dt: float = 0.0

dt_max: float = 0.0

adapt_dt: bool = True

dt_tune: float = 1.5

dt_step: int = 10

robert: float = 0.01

filter_type: str = 'hyperviscous'

filter_exp: float = 8.0

k_cut: float = 0.0

dealiasing: str = 'isotropic'

filt_tune: float = 1.0

bot_drag: float = 0.0

top_drag: float = 0.0

therm_drag: float = 0.0

quad_drag: float = 0.0

qd_angle: float = 0.0

use_forcing: bool = False

norm_forcing: bool = False

forc_coef: float = 0.0

forc_corr: float = 0.0

kf_min: float = 0.0

kf_max: float = 0.0

use_topo: bool = False

use_tracer: bool = False

linear: bool = False

idum: int = -7

pi: float = 3.141592653589793

q: ndarray | None = None

q_o: ndarray | None = None

rhs: ndarray | None = None

psi: ndarray | None = None

psi_o: ndarray | None = None

cntr: int = 0

time: float = 0.0

initialize(psi_init=None, dz=None, rho=None, ubar=None, vbar=None, hb=None)

Set up spectral grid, stratification, and initial fields.

Parameters:

• psi_init (ndarray (nz, nky, nkx) complex, optional) – Initial streamfunction in spectral space.

• dz ((nz,) layer thicknesses (should sum to 1))

• rho ((nz,) layer densities)

• ubar ((nz,) mean zonal / meridional velocity profiles)

• vbar ((nz,) mean zonal / meridional velocity profiles)

• hb ((nky, nkx) complex bottom topography in spectral space)

get_pv(psi)

Compute PV q from streamfunction ψ (spectral).

q = -|k|²ψ + S·ψ

Translates Fortran Get_pv in qg_run_tools.f90. psi : (nz, nky, nkx) or (nv, nky, nkx) for surf_buoy Returns (nz, nky, nkx)

invert_pv()

Invert PV q → ψ by solving the tridiagonal system per wavenumber.

Translates Fortran Invert_pv in qg_run_tools.f90. Returns psi of shape (nz, nky, nkx).

step()

Advance the model one time step.

Translates the main time loop body in qg_driver.f90.

run(n_steps)

Run the model for n_steps time steps.

get_psi_grid()

Return streamfunction in physical space, shape (nz, ny, nx).

get_q_grid()

Return PV in physical space, shape (nz, ny, nx).

get_u_grid()

Return zonal velocity, shape (nz, ny, nx).

get_v_grid()

Return meridional velocity, shape (nz, ny, nx).

property nx

property ny

property nkx

property nky