Data assimilation system for the PyQG model
This repository provides data assimilation (DA) capability for the PyQG model, as documented in the Lorenz-96 notebooks.
DA_core.py: data assimilation related functionsDA_QG2L.ipynb: notebook for data assimilation examples in 2-layer QG modelDA_analysis.ipynb: notebook to analyze DA resultsENKF.py,ENKF_training.py: Python script to run EnKF DA experiments in bulk
ML_core.py: machine learning related functionsB_UNet.ipynb: notebook to train U-Net to predict ensemble covariancesUNet_analysis.ipynb: notebook to analyze trained U-Netstraining.py: Python script to train U-NetUNetKF.py: Python script to run UNetKF DA experiments
For the code clinic, I would like to optimize the code related to the training and inference, especially in the case of using GPU.
The scheme of the data pipline is as follows:
At each time step, the full q datasets have the size of (level,Ny,Nx), so the full covariance matrix of q would have size of (level,Ny,Nx,level,Ny,Nx).
In this current PyQG implementation of EnKF, we use the full covariance matrices during the data assimilation step. However, the full covariance matrices at all time steps are prohibitively big to save for training U-Nets.
Since we normally use covariance localization in EnKF applications, only part of the full covariance matrix is used (usually based on physical distance), we can save only localized matrices.
As a result, the saved q datasets have the size of (time,level,Ny,Nx), while the saved covariance matrices B have the size of (time,level,Nx,Nx,level,Ny_local,Nx_local), where Ny_local and Nx_local are significantly smaller than Ny and Nx.
During training, each data sample consists of a localized q matrix and a localized B matrix. The localized B with size of (level,Ny_local,Nx_local) would simply be a subset of the full dataset, while the localized q is taken as subset of the full matrix at runtime.
The same process happens during inference. When the U-Net is applied in the DA processs, a localized q matrix is constructed around each model gridpoint.
Because the training samples are taken as moving subsets of the full datasets, after reading the input data, the "get item" function takes one sample at a time to generate the data batches.