Source code for pycif.plugins.datastreams.fluxes.chimere.fetch

import os
import datetime
import pandas as pd
import xarray as xr
import itertools
from .....utils import path
from .....utils.classes.setup import Setup
from .....utils.check.errclass import CifError
from .....utils.hdf5 import _hdf5_lock




def fetch(ref_dir, ref_file, input_interval, target_dir, tracer=None, **kwargs):
    """Fetch CHIMERE AEMISSIONS/BEMISSIONS files covering the interval.

    Builds a date range at ``tracer.file_freq``, links each existing
    per-period file into ``target_dir`` with hourly sub-dates. If the linked
    file declares a ``bottom_top`` dimension different from the domain's
    current number of levels, ``tracer.domain`` is rebuilt with the file's
    vertical level count (horizontal geometry copied as-is, sigma
    coefficients truncated to the new number of levels).

    Args:
        ref_dir (str): directory where the original files are found.
        ref_file (str): (template) name of the original files.
        input_interval (list): simulation interval, as a list of the two
            bounding dates.
        target_dir (str): directory where links to the original files are
            created.
        tracer: the tracer Plugin, giving access to ``file_freq``,
            ``name`` and ``domain``.
        **kwargs: unused, kept for interface compatibility.

    Returns:
        (dict, dict): ``list_files`` and ``list_dates``.

        list_files: for each date that begins a period, a list containing
            the names of the files that are available for the dates within
            this period.
        list_dates: for each date that begins a period, a list containing
            the date intervals matching the files listed in ``list_files``.

    Raises:
        CifError: if ``tracer.file_freq`` is not set, or if no file matching
            ``ref_dir``/``ref_file`` is found over the requested interval.
    """
    if tracer.file_freq == "":
        raise CifError(
            f"File frequency is not set for tracer {tracer.name} "
            f"while it is required to fetch files from {file}"
        )

    list_period_dates = \
        pd.date_range(input_interval[0], input_interval[1], freq=tracer.file_freq)
    list_dates = {}
    list_files = {}
    for dd in list_period_dates:
        file = dd.strftime(f"{ref_dir}/{ref_file}")
        if os.path.isfile(file):
            file_hours = pd.date_range(
                dd, dd + pd.to_timedelta(tracer.file_freq), freq="1h")
            list_dates[dd] = [[hh, hh + datetime.timedelta(hours=1)]
                              for hh in file_hours]
            list_files[dd] = (len(file_hours) * [file])

            # Fetching
            target_file = f"{target_dir}/{os.path.basename(file)}"
            path.link(file, target_file)

    # Raise exception if no file was found
    if not list(itertools.chain(*list_files.values())):
        raise CifError(
            f"Could not find any file fitting the requested structure:\nPlease check your informations: \n   - ref_dir: {ref_dir}\n   - ref_file: {ref_file}\n   - date_range: {input_interval}\n   - file_freq: {tracer.file_freq}\nThis set-up leads to the following dates and files being looped on:\n"
            + "\n".join([dd.strftime(f"    - {dd}: {ref_dir}/{ref_file}")
                         for dd in list_period_dates])
        )

    # Alter vertical domain depending on input files
    # if bottom_top dimension is available
    file_ref = list(itertools.chain(*list_files.values()))[0]

    nlevemis = 0
    try:
        with _hdf5_lock:
            nlevemis = xr.open_dataset(file_ref).sizes["bottom_top"]
    except KeyError:
        pass

    if nlevemis > 0:
        domain_in = tracer.domain
        if domain_in.nlev != nlevemis:
            domain_out = Setup.load_registered(
                domain_in.plugin.name, domain_in.plugin.version,
                "domain", plg_orig=domain_in
            )

            domain_out.nlon = domain_in.nlon
            domain_out.nlat = domain_in.nlat
            domain_out.zlon = domain_in.zlon
            domain_out.zlat = domain_in.zlat
            domain_out.zlonc = domain_in.zlonc
            domain_out.zlatc = domain_in.zlatc

            domain_out.nlon_side = domain_in.nlon_side
            domain_out.nlat_side = domain_in.nlat_side
            domain_out.zlonc_side = domain_in.zlonc_side
            domain_out.zlatc_side = domain_in.zlatc_side
            domain_out.zlon_side = domain_in.zlon_side
            domain_out.zlat_side = domain_in.zlat_side

            # Read nlevemis from one of the available files
            domain_out.pressure_unit = domain_in.pressure_unit
            domain_out.nlev = nlevemis
            domain_out.sigma_a = domain_in.sigma_a[:domain_out.nlev + 1]
            domain_out.sigma_b = domain_in.sigma_b[:domain_out.nlev + 1]
            domain_out.sigma_a_mid = domain_in.sigma_a_mid[:domain_out.nlev]
            domain_out.sigma_b_mid = domain_in.sigma_b_mid[:domain_out.nlev]

            tracer.domain = domain_out

    return list_files, list_dates