Source code for pycif.plugins.transforms.basic.vertical_interpolation.utils.array.linear

import numpy as np
import pandas as pd
import xarray as xr




def linear_fwd(transf, pres_in, pres_out,
               mode, inout_datastore, trid, ddi):
    nlev_in = pres_in.size
    nlev_out = pres_out.size
    xmod = inout_datastore["inputs"]

    pres_tmp = np.sort(np.unique(np.append(pres_in, pres_out)))
    df_pres = pd.DataFrame(range(len(pres_in)), index=pres_in)
    df_pres = df_pres.reindex(pres_tmp).interpolate(method="index")
    if transf.fill_nans:
        df_pres = df_pres.fillna(method="bfill").fillna(method="ffill")
    df_pres = df_pres.reindex(pres_out)
    data_id = ["spec"]
    if mode == "tl" and "incr" in xmod[trid][ddi]:
        data_id += ["incr"]
    inout_datastore["outputs"][trid][ddi] = {}
    for did in data_id:
        var_in = xmod[trid][ddi][did]

        ntimes, nlev, nlat, nlon = var_in.shape
        var_out = np.zeros((ntimes, nlev_out, nlat, nlon), dtype=var_in.dtype)

        for k, dd in enumerate(pres_out):
            ind = df_pres.iloc[k, 0]
            if np.isnan(ind):
                continue
            dmin = np.floor(ind).astype(int)
            wgt = ind - dmin
            try:
                var_out[:, k] = (
                    var_in[:, dmin] * (1 - wgt)
                    + var_in[:, min(dmin + 1, nlev_in - 1)] * wgt
                ).values
            except:
                print(__file__)
                import code
                code.interact(local=dict(locals(), **globals()))
        inout_datastore["outputs"][trid][ddi][did] = xr.DataArray(
            var_out,
            coords={"time": var_in.time},
            dims=("time", "lev", "lat", "lon"),
        )