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

import copy
import pandas as pd
import numpy as np
from .......utils.datastores.empty import init_empty
from .......utils.parallel import thread
from .......utils.check.errclass import CifError




def forward(
    transf,
    ddi,
    inputs,
    inout_datastore,
    onlyinit, mapper, sampled_out, sparse_out,
    save_debug, nthreads=1
):

    if onlyinit or ddi not in mapper["interpol_indexes"]:
        return

    interpol_indexes = mapper["interpol_indexes"][ddi]

    # Load input and output dates from transform
    trid_ref = list(mapper["inputs"].keys())[0]
    in_dates = mapper["inputs"][trid_ref]["input_dates"]

    out_dates = mapper["outputs"][trid_ref]["input_dates"]
    if np.size(out_dates[ddi]) == 1:
        raise CifError("Output dates should not be single for adjoint interpolation")
    else:
        out_dates_start = out_dates[ddi]["start_date"]
        out_dates_end = out_dates[ddi]["end_date"]

    # Output dates from output datastore
    ds_out_dates_start = out_dates_start
    ds_out_dates_end = out_dates_end

    # Threading adjoint propogation
    nthreads_tmp = min(nthreads, len(mapper["inputs"]))
    thread_intervals = np.linspace(
        0, len(mapper["inputs"]), nthreads_tmp + 1
    ).astype(int)
    list_trids = copy.deepcopy(list(mapper["inputs"].keys()))

    @thread
    def thread_function(ithread):

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

            # Initialize empty output datastore
            outputs = None

            # Loop over target dates
            for ddtarget in interpol_indexes:
                # Fetch adjoint metadata
                wgt = transf.metadata[ddi][ddtarget]["weights"]
                ds_tmp_mask = transf.metadata[ddi][ddtarget]["target_mask"]
                data_in = copy.deepcopy(inputs[trid][ddtarget])
                if not isinstance(data_in, pd.DataFrame):
                    continue

                # Initiate outputs if still None
                if outputs is None:
                    outputs = init_empty(nlines=len(ds_tmp_mask))

                # Put new data into respective outputs
                if sparse_out or sampled_out:
                    # input_weights = outputs_tmp['interpol_indexes']
                    # Fill empty columns
                    columns = [c for c in ["spec", "incr"]
                               if c in data_in["maindata"]]
                    for c in columns:
                        if c not in outputs["maindata"]:
                            outputs[("maindata", c)] = 0.

                    # Group once for all
                    data_out = \
                        (data_in["maindata"].loc[:, columns]
                            * wgt[("interpol_indexes", "weights")].values[:, np.newaxis]) \
                        .groupby(wgt[("target_index", "")].values).sum()
                    output_index = wgt.groupby(wgt[("target_index", "")].values).first()[
                        ("target_index", "")].values
                    for c in columns:
                        outputs.loc[output_index, ("maindata", c)] = \
                            outputs.loc[output_index, ("maindata", c)].values \
                            + data_out[c].values

                    # # Propagate info from previous transforms
                    # # if saving debug info
                    # if save_debug:
                    #     outputs_tmp = copy.deepcopy(outputs.loc[ds_tmp_mask])
                    #     data_out = \
                    #         data_in["metadata"].groupby(
                    #             outputs_tmp.loc[:, "ref_index"]).first()
                    #     for c in data_in.columns:
                    #         if c[0] == "maindata":
                    #             continue

                    #         if c[1] not in data_out:
                    #             outputs.loc[:, c] = np.nan

                    #         outputs.loc[ds_tmp_mask, c] = data_out[c[1]].values

                    #     debug_cols = pd.MultiIndex.from_product(
                    #         [[transf.transform_id], out_index.columns]
                    #     )
                    #     group_index = out_index.groupby("out_index")
                    #     df_debug = pd.DataFrame(
                    #         {c: group_index[c[1]].apply(
                    #             list) for c in debug_cols},
                    #         index=outputs.index,
                    #         columns=debug_cols)

                    #     for c in df_debug.columns:
                    #         outputs.loc[:, c] = df_debug[c].values

                else:
                    mask = data_in["metadata"]["tstep"].isin(
                        out_index["indexes"])
                    data_out = copy.deepcopy(data_in.loc[mask])
                    data_out.loc[:, ("metadata", "tstep")] = \
                        out_index.set_index("indexes").loc[
                            data_out["metadata"]["tstep"], "out_index"].values
                    outputs = pd.concat([outputs, data_out])

            # Putting in outputs for sparse_out
            if (not sparse_out or sampled_out) and outputs is not None:
                inout_datastore["outputs"][trid][ddi] = outputs

    # Apply threaded function
    thread_function(range(nthreads_tmp))