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

import numpy as np
import xarray as xr

from logging import warning

from .....utils.datastores.empty import init_empty




def forward(
        transform,
        inout_datastore,
        controlvect,
        obsvect,
        mapper,
        di,
        df,
        mode,
        runsubdir,
        workdir,
        onlyinit=False,
        **kwargs
):
    """Extract point-observation values from a gridded model output field.

    For each (input tracer, output tracer) pair in the mapper, reads the
    gridded field ``xmod_in["spec"]`` and uses the pre-computed grid
    indices ``(i, j, level, tstep)`` stored in the output observation
    DataFrame's metadata to extract the corresponding model values via
    xarray label-based indexing.

    In TL mode (``mode = 'tl'``) the ``'incr'`` field is sampled in the
    same way.

    Single-level fields are handled by setting ``lev = 0`` for all
    observations.

    Args:
        transform (Plugin): array2sampled instance (unused directly).
        inout_datastore (dict): mutable datastore; ``'inputs'`` has the
            gridded field, ``'outputs'`` has the sparse observation
            DataFrame to be populated.
        controlvect: unused.
        obsvect: unused.
        mapper (dict): transform mapper.
        di (datetime): sub-simulation start date.
        df (datetime): sub-simulation end date.
        mode (str): ``'fwd'`` or ``'tl'``.
        runsubdir (str): unused.
        workdir (str): unused.
        onlyinit (bool): if ``True``, return immediately.
        **kwargs: unused.
    """
    if onlyinit:
        return

    ddi = min(di, df)
    for trid_in, trid_out in zip(mapper["inputs"], mapper["outputs"]):
        xmod_in = inout_datastore["inputs"][trid_in][ddi]

        # Fetch metadata from adjoint
        metadata = transform.metadata[ddi][trid_in]
        t = metadata["tstep"]
        i = metadata["i"]
        j = metadata["j"]
        lev = metadata["level"]

        # Initialize output with size from metadata
        xmod_out = init_empty(nlines=len(t))

        columns = ["spec", "incr"]
        for c in columns:
            if c not in xmod_in:
                xmod_out[("maindata", c)] = np.nan
                continue

            xmod_out[("maindata", c)] = xmod_in[c].isel(
                lat=xr.DataArray(i),
                lon=xr.DataArray(j),
                time=xr.DataArray(t),
                lev=xr.DataArray(lev)).data

        inout_datastore["outputs"][trid_out][ddi] = xmod_out