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

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




def adjoint(
    transf, ddi, mapper, outputs, inout_datastore, onlyinit,
    nthreads=1
):

    if 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 = outputs[trid_ref][ddi]["metadata"]["date"]
    ds_out_dates_end = outputs[trid_ref][ddi]["metadata"]["enddate"]

    # Initializing inputs
    inout_datastore["inputs"][trid_ref] = {}

    # Initialize metadata if not already done
    transf.metadata = getattr(transf, "metadata", {})
    transf.metadata[ddi] = {}

    # Now loop over target dates
    for ddtarget in interpol_indexes:
        # Crop out_dates to the date interval covered by input dates
        in_dates_all = in_dates[ddtarget].stack().reset_index(
            drop=True).drop_duplicates()
        tmp_mask = (
            (out_dates_start <= in_dates_all.max())
            & (out_dates_end >= in_dates_all.min())
        )

        # Now manage for redundant dates
        out_dates_end_tmp = out_dates_end[tmp_mask]
        out_dates_start_tmp = out_dates_start[tmp_mask]

        out_dates_tmp = pd.DataFrame(
            data={'date_start': out_dates_start_tmp,
                  'date_end': out_dates_end_tmp,
                  'index': range(len(out_dates_end_tmp))}
        )
        unique_index = out_dates_tmp.drop_duplicates(
            ['date_start', 'date_end'])['index'].values
        out_dates_ref = pd.DataFrame(
            index=pd.MultiIndex.from_frame(
                out_dates_tmp.iloc[unique_index].loc[:, ['date_start', 'date_end']]),
            data={'index': unique_index}
        )

        # Find correspondance between reference dates and present datastore
        ds_tmp_mask = (
            (ds_out_dates_start <= in_dates_all.max())
            & (ds_out_dates_end >= in_dates_all.min())
        )

        ds_out_dates_end_tmp = ds_out_dates_end[ds_tmp_mask]
        ds_out_dates_start_tmp = ds_out_dates_start[ds_tmp_mask]

        target = pd.MultiIndex.from_frame(pd.DataFrame({
            'date_start': ds_out_dates_start_tmp,
            'date_end': ds_out_dates_end_tmp}))
        out_unique_index = out_dates_ref.loc[target, 'index'].values

        # Merge outputs with indexes
        inter_index = copy.deepcopy(interpol_indexes[ddtarget])
        iterables = [["interpol_indexes"], inter_index.columns]
        inter_index.columns = pd.MultiIndex.from_product(iterables)

        # 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]

                outputs_tmp = outputs[trid][ddi].loc[ds_tmp_mask]
                outputs_tmp.loc[:, "target_index"] = np.arange(
                    len(outputs[trid][ddi])
                )[ds_tmp_mask]
                outputs_tmp.index = out_unique_index
                data_out = pd.merge(outputs_tmp,
                                    inter_index,
                                    left_index=True, right_index=True)

                data_out[("metadata", "tstep")] = \
                    data_out["interpol_indexes"]["indexes"].values.astype(int)
                data_out[("metadata", "dtstep")] = 1

                # Put new data into respective input
                # for data_id in inout_datastore["inputs"]:
                #     if data_id != trid:
                #         continue

                inout_datastore["inputs"][trid][ddtarget] = \
                    data_out

                wgt = data_out.loc[:, [
                    ("interpol_indexes", "weights"), ("target_index", "")]]
                del inout_datastore["inputs"][trid][ddtarget]["interpol_indexes"]

                # Update metadata for later forward run
                transf.metadata[ddi][ddtarget] = {
                    "weights": wgt,
                    "target_mask": ds_tmp_mask
                }

                # Stop here if only init
                if onlyinit:
                    continue

                inout_datastore["inputs"][trid][
                    ddtarget].loc[:, ("maindata", "adj_out")] *= \
                    wgt.loc[:, ("interpol_indexes", "weights")].values

        # Apply threaded function
        thread_function(range(nthreads_tmp))