Source code for pycif.plugins.datastreams.fields.grib2_ecmwf.read

from logging import debug, warning

import numpy as np
import xarray as xr

from .utils import GribDataset, get_grid_type, get_jscan




def read(
    self,
    name,
    varnames,
    dates,
    files,
    interpol_flx=False,
    comp_type=None,
    tracer=None,
    **kwargs,
):
    xout = []
    outdates = []
    for dd, dd_file in zip(dates, files):
        spec_cum = []
        for ddi_file in dd_file:
            ds = GribDataset(
                ddi_file,
                filter_by_keys=tracer.filter_by_keys_dict,
            )

            speci = ds[varnames]
            grid_type = get_grid_type(ddi_file, tracer.filter_by_keys_dict)

            if grid_type == "regular":
                # Flip latitudes if jscan = 0
                jscan = get_jscan(ddi_file, tracer.filter_by_keys_dict)
                if jscan == 0:
                    speci = speci[..., ::-1, :]

            # Flip levels if not surface field
            if not tracer.surface:
                speci = speci[::-1, ...]

            spec_cum.append(speci)

        # decumulation if decumul
        if tracer.decumul:
            debug(f"Decumulation done for date {dd[0]} and file {dd_file}")
            if len(spec_cum) > 1:
                spec = spec_cum[1] - spec_cum[0]
            else:
                warning(
                    f"impossible DECUMULATION ? for date {dd[0]} and file {dd_file}"
                )
                spec = spec_cum[0]
        else:
            spec = spec_cum[0]

        xout.append(spec)
        outdates.append(dd[0])

    xout = np.array(xout)
    if grid_type != "regular":
        xout = xout[..., np.newaxis, :]

    if tracer.surface:
        xout = xout[:, np.newaxis, ...]
    else:
        xout = xout

    # Expand surface
    if tracer.expand_psurf:
        domain = tracer.domain
        sigma_a = domain.sigma_a[np.newaxis, :, np.newaxis, np.newaxis]
        sigma_b = domain.sigma_b[np.newaxis, :, np.newaxis, np.newaxis]

        pinterface = np.exp(xout) * sigma_b + sigma_a
        xout = 0.5 * (pinterface[:, 1:] + pinterface[:, :-1])

        # Pressure difference
        if tracer.pressure_thickness:
            xout = np.diff(pinterface, axis=1)

    xmod = xr.DataArray(
        xout,
        coords={"time": outdates},
        dims=("time", "lev", "lat", "lon"),
    )

    return xmod