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

import datetime
import os

import numpy as np
import pandas as pd
import xarray as xr

from .....utils import path
from .....utils.classes.setup import Setup




def fetch(ref_dir, ref_file, date_interval, target_dir, tracer=None, **kwargs):

    # Reshape input interval to include full years
    datei, datef = date_interval
    datei = datetime.datetime(year=datei.year, month=1, day=1)
    if datetime.datetime(year=datef.year, month=1, day=1) < datef:
        datef += pd.DatetimeIndex([datef]).days_in_month[0] * datetime.timedelta(days=1)

    list_period_dates = pd.date_range(datei, datef, freq=tracer.file_freq)
    list_dates = {}
    list_files = {}
    valid_files = []
    for dd in list_period_dates:
        file = dd.strftime("{}/{}".format(ref_dir, ref_file))
        if not os.path.isfile(file) or file in valid_files:
            continue

        # Read times
        dates = xr.open_dataset(file)["time"][:]

        # Round dates
        dates = dates.to_dataframe()["time"].dt.round("H").values

        # Check if monthly resolution
        freq = np.unique(np.diff(dates))
        dates = pd.DatetimeIndex(dates)
        if len(freq) != 1 and np.all(pd.TimedeltaIndex(np.diff(dates)).days >= 28):
            dates -= (dates.day - 1) * pd.Timedelta("1D") + dates.hour * pd.Timedelta(
                "1h"
            )
            freq = dates.days_in_month[-1] * pd.Timedelta("1D")
            freq = freq.to_numpy()

        dates = pd.DatetimeIndex(np.append(dates.values, dates.values[-1] + freq))

        # Shift year if using a fixed file
        year_ref = dates.year.min()
        if dd.year != year_ref:
            dates = dates + pd.DateOffset(years=dd.year - year_ref)

        # Put dates into correct format
        list_dates[dd] = [[hh0, hh1] for hh0, hh1 in zip(dates[:-1], dates[1:])]

        # Fetching
        target_file = "{}/{}".format(target_dir, os.path.basename(file))
        path.link(file, target_file)
        list_files[dd] = len(list_dates[dd]) * [target_file]

        # Keep in memory processed files
        valid_files.append(file)

    # Alter the domain from LMDZ by keeping only the surface layer
    domain_in = tracer.domain
    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 = 1
    domain_out.sigma_a_mid = np.array([0])
    domain_out.sigma_b_mid = np.array([1])
    domain_out.sigma_a = np.array([0])
    domain_out.sigma_b = np.array([1])

    tracer.domain = domain_out

    return list_files, list_dates