High-level python script for computing geopotential, geopotential height, and geometric height on ECMWF model levels from a GRIB file

comp_ecmwf_gepot_at_ml.py

import cfgrib
import numpy
from typing import NamedTuple
import gc

"""
This script produces output within +/- 0.5 m2.s-2 of the ECMWF reference output
It is not optimised, so it is slower than the reference script

Variables naming follows the nomenclature used in the ECMWF documentation:

ref_ds - reference dataset with surface data
grib_data - dataset with all the variables and output

pv - array of a&b coefficients for the hybrid vertical coordinate
coeffs - namedtuple of a&b coefficients for the hybrid vertical coordinate

ph - array of pressure values at (positive) half levels
pf - array of pressure values at full levels
alpha - array of alpha values at full levels
zh - array of geopotential values at (positive) half levels

GRIB variables:
t - temperature
q - specific humidity
lnsp - natural logarithm of surface pressure
pres - pressure
z - geopotential
u - zonal wind
v - meridional wind
w - vertical wind
d - divergence
vtmp - virtual temperature
gh - geopotential height
h - geometric height
"""

G = 9.80665  # m*s-2 gravitational acceleration
R_DRY = 287.0597  # J*K-1*kg-1 dry air gas constant
R_VAP = 461.5250  # J*K-1*kg-1 water vapour gas constant
RE = 6371229.0  # m Earth radius
N_LVLS = 137  # number of vertical levels


class Coeffs(NamedTuple):
    a: numpy.ndarray
    b: numpy.ndarray


gc.enable()

ref_ds, grib_data = cfgrib.open_datasets(
    "levels-ll.grib", backend_kwargs={"read_keys": ["pv"]}
)

grib_data = grib_data.assign(
    vtmp=grib_data.t * (1 + (((R_VAP / R_DRY) - 1) * grib_data.q))
)
grib_data.vtmp.attrs["units"] = "K"
grib_data.vtmp.attrs["long_name"] = "Virtual temperature"

pv = ref_ds.z.attrs["GRIB_pv"]
assert pv.shape == ((N_LVLS + 1) * 2,)
coeffs = Coeffs(*numpy.split(pv, 2))

ph = numpy.empty(
    (
        grib_data.hybrid.size + 1,
        grib_data.latitude.size,
        grib_data.longitude.size,
    )
)
ph = coeffs.a[:, None, None] + (
    coeffs.b[:, None, None] * numpy.exp(ref_ds.lnsp.values)[None, :, :]
)

pf = 0.5 * (ph[1:] + ph[:-1])

alpha = 1.0 - (ph[1:-1] / (ph[2:] - ph[1:-1])) * numpy.log(ph[2:] / ph[1:-1])
alpha = numpy.concatenate((numpy.full_like(alpha[None, 0], numpy.log(2)), alpha))

zh = numpy.empty_like(pf)

for i in range(zh.shape[0]):
    print(i)

    j = i + 2  # index of sigma
    lvl = ref_ds.z.values + numpy.sum(
        (R_DRY * grib_data.vtmp[j - 1 :].values * numpy.log(ph[j:] / ph[j - 1 : -1])),
        axis=0,
    )
    zh[i] = lvl

    gc.collect()

zf = zh + alpha * R_DRY * grib_data.vtmp.values

grib_data = grib_data.assign(z=(("hybrid", "latitude", "longitude"), zf))
grib_data = grib_data.assign(pres=(("hybrid", "latitude", "longitude"), pf))
grib_data = grib_data.assign(gh=grib_data.z / G)
grib_data = grib_data.assign(h=(grib_data.gh * RE) / (RE - grib_data.gh))

grib_data.z.attrs["units"] = "m2.s-2"
grib_data.z.attrs["long_name"] = "Geopotential"
grib_data.pres.attrs["units"] = "Pa"
grib_data.pres.attrs["long_name"] = "Pressure"
grib_data.gh.attrs["units"] = "gpm"
grib_data.gh.attrs["long_name"] = "Geopotential height"
grib_data.h.attrs["units"] = "m"
grib_data.h.attrs["long_name"] = "Geometric height"

grib_data.to_netcdf("full_data.nc")