Source code for pycif.plugins.transforms.basic.vertical_interpolation.adjoint

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

try:
    import cPickle as pickle
except ImportError:
    import pickle

from .utils.sparse.adjoint import sparse_adjoint
from .utils.array.adjoint import array_adjoint




def adjoint(
    transf,
    inout_datastore,
    controlvect,
    obsvect,
    mapper,
    di,
    df,
    mode,
    runsubdir,
    workdir,
    onlyinit=False,
    **kwargs
):
    """Adjoint vertical interpolation: scatter observation sensitivities back to model levels.

    Dispatches to the sparse or array adjoint implementation.  For sparse
    output (observation-indexed data), scatters ``adj_out`` from each
    observation back to the model level(s) it corresponds to.  For array
    output, applies the transposed interpolation weights.

    Args:
        transf (Plugin): vertical_interpolation instance (carries
            ``sparse_out``, ``sampled_out``).
        inout_datastore (dict): mutable datastore.
        controlvect: unused.
        obsvect: unused.
        mapper (dict): transform mapper.
        di (datetime): sub-simulation start date.
        df (datetime): sub-simulation end date.
        mode (str): ``'adj'``.
        runsubdir (str): unused.
        workdir (str): unused.
        onlyinit (bool): if ``True`` (array path), return immediately.
        **kwargs: forwarded to the sparse/array implementations.
    """

    ddi = min(di, df)

    is_sparse_out = transf.sparse_out
    is_sampled_out = transf.sampled_out

    # Force sparse or sampled, process batch computation altogether,
    # Otherwise, loop over trid
    if is_sparse_out or is_sampled_out:
        sparse_adjoint(transf, mapper, inout_datastore,
                       ddi, onlyinit, **kwargs)

    else:
        # Do nothing for non-sparse data and onlyinit mode
        if onlyinit:
            return

        for trid in mapper["outputs"]:
            # Deal with full data
            array_adjoint(
                transf, mapper, inout_datastore,
                trid, ddi, onlyinit, **kwargs
            )