Construct a model

In this example, we showcase how to construct a pyshqg.core.model.QGModel instance.

import yaml

from pyshqg.backend.numpy_backend import NumpyBackend as Backend
from pyshqg.preprocessing.reference_data import load_reference_data, interpolate_data
from pyshqg.preprocessing.reference_data import load_test_data
from pyshqg.preprocessing.vertical_parametrisation import VerticalParametrisation
from pyshqg.preprocessing.orography import Orography
from pyshqg.core.spectral_transformations import SpectralTransformations
from pyshqg.core.poisson import QGPoissonSolver
from pyshqg.core.dissipation import QGEkmanDissipation, QGSelectiveDissipation
from pyshqg.core.dissipation import QGThermalDissipation, QGDissipation
from pyshqg.core.source import QGForcing
from pyshqg.core.model import QGModel
from pyshqg.core.constructors import construct_model

Construct the backend

In all cases, a backend is required to create a pyshqg.core.model.QGModel instance.

For the present example, let us use the standard numpy-based backend with double precision, i.e. with floatx = 'float64'.

backend = Backend('float64')

Step-by-step manual construction

Beyond the backend, a pyshqg.core.model.QGModel instance is constructed using six objects:

1. the spectral transformations;

2. the vertical parametrisation;

3. the orography;

4. the poisson solver;

5. the dissipation processes;

6. the forcing processes.

Note that, since these objects are inter-connected, they should be created in that specific order.

Spectral transformations

The spectral transformations object is a pyshqg.core.spectral_transformations.SpectralTransformations instance supposed to handle all spectral transformations. It is constructed using

• backend: the backend;

• T: the internal spectral truncature;

• T_grid: the grid spectral truncature;

• planet_radius: the planetary radius;

• planet_omega: the planetary rotation speed.

Note that the size of the spectral and grid spaces are directly determined from T and T_grid: the size of the spectral space is (2, T+1, T+1) and the size of the grid space is (Nlat, Nlon) where Nlat = T_grid+1 and Nlon = 2*N_lat.

For the present example, let us use T = 21, T_grid = 31, planet_radius = 6371000 (which is approximately the Earth radius in meter), and planet_omega = 7.292e-05 (which is approximately the Earth rotation speed in radian per second). Note in particular that the size of the spectral space is (2, 22, 22) and the size of the grid space is (32, 64).

spectral_transformations = SpectralTransformations(
    backend=backend,
    T=21,
    T_grid=31,
    planet_radius=6371000,
    planet_omega=7.292e-05,
)

Vertical parametrisation

The vertical parametrisation object is a pyshqg.preprocessing.vertical_parametrisation.VerticalParametrisation instance supposed to handle the vertical parametrisation of the model. It is constructed using

• rossby_radius_list: the list of Rossby radius for each level interface.

Note that the number of vertical levels in the model is directly determined from the length of the list: Nz = len(rossby_radius_list)+1. Furthermore, the model implicitly assumes that the first model level is the highest level. Therefore, the first number in this list is the Rossby radius corresponding to the interface between the first two model levels.

For the present example, let us use three vertical levels, with Rossby radii 700000 and 450000 (in meter) for the interface between model levels.

vertical_parametrisation = VerticalParametrisation(
    rossby_radius_list=[700000, 450000],
)

Orography

The orography object is a pyshqg.preprocessing.orography.Orography instance supposed to handle the orography of the model. It is constructed using

• land_sea_mask: the land-sea mask coefficients in grid space;

• orography: the orography coefficients in grid space.

These coefficients can be accessed from the reference dataset using the pyshqg.preprocessing.reference_data.load_reference_data() function. Note however that this reference dataset is available in the T63 grid, which means that they may need to be interpolated in the model grid, for example using the pyshqg.preprocessing.reference_data.interpolate_data() function.

Let us start by loading the reference data. For this, we use:

• grid_truncature = 63: the grid truncature of the reference data (only 63 is available at the moment);

• load = True: indicates that we want to load the data into memory;

• padding = False: indicates that we do not want to apply padding in the latitude direction.

Note that padding is only necessary when the reference data must be interpolated in a grid larger than T63.

ds_reference = load_reference_data(
    grid_truncature=63,
    load=True,
    padding=False,
)

We now can interpolate the reference data using

• ds = ds_reference: the reference data to interpolate;

• lat = spectral_transformations.lats: the latitude nodes;

• lon = spectral_transformations.lons: the longitude nodes;

• methods: the list of interpolation methods for each variable in the dataset (always 'quintic' here).

Note that we provide an interpolation for the land-sea mask and the orography field, but also for the forcing. Even though the forcing is not required to construct the orography object, it will be used later on to construct the forcing object.

ds_interpolated = interpolate_data(
    ds=ds_reference,
    lat=spectral_transformations.lats,
    lon=spectral_transformations.lons,
    methods={'forcing': 'quintic', 'land_sea_mask': 'quintic', 'orography': 'quintic'},
)

We now can construct the orography object.

orography = Orography(
    land_sea_mask=ds_interpolated.land_sea_mask.to_numpy(),
    orography=ds_interpolated.orography.to_numpy(),
)

Poisson solver

