This notebook shows how to use the DiffEqPy package to integrate the qgs model and compare with the qgs Runge-Kutta integrator.
You must finish the install of DiffEqPy manually before running this notebook!
The DiffEqPy package first installation step is done by Anaconda in the qgs environment but then you must install Julia and follow the final manual installation instruction found in the link above.
These can be summed up as opening a terminal and doing:
conda activate qgs
python
and then inside the Python command line interface do:
>>> import diffeqpy
>>> diffeqpy.install()
which will then finalize the installation.
This notebook use the model version with a simple 2-layer channel QG atmosphere truncated at wavenumber 2 on a beta-plane with a simple orography (a montain and a valley).
More detail can be found in the articles:
or in the documentation and on readthedocs.
First, setting the path and loading of some modules
import sys, os
sys.path.extend([os.path.abspath('../')])
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
Initializing the random number generator (for reproducibility). -- Disable if needed.
np.random.seed(210217)
Importing the model's modules
from qgs.params.params import QgParams
from qgs.integrators.integrator import RungeKuttaIntegrator
from qgs.functions.tendencies import create_tendencies
from qgs.plotting.util import std_plot
Importing Julia DifferentialEquations.jl package
from julia.api import Julia
jl = Julia(compiled_modules=False)
from diffeqpy import de
Importing also Numba
from numba import jit
Ignoring Numba performance warnings. -- Comment to disable if you want to see them.
import warnings
warnings.filterwarnings("ignore")
General parameters
# Time parameters
dt = 0.1
# Saving the model state n steps
write_steps = 5
number_of_trajectories = 1
number_of_perturbed_trajectories = 10
Setting some model parameters
# Model parameters instantiation with some non-default specs
model_parameters = QgParams({'phi0_npi': np.deg2rad(50.)/np.pi, 'hd':0.3})
# Mode truncation at the wavenumber 2 in both x and y spatial coordinate
model_parameters.set_atmospheric_channel_fourier_modes(2, 2)
# Changing (increasing) the orography depth and the meridional temperature gradient
model_parameters.ground_params.set_orography(0.4, 1)
model_parameters.atemperature_params.set_thetas(0.2, 0)
# Printing the model's parameters
model_parameters.print_params()
Creating the tendencies function
%%time
f, Df = create_tendencies(model_parameters)
Defining an integrator
integrator = RungeKuttaIntegrator()
integrator.set_func(f)
Start on a random initial condition and integrate over a transient time to obtain an initial condition on the attractors
%%time
ic = np.random.rand(model_parameters.ndim)*0.1
integrator.integrate(0., 200000., dt, ic=ic, write_steps=0)
time, ic = integrator.get_trajectories()
Now integrate to obtain with the qgs RK4 integrator
%%time
integrator.integrate(0., 200., dt, ic=ic, write_steps=write_steps)
time, traj = integrator.get_trajectories()
And also with the DifferentialEquations ODE Solver
# defining a function with the correct header
@jit
def f_jl(x,p,t):
u = f(t,x)
return u
%%time
# Defining the problem and integrating
prob = de.ODEProblem(f_jl, ic, (0., 200.))
sol = de.solve(prob, saveat=write_steps * dt)
Plotting the result
var = 10
jtraj = np.array(sol.u).T
plt.figure(figsize=(10, 8))
plt.plot(model_parameters.dimensional_time*time, traj[var], label="Qgs RK4 integrator")
plt.plot(model_parameters.dimensional_time*time, jtraj[var], label="DiffEqPy default integrator")
plt.legend()
plt.xlabel('time (days)')
plt.ylabel('$'+model_parameters.latex_var_string[var]+'$');