<jupyter_start><jupyter_text>This notebook demonstrates how to add metadata and units consistency to xarray fields loaded from grid using cfgrib (https://github.com/ecmwf/cfgrib)This is useful for dimensional analysis and physical type correctness of computations in post-processing of NWP simulationshttps://arxiv.org/pdf/1807.07643.pdf First we need to setup the installation of cfgrib for loading grib files into xarray datasets<jupyter_code>!pip install -q condacolab
import condacolab
condacolab.install()
!conda install -c conda-forge cfgrib
!python -m cfgrib selfcheck<jupyter_output><empty_output><jupyter_text>Download a sample grib file and small dataset classes that add pint units to xarray arrays<jupyter_code>from getpass import getpass
secret = getpass('Enter the secret value: ')
import subprocess
!rm -rf grib_files
subprocess.run(['git', 'clone', 'https://'+secret+'@github.com/cosunae/grib_files.git'], check=True)
!rm -rf postproc_pt1
subprocess.run(['git', 'clone', 'https://'+secret+'@github.com/cosunae/postproc_pt1'], check=True)

!pip install postproc_pt1/
!python -m cfgrib selfcheck
from dataset import open_datasets

dss = open_datasets('grib_files/cosmo-eu/lfff00000000_2014010400.gb2', engine='cfgrib', backend_kwargs={'filter_by_keys': {'typeOfLevel': 'generalVerticalLayer'}})<jupyter_output><empty_output><jupyter_text>Since there u and v are staggered fields (i.e. have different lon,lat coordinates), not all fields can be inserted into the same hypercube. Therefore 3 different datasets are generated.<jupyter_code>mass_ds = dss[0]
u_ds = dss[1]
v_ds = dss[2]<jupyter_output><empty_output><jupyter_text>Inspect the datasets. Click "show data repr" of latitude/longitude coordinates to see the different coordinates system between mass_ds (mass point dataset) and the u_ds (u point dataset).<jupyter_code>mass_ds<jupyter_output><empty_output><jupyter_text>Unpack the individual fields<jupyter_code>t = mass_ds['t']
q = mass_ds['q']
QI = mass_ds['QI']
pres = mass_ds['pres']
u = u_ds['u']
v = v_ds['v']
import xarray as xr
xr.set_options(keep_attrs = True)
# Following are compatible since dimensionless
QQ = q+QI*0.01
print(QQ.units)
# exception trigger due to incompatibility of dimensions
try:
    f = t+q
except RuntimeError as e:
    print("Testing error (as it should be)",e)

# compute T at half levels, and check it preserves metadata and units 
def half_levels(f):
  cf = f.cumsum(dim='generalVerticalLayer')
  return (cf[2:] - cf[:-2])/float(2)

t_half = half_levels(t)

# Add t0
t0 = t.isel(generalVerticalLayer=0)
t_half = t_half+t0

# consistency check
try: 
    f = t_half + q
except RuntimeError as e:
    print("Testing error (as it should be)",e)

# U & V are compatible. They both are defined in the same indexing x,y although have different lon,lat coordinates
(u+v).isel(generalVerticalLayer=0).plot()<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup