Source code for pycif.plugins.datastreams.meteos.chimere_meteo.write

import copy
import os

from netCDF4 import Dataset
import numpy as np
import pandas as pd
import xarray as xr

from .....utils.classes.fluxes import Flux
from logging import info



def write(self, name, meteo_file, meteo_data, mode="a", **kwargs):
    """Write flux to METEO CHIMERE compatible files.

    Args:
        self (Fluxes): the Fluxes plugin
        flx_file (str): the file where to write fluxes
        flx (xarray.DataArray): fluxes data to write
        mode (str): 'w' to overwrite, 'a' to append
        """

    # If mode is 'a' but file does not exit, switch to mode 'w'
    if mode == "a" and not os.path.isfile(meteo_file):
        mode = "w"

    # If only one level, drop dimension
    if meteo_data.dims.get("lev", -1) == 1:
        meteo_data = meteo_data.squeeze("lev")

    # Rename dimensions
    rename_dims = {"time": "Time", "lev": "bottom_top",
                   "lon": "west_east", "lat": "south_north"}
    for dim in rename_dims:
        if dim in meteo_data.dims:
            meteo_data = meteo_data.rename({dim: rename_dims[dim]})

    # Transform times to CHIMERE strings
    str_dates = list(
        meteo_data["Time"].dt.strftime("%Y-%m-%d_%H:%M:00")
        .to_pandas().values
    )
    dtype = np.dtype(('S', 19))
    meteo_data["Times"] = xr.DataArray(
        str_dates, dims=["Time"]).astype(dtype)

    meteo_data = meteo_data.drop_vars("Time")

    # Add lon/lat/sigma_a/sigma_b
    domain = self.domain
    meteo_data["lon"] = (("south_north", "west_east"), domain.zlon)
    meteo_data["lat"] = (("south_north", "west_east"), domain.zlat)
    meteo_data["a_vcoord"] = (("bottom_top"), domain.sigma_a[1:])
    meteo_data["b_vcoord"] = (("bottom_top"), domain.sigma_b[1:])

    # Duplicate time axis to accelerate reading of nphourm
    encoding = {'Times': {'char_dim_name': 'DateStrLen'}}
    if name == "nphourm":
        meteo_data["Times2"] = (
            ("Time2"), copy.deepcopy(meteo_data["Times"].values)
        )
        meteo_data["nphourm2"] = (
            ("Time2"), copy.deepcopy(meteo_data["nphourm"].values)
        )
        encoding["Times2"] = {'char_dim_name': 'DateStrLen'}

    # Dump to METEO.nc file
    if self.ncformat=='NETCDF4' :
        comp = dict(zlib=True, complevel=5)
        for var in meteo_data.data_vars: 
            meteo_data[var].encoding.update(comp)
        meteo_data = meteo_data.chunk(chunks={"Time":1,
                                              "bottom_top": domain.sigma_a[1:].size, 
                                              "south_north":domain.zlat.shape[0],
                                              "west_east":domain.zlat.shape[1]})
    meteo_data.to_netcdf(
         meteo_file, mode, format=self.ncformat,
         encoding=encoding, unlimited_dims={'Time': True})