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

import numpy as np
import xarray as xr




def write(self, name, flx_file, flx, mode="a", metadata=None, **kwargs):
    """Write flux to LMDZ-DISPERSION compatible files.
    The shape follows the LMDZ physical vectorial shape grid.

    Args:
        self (Fluxes): the Fluxes plugin
        flx_file (str): the file where to write fluxes
        flx (xarray.Dataset): fluxes data to write
    """
    
    # Check that domain info is in metadata
    if metadata is None:
        raise Exception("Trying to write LMDZ sflx files without metadata")
    
    if "domain" not in metadata:
        raise Exception("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 Exception("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
    ds_out.to_netcdf(flx_file)