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

from logging import warning
import numpy as np
import copy
import pandas as pd
from .vcoordfromfile import vcoordfromfile
from .......utils.check.errclass import CifWarning, CifNotImplementedError




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

    # Initialize outputs
    xmod_in = inout_datastore["inputs"]
    inout_datastore["outputs"] = {
        trid: {ddi: {}}
        for trid in xmod_in
    }

    # Stop here if onlyinit
    if onlyinit:
        return

    # Raise error for unrecognized method
    if transf.method not in transf.input_arguments["method"]["accepted"]:
        raise CifWarning(f"The method {transf.method} is not recognized!")

    # Just forward propagation for static-levels
    if transf.method == "static-levels":
        inout_datastore["outputs"] = {
            trid: {ddi: xmod_in[trid][ddi]} for trid in xmod_in
        }
        return

    # Otherwise loop over trids
    for trid in xmod_in[trid][ddi]:

        raise CifNotImplementedError(
            f"Method {transf.method} is not yet implemented for vertical interpolation"
        )

        # Reshaping the dataframe to unfold data
        xmod_out = inout_datastore["outputs"][trid][ddi]
        nobs = len(xmod_out)
        nlev_inputs = mapper["inputs"][trid]["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)

        coord_out = mapper["outputs"][trid].get("coord_out", "height")

        spec = xmod_in[trid][ddi][("maindata", "spec")].values.reshape((nobs, -1))
        incr = xmod_in[trid][ddi][("maindata", "incr")].values.reshape((nobs, -1))
        hthick = xmod_in[("hlay", trid[1])][ddi][
            ("maindata", "spec")].values.reshape((nobs, -1))
        hlay = np.concatenate([
            np.zeros((nobs, 1)),
            hthick.cumsum(axis=1),
            np.inf * np.ones((nobs, 1))
        ], axis=1)

        # Use orography or not
        # oro = xmod_in[("oro", trid[1])][ddi][
        #     ("maindata", "spec")].values.reshape((nobs, -1))

        # Find where the data is in the input mesh
        alt_low = xmod_out["metadata"]["alt"].values
        alt_high = xmod_out["metadata"]["alt"] \
            + xmod_out["metadata"].get("dalt", 0) + 250
        alt_high = alt_high.values
        alt_mid = 0.5 * (alt_high + alt_low)

        levmeshin = np.arange(0, nlev_inputs + 2)
        levout_low = np.zeros(nobs)
        levout_high = np.zeros(nobs)
        levout_mid = np.zeros(nobs)
        weights = {"i": [], "wgt": []}
        for i in range(nobs):
            levout_low[i] = np.interp(alt_low[i], hlay[i], levmeshin)
            levout_high[i] = np.interp(alt_high[i], hlay[i], levmeshin)
            levout_mid[i] = np.interp(alt_mid[i], hlay[i], levmeshin)

        # Compute weights per layers
        iout_mid = np.floor(levout_mid).astype(int)
        if transf.method == "match-layer":
            xmod_out.loc[:, ("metadata", "level")] = iout_mid
            xmod_out.loc[:, ("maindata", "spec")] = spec[range(nobs), iout_mid]
            if mode == "tl":
                xmod_out.loc[:, ("maindata", "incr")] = incr[range(nobs), iout_mid]

            weights["i"] = iout_mid
            weights["wgt"] = np.ones(nobs)

        # elif transf.method == "linear":

        # Save weights for later adjoint
        transf.fwd_weights = weights


        # print(__file__)
        # import code
        # code.interact(local=dict(locals(), **globals()))