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

import copy
import xarray as xr
import numpy as np




def adjoint(
        transform,
        inout_datastore,
        controlvect,
        obsvect,
        mapper,
        di,
        df,
        mode,
        runsubdir,
        workdir,
        onlyinit=False,
        **kwargs
):
    """Propagate adjoint sensitivities through the clip/crop mask.

    The adjoint of :func:`forward`: copies ``adj_out`` from the output
    datastore to the input datastore, then zeroes the same spatial mask
    (sensitivities at zeroed cells in the forward pass are zero in the
    adjoint too).

    Args:
        transform (Plugin): clipcrop instance.
        inout_datastore (dict): mutable datastore.
        controlvect: unused.
        obsvect: unused.
        mapper (dict): transform mapper.
        di (datetime): sub-simulation start date.
        df (datetime): unused.
        mode (str): ``'adj'``.
        runsubdir (str): unused.
        workdir (str): unused.
        onlyinit (bool): if ``True``, return immediately.
        **kwargs: unused.
    """

    if onlyinit:
        return

    ddi = min(di, df)

    xmod_in = inout_datastore["inputs"]
    xmod_out = inout_datastore["outputs"]

    parameter = transform.parameter
    component = transform.component
    trid = (component, parameter)

    # Find pixels to clip/crop
    domain = mapper["inputs"][trid]["domain"]
    mask = (domain.zlon >= transform.xmin) & (domain.zlon <= transform.xmax) \
        & (domain.zlat >= transform.ymin) & (domain.zlat <= transform.ymax)
    if transform.mode == "crop":
        mask = ~mask

    # Propagate adj_out
    for trid in xmod_out:
        if "adj_out" not in xmod_out[trid][di]:
            continue

        xmod_in[trid][di]["adj_out"] = \
            copy.deepcopy(xmod_out[trid][di]["adj_out"])

        xmod_in[trid][di]["adj_out"].values[..., mask] = 0.