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

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




def closest_fwd(transf, pres_in, pres_out,
                mode, inout_datastore, trid, ddi, mapper):

    nlev_in = pres_in.size
    nlev_out = pres_out.size
    xmod = inout_datastore["inputs"]

    is_sparse_in = mapper["inputs"][trid].get("sparse_data", False)
    if not is_sparse_in:
        data_id = ["spec"] if mode == "fwd" else ["spec", "incr"]
        inout_datastore["outputs"][trid][ddi] = {}
        for did in data_id:
            if did not in xmod[trid][ddi]:
                continue

            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)
            dpres = np.abs(pres_out[:, np.newaxis] - pres_in)
            for kin, kout in enumerate(np.argmin(dpres, axis=0)):
                var_out[:, kout] += var_in[:, kin]

            inout_datastore["outputs"][trid][ddi][did] = xr.DataArray(
                var_out,
                coords={"time": var_in.time},
                dims=("time", "lev", "lat", "lon"),
            )

    else:
        var_in = xmod[trid][ddi]
        var_out = copy.deepcopy(var_in)
        dpres = (pres_out[:, np.newaxis] - pres_in).T
        min_in = np.abs(dpres).argmin(axis=1)
        var_out.loc[:, ("metadata", "level")] = min_in

        inout_datastore["outputs"][trid][ddi] = var_out