Source code for pycif.plugins.transforms.system.loadfromoutputs.adjoint

from logging import info
import copy
import pandas as pd
import numpy as np
from .....utils.datastores.empty import init_empty




def adjoint(
        transform,
        inout_datastore,
        controlvect,
        obsvect,
        mapper,
        di,
        df,
        mode,
        runsubdir,
        workdir,
        onlyinit=False,
        do_simu=True,
        check_transforms=False,
        **kwargs
):
    """Write model-space adjoint sensitivities to native model input format.

    The adjoint of :func:`forward`: for each input tracer, collects the
    adjoint sensitivity arrays from the output datastore (populated by
    ``array2sampled`` or ``sparse2sample``), initialises empty datastores
    where data is absent, and calls the model's ``outputs2native_adj``
    function to write the sensitivities to the model-native format.

    Also performs a deep copy of the ``'inputs'`` datastore to prevent
    later in-place modifications from corrupting the adjoint pass.

    Batch-sampling tracers (``__sample#N`` naming) are grouped and
    processed together via ``trids_in_ref`` to avoid redundant I/O.

    Args:
        transform (Plugin): loadfromoutputs instance (carries
            ``outputs2native_adj``).
        inout_datastore (dict): mutable datastore; ``'outputs'`` holds
            the sensitivities; ``'inputs'`` is updated in-place.
        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): passed to ``outputs2native_adj``.
        workdir (str): unused.
        onlyinit (bool): forwarded to ``outputs2native_adj``.
        do_simu (bool): forwarded to ``outputs2native_adj``.
        check_transforms (bool): forwarded to ``outputs2native_adj``.
        **kwargs: unused.
    """

    datastore = inout_datastore["outputs"]
    trids_in = list(mapper["inputs"].keys())
    trids_in_ref = list(set(
        [(trid[0], trid[1].split("__sample#")[0]) for trid in trids_in]))
    for trid in trids_in_ref:
        input_type = trid[0]

        # Batch dump if any
        trids2dump = [
            tr for tr in trids_in
            if tr[1].split("__sample#")[0] == trid[1]
        ]

        # If input parameter is '',
        # dumps all available parameters of this component
        if trid[1] == "":
            trids2dump = [t for t in datastore if t[0] == trid[0]]

        # Initialize empty datastore if empty
        data2dump = {}
        for t in trids2dump:
            data2dump[t] = {}
            for di in datastore.get(t, {}):
                if len(datastore[t][di]) == 0:
                    data2dump[t][di] = init_empty()
                    continue

                data2dump[t][di] = datastore[t][di]

        # Store to input datastore
        inout_datastore["inputs"].update(data2dump)

        # Dump to model inputs
        transform.outputs2native_adj(
            data2dump, input_type, di, df, runsubdir, mode,
            onlyinit=onlyinit, do_simu=do_simu,
            check_transforms=check_transforms
        )

    # Hard save input
    inout_datastore["inputs"] = copy.deepcopy(inout_datastore["inputs"])