Source code for pycif.plugins.transforms.basic.time_interpolation.utils.array.forward

import numpy as np
import xarray as xr
import pandas as pd
import copy
from .......utils.parallel import thread
from .......utils.check.errclass import CifError




def forward(
    ddi, mapper, inout_datastore, inputs, nthreads=1
):
    list_trids = copy.deepcopy(list(mapper["inputs"].keys()))

    # If the period has no interpolation indexes, just pass
    if ddi not in mapper["interpol_indexes"]:
        return

    if not mapper["do_interpolation"][ddi]:
        # Re-order periods if necessary
        if np.any(mapper["reorder_periods"][ddi]):
            interpol_indexes = mapper["interpol_indexes"][ddi]
            target_period = list(interpol_indexes.keys())[0]

            for trid in list_trids:
                inout_datastore["outputs"][trid][ddi] = \
                    inputs[trid][target_period]

        # Otherwise, just forward the inputs
        else:
            for trid in list_trids:
                inout_datastore["outputs"][trid][ddi] = \
                    inputs[trid].get(ddi, {})

        return

    # Handling output dates
    trid_ref = list(mapper["inputs"].keys())[0]
    out_dates = mapper["outputs"][trid_ref]["input_dates"][ddi]
    if np.size(out_dates) == 0:
        raise CifError("This should not happen")
    elif np.size(out_dates) == 1:
        raise CifError("Output dates should not be single for forward interpolation")
    else:
        out_dates_start = out_dates["start_date"]
        out_dates_end = out_dates["end_date"]

    # Threading filling of outputs
    thread_intervals = np.linspace(
        0, len(mapper["inputs"]), nthreads + 1
    ).astype(int)
    list_trids = copy.deepcopy(list(mapper["inputs"].keys()))
    all_interpol_indexes = mapper["interpol_indexes"][ddi]

    @thread
    def thread_function(ithread):
        for itrid in range(thread_intervals[ithread], thread_intervals[ithread + 1]):
            trid = list_trids[itrid]

            # Loop on sub-periods
            outputs = {}
            for ddtarget in all_interpol_indexes:
                interpol_indexes = all_interpol_indexes[ddtarget]
                interpol_indexes.loc[
                    interpol_indexes["weights"] == 0, "weights"] = 1
                data_in = inputs[trid][ddtarget]
                for did in data_in:
                    # Skip adj_out
                    if did == "adj_out":
                        continue

                    if did not in outputs:
                        outputs[did] = xr.DataArray(
                            np.zeros(
                                (len(out_dates_start), *data_in[did].shape[1:])),
                            coords={"time": out_dates_start},
                            dims=("time", "lev", "lat", "lon"),
                        )
                    data_out = xr.DataArray(
                        data_in[did].values[
                            interpol_indexes["indexes"].astype(int).values
                        ]
                        * interpol_indexes["weights"].values[
                            :, np.newaxis, np.newaxis, np.newaxis
                        ],
                        coords={"time": out_dates_start[interpol_indexes.index]},
                        dims=("time", "lev", "lat", "lon"),
                    )

                    # Account for filter on crop dates
                    in_dates = mapper["inputs"][trid]["input_dates"][ddtarget]
                    in_dates_all = pd.DatetimeIndex(
                        in_dates.stack().drop_duplicates().sort_values())
                    ddi_mask = (
                        (out_dates_start <= in_dates_all.max())
                        & (out_dates_end >= in_dates_all.min())
                    )
                    target_indexes = np.where(ddi_mask)[0]

                    np.add.at(
                        outputs[did].values,
                        target_indexes[interpol_indexes.index.values],
                        data_out.values
                    )

            inout_datastore["outputs"][trid][ddi] = outputs

    # Apply threaded function
    thread_function(range(nthreads))