Source code for pycif.plugins.transforms.system.sparse2sample.forward

import numpy as np
import xarray as xr

from logging import warning




def forward(
        transform,
        inout_datastore,
        controlvect,
        obsvect,
        mapper,
        di,
        df,
        mode,
        runsubdir,
        workdir,
        onlyinit=False,
        **kwargs
):
    ddi = min(di, df)
    for trid in mapper["inputs"]:
        xmod_in = inout_datastore["inputs"][trid][ddi]
        xmod_out = inout_datastore["outputs"][trid][ddi]

        t = xmod_in["metadata"]["tstep"].astype(int).values
        lev = xmod_in["metadata"]["level"].astype(int).values
        i = xmod_in["metadata"]["i"].astype(int).values
        j = xmod_in["metadata"]["j"].astype(int).values

        # Output shape
        nlon = mapper["outputs"][trid]["domain"].nlon
        nlat = mapper["outputs"][trid]["domain"].nlat
        nlev = mapper["outputs"][trid]["domain"].nlev
        ntimes = len(mapper["outputs"][trid]["input_dates"][ddi])

        columns = ["spec"] if mode == "fwd" else ["spec", "incr"]
        for c in columns:
            if c not in xmod_in["maindata"]:
                continue

            var_out = np.zeros((ntimes, nlev, nlat, nlon))
            np.add.at(var_out, (t, lev, i, j), xmod_in[("maindata", c)].values)

            xmod_out[c] = xr.DataArray(
                var_out,
                coords={"time": mapper["outputs"][trid]["input_dates"][ddi][:, 0]},
                dims=("time", "lev", "lat", "lon"),
            )