Source code for pycif.plugins.datastreams.backgrounds.tm5_background.fetch

import os
import datetime
import pandas as pd
import numpy as np
from netCDF4 import Dataset
from .....utils import path
from .....utils.hdf5 import _hdf5_lock




def fetch(ref_dir, ref_file, date_interval, target_dir, tracer=None, **kwargs):
    """Locate TM5-4DVAR annual background files and link them to the working directory.

    The requested date interval is expanded to cover full calendar years,
    and one yearly-formatted file name (``ref_file`` formatted with the
    year's start date) is checked per year. Each existing file is symlinked
    into ``target_dir`` and opened to read its time stamps
    (``tracer.time_varname``), from which hourly ``[t, t + 1h]`` date pairs
    are built.

    Args:
        ref_dir: directory holding the TM5-4DVAR background files.
        ref_file: strftime-style file name pattern (relative to
            ``ref_dir``), formatted once per calendar year.
        date_interval: 2-element sequence ``(datei, datef)`` giving the
            requested date range.
        target_dir: directory where matching files are symlinked.
        tracer: the background datastream Plugin, used for
            ``tracer.time_varname``.

    Returns:
        tuple: ``(list_files, list_dates)``, dicts keyed by each year's
        start date. ``list_files`` maps each key to a list of the (single,
        repeated) linked file path, one entry per time stamp found in the
        file; ``list_dates`` maps each key to the corresponding list of
        ``[t, t + 1h]`` date pairs. Years without an existing file are
        simply absent from both dicts.
    """

    # Reshape input interval to include full years
    datei, datef = date_interval
    datei = datetime.datetime(year=datei.year, month=1, day=1)
    datef = datetime.datetime(year=datef.year + 1, month=1, day=1)
    
    list_dates_files = pd.date_range(datei, datef, freq="1YS")
    
    list_dates = {}
    list_files = {}
    for dd in list_dates_files:
        file = dd.strftime(f"{ref_dir}/{ref_file}")
        if os.path.isfile(file):
            # Fetching
            target_file = f"{target_dir}/{os.path.basename(file)}"
            path.link(file, target_file)

            # Time stamps
            with _hdf5_lock:
                with Dataset(file) as bkg:
                    ts = np.array([
                        datetime.datetime(*t) for t in bkg[tracer.time_varname][:]
                    ])
                
            list_dates[dd] = [
                [hh0, hh0 + datetime.timedelta(hours=1)]
                for hh0 in ts]

            list_files[dd] = (len(list_dates[dd]) * [target_file])
    
    return list_files, list_dates