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

from logging import warning, info
import numpy as np
import copy
from .vcoordfromfile import vcoordfromfile
from .......utils.check.errclass import CifWarning




def sparse_adjoint(transf, mapper, inout_datastore,
                   ddi, onlyinit, **kwargs):

    xmod_out = inout_datastore["outputs"]
    inout_datastore["inputs"] = {}
    trid_ref = list(xmod_out)[0]
    if transf.method not in transf.input_arguments["method"]["accepted"]:
        raise CifWarning(f"The method {transf.method} is not recognized!")

    if transf.method == "static-levels":
        # First propagate datastores to inputs
        for trid in inout_datastore["outputs"]:
            if trid not in inout_datastore["inputs"]:
                inout_datastore["inputs"][trid] = {}

            inout_datastore["inputs"][trid][ddi] = xmod_out[trid][ddi]

        # if no statlev file or ignore_level, just keep same levels
        # WARNING: this can be dangerous!
        file_statlev = transf.file_statlev
        if file_statlev == "" or transf.ignore_level:
            info(
                f"Kept levels as defined in original datastore for {trid_ref}"
            )
            return

        # Otherwise, fetch levels from pre-defined file
        levels = vcoordfromfile(
            inout_datastore["outputs"][trid_ref][ddi],
            file_statlev,
            **kwargs
        )

        for trid in inout_datastore["outputs"]:
            inout_datastore["inputs"][trid][ddi][(
                "metadata", "level")] = levels

        return

    # Extract limited layers if already computed
    if hasattr(transf, "fwd_weights"):
        weights = transf.fwd_weights
        nobs = len(xmod_out[trid_ref][ddi])
        for trid_in in inout_datastore["outputs"]:
            if trid_in not in inout_datastore["inputs"]:
                inout_datastore["inputs"][trid_in] = {}

            inout_datastore["inputs"][trid_in][ddi] = \
                xmod_out[xmod_out][ddi].iloc[range(nobs)]
            inout_datastore["inputs"][trid_in][ddi].loc[
                :, ("metadata", "level")] = weights["i"]

            if not onlyinit:
                inout_datastore["inputs"][trid_in][ddi].loc[
                    :, ("maindata", "adj_out")] *= weights["wgt"]

        return

    # Other methods require the full extraction of the input column
    # If not already done in previous forward run
    warning(f"Doing vertical interpolation for {trid_ref} with method '"
            f"{transf.method}'. \n Need all input levels to be extracted. \n"
            f"This can be computationally demanding and should be used only once"
            f" to determine the matching layers, then reused with pre-computed "
            f"levels")

    # Reshaping the dataframe to unfold data
    nobs = len(xmod_out[trid_ref][ddi])
    nlev_inputs = mapper["inputs"][trid_ref]["domain"].nlev
    dlev = np.ones(nobs, dtype=int) * nlev_inputs

    # Index in the original data of the level-extended dataframe
    native_inds_main = np.append([0], dlev.cumsum())

    # Output index
    idx = np.zeros((native_inds_main[-1]), dtype=int)
    idx[native_inds_main[:-1]] = np.arange(nobs)
    np.maximum.accumulate(idx, out=idx)

    # Copy output datastore to inputs
    for trid_in in inout_datastore["outputs"]:
        inout_datastore["inputs"][trid_in][ddi] = \
            copy.deepcopy(xmod_out[trid_ref][ddi].iloc[idx])
        inout_datastore["inputs"][trid_in][ddi].loc[:, ("metadata", "level")] = \
            np.arange(nobs * nlev_inputs) - idx * nlev_inputs

    # Propagate adjoint to correct indexes
    if onlyinit:
        return