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

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 adjoint(
    ddi, mapper, inout_datastore, outputs, onlyinit,
    nthreads=1
):

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

    if onlyinit:
        return

    # Initializing inputs
    list_trids = copy.deepcopy(list(mapper["inputs"].keys()))
    inout_datastore["inputs"] = {
        trid: {} for trid in list_trids
    }

    if not mapper["do_interpolation"][ddi]:
        if sum(mapper["reorder_periods"][ddi]) != 0:
            interpol_indexes = mapper["interpol_indexes"][ddi]
            target_period = list(interpol_indexes.keys())[0]

            for trid in list_trids:
                inout_datastore["inputs"][trid][target_period] = \
                    outputs[trid][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 adjoint interpolation")
    else:
        out_dates_start = out_dates["start_date"]
        out_dates_end = out_dates["end_date"]

    all_interpol_indexes = mapper["interpol_indexes"][ddi]
    data_out = outputs[trid_ref][ddi]["adj_out"]

    for ddtarget in all_interpol_indexes:
        dates_in = mapper["inputs"][trid_ref]["input_dates"][ddtarget]["start_date"]
        interpol_indexes = all_interpol_indexes[ddtarget]
        interpol_indexes.loc[
            interpol_indexes["weights"] == 0, "weights"] = 1

        # Account for filter on crop dates
        in_dates = mapper["inputs"][trid_ref]["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]

        mesh = np.meshgrid(interpol_indexes["indexes"].astype(int),
                           np.arange(data_out.shape[1]),
                           np.arange(data_out.shape[2]),
                           np.arange(data_out.shape[3]),
                           indexing="ij")

        # Threading filling of outputs
        thread_intervals = np.linspace(
            0, len(mapper["inputs"]), nthreads + 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]

                data_out_tmp = outputs[trid][ddi]["adj_out"]

                data_in = np.zeros((len(dates_in), *data_out_tmp.shape[1:]))
                np.add.at(
                    data_in,
                    tuple(mesh),
                    data_out_tmp.values[target_indexes[interpol_indexes.index.values]]
                    * interpol_indexes["weights"].values[
                        :, np.newaxis, np.newaxis, np.newaxis]
                )

                data_in = \
                    xr.DataArray(data_in, coords={"time": dates_in},
                                 dims=("time", "lev", "lat", "lon"))

                if ddtarget not in inout_datastore["inputs"][trid]:
                    inout_datastore["inputs"][trid][ddtarget] = {}

                if "adj_out" not in inout_datastore["inputs"][trid][ddtarget]:
                    inout_datastore["inputs"][trid][ddtarget]["adj_out"] = \
                        data_in

                else:
                    inout_datastore["inputs"][trid][ddtarget]["adj_out"] += \
                        data_in

        # Apply threaded function
        thread_function(range(nthreads))