qg-python 1.0.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
qg_model.py
ADDED
|
@@ -0,0 +1,753 @@
|
|
|
1
|
+
"""
|
|
2
|
+
QG-Python: 1.5-layer Quasi-Geostrophic Model
|
|
3
|
+
=============================================
|
|
4
|
+
Faithful Python port of the QG-C model (Pavel Sakov, Bureau of Meteorology).
|
|
5
|
+
|
|
6
|
+
PHYSICS:
|
|
7
|
+
Potential vorticity equation:
|
|
8
|
+
dq/dt = -r·J(psi,q) - rkb·zeta + rkh·nabla²zeta - rkh2·nabla⁴zeta + curl(tau) - dpsi/dy
|
|
9
|
+
q = nabla²psi - F·psi (potential vorticity)
|
|
10
|
+
with Dirichlet boundary conditions: psi = 0
|
|
11
|
+
|
|
12
|
+
TIME-STEPPING SCHEMES:
|
|
13
|
+
order1 - explicit Euler (1st order)
|
|
14
|
+
order2 - midpoint Runge-Kutta (2nd order)
|
|
15
|
+
rk4 - classical Runge-Kutta (4th order) [default]
|
|
16
|
+
dp5 - Dormand-Prince (5th order)
|
|
17
|
+
leapfrog - Leapfrog with Robert-Asselin (RA) or Robert-Asselin-Williams (RAW) filter
|
|
18
|
+
|
|
19
|
+
LEAPFROG FILTER OPTIONS:
|
|
20
|
+
--ra-alpha FLOAT Robert-Asselin filter parameter (default: 0.1)
|
|
21
|
+
--raw Use RAW (Robert-Asselin-Williams) filter instead of RA
|
|
22
|
+
RAW is 3rd-order accurate in the filter correction vs
|
|
23
|
+
1st-order for RA, and reduces spurious amplitude damping.
|
|
24
|
+
|
|
25
|
+
OUTPUT:
|
|
26
|
+
If netCDF4 is installed -> .nc (compatible with the original QG-C)
|
|
27
|
+
Otherwise -> .npz (NumPy compressed, readable with np.load)
|
|
28
|
+
|
|
29
|
+
USAGE:
|
|
30
|
+
python qg_model.py # default parameters
|
|
31
|
+
python qg_model.py my_config.prm # parameter file
|
|
32
|
+
python qg_model.py --tend 1000 --scheme dp5 # command-line options
|
|
33
|
+
python qg_model.py --scheme leapfrog --ra-alpha 0.1 # Leapfrog + RA filter
|
|
34
|
+
python qg_model.py --scheme leapfrog --raw --ra-alpha 0.1 # Leapfrog + RAW filter
|
|
35
|
+
|
|
36
|
+
From Python:
|
|
37
|
+
from qg_model import QGModel, QGParams
|
|
38
|
+
prm = QGParams(tend=5000, dtout=50, scheme='leapfrog',
|
|
39
|
+
ra_alpha=0.1, raw_filter=False)
|
|
40
|
+
m = QGModel(prm)
|
|
41
|
+
psi, q = m.run('restart', fname='qg_fourier.npz')
|
|
42
|
+
|
|
43
|
+
DEPENDENCIES:
|
|
44
|
+
numpy, scipy (required)
|
|
45
|
+
netCDF4 (optional - for .nc output)
|
|
46
|
+
"""
|
|
47
|
+
|
|
48
|
+
__version__ = "1.0.0"
|
|
49
|
+
|
|
50
|
+
import numpy as np
|
|
51
|
+
from scipy.sparse import diags
|
|
52
|
+
from scipy.sparse.linalg import factorized
|
|
53
|
+
from dataclasses import dataclass, field
|
|
54
|
+
from typing import Optional
|
|
55
|
+
import time
|
|
56
|
+
import os
|
|
57
|
+
|
|
58
|
+
try:
|
|
59
|
+
import netCDF4 as _nc
|
|
60
|
+
HAS_NC = True
|
|
61
|
+
except ImportError:
|
|
62
|
+
HAS_NC = False
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
# ============================================================================
|
|
66
|
+
# Parameters
|
|
67
|
+
# ============================================================================
|
|
68
|
+
|
|
69
|
+
@dataclass
|
|
70
|
+
class QGParams:
|
|
71
|
+
# --- Grid ---
|
|
72
|
+
mrefin: int = 7
|
|
73
|
+
nx1: int = 3
|
|
74
|
+
ny1: int = 3
|
|
75
|
+
lx: float = 1.0
|
|
76
|
+
|
|
77
|
+
# --- Physics ---
|
|
78
|
+
rkb: float = 3.0e-4
|
|
79
|
+
rkh: float = -5.0e-8
|
|
80
|
+
rkh2: float = 1.0e-12
|
|
81
|
+
f: float = 1600.0
|
|
82
|
+
r: float = 1.0e-5
|
|
83
|
+
a: float = 6.29
|
|
84
|
+
k: float = 1.0
|
|
85
|
+
|
|
86
|
+
# --- Time integration ---
|
|
87
|
+
scheme: str = 'rk4'
|
|
88
|
+
dt: float = 1.0
|
|
89
|
+
tend: float = 50000.0
|
|
90
|
+
dtout: float = 20.0
|
|
91
|
+
|
|
92
|
+
# --- Leapfrog filter options ---
|
|
93
|
+
ra_alpha: float = 0.1 # Robert-Asselin filter strength (0 = off)
|
|
94
|
+
raw_filter: bool = False # True -> use RAW (Williams 2009) instead of RA
|
|
95
|
+
|
|
96
|
+
# --- I/O ---
|
|
97
|
+
outfname: str = 'qg.nc'
|
|
98
|
+
save_q: bool = False
|
|
99
|
+
verbose: bool = True
|
|
100
|
+
|
|
101
|
+
m: int = field(init=False)
|
|
102
|
+
n: int = field(init=False)
|
|
103
|
+
|
|
104
|
+
def __post_init__(self):
|
|
105
|
+
ksq = 2 ** (self.mrefin - 1)
|
|
106
|
+
self.m = self.ny1 * ksq + 1
|
|
107
|
+
self.n = self.nx1 * ksq + 1
|
|
108
|
+
|
|
109
|
+
@classmethod
|
|
110
|
+
def from_file(cls, fname: str) -> 'QGParams':
|
|
111
|
+
prm = cls()
|
|
112
|
+
_conv = {
|
|
113
|
+
'mrefin': ('mrefin', int),
|
|
114
|
+
'nx1': ('nx1', int),
|
|
115
|
+
'ny1': ('ny1', int),
|
|
116
|
+
'lx': ('lx', float),
|
|
117
|
+
'rkb': ('rkb', float),
|
|
118
|
+
'rkh': ('rkh', float),
|
|
119
|
+
'rkh2': ('rkh2', float),
|
|
120
|
+
'f': ('f', float),
|
|
121
|
+
'r': ('r', float),
|
|
122
|
+
'a': ('a', float),
|
|
123
|
+
'k': ('k', float),
|
|
124
|
+
'scheme': ('scheme', str),
|
|
125
|
+
'dt': ('dt', float),
|
|
126
|
+
'tend': ('tend', float),
|
|
127
|
+
'dtout': ('dtout', float),
|
|
128
|
+
'ra_alpha': ('ra_alpha', float),
|
|
129
|
+
'raw_filter': ('raw_filter', lambda x: x.strip() in ('1','yes','true','Yes')),
|
|
130
|
+
'outfname': ('outfname', str),
|
|
131
|
+
'save_q': ('save_q', lambda x: x.strip() in ('1','yes','true','Yes')),
|
|
132
|
+
}
|
|
133
|
+
with open(fname) as fh:
|
|
134
|
+
for line in fh:
|
|
135
|
+
line = line.strip()
|
|
136
|
+
if not line or line.startswith('#'):
|
|
137
|
+
continue
|
|
138
|
+
if '=' not in line:
|
|
139
|
+
continue
|
|
140
|
+
key, _, val = line.partition('=')
|
|
141
|
+
key = key.strip().lower()
|
|
142
|
+
val = val.strip()
|
|
143
|
+
if key in _conv:
|
|
144
|
+
attr, conv = _conv[key]
|
|
145
|
+
setattr(prm, attr, conv(val))
|
|
146
|
+
prm.__post_init__()
|
|
147
|
+
return prm
|
|
148
|
+
|
|
149
|
+
def describe(self):
|
|
150
|
+
print(f" grid : {self.m} x {self.n} (mrefin={self.mrefin}, nx1={self.nx1}, ny1={self.ny1})")
|
|
151
|
+
print(f" lx : {self.lx}")
|
|
152
|
+
print(f" rkb/rkh/rkh2: {self.rkb} / {self.rkh} / {self.rkh2}")
|
|
153
|
+
print(f" F, r, A, k : {self.f}, {self.r}, {self.a}, {self.k}")
|
|
154
|
+
print(f" scheme : {self.scheme}")
|
|
155
|
+
if self.scheme == 'leapfrog':
|
|
156
|
+
filt = 'RAW' if self.raw_filter else 'RA'
|
|
157
|
+
print(f" LF filter : {filt} alpha={self.ra_alpha}")
|
|
158
|
+
print(f" dt / tend : {self.dt} / {self.tend}")
|
|
159
|
+
print(f" dtout : {self.dtout}")
|
|
160
|
+
print(f" output : {self.outfname} (save_q={self.save_q})")
|
|
161
|
+
|
|
162
|
+
|
|
163
|
+
# ============================================================================
|
|
164
|
+
# Numerical operators
|
|
165
|
+
# ============================================================================
|
|
166
|
+
|
|
167
|
+
def arakawa(h: float, psi: np.ndarray, q: np.ndarray) -> np.ndarray:
|
|
168
|
+
c = 12.0 * h * h
|
|
169
|
+
J = np.zeros_like(psi)
|
|
170
|
+
|
|
171
|
+
pjm = psi[:-2, 1:-1]; pjp = psi[2:, 1:-1]
|
|
172
|
+
pim = psi[1:-1, :-2]; pip = psi[1:-1, 2:]
|
|
173
|
+
pmm = psi[:-2, :-2]; pmp = psi[:-2, 2:]
|
|
174
|
+
ppm = psi[2:, :-2]; ppp = psi[2:, 2:]
|
|
175
|
+
|
|
176
|
+
qjm = q[:-2, 1:-1]; qjp = q[2:, 1:-1]
|
|
177
|
+
qim = q[1:-1, :-2]; qip = q[1:-1, 2:]
|
|
178
|
+
qmm = q[:-2, :-2]; qmp = q[:-2, 2:]
|
|
179
|
+
qpm = q[2:, :-2]; qpp = q[2:, 2:]
|
|
180
|
+
|
|
181
|
+
J[1:-1, 1:-1] = (
|
|
182
|
+
(pjm - pim ) * qmm
|
|
183
|
+
+ (pmm + pjm - ppm - pjp ) * qim
|
|
184
|
+
+ (pim - pjp ) * qpm
|
|
185
|
+
+ (pmp + pip - pmm - pim ) * qjm
|
|
186
|
+
+ (pim + ppm - pip - ppp ) * qjp
|
|
187
|
+
+ (pip - pjm ) * qmp
|
|
188
|
+
+ (pjp + ppp - pjm - pmp ) * qip
|
|
189
|
+
+ (pjp - pip ) * qpp
|
|
190
|
+
) / c
|
|
191
|
+
return J
|
|
192
|
+
|
|
193
|
+
|
|
194
|
+
def laplacian(h: float, A: np.ndarray) -> np.ndarray:
|
|
195
|
+
h2 = 1.0 / (h * h)
|
|
196
|
+
L = np.zeros_like(A)
|
|
197
|
+
L[1:-1, 1:-1] = (
|
|
198
|
+
A[:-2, 1:-1] + A[2:, 1:-1] + A[1:-1, :-2] + A[1:-1, 2:]
|
|
199
|
+
- 4.0 * A[1:-1, 1:-1]
|
|
200
|
+
) * h2
|
|
201
|
+
return L
|
|
202
|
+
|
|
203
|
+
|
|
204
|
+
# ============================================================================
|
|
205
|
+
# Helmholtz solver
|
|
206
|
+
# ============================================================================
|
|
207
|
+
|
|
208
|
+
def _build_helmholtz_matrix(m: int, n: int, h: float, F: float):
|
|
209
|
+
mn = m * n
|
|
210
|
+
h2 = 1.0 / (h * h)
|
|
211
|
+
|
|
212
|
+
diag = np.full(mn, -(4.0 * h2 + F))
|
|
213
|
+
off = np.full(mn, h2)
|
|
214
|
+
|
|
215
|
+
for j in range(m):
|
|
216
|
+
for i in range(n):
|
|
217
|
+
if j == 0 or j == m-1 or i == 0 or i == n-1:
|
|
218
|
+
diag[j*n + i] = 1.0
|
|
219
|
+
|
|
220
|
+
off_e = off.copy()
|
|
221
|
+
off_w = off.copy()
|
|
222
|
+
for j in range(m):
|
|
223
|
+
k = j*n + (n-1)
|
|
224
|
+
off_e[k] = 0.0
|
|
225
|
+
off_w[k] = 0.0
|
|
226
|
+
|
|
227
|
+
border = np.zeros(mn, dtype=bool)
|
|
228
|
+
border[:n] = True
|
|
229
|
+
border[-n:] = True
|
|
230
|
+
border[::n] = True
|
|
231
|
+
border[n-1::n] = True
|
|
232
|
+
|
|
233
|
+
off_e[border] = 0.0
|
|
234
|
+
off_w[border] = 0.0
|
|
235
|
+
off_n = off.copy(); off_n[border] = 0.0
|
|
236
|
+
off_s = off.copy(); off_s[border] = 0.0
|
|
237
|
+
|
|
238
|
+
A = diags(
|
|
239
|
+
[off_s[:mn-n], off_w[:mn-1], diag, off_e[:mn-1], off_n[:mn-n]],
|
|
240
|
+
[-n, -1, 0, 1, n],
|
|
241
|
+
shape=(mn, mn),
|
|
242
|
+
format='csr'
|
|
243
|
+
)
|
|
244
|
+
return A
|
|
245
|
+
|
|
246
|
+
|
|
247
|
+
class HelmholtzSolver:
|
|
248
|
+
def __init__(self, m: int, n: int, h: float, F: float):
|
|
249
|
+
self.m = m
|
|
250
|
+
self.n = n
|
|
251
|
+
A = _build_helmholtz_matrix(m, n, h, F).tocsc() # CSC: what SuperLU wants
|
|
252
|
+
self._lu = factorized(A)
|
|
253
|
+
|
|
254
|
+
def solve(self, q: np.ndarray) -> np.ndarray:
|
|
255
|
+
rhs = q.ravel().copy()
|
|
256
|
+
m, n = self.m, self.n
|
|
257
|
+
mn = m * n
|
|
258
|
+
border_idx = (
|
|
259
|
+
list(range(n))
|
|
260
|
+
+ list(range(mn-n, mn))
|
|
261
|
+
+ list(range(0, mn, n))
|
|
262
|
+
+ list(range(n-1, mn, n))
|
|
263
|
+
)
|
|
264
|
+
rhs[border_idx] = 0.0
|
|
265
|
+
return self._lu(rhs).reshape(m, n)
|
|
266
|
+
|
|
267
|
+
|
|
268
|
+
# ============================================================================
|
|
269
|
+
# QG Model
|
|
270
|
+
# ============================================================================
|
|
271
|
+
|
|
272
|
+
class QGModel:
|
|
273
|
+
|
|
274
|
+
def __init__(self, prm: QGParams):
|
|
275
|
+
self.prm = prm
|
|
276
|
+
self.m = prm.m
|
|
277
|
+
self.n = prm.n
|
|
278
|
+
self.h = prm.lx / (prm.n - 1)
|
|
279
|
+
self.t = 0.0
|
|
280
|
+
|
|
281
|
+
self.psi = np.zeros((prm.m, prm.n))
|
|
282
|
+
self.q = np.zeros((prm.m, prm.n))
|
|
283
|
+
|
|
284
|
+
# Leapfrog requires q at the previous time level
|
|
285
|
+
self._q_prev = None # set to None; initialised at first LF step
|
|
286
|
+
|
|
287
|
+
i_idx = np.arange(prm.n)
|
|
288
|
+
tmp = np.sin(prm.k * 2.0 * np.pi * i_idx / prm.n)
|
|
289
|
+
self.curlt = -prm.a * tmp * np.abs(tmp)
|
|
290
|
+
|
|
291
|
+
if prm.verbose:
|
|
292
|
+
print(f" Building Helmholtz solver ({prm.m}x{prm.n})...", end=' ', flush=True)
|
|
293
|
+
t0 = time.time()
|
|
294
|
+
self.helm = HelmholtzSolver(prm.m, prm.n, self.h, prm.f)
|
|
295
|
+
if prm.verbose:
|
|
296
|
+
print(f"done ({time.time()-t0:.2f}s)")
|
|
297
|
+
|
|
298
|
+
self._out_t = []
|
|
299
|
+
self._out_psi = []
|
|
300
|
+
self._out_q = []
|
|
301
|
+
|
|
302
|
+
def _calc_psi(self, q):
|
|
303
|
+
return self.helm.solve(q)
|
|
304
|
+
|
|
305
|
+
def _rhs(self, q, psi):
|
|
306
|
+
prm = self.prm
|
|
307
|
+
h = self.h
|
|
308
|
+
|
|
309
|
+
J = arakawa(h, psi, q)
|
|
310
|
+
zeta = laplacian(h, psi)
|
|
311
|
+
zeta2 = laplacian(h, zeta)
|
|
312
|
+
zeta4 = laplacian(h, zeta2)
|
|
313
|
+
|
|
314
|
+
dpsi_dy = np.zeros_like(psi)
|
|
315
|
+
dpsi_dy[1:-1, 1:-1] = (psi[2:, 1:-1] - psi[:-2, 1:-1]) / (2.0 * h)
|
|
316
|
+
|
|
317
|
+
F = np.zeros_like(q)
|
|
318
|
+
F[1:-1, 1:-1] = (
|
|
319
|
+
- prm.r * J[1:-1, 1:-1]
|
|
320
|
+
- prm.rkb * zeta[1:-1, 1:-1]
|
|
321
|
+
+ prm.rkh * zeta2[1:-1, 1:-1]
|
|
322
|
+
- prm.rkh2 * zeta4[1:-1, 1:-1]
|
|
323
|
+
+ self.curlt[1:-1]
|
|
324
|
+
- dpsi_dy[1:-1, 1:-1]
|
|
325
|
+
)
|
|
326
|
+
return F
|
|
327
|
+
|
|
328
|
+
# ── existing schemes ─────────────────────────────────────────────────────
|
|
329
|
+
|
|
330
|
+
def _step_order1(self):
|
|
331
|
+
dt = self.prm.dt
|
|
332
|
+
psi1 = self._calc_psi(self.q)
|
|
333
|
+
k1 = self._rhs(self.q, psi1)
|
|
334
|
+
self.q += dt * k1
|
|
335
|
+
self.psi = self._calc_psi(self.q)
|
|
336
|
+
self.t += dt
|
|
337
|
+
|
|
338
|
+
def _step_order2(self):
|
|
339
|
+
dt = self.prm.dt
|
|
340
|
+
q0 = self.q.copy()
|
|
341
|
+
psi0 = self._calc_psi(q0)
|
|
342
|
+
k1 = self._rhs(q0, psi0)
|
|
343
|
+
qm = q0 + 0.5 * dt * k1
|
|
344
|
+
psim = self._calc_psi(qm)
|
|
345
|
+
k2 = self._rhs(qm, psim)
|
|
346
|
+
self.q = q0 + dt * k2
|
|
347
|
+
self.psi = self._calc_psi(self.q)
|
|
348
|
+
self.t += dt
|
|
349
|
+
|
|
350
|
+
def _step_rk4(self):
|
|
351
|
+
dt = self.prm.dt
|
|
352
|
+
q0 = self.q.copy()
|
|
353
|
+
|
|
354
|
+
psi1 = self._calc_psi(q0); k1 = self._rhs(q0, psi1)
|
|
355
|
+
psi2 = self._calc_psi(q0 + 0.5*dt*k1); k2 = self._rhs(q0+0.5*dt*k1, psi2)
|
|
356
|
+
psi3 = self._calc_psi(q0 + 0.5*dt*k2); k3 = self._rhs(q0+0.5*dt*k2, psi3)
|
|
357
|
+
psi4 = self._calc_psi(q0 + dt*k3); k4 = self._rhs(q0+dt*k3, psi4)
|
|
358
|
+
|
|
359
|
+
self.q = q0 + (dt / 6.0) * (k1 + 2.0*k2 + 2.0*k3 + k4)
|
|
360
|
+
self.psi = self._calc_psi(self.q)
|
|
361
|
+
self.t += dt
|
|
362
|
+
|
|
363
|
+
def _step_dp5(self):
|
|
364
|
+
dt = self.prm.dt
|
|
365
|
+
q0 = self.q.copy()
|
|
366
|
+
|
|
367
|
+
def f(q): return self._rhs(q, self._calc_psi(q))
|
|
368
|
+
|
|
369
|
+
k1 = f(q0)
|
|
370
|
+
k2 = f(q0 + dt * 0.2 * k1)
|
|
371
|
+
k3 = f(q0 + dt * (3/40*k1 + 9/40*k2))
|
|
372
|
+
k4 = f(q0 + dt * (44/45*k1 - 56/15*k2 + 32/9*k3))
|
|
373
|
+
k5 = f(q0 + dt * (19372/6561*k1 - 25360/2187*k2
|
|
374
|
+
+ 64448/6561*k3 - 212/729*k4))
|
|
375
|
+
k6 = f(q0 + dt * (9017/3168*k1 - 355/33*k2
|
|
376
|
+
+ 46732/5247*k3 + 49/176*k4 - 5103/18656*k5))
|
|
377
|
+
|
|
378
|
+
self.q = q0 + dt * (35/384*k1 + 500/1113*k3 + 125/192*k4
|
|
379
|
+
- 2187/6784*k5 + 11/84*k6)
|
|
380
|
+
self.psi = self._calc_psi(self.q)
|
|
381
|
+
self.t += dt
|
|
382
|
+
|
|
383
|
+
# ── Leapfrog + Robert-Asselin / RAW filter ────────────────────────────────
|
|
384
|
+
#
|
|
385
|
+
# Standard Leapfrog (2nd-order):
|
|
386
|
+
# q^{n+1} = q^{n-1} + 2*dt * F(q^n, psi^n)
|
|
387
|
+
#
|
|
388
|
+
# Robert-Asselin (RA) filter applied to q^n after each step:
|
|
389
|
+
# q^n <- q^n + alpha/2 * (q^{n-1} - 2*q^n + q^{n+1})
|
|
390
|
+
#
|
|
391
|
+
# Robert-Asselin-Williams (RAW) filter (Williams 2009):
|
|
392
|
+
# Splits the RA correction between q^n and q^{n+1} using parameter beta,
|
|
393
|
+
# preserving 2nd-order accuracy of the physical mode while still
|
|
394
|
+
# suppressing the computational mode. With beta=0.5 (default):
|
|
395
|
+
# q^n <- q^n + alpha*(1-beta)/2 * (q^{n-1} - 2*q^n + q^{n+1})
|
|
396
|
+
# q^{n+1}<- q^{n+1} - alpha*beta/2 * (q^{n-1} - 2*q^n + q^{n+1})
|
|
397
|
+
#
|
|
398
|
+
# First leapfrog step (n=0) uses a single RK4 step to obtain q^1,
|
|
399
|
+
# ensuring a clean start without the computational mode.
|
|
400
|
+
|
|
401
|
+
def _step_leapfrog(self):
|
|
402
|
+
dt = self.prm.dt
|
|
403
|
+
alpha = self.prm.ra_alpha
|
|
404
|
+
raw = self.prm.raw_filter
|
|
405
|
+
beta = 0.5 # RAW splitting parameter (Williams 2009 recommends 0.5)
|
|
406
|
+
|
|
407
|
+
# ── bootstrap: first step uses RK4 to avoid exciting computational mode
|
|
408
|
+
if self._q_prev is None:
|
|
409
|
+
self._q_prev = self.q.copy()
|
|
410
|
+
self._step_rk4() # advances self.q and self.t by dt
|
|
411
|
+
return
|
|
412
|
+
|
|
413
|
+
# ── standard leapfrog leap
|
|
414
|
+
q_nm1 = self._q_prev # q^{n-1}
|
|
415
|
+
q_n = self.q.copy() # q^n
|
|
416
|
+
psi_n = self._calc_psi(q_n) # psi^n (one Helmholtz solve)
|
|
417
|
+
|
|
418
|
+
F_n = self._rhs(q_n, psi_n)
|
|
419
|
+
q_np1 = q_nm1 + 2.0 * dt * F_n # leapfrog step: q^{n+1}
|
|
420
|
+
|
|
421
|
+
# ── filter correction
|
|
422
|
+
delta = q_nm1 - 2.0 * q_n + q_np1 # 2nd-difference (centred on n)
|
|
423
|
+
|
|
424
|
+
if raw:
|
|
425
|
+
# RAW filter (Williams 2009) — less dissipative than plain RA
|
|
426
|
+
q_n_filtered = q_n + 0.5 * alpha * (1.0 - beta) * delta
|
|
427
|
+
q_np1_filtered = q_np1 - 0.5 * alpha * beta * delta
|
|
428
|
+
else:
|
|
429
|
+
# Classic Robert-Asselin filter
|
|
430
|
+
q_n_filtered = q_n + 0.5 * alpha * delta
|
|
431
|
+
q_np1_filtered = q_np1 # RA does not modify q^{n+1}
|
|
432
|
+
|
|
433
|
+
# ── advance state
|
|
434
|
+
self._q_prev = q_n_filtered # filtered q^n becomes q^{n-1}
|
|
435
|
+
self.q = q_np1_filtered # filtered q^{n+1} is new state
|
|
436
|
+
self.psi = self._calc_psi(self.q)
|
|
437
|
+
self.t += dt
|
|
438
|
+
|
|
439
|
+
# ── dispatcher ───────────────────────────────────────────────────────────
|
|
440
|
+
|
|
441
|
+
def _do_step(self):
|
|
442
|
+
s = self.prm.scheme
|
|
443
|
+
if s == 'order1': self._step_order1()
|
|
444
|
+
elif s == 'order2': self._step_order2()
|
|
445
|
+
elif s == 'rk4': self._step_rk4()
|
|
446
|
+
elif s == 'dp5': self._step_dp5()
|
|
447
|
+
elif s == 'leapfrog': self._step_leapfrog()
|
|
448
|
+
else: raise ValueError(f"Unknown scheme: '{s}'")
|
|
449
|
+
|
|
450
|
+
if not np.isfinite(self.psi).all():
|
|
451
|
+
raise RuntimeError(
|
|
452
|
+
f"\n\n *** INSTABILITY (NaN/Inf) at t={self.t:.1f} ***\n"
|
|
453
|
+
f" Suggestions:\n"
|
|
454
|
+
f" 1. Reduce dt (current={self.prm.dt}) -> try dt <= {self.prm.dt/4:.4g}\n"
|
|
455
|
+
f" 2. Increase RA filter: --ra-alpha 0.2\n"
|
|
456
|
+
f" 3. Try RAW filter: --raw\n"
|
|
457
|
+
f" 4. Disable neg. fric.: --rkh 0\n"
|
|
458
|
+
f" 5. More dissipation: --rkh2 1e-11\n"
|
|
459
|
+
)
|
|
460
|
+
|
|
461
|
+
def _q_from_psi(self):
|
|
462
|
+
self.q = laplacian(self.h, self.psi) - self.prm.f * self.psi
|
|
463
|
+
|
|
464
|
+
# ── initial conditions ───────────────────────────────────────────────────
|
|
465
|
+
|
|
466
|
+
def set_ic(self, kind: str = 'zero', **kwargs):
|
|
467
|
+
m, n = self.m, self.n
|
|
468
|
+
self.psi = np.zeros((m, n))
|
|
469
|
+
self.q = np.zeros((m, n))
|
|
470
|
+
self.t = 0.0
|
|
471
|
+
self._q_prev = None # reset leapfrog memory
|
|
472
|
+
|
|
473
|
+
if kind == 'zero':
|
|
474
|
+
pass
|
|
475
|
+
|
|
476
|
+
elif kind == 'fourier':
|
|
477
|
+
amplitude = kwargs.get('amplitude', 1.0)
|
|
478
|
+
kmax = kwargs.get('kmax', 5)
|
|
479
|
+
seed = kwargs.get('seed', 42)
|
|
480
|
+
rng = np.random.default_rng(seed)
|
|
481
|
+
|
|
482
|
+
j_idx = np.arange(m)
|
|
483
|
+
i_idx = np.arange(n)
|
|
484
|
+
psi = np.zeros((m, n))
|
|
485
|
+
for kj in range(1, kmax + 1):
|
|
486
|
+
for ki in range(1, kmax + 1):
|
|
487
|
+
amp = rng.uniform(-1, 1) * amplitude / (kj**2 + ki**2)
|
|
488
|
+
psi += amp * np.outer(
|
|
489
|
+
np.sin(kj * np.pi * j_idx / (m - 1)),
|
|
490
|
+
np.sin(ki * np.pi * i_idx / (n - 1))
|
|
491
|
+
)
|
|
492
|
+
self.psi = psi
|
|
493
|
+
self._q_from_psi()
|
|
494
|
+
|
|
495
|
+
elif kind == 'vortex':
|
|
496
|
+
vortices = kwargs.get('vortices', [
|
|
497
|
+
{'x': 0.5, 'y': 0.5, 'r': 0.1, 'amp': 1.0}
|
|
498
|
+
])
|
|
499
|
+
lx = self.prm.lx
|
|
500
|
+
x = np.linspace(0, lx, n)
|
|
501
|
+
y = np.linspace(0, lx, m)
|
|
502
|
+
X, Y = np.meshgrid(x, y)
|
|
503
|
+
|
|
504
|
+
psi = np.zeros((m, n))
|
|
505
|
+
for v in vortices:
|
|
506
|
+
x0 = v.get('x', 0.5) * lx
|
|
507
|
+
y0 = v.get('y', 0.5) * lx
|
|
508
|
+
r = v.get('r', 0.1) * lx
|
|
509
|
+
amp = v.get('amp', 1.0)
|
|
510
|
+
psi += amp * np.exp(-((X-x0)**2 + (Y-y0)**2) / (2.0 * r**2))
|
|
511
|
+
|
|
512
|
+
psi[0, :] = 0.0
|
|
513
|
+
psi[-1, :] = 0.0
|
|
514
|
+
psi[:, 0] = 0.0
|
|
515
|
+
psi[:, -1] = 0.0
|
|
516
|
+
self.psi = psi
|
|
517
|
+
self._q_from_psi()
|
|
518
|
+
|
|
519
|
+
elif kind == 'restart':
|
|
520
|
+
fname = kwargs.get('fname')
|
|
521
|
+
record = kwargs.get('record', -1)
|
|
522
|
+
if fname is None:
|
|
523
|
+
raise ValueError("set_ic('restart', ...) requires fname=<path>")
|
|
524
|
+
|
|
525
|
+
if fname.endswith('.npz'):
|
|
526
|
+
data = np.load(fname)
|
|
527
|
+
self.psi = data['psi'][record].astype(float)
|
|
528
|
+
if 'q' in data:
|
|
529
|
+
self.q = data['q'][record].astype(float)
|
|
530
|
+
else:
|
|
531
|
+
self._q_from_psi()
|
|
532
|
+
if 't' in data:
|
|
533
|
+
self.t = float(data['t'][record])
|
|
534
|
+
|
|
535
|
+
elif fname.endswith('.nc') or fname.endswith('.nc4'):
|
|
536
|
+
if not HAS_NC:
|
|
537
|
+
raise ImportError("netCDF4 is not installed; use a .npz file instead")
|
|
538
|
+
ds = _nc.Dataset(fname, 'r')
|
|
539
|
+
self.psi = np.array(ds.variables['psi'][record], dtype=float)
|
|
540
|
+
if 'q' in ds.variables:
|
|
541
|
+
self.q = np.array(ds.variables['q'][record], dtype=float)
|
|
542
|
+
else:
|
|
543
|
+
self._q_from_psi()
|
|
544
|
+
if 't' in ds.variables:
|
|
545
|
+
self.t = float(ds.variables['t'][record])
|
|
546
|
+
ds.close()
|
|
547
|
+
else:
|
|
548
|
+
raise ValueError(f"Unrecognised file format: {fname!r} (use .npz or .nc)")
|
|
549
|
+
|
|
550
|
+
if self.prm.verbose:
|
|
551
|
+
print(f" restart from '{fname}', record={record}, t={self.t:.1f}")
|
|
552
|
+
|
|
553
|
+
else:
|
|
554
|
+
raise ValueError(
|
|
555
|
+
f"Unknown IC kind='{kind}'. "
|
|
556
|
+
"Valid: 'zero', 'fourier', 'vortex', 'restart'"
|
|
557
|
+
)
|
|
558
|
+
|
|
559
|
+
# ── output helpers ───────────────────────────────────────────────────────
|
|
560
|
+
|
|
561
|
+
def _save_record(self):
|
|
562
|
+
self._out_t.append(float(self.t))
|
|
563
|
+
self._out_psi.append(self.psi.copy().astype(np.float32))
|
|
564
|
+
if self.prm.save_q:
|
|
565
|
+
self._out_q.append(self.q.copy().astype(np.float32))
|
|
566
|
+
|
|
567
|
+
def _write_output(self):
|
|
568
|
+
prm = self.prm
|
|
569
|
+
fname = prm.outfname
|
|
570
|
+
t_arr = np.array(self._out_t)
|
|
571
|
+
p_arr = np.stack(self._out_psi)
|
|
572
|
+
|
|
573
|
+
if HAS_NC:
|
|
574
|
+
if not fname.endswith('.nc'):
|
|
575
|
+
fname += '.nc'
|
|
576
|
+
ds = _nc.Dataset(fname, 'w', format='NETCDF4')
|
|
577
|
+
ds.setncattr('dt', prm.dt); ds.setncattr('lx', prm.lx)
|
|
578
|
+
ds.setncattr('rkb', prm.rkb); ds.setncattr('rkh', prm.rkh)
|
|
579
|
+
ds.setncattr('rkh2', prm.rkh2); ds.setncattr('F', prm.f)
|
|
580
|
+
ds.setncattr('r', prm.r); ds.setncattr('A', prm.a)
|
|
581
|
+
ds.setncattr('k', prm.k); ds.setncattr('scheme',prm.scheme)
|
|
582
|
+
ds.setncattr('ra_alpha', prm.ra_alpha)
|
|
583
|
+
ds.setncattr('raw_filter', int(prm.raw_filter))
|
|
584
|
+
ds.setncattr('comment', 'qg-python output')
|
|
585
|
+
ds.createDimension('record', None)
|
|
586
|
+
ds.createDimension('j', prm.m)
|
|
587
|
+
ds.createDimension('i', prm.n)
|
|
588
|
+
v = ds.createVariable('t', 'f8', ('record',)); v[:] = t_arr
|
|
589
|
+
v = ds.createVariable('psi', 'f4', ('record','j','i')); v[:] = p_arr
|
|
590
|
+
if prm.save_q:
|
|
591
|
+
v = ds.createVariable('q', 'f4', ('record','j','i'))
|
|
592
|
+
v[:] = np.stack(self._out_q)
|
|
593
|
+
ds.close()
|
|
594
|
+
else:
|
|
595
|
+
if not fname.endswith('.npz'):
|
|
596
|
+
fname = os.path.splitext(fname)[0] + '.npz'
|
|
597
|
+
kw = dict(t=t_arr, psi=p_arr,
|
|
598
|
+
lx=prm.lx, dt=prm.dt, rkb=prm.rkb,
|
|
599
|
+
rkh=prm.rkh, rkh2=prm.rkh2, F=prm.f,
|
|
600
|
+
r=prm.r, A=prm.a, k=prm.k,
|
|
601
|
+
ra_alpha=prm.ra_alpha,
|
|
602
|
+
raw_filter=int(prm.raw_filter))
|
|
603
|
+
if prm.save_q:
|
|
604
|
+
kw['q'] = np.stack(self._out_q)
|
|
605
|
+
np.savez_compressed(fname, **kw)
|
|
606
|
+
|
|
607
|
+
if prm.verbose:
|
|
608
|
+
print(f"\n Output saved: {fname} ({len(t_arr)} records)")
|
|
609
|
+
|
|
610
|
+
# ── main loop ────────────────────────────────────────────────────────────
|
|
611
|
+
|
|
612
|
+
def run(self, ic: str = 'zero', **ic_kwargs):
|
|
613
|
+
prm = self.prm
|
|
614
|
+
self.set_ic(ic, **ic_kwargs)
|
|
615
|
+
self._save_record()
|
|
616
|
+
|
|
617
|
+
nstep = int((prm.tend - self.t) / prm.dt)
|
|
618
|
+
dnout = max(1, int(prm.dtout / prm.dt))
|
|
619
|
+
|
|
620
|
+
if prm.verbose:
|
|
621
|
+
print(f" nstep = {nstep}, nrecord ~= {nstep // dnout}")
|
|
622
|
+
print(f" integrating [{prm.scheme}]:", end='', flush=True)
|
|
623
|
+
|
|
624
|
+
t0 = time.time()
|
|
625
|
+
for step in range(1, nstep + 1):
|
|
626
|
+
self._do_step()
|
|
627
|
+
if step % dnout == 0:
|
|
628
|
+
self._save_record()
|
|
629
|
+
if prm.verbose:
|
|
630
|
+
print(':', end='', flush=True)
|
|
631
|
+
|
|
632
|
+
elapsed = time.time() - t0
|
|
633
|
+
if prm.verbose:
|
|
634
|
+
print(f"\n done in {elapsed:.1f}s ({nstep/elapsed:.0f} steps/s)")
|
|
635
|
+
|
|
636
|
+
self._write_output()
|
|
637
|
+
return self.psi, self.q
|
|
638
|
+
|
|
639
|
+
|
|
640
|
+
# ============================================================================
|
|
641
|
+
# Command-line interface
|
|
642
|
+
# ============================================================================
|
|
643
|
+
|
|
644
|
+
def main(argv=None):
|
|
645
|
+
"""Command-line entry point (``qg-python`` once installed from PyPI)."""
|
|
646
|
+
import argparse
|
|
647
|
+
|
|
648
|
+
parser = argparse.ArgumentParser(
|
|
649
|
+
description='QG-Python: 1.5-layer quasi-geostrophic model',
|
|
650
|
+
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
651
|
+
epilog="""
|
|
652
|
+
Examples — restart from equilibrium state (recommended for analysis):
|
|
653
|
+
python qg_model.py --scheme rk4 --dt 1 --tend 110000 --dtout 40 --ic restart --restart-file qg_fourier.npz --out qg_rk4_ref.npz
|
|
654
|
+
python qg_model.py --scheme leapfrog --dt 1 --tend 110000 --dtout 40 --ic restart --restart-file qg_fourier.npz --ra-alpha 0.1 --out qg_lf_ra.npz
|
|
655
|
+
python qg_model.py --scheme leapfrog --dt 1 --tend 110000 --dtout 40 --ic restart --restart-file qg_fourier.npz --ra-alpha 0.1 --raw --out qg_lf_raw.npz
|
|
656
|
+
|
|
657
|
+
Stability sweep (vary dt):
|
|
658
|
+
python qg_model.py --scheme leapfrog --dt 0.5 --tend 110000 --dtout 40 --ic restart --restart-file qg_fourier.npz --ra-alpha 0.1 --out qg_lf_dt05.npz
|
|
659
|
+
python qg_model.py --scheme leapfrog --dt 2.0 --tend 110000 --dtout 40 --ic restart --restart-file qg_fourier.npz --ra-alpha 0.1 --out qg_lf_dt2.npz
|
|
660
|
+
""")
|
|
661
|
+
parser.add_argument('prmfile', nargs='?', help='parameter file (.prm)')
|
|
662
|
+
|
|
663
|
+
# Time
|
|
664
|
+
parser.add_argument('--scheme', choices=['order1','order2','rk4','dp5','leapfrog'])
|
|
665
|
+
parser.add_argument('--tend', type=float)
|
|
666
|
+
parser.add_argument('--dt', type=float)
|
|
667
|
+
parser.add_argument('--dtout', type=float)
|
|
668
|
+
parser.add_argument('--out', metavar='FILE')
|
|
669
|
+
parser.add_argument('--save-q', action='store_true')
|
|
670
|
+
parser.add_argument('--quiet', action='store_true')
|
|
671
|
+
|
|
672
|
+
# Grid
|
|
673
|
+
parser.add_argument('--mrefin', type=int)
|
|
674
|
+
parser.add_argument('--nx1', type=int)
|
|
675
|
+
parser.add_argument('--ny1', type=int)
|
|
676
|
+
|
|
677
|
+
# Physics
|
|
678
|
+
parser.add_argument('--rkb', type=float)
|
|
679
|
+
parser.add_argument('--rkh', type=float)
|
|
680
|
+
parser.add_argument('--rkh2', type=float)
|
|
681
|
+
parser.add_argument('--a', type=float)
|
|
682
|
+
parser.add_argument('--f', type=float)
|
|
683
|
+
parser.add_argument('--r', type=float)
|
|
684
|
+
|
|
685
|
+
# Leapfrog filter
|
|
686
|
+
parser.add_argument('--ra-alpha', type=float, default=None,
|
|
687
|
+
metavar='ALPHA',
|
|
688
|
+
help='Robert-Asselin filter strength (default: 0.1). '
|
|
689
|
+
'Used with --scheme leapfrog.')
|
|
690
|
+
parser.add_argument('--raw', action='store_true',
|
|
691
|
+
help='Use RAW (Robert-Asselin-Williams) filter instead of classic RA. '
|
|
692
|
+
'Less dissipative; recommended for long runs.')
|
|
693
|
+
|
|
694
|
+
# IC
|
|
695
|
+
parser.add_argument('--ic', default='zero',
|
|
696
|
+
choices=['zero','fourier','vortex','restart'])
|
|
697
|
+
parser.add_argument('--amplitude', type=float, default=1.0)
|
|
698
|
+
parser.add_argument('--kmax', type=int, default=5)
|
|
699
|
+
parser.add_argument('--seed', type=int, default=42)
|
|
700
|
+
parser.add_argument('--restart-file', metavar='FILE')
|
|
701
|
+
parser.add_argument('--restart-record', type=int, default=-1)
|
|
702
|
+
|
|
703
|
+
args = parser.parse_args(argv)
|
|
704
|
+
|
|
705
|
+
prm = QGParams.from_file(args.prmfile) if args.prmfile else QGParams()
|
|
706
|
+
|
|
707
|
+
# Time
|
|
708
|
+
if args.scheme: prm.scheme = args.scheme
|
|
709
|
+
if args.tend: prm.tend = args.tend
|
|
710
|
+
if args.dt: prm.dt = args.dt
|
|
711
|
+
if args.dtout: prm.dtout = args.dtout
|
|
712
|
+
if args.out: prm.outfname = args.out
|
|
713
|
+
if args.save_q: prm.save_q = True
|
|
714
|
+
if args.quiet: prm.verbose = False
|
|
715
|
+
|
|
716
|
+
# Grid
|
|
717
|
+
if args.mrefin or args.nx1 or args.ny1:
|
|
718
|
+
if args.mrefin: prm.mrefin = args.mrefin
|
|
719
|
+
if args.nx1: prm.nx1 = args.nx1
|
|
720
|
+
if args.ny1: prm.ny1 = args.ny1
|
|
721
|
+
prm.__post_init__()
|
|
722
|
+
|
|
723
|
+
# Physics
|
|
724
|
+
if args.rkb is not None: prm.rkb = args.rkb
|
|
725
|
+
if args.rkh is not None: prm.rkh = args.rkh
|
|
726
|
+
if args.rkh2 is not None: prm.rkh2 = args.rkh2
|
|
727
|
+
if args.a is not None: prm.a = args.a
|
|
728
|
+
if args.f is not None: prm.f = args.f
|
|
729
|
+
if args.r is not None: prm.r = args.r
|
|
730
|
+
|
|
731
|
+
# Leapfrog filter
|
|
732
|
+
if args.ra_alpha is not None: prm.ra_alpha = args.ra_alpha
|
|
733
|
+
if args.raw: prm.raw_filter = True
|
|
734
|
+
|
|
735
|
+
print(" QG-Python")
|
|
736
|
+
prm.describe()
|
|
737
|
+
print(f" netCDF4 available : {HAS_NC}")
|
|
738
|
+
print(f" initial condition : {args.ic}")
|
|
739
|
+
|
|
740
|
+
ic_kwargs: dict = {}
|
|
741
|
+
if args.ic == 'fourier':
|
|
742
|
+
ic_kwargs = dict(amplitude=args.amplitude, kmax=args.kmax, seed=args.seed)
|
|
743
|
+
elif args.ic == 'restart':
|
|
744
|
+
if not args.restart_file:
|
|
745
|
+
parser.error("--ic restart requires --restart-file FILE")
|
|
746
|
+
ic_kwargs = dict(fname=args.restart_file, record=args.restart_record)
|
|
747
|
+
|
|
748
|
+
model = QGModel(prm)
|
|
749
|
+
model.run(ic=args.ic, **ic_kwargs)
|
|
750
|
+
|
|
751
|
+
|
|
752
|
+
if __name__ == '__main__':
|
|
753
|
+
main()
|
|
@@ -0,0 +1,358 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: qg-python
|
|
3
|
+
Version: 1.0.0
|
|
4
|
+
Summary: Reproducible 1.5-layer quasi-geostrophic model (port of QG-C) with a direct sparse Helmholtz solver, RK4 / dp5 / Leapfrog (RA, RAW) integrators and stability diagnostics
|
|
5
|
+
Author-email: "Elias D. Nino-Ruiz" <enino@uninorte.edu.co>
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://aml-cs.github.io/software/qg-python/
|
|
8
|
+
Project-URL: Repository, https://github.com/enino84/qg-leapfrog-stability
|
|
9
|
+
Project-URL: Paper, https://doi.org/10.1002/gdj3.70107
|
|
10
|
+
Project-URL: Dataset, https://doi.org/10.5281/zenodo.22168803
|
|
11
|
+
Keywords: quasi-geostrophic,data assimilation,leapfrog,robert-asselin-williams,geophysical fluid dynamics
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Intended Audience :: Science/Research
|
|
14
|
+
Classifier: Topic :: Scientific/Engineering :: Atmospheric Science
|
|
15
|
+
Classifier: Operating System :: OS Independent
|
|
16
|
+
Requires-Python: >=3.9
|
|
17
|
+
Description-Content-Type: text/markdown
|
|
18
|
+
License-File: LICENSE
|
|
19
|
+
Requires-Dist: numpy>=1.24
|
|
20
|
+
Requires-Dist: scipy>=1.10
|
|
21
|
+
Provides-Extra: plots
|
|
22
|
+
Requires-Dist: matplotlib>=3.7; extra == "plots"
|
|
23
|
+
Provides-Extra: netcdf
|
|
24
|
+
Requires-Dist: netCDF4; extra == "netcdf"
|
|
25
|
+
Provides-Extra: all
|
|
26
|
+
Requires-Dist: matplotlib>=3.7; extra == "all"
|
|
27
|
+
Requires-Dist: netCDF4; extra == "all"
|
|
28
|
+
Dynamic: license-file
|
|
29
|
+
|
|
30
|
+
# qg-leapfrog-stability
|
|
31
|
+
|
|
32
|
+
[](https://doi.org/10.5281/zenodo.22168803)
|
|
33
|
+
[](LICENSE)
|
|
34
|
+
|
|
35
|
+
**QG-Python: a reproducible 1.5-layer quasi-geostrophic modelling service
|
|
36
|
+
with spectral stability and filter sensitivity diagnostics**
|
|
37
|
+
|
|
38
|
+
This repository contains the QG-Python model, the analysis scripts, and the
|
|
39
|
+
shell drivers associated with the Data Services Article published in
|
|
40
|
+
*Geoscience Data Journal* (Royal Meteorological Society / Wiley):
|
|
41
|
+
|
|
42
|
+
> Nino-Ruiz, E. D. (2026). **QG-Python: A Reproducible 1.5-Layer
|
|
43
|
+
> Quasi-Geostrophic Modelling Service with Direct Sparse Solvers, Spectral
|
|
44
|
+
> Stability Diagnostics, and Filter Sensitivity Data**.
|
|
45
|
+
> *Geoscience Data Journal*. doi:[10.1002/gdj3.70107](https://doi.org/10.1002/gdj3.70107)
|
|
46
|
+
|
|
47
|
+
The archived snapshot of this repository **together with the simulation
|
|
48
|
+
outputs** is the citable version of record:
|
|
49
|
+
doi:[10.5281/zenodo.22168803](https://doi.org/10.5281/zenodo.22168803)
|
|
50
|
+
|
|
51
|
+
---
|
|
52
|
+
|
|
53
|
+
## Install
|
|
54
|
+
|
|
55
|
+
```bash
|
|
56
|
+
pip install qg-python # model only (numpy, scipy)
|
|
57
|
+
pip install "qg-python[all]" # + matplotlib and netCDF4 for the viewer and .nc output
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
```python
|
|
61
|
+
from qg_model import QGModel, QGParams
|
|
62
|
+
|
|
63
|
+
prm = QGParams(mrefin=6, tend=1000, dtout=20, scheme='leapfrog',
|
|
64
|
+
ra_alpha=0.1, raw_filter=True, outfname='qg_run.npz')
|
|
65
|
+
psi, q = QGModel(prm).run('fourier', seed=42)
|
|
66
|
+
```
|
|
67
|
+
|
|
68
|
+
The same model is available on the command line as `qg-python` (run `qg-python --help`).
|
|
69
|
+
To reproduce the experiments of the paper, clone this repository and run `bash run_all.sh`.
|
|
70
|
+
|
|
71
|
+
---
|
|
72
|
+
|
|
73
|
+
## Overview
|
|
74
|
+
|
|
75
|
+
We implement the Leapfrog scheme with Robert-Asselin (RA) and
|
|
76
|
+
Robert-Asselin-Williams (RAW) filters in a 1.5-layer quasi-geostrophic (QG)
|
|
77
|
+
model, and carry out a systematic numerical analysis covering:
|
|
78
|
+
|
|
79
|
+
- Replacement of the original Fortran multigrid Helmholtz solver with a direct
|
|
80
|
+
sparse LU factorization (SuperLU via SciPy), enabling unambiguous stability analysis
|
|
81
|
+
- Eigenspectrum of the linearized operator $J_\mathcal{F}(q_\mathrm{eq})$ at two
|
|
82
|
+
grid resolutions (49×49 and 97×97), revealing resolution-dependent deviation
|
|
83
|
+
from the classical von Neumann bound
|
|
84
|
+
- Maximum stable time step $\Delta t_{\max}$ as a function of grid spacing $h$
|
|
85
|
+
- First estimate of the maximum Lyapunov exponent for this model
|
|
86
|
+
($\sigma_1 = 0.0318$, $\tau_L = 31.5$ time units)
|
|
87
|
+
- Statistical uncertainty propagation: kinetic energy and enstrophy PDFs,
|
|
88
|
+
energy spectra, and KL divergences across schemes
|
|
89
|
+
- Sensitivity of filter-induced dissipation to the RA parameter
|
|
90
|
+
$\alpha \in \{0.05, 0.10, 0.20, 0.40\}$
|
|
91
|
+
|
|
92
|
+
The QG model is a faithful Python port of the QG-C model by
|
|
93
|
+
[Pavel Sakov (Bureau of Meteorology)](https://github.com/sakov/qg-c),
|
|
94
|
+
extended with the Leapfrog family of integrators.
|
|
95
|
+
|
|
96
|
+
---
|
|
97
|
+
|
|
98
|
+
## Repository structure
|
|
99
|
+
|
|
100
|
+
```
|
|
101
|
+
qg-leapfrog-stability/
|
|
102
|
+
│
|
|
103
|
+
├── qg_model.py # QG model — all time-stepping schemes
|
|
104
|
+
├── qg_viewer.py # Interactive field viewer (.npz files)
|
|
105
|
+
├── run_all.sh # Single script to reproduce all results
|
|
106
|
+
│
|
|
107
|
+
├── analysis/
|
|
108
|
+
│ ├── plot_style.py # Shared publication style settings
|
|
109
|
+
│ ├── jacobian_spectrum.py # Build J_F and compute eigenspectrum
|
|
110
|
+
│ ├── stability_regions.py # Plot eigenspectrum + stability regions
|
|
111
|
+
│ ├── lyapunov.py # Benettin algorithm for Lyapunov exponent
|
|
112
|
+
│ ├── uncertainty.py # PDFs, power spectra, KL divergence
|
|
113
|
+
│ ├── plot_fields.py # Spatial snapshots of streamfunction
|
|
114
|
+
│ └── plot_alpha_sensitivity.py # RA filter sensitivity figures
|
|
115
|
+
│
|
|
116
|
+
├── experiments/
|
|
117
|
+
│ ├── run_reference.sh # Spin-up + RK4 reference run
|
|
118
|
+
│ ├── run_leapfrog_long.sh # LF+RA and LF+RAW main runs
|
|
119
|
+
│ ├── run_leapfrog_sweep.sh # RA filter alpha sweep
|
|
120
|
+
│ └── run_analysis.sh # Eigenspectrum + Lyapunov diagnostics
|
|
121
|
+
│
|
|
122
|
+
├── data/
|
|
123
|
+
│ ├── README.md # Exact command line behind every run
|
|
124
|
+
│ ├── blowup_times.txt # dt=1.0 instability onset, LF+RA / LF+RAW
|
|
125
|
+
│ ├── lyapunov_results.txt # Benettin MLE summary
|
|
126
|
+
│ └── uncertainty_results.txt # KE / enstrophy / KL divergence summary
|
|
127
|
+
│
|
|
128
|
+
├── paper/
|
|
129
|
+
│ └── README.md # Pointer to the published article
|
|
130
|
+
│
|
|
131
|
+
├── figures/ # Generated locally — not tracked by git
|
|
132
|
+
├── requirements.txt
|
|
133
|
+
├── CITATION.cff
|
|
134
|
+
├── LICENSE
|
|
135
|
+
└── README.md
|
|
136
|
+
```
|
|
137
|
+
|
|
138
|
+
---
|
|
139
|
+
|
|
140
|
+
## Dependencies
|
|
141
|
+
|
|
142
|
+
```
|
|
143
|
+
numpy >= 1.24
|
|
144
|
+
scipy >= 1.10
|
|
145
|
+
matplotlib >= 3.7
|
|
146
|
+
```
|
|
147
|
+
|
|
148
|
+
Install with:
|
|
149
|
+
|
|
150
|
+
```bash
|
|
151
|
+
pip install -r requirements.txt
|
|
152
|
+
```
|
|
153
|
+
|
|
154
|
+
No additional compiled extensions are required. The Helmholtz solver uses
|
|
155
|
+
`scipy.sparse.linalg.factorized` (SuperLU backend).
|
|
156
|
+
|
|
157
|
+
---
|
|
158
|
+
|
|
159
|
+
## Reproducing the results
|
|
160
|
+
|
|
161
|
+
All steps are collected in a single script. Run from the root of the repository:
|
|
162
|
+
|
|
163
|
+
```bash
|
|
164
|
+
bash run_all.sh
|
|
165
|
+
```
|
|
166
|
+
|
|
167
|
+
The script executes the following steps in order:
|
|
168
|
+
|
|
169
|
+
| Step | Description | Runtime (approx.) |
|
|
170
|
+
|------|-------------|-------------------|
|
|
171
|
+
| 1 | Spin-up from rest (mrefin=6, tend=100000, RK4, dt=1) | ~45 min |
|
|
172
|
+
| 2 | RK4 reference run (dt=1, t=100000→110000) | ~25 min |
|
|
173
|
+
| 3 | LF+RA main run (α=0.10, dt=0.5) | ~12 min |
|
|
174
|
+
| 4 | LF+RAW main run (α=0.10, dt=0.5) | ~12 min |
|
|
175
|
+
| 5 | Alpha sensitivity sweep (α=0.05, 0.20, 0.40) | ~36 min |
|
|
176
|
+
| 6 | Jacobian eigenspectrum, mrefin=6 (97×97) | ~2 min |
|
|
177
|
+
| 7a | Spin-up from rest, mrefin=5 (dt=0.1) | ~45 min |
|
|
178
|
+
| 7b | Jacobian eigenspectrum, mrefin=5 (49×49) | ~1 min |
|
|
179
|
+
| 8 | Lyapunov exponent (Benettin, 50000 steps) | ~90 min |
|
|
180
|
+
| 9 | All figures | ~1 min |
|
|
181
|
+
|
|
182
|
+
Total runtime on a modern laptop: approximately 4–5 hours.
|
|
183
|
+
|
|
184
|
+
> **Note:** The mrefin=5 spin-up (Step 7a) requires `dt=0.1` — larger time steps
|
|
185
|
+
> cause blow-up during spin-up at this resolution.
|
|
186
|
+
|
|
187
|
+
### Running individual steps
|
|
188
|
+
|
|
189
|
+
The pipeline is also split into four self-contained drivers under
|
|
190
|
+
`experiments/`, which can be run in isolation for partial reproduction:
|
|
191
|
+
|
|
192
|
+
```bash
|
|
193
|
+
bash experiments/run_reference.sh # spin-up + RK4 reference
|
|
194
|
+
bash experiments/run_leapfrog_long.sh # LF+RA and LF+RAW main runs
|
|
195
|
+
bash experiments/run_leapfrog_sweep.sh # alpha sensitivity sweep
|
|
196
|
+
bash experiments/run_analysis.sh # eigenspectrum + Lyapunov
|
|
197
|
+
```
|
|
198
|
+
|
|
199
|
+
Each command in `run_all.sh` is likewise self-contained. For example, to
|
|
200
|
+
recompute only the eigenspectrum:
|
|
201
|
+
|
|
202
|
+
```bash
|
|
203
|
+
python analysis/jacobian_spectrum.py --restart qg_zero.npz --neigs 200 --no-matrix --out jacobian_spectrum.npz
|
|
204
|
+
```
|
|
205
|
+
|
|
206
|
+
---
|
|
207
|
+
|
|
208
|
+
## Key results
|
|
209
|
+
|
|
210
|
+
| Quantity | Value |
|
|
211
|
+
|----------|-------|
|
|
212
|
+
| Grid (main experiments) | 97×97 (mrefin=6), $h=0.01042$ |
|
|
213
|
+
| Analytical $\Delta t_{\max}^{LF}$ (mrefin=6) | 0.090 |
|
|
214
|
+
| Numerical $\Delta t_{\max}^{LF}$ (mrefin=6) | 1.236 (factor **13.7×** larger) |
|
|
215
|
+
| Analytical $\Delta t_{\max}^{LF}$ (mrefin=5) | 0.362 |
|
|
216
|
+
| Numerical $\Delta t_{\max}^{LF}$ (mrefin=5) | 0.235 (factor **1.54×** smaller) |
|
|
217
|
+
| Maximum Lyapunov exponent $\sigma_1$ | 0.0318 |
|
|
218
|
+
| Lyapunov time $\tau_L$ | 31.5 time units |
|
|
219
|
+
| LF+RA kinetic energy bias (α=0.10) | −21% |
|
|
220
|
+
| LF+RA enstrophy bias (α=0.10) | −46% |
|
|
221
|
+
| LF+RAW kinetic energy bias | −2% |
|
|
222
|
+
| LF+RAW enstrophy bias | +4% |
|
|
223
|
+
|
|
224
|
+
---
|
|
225
|
+
|
|
226
|
+
## Time-stepping schemes
|
|
227
|
+
|
|
228
|
+
| Scheme | Order | RHS evals / step | Notes |
|
|
229
|
+
|--------|-------|-------------------|-------|
|
|
230
|
+
| `order1` | 1 | 1 | Explicit Euler |
|
|
231
|
+
| `order2` | 2 | 2 | Midpoint RK |
|
|
232
|
+
| `rk4` | 4 | 4 | Default in QG-C |
|
|
233
|
+
| `dp5` | 5 | 6 | Dormand-Prince |
|
|
234
|
+
| `leapfrog` | 2 | 1 | + RA or RAW filter |
|
|
235
|
+
|
|
236
|
+
The first Leapfrog step uses a single RK4 step as bootstrap to avoid
|
|
237
|
+
exciting the computational mode from the start.
|
|
238
|
+
|
|
239
|
+
### Leapfrog filter options
|
|
240
|
+
|
|
241
|
+
```bash
|
|
242
|
+
# Robert-Asselin filter (classic, Asselin 1972)
|
|
243
|
+
python qg_model.py --scheme leapfrog --ra-alpha 0.1 ...
|
|
244
|
+
|
|
245
|
+
# Robert-Asselin-Williams filter (less dissipative, Williams 2009)
|
|
246
|
+
python qg_model.py --scheme leapfrog --ra-alpha 0.1 --raw ...
|
|
247
|
+
```
|
|
248
|
+
|
|
249
|
+
---
|
|
250
|
+
|
|
251
|
+
## Data availability
|
|
252
|
+
|
|
253
|
+
The simulation output files (`.npz`) are several hundred megabytes and are
|
|
254
|
+
therefore not tracked by git. They are archived, together with a snapshot of
|
|
255
|
+
this repository, on Zenodo:
|
|
256
|
+
|
|
257
|
+
> **DOI**: [10.5281/zenodo.22168803](https://doi.org/10.5281/zenodo.22168803)
|
|
258
|
+
> **Record**: <https://doi.org/10.5281/zenodo.22168803>
|
|
259
|
+
> **Development version**: <https://github.com/enino84/qg-leapfrog-stability>
|
|
260
|
+
|
|
261
|
+
The Zenodo record is the authoritative citation target: it holds both the
|
|
262
|
+
tagged source snapshot (`qg-python-v1.0.zip`) and the simulation outputs. This
|
|
263
|
+
GitHub repository is the development mirror and tracks the source only.
|
|
264
|
+
|
|
265
|
+
The archive contains the equilibrium restart fields at `t = 100000` for both
|
|
266
|
+
resolutions, the eigenvalue clouds of the linearised operator, the Benettin
|
|
267
|
+
trace of the maximum Lyapunov exponent, and the time-resolved kinetic energy,
|
|
268
|
+
enstrophy and isotropic spectra for RK4, LF+RA and LF+RAW across
|
|
269
|
+
alpha in {0.05, 0.10, 0.20, 0.40}. `data/README.md` documents the exact
|
|
270
|
+
command line behind each file, so every output can also be regenerated from
|
|
271
|
+
scratch.
|
|
272
|
+
|
|
273
|
+
---
|
|
274
|
+
|
|
275
|
+
## Reference model
|
|
276
|
+
|
|
277
|
+
This work builds on:
|
|
278
|
+
|
|
279
|
+
> Sakov, P. (2024). *QG-C: A quasi-geostrophic model in C*.
|
|
280
|
+
> Bureau of Meteorology, Australia.
|
|
281
|
+
> https://github.com/sakov/qg-c
|
|
282
|
+
|
|
283
|
+
Key numerical methods:
|
|
284
|
+
|
|
285
|
+
> Arakawa, A. (1966). Computational design for long-term numerical integration
|
|
286
|
+
> of the equations of fluid motion.
|
|
287
|
+
> *Journal of Computational Physics*, 1(1), 119–143.
|
|
288
|
+
|
|
289
|
+
> Asselin, R. (1972). Frequency filter for time integrations.
|
|
290
|
+
> *Monthly Weather Review*, 100(6), 487–490.
|
|
291
|
+
|
|
292
|
+
> Williams, P. D. (2009). A proposed modification to the Robert-Asselin
|
|
293
|
+
> time filter. *Monthly Weather Review*, 137(8), 2538–2546.
|
|
294
|
+
|
|
295
|
+
> Li, X. S., & Demmel, J. W. (2003). SuperLU_DIST: A scalable distributed-memory
|
|
296
|
+
> sparse direct solver for unsymmetric linear systems.
|
|
297
|
+
> *ACM Transactions on Mathematical Software*, 29(2), 110–140.
|
|
298
|
+
|
|
299
|
+
---
|
|
300
|
+
|
|
301
|
+
## License
|
|
302
|
+
|
|
303
|
+
The source code in this repository is released under the MIT License — see
|
|
304
|
+
[LICENSE](LICENSE). The simulation output files distributed through the Zenodo
|
|
305
|
+
archive are released under
|
|
306
|
+
[CC-BY-4.0](https://creativecommons.org/licenses/by/4.0/).
|
|
307
|
+
|
|
308
|
+
---
|
|
309
|
+
|
|
310
|
+
## Citation
|
|
311
|
+
|
|
312
|
+
Please cite both the article and the archived software release. Machine-readable
|
|
313
|
+
metadata is in [CITATION.cff](CITATION.cff).
|
|
314
|
+
|
|
315
|
+
```bibtex
|
|
316
|
+
@article{NinoRuiz2026qgpython,
|
|
317
|
+
author = {Nino-Ruiz, Elias D.},
|
|
318
|
+
title = {{QG-Python}: A Reproducible 1.5-Layer Quasi-Geostrophic Modelling
|
|
319
|
+
Service with Direct Sparse Solvers, Spectral Stability Diagnostics,
|
|
320
|
+
and Filter Sensitivity Data},
|
|
321
|
+
journal = {Geoscience Data Journal},
|
|
322
|
+
year = {2026},
|
|
323
|
+
doi = {10.1002/gdj3.70107}
|
|
324
|
+
}
|
|
325
|
+
|
|
326
|
+
@software{NinoRuiz2026qgpythonsoftware,
|
|
327
|
+
author = {Nino-Ruiz, Elias D.},
|
|
328
|
+
title = {{QG-Python}: a reproducible 1.5-layer quasi-geostrophic modelling
|
|
329
|
+
service with spectral stability and filter sensitivity diagnostics},
|
|
330
|
+
version = {1.0},
|
|
331
|
+
year = {2026},
|
|
332
|
+
publisher = {Zenodo},
|
|
333
|
+
doi = {10.5281/zenodo.22168803}
|
|
334
|
+
}
|
|
335
|
+
```
|
|
336
|
+
|
|
337
|
+
---
|
|
338
|
+
|
|
339
|
+
## A note on the enstrophy diagnostic
|
|
340
|
+
|
|
341
|
+
Enstrophy is reported throughout as the domain average of
|
|
342
|
+
`(nabla^2 psi)^2`, evaluated with the standard 5-point Laplacian on the
|
|
343
|
+
**interior nodes only**. This is the definition in
|
|
344
|
+
`analysis/uncertainty.py::enstrophy`, and it is the one every other script
|
|
345
|
+
imports.
|
|
346
|
+
|
|
347
|
+
The wall ring must be excluded from the average. `qg_model.py` imposes
|
|
348
|
+
homogeneous Dirichlet conditions `psi = 0` on all four walls, so at a wall node
|
|
349
|
+
the 5-point stencil collapses to `psi_adjacent / h^2`; with `h ~ 0.01` that is
|
|
350
|
+
`O(1e5)`, and squaring it lets a one-node ring — about 4% of the points —
|
|
351
|
+
dominate the domain average and inflate its variance by considerably more,
|
|
352
|
+
while leaving the *relative* biases between time-stepping schemes almost
|
|
353
|
+
unchanged.
|
|
354
|
+
|
|
355
|
+
`analysis/make_tables.py` regenerates the LaTeX rows of both tables in the
|
|
356
|
+
article directly from the `.npz` runs, so the statistical summary and the alpha
|
|
357
|
+
sweep cannot drift apart. Pass `--compare` to also print the wall-inclusive
|
|
358
|
+
average.
|
|
@@ -0,0 +1,7 @@
|
|
|
1
|
+
qg_model.py,sha256=zQDizQBEUckoHmfZK0NzNIlA8k7xmcGJu5tnK9UVNqU,28273
|
|
2
|
+
qg_python-1.0.0.dist-info/licenses/LICENSE,sha256=ESYyLizI0WWtxMeS7rGVcX3ivMezm-HOd5WdeOh-9oU,1056
|
|
3
|
+
qg_python-1.0.0.dist-info/METADATA,sha256=LAcPMpoqZR6KiOpOU0ImZYe5CgCHokzcvC39uUWcFGU,13353
|
|
4
|
+
qg_python-1.0.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
|
|
5
|
+
qg_python-1.0.0.dist-info/entry_points.txt,sha256=j1_sCopSmMJgdJUqK4vv854GBIq8YP4ff_GeN5rXNFw,44
|
|
6
|
+
qg_python-1.0.0.dist-info/top_level.txt,sha256=V_Kos62ZmivnfoqubVC0_ppGPcsurn-a_44QxNU7vVU,9
|
|
7
|
+
qg_python-1.0.0.dist-info/RECORD,,
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
qg_model
|