The poisson solver object is a pyshqg.core.poisson.QGPoissonSolver instance supposed to handle the transformation between vorticity and stream function. It is constructed using

• backend: the backend;

• spectral_transformations: the spectral transformations object;

• vertical_parametrisation: the vertical parametrisation object;

• orography: the orography object;

• orography_scale: the vertical length scale for the orography in the Poisson-like equation between \(q\) and \(\psi\).

For the present example, we use orography_scale = 9000 (in meter).

poisson_solver = QGPoissonSolver(
    backend=backend,
    spectral_transformations=spectral_transformations,
    vertical_parametrisation=vertical_parametrisation,
    orography=orography,
    orography_scale=9000,
)

Ekman dissipation

The Ekman dissipation object is a pyshqg.core.dissipation.QGEkmanDissipation instance supposed to handle the Ekman dissipation processes. It is constructed using

• backend: the backend;

• spectral_transformations: the spectral transformations object;

• orography: the orography object;

• orography_scale: the vertical length scale for the orography in \(\mu\);

• tau: the \(\tau\) parameter of the friction coefficient \(\mu\);

• weight_land_sea_mask: the weight of the land-sea mask contribution in \(\mu\);

• weight_orography: the weight of the orography contribution in \(\mu\).

For the present example, we use orography_scale = 1000 (in meter), tau =  259200 (in second), weight_land_sea_mask = 0.5, and weight_orography = 0.5.

ekman=QGEkmanDissipation(
    backend=backend,
    spectral_transformations=spectral_transformations,
    orography=orography,
    orography_scale=1000,
    tau=259200,
    weight_land_sea_mask=0.5,
    weight_orography=0.5,
)

Selective dissipation

The selective dissipation object is a pyshqg.core.dissipation.QGSelectiveDissipation instance supposed to handle the selective dissipation processes. It is constructed using

• backend: the backend;

• spectral_transformations: the spectral transformations object;

• tau: the \(\tau\) parameter of the selective dissipation.

For the present example, we use tau = 8640 (in second).

selective=QGSelectiveDissipation(
    backend=backend,
    spectral_transformations=spectral_transformations,
    tau=8640,
)

Thermal dissipation

The thermal dissipation object is a pyshqg.core.dissipation.QGThermalDissipation instance supposed to handle the thermal dissipation processes. It is constructed using

• backend: the backend;

• vertical_parametrisation: the vertical parametrisation object;

• tau: the \(\tau\) parameter of the thermal dissipation.

For the present example, we use tau = 2160000 (in second).

thermal=QGThermalDissipation(
    backend,
    vertical_parametrisation=vertical_parametrisation,
    tau=2160000,
),

Dissipation

The dissipation object is a pyshqg.core.dissipation.QGDissipation instance which is simply a convenient wrapper to hold the Ekman, selective, and thermal dissipation.

dissipation = QGDissipation(
    ekman=ekman,
    selective=selective,
    thermal=thermal,
)

Forcing

The forcing object is a pyshqg.core.source.QGForcing instance supposed to handle the forcing processes. It is constructed using

• backend: the backend;

• forcing: the forcing coefficients in grid space.

Note that the forcing coefficients have already been computed alongside the orography.

forcing = QGForcing(
    backend,
    forcing=ds_interpolated.forcing.to_numpy(),
)

Model

We now have all the pieces to construct the model. Note that the vertical parametrisation and orography objects are not required here, because they contain only preprocessing computations that have already been performed at this point.

model = QGModel(
    backend=backend,
    spectral_transformations=spectral_transformations,
    poisson_solver=poisson_solver,
    dissipation=dissipation,
    forcing=forcing,
)

Construction using a configuration dictionary

Alternatively, one can use the pyshqg.core.constructors.construct_model() function to construct the model using a configuration object stored as a dictionary.

For example, the test data contains a dictionary that will produce the model used in the unit tests. Let us see how this works.

Let us start by loading the test data. For the present example, we choose the test data at spectral resolution T21 and grid resolution T31.

ds_test = load_test_data(
    internal_truncature=21,
    grid_truncature=31,
)

The model configuration is stored in ds_test.config. Note that this is the exact same configuration that we have used to manually construct the model, except for the fact that it contains additional sections to create the integrators.

print(yaml.dump(ds_test.config))

abm_integration:
  dt: 1800
  method: abm
data_interpolation:
  forcing: quintic
  land_sea_mask: quintic
  orography: quintic
dissipation:
  ekman:
    orography_scale: 1000.0
    tau: 259200.0
    weight_land_sea_mask: 0.5
    weight_orography: 0.5
  selective:
    tau: 8640.0
  thermal:
    tau: 2160000.0
ee_integration:
  dt: 1800
  method: ee
poisson_solver:
  orography_scale: 9000.0
reference_data:
  grid_truncature: 63
  load: true
  padding: true
rk2_integration:
  dt: 1800
  method: rk2
rk4_integration:
  dt: 1800
  method: rk4
spectral_transformations:
  T: 21
  T_grid: 31
  planet_omega: 7.292e-05
  planet_radius: 6371000.0
vertical_parametrisation:
  rossby_radius_list:
  - 700000.0
  - 450000.0

With this config, we can simply construct the model as follows.

model = construct_model(backend, ds_test.config)