Source code for pycif.plugins.datastreams.fluxes.lmdz_sflx.write

import numpy as np
import xarray as xr
from .....utils.check.errclass import CifError
from .....utils.hdf5 import _hdf5_lock




def write(self, name, flx_file, flx, mode="a", metadata=None, **kwargs):
    """Write a single-level flux DataArray to an LMDZ sflx compatible NetCDF file.

    Flattens the gridded ``(lat, lon)`` flux into the LMDZ vector
    convention: North/South pole rows are averaged zonally into single
    values, the remaining rows are flattened row-major (dropping the
    wrap-around longitude at the date line), and pole values are placed at
    the start/end of the resulting ``vector`` dimension.

    Args:
        self (Fluxes): the Fluxes plugin
        name (str): name of the flux variable to write
        flx_file (str): the file where to write fluxes
        flx (xarray.DataArray): fluxes data to write, with a single
            vertical level
        mode (str): Unused, kept for interface consistency; the file is
            always (over)written.
        metadata (dict): Must contain a ``'domain'`` key with the `Domain`
            object used to derive latitude ordering and grid size.

    Raises:
        CifError: If `metadata` (or its ``'domain'`` key) is missing, or if
            `flx` has more than one vertical level.
    """
    
    # Check that domain info is in metadata
    if metadata is None:
        raise CifError("Trying to write LMDZ sflx files without metadata")
    
    if "domain" not in metadata:
        raise CifError("Trying to write LMDZ sflx files without domain information")
    
    domain = metadata["domain"]
    lon, lat = domain.zlon[0], domain.zlat[:, 0]
    time = flx.time.values
    ntime = time.size
    
    # Cannot dump multi-level emissions
    nlev = flx.lev.size
    if nlev != 1:
        raise CifError("Can't dump emissions with several vertical levels")
    
    # From North to South
    flx_data = flx.values[:, 0, :, :]
    if lat[0] < lat[-1]:
        flx_data = np.flip(flx_data, 1)
    
    # Now flatten the fluxes
    values_north = flx_data[:, 0, :].mean(1)  # longitude average
    values_south = flx_data[:, -1, :].mean(1)
    values_north = values_north[:, np.newaxis]
    values_south = values_south[:, np.newaxis]
    
    # Flatten the rest of the glob (remove the last longitude value which is a
    # repetition of the first one at the date line at -180)
    flat_dim = (domain.nlat - 2) * (domain.nlon - 1)
    flx_data = flx_data[:, 1:-1, :-1]
    flx_data = np.reshape(flx_data, (ntime, flat_dim))
    flx_data = np.append(values_north, flx_data, axis=1)
    flx_data = np.append(flx_data, values_south, axis=1)

    vector = np.linspace(1, flat_dim + 2, flat_dim + 2)
    ds_out = xr.Dataset({name: (['time', 'vector'], flx_data)},
                        coords={'vector': (['vector'], vector),
                                'time': (['time'], time)})
    
    # Dump to file
    with _hdf5_lock:
        ds_out.to_netcdf(flx_file)