Source code for snowtools.utils.xarray_snowtools_preprocess

# -*- coding: utf-8 -*-

import xarray

# The following dictionnaries are used to control default variable and dimension names mapping
# TODO : Réfléchir à un mapping plus cohérent nettoyer les dicts actuels (cf ticket #295)
dimension_map = {'x': 'xx', 'y': 'yy', 'location': 'Number_of_points', 'Number_of_patches': 'tile',
        'valid_time': 'time'}
variables_map = {'Rainf_ds': 'Rainf', 'Snowf_ds': 'Snowf', 'band_data': 'ZS', 'prec': 'Precipitation',
        'rr': 'Precipitation', 'massif_number': 'massif_num'}




def preprocess(ds, mapping=dict(), decode_time=True, transpose=False, sort_dims=True):
    """

    This is the main method to call when opening a SURFEX IO NetCDF file.

    It deals with the following actions :

    1. Update variable and dimension names to ensure that standard names can be used in the following processes

    2. Ensure that 'missing_value' and "_FillValue" attributes do not differ to avoid crashing when trying to write data

    3. Decode time dimension properly if necessary (i.e ensure that the associated coordinates are datetime objects)

    Direct usage example:

    .. code-block:: python

        import xarray as xr

        from snowtools.utils.xarray_snowtools import preprocess

        ds = xr.open_dataset('INPUT.nc', decode_times=False)
        ds = preprocess(ds)


    :param ds: xarray object to preprocess
    :type ds: xarray Dataset or Dataarray
    :param mapping: User-defined dictionnary to map variable or dimension names (it will be used as a complement to
                    the default mapping dictionnaries).
    :type mapping: dict
    :param decode_time: Manually decode the time variable when xarray fails to do it properly (SURFEX outputs)
    :type decode_time: bool
    :param transpose: Put time dimension as first dimension in case of data processing through numpy arrays
    :type transpose: bool
    """

    # Update variable and dimension names to ensure that standard names can be used from now on
    ds = update_names(ds, mapping)

    # Ensure that 'missing_value' and "_FillValue" attributes do not differ to avoid crashing when trying to write data
    ds = check_encoding(ds)

    # Decode time dimension properly (if necessary)
    if decode_time:
        ds = decode_time_dimension(ds)

    # Ensure that the time dimension is in the first one for numpy-based tools backward compatibility
    if transpose:
        ds = transpose(ds)

    if sort_dims:
        ds = sort_dimensions(ds)

    return ds



def update_names(ds, mapping):
    """
    Map variable and dimension names to snowtools-standard ones.

    :param ds: xarray object to preprocess
    :type ds: xarray Dataset or Dataarray
    :param mapping: User-defined dictionnary to map variable or dimension names (it will be used as a complement to
                    the default mapping dictionnaries).
    :type mapping: dict
    """
    # Do not directly modify *variables_map* in case several calls to "preprocess" are made from
    # the same session
    default_map = variables_map.copy()
    default_map.update(dimension_map)

    if isinstance(ds, xarray.core.dataarray.DataArray):
        list_entries = list(ds.dims)
    else:
        list_entries = list(ds.keys()) + list(ds.dims)

    update_dict = {key: default_map[key] for key in list_entries if key in default_map.keys()}
    # The optionnal user-defined mapping dictionnary overwrites the default renaming
    user_mapping = {key: mapping[key] for key in list_entries if key in mapping.keys()}

    update_dict.update(user_mapping)

    ds = ds.rename(update_dict)

    # Set orginal names as attribute for backtracking
    if len(update_dict) > 0:
        backtrack = ', '.join([f'{value}: {key}' for key, value in update_dict.items()])
        ds = ds.assign_attrs(original_name=backtrack)

    return ds


def check_encoding(ds):
    """
    Ensure that "missing_value" and "_FillValue" attributes do not differ to avoid crashing when trying to write data.

    See snowtools ticket #282 or xarray ticket #7722 (https://github.com/pydata/xarray/issues/7722) for more
    information.

    :param ds: xarray object to preprocess
    :type ds: xarray Dataset or Dataarray
    """
    if isinstance(ds, xarray.core.dataarray.Dataset):
        for var in ds.keys():
            if 'missing_value' in ds[var].encoding.keys():
                ds[var].encoding['missing_value'] = ds[var].encoding['_FillValue']

    elif isinstance(ds, xarray.core.dataarray.DataArray):
        if 'missing_value' in ds.encoding.keys():
            ds.encoding['missing_value'] = ds.encoding['_FillValue']

    return ds


def decode_time_dimension(ds):
    """
    Manually decode time variable when xarray fails to do it properly (SURFEX outputs)

    :param ds: xarray object to preprocess
    :type ds: xarray Dataset or Dataarray

    """
    if 'time' in list(ds.coords):
        time = xarray.Dataset({"time": ds.time})
        time = xarray.decode_cf(time)
        ds['time'] = time.time
    return ds


def sort_dimensions(ds):
    """
    Sort dimension values.

    :param ds: xarray object to preprocess
    :type ds: xarray Dataset or Dataarray

    """
    ds = ds.sortby(list(ds.dims))
    return ds