Running Parameterized FCNNs

Running Parameterized FCNNs#

Now that we have taken time to explore and understand FCNNs, we will now run a parameterization of our own.

# Import required modules
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
import pyqg
import json
import fsspec
from pyqg_parameterization_benchmarks.neural_networks import FCNNParameterization, FullyCNN

%matplotlib inline

Below we are loading a sample dataset, as training data for our parameterization, of a high-resolution simulation in the eddy configuration started from random noise and run over a span of 10 simulation years that has been coarsened using Operator 2 (a spectral truncation + softer Gaussian filter).

The parameterization we are using for this example is an FCNN trained on the values of PV and velocity to estimate PV subgrid forcing computed with Operator 2.

param = FCNNParameterization.train_on(training_ds, '/home/jovyan/models/fcnn_qu_to_Sq2', 
                                      inputs=['q', 'u', 'v'], 
                                      targets = ['q_subgrid_forcing']) 

Now that we have a trained parameterization we can begin testing it. We call test_offline() on the parameterization instance to begin evaluating the parameterization on an offline dataset.

# Helpers methods to help visualize offline metrics/performance
def imshow(arr):
    plt.imshow(arr, vmin=0, vmax=1, cmap='inferno')
    mean = arr.mean().data
    plt.text(32, 32, f"{mean:.2f}", color=('white' if mean<0.75 else 'black'),
             fontweight='bold', ha='center', va='center', fontsize=16)
    plt.xticks([]); plt.yticks([])
    
def colorbar(label):
    plt.colorbar().set_label(label, fontsize=16,rotation=0,ha='left',va='center')

Upon calling test_offline(), an xarray.Dataset object describing the predictions made by the parameterization is returned. Included within this dataset are a number of computed metrics.

plt.figure(figsize=(7.6,6)).suptitle("Offline performance")

for z in [0,1]:
    plt.subplot(2,2,z+1,title=f"{['Upper','Lower'][z]} layer")
    imshow(preds1.q_subgrid_forcing_spatial_correlation.isel(lev=z))
    if z: colorbar("Correlation")
    
for z in [0,1]:
    plt.subplot(2,2,z+3)
    imshow(preds1.q_subgrid_forcing_spatial_skill.isel(lev=z))
    if z: colorbar("$R^2$")

../_images/cf01dfcf19e0c494698bc739114738e98d4ca5ebd0c51eebccee096b50bcfb02.png

We can improve these offline metrics by training our FCNN on more datasets. Below, we feed another dataset into each model, extract the input and target features from the dataset, and standardize by calling fit() which also then initiates a training session on the inputs and targets.

training_ds2 = xr.open_zarr('/home/jovyan/datasets/eddy/op2/ds3').load()
for m in param.models:
    new_inputs = m.extract_inputs(training_ds2)
    new_targets = m.extract_targets(training_ds2)
    m.fit(new_inputs, new_targets, rescale=True)

We then run predictions again with the same, initial test dataset on the parameterization after its been trained on another dataset.

plt.figure(figsize=(7.6,6)).suptitle("Offline performance")

for z in [0,1]:
    plt.subplot(2,2,z+1,title=f"{['Upper','Lower'][z]} layer")
    imshow(preds2.q_subgrid_forcing_spatial_correlation.isel(lev=z))
    if z: colorbar("Correlation")
    
for z in [0,1]:
    plt.subplot(2,2,z+3)
    imshow(preds2.q_subgrid_forcing_spatial_skill.isel(lev=z))
    if z: colorbar("$R^2$")

../_images/4d40a338cff401a9e706aa2e2b9af178fc2f753d2979fbc7d00e5f3bb9f69e56.png

As shown, our computed metrics improved with additional training. We then save our newly trained models to the model directory path for our parameterization in order to keep the updated trained state of the models.

param.models[0].save('/home/jovyan/models/fcnn_qu_to_Sq2/models/0')
param.models[1].save('/home/jovyan/models/fcnn_qu_to_Sq2/models/1')

We can then reload the FCNN parameterization in future use cases by instantiating a new FCNNParameterization instance and passing in the directory path at which the parameterization is saved, as follows:

param = FCNNParameterization('/home/jovyan/models/fcnn_qu_to_Sq2')