Source code for pycif.plugins.transforms.basic.families.forward

import numpy as np
import copy




def forward(
        transform,
        inout_datastore,
        controlvect,
        obsvect,
        mapper,
        di,
        df,
        mode,
        runsubdir,
        workdir,
        onlyinit=False,
        **kwargs
):

    trids_in = list(mapper["inputs"].keys())
    trids_out = list(mapper["outputs"].keys())

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

    # Retrieve reference configuration if batch sampling
    trids_in_ref = list(set(
        [(trid[0], trid[1].split("__sample#")[0]) for trid in trids_in]))
    trids_out_ref = list(set(
        [(trid[0], trid[1].split("__sample#")[0]) for trid in trids_out]))

    # Build matching dictionary between inputs and outputs
    # depending on batch sampling or not
    inputs2outputs_map = {}
    for trid_out in trids_out:
        # If trid_out in trids_out_ref, means that unperturbed
        if trid_out in trids_out_ref:
            inputs2outputs_map[trid_out] = trids_in_ref[:]
            continue

        sample_id = trid_out[1].split("__sample#")[1]
        inputs2outputs_map[trid_out] = [
            trid_in if trid_in in trids_in
            else (trid_in[0], f"{trid_in[1]}__sample#{sample_id}")
            for trid_in in trids_in_ref
        ]

    keys = ["spec"]
    if mode == "tl":
        keys.append("incr")

    # Run the main code for all group of samples
    for trid_out in trids_out:
        trids_in_tmp = inputs2outputs_map[trid_out]

        # If sparse data, different treatment
        if mapper["outputs"][trid_out].get("sparse_data", False):
            if onlyinit:
                return

            inif = "yes"

            # Copy outputs from first inputs
            inout_datastore["outputs"][trid_out][di] = copy.deepcopy(
                xmod_in[trids_in_tmp[0]][di]
            )

            for k in keys:
                inout_datastore["outputs"][trid_out][di].loc[:, ("maindata", k)] = \
                    np.sum([
                        xmod_in[trid_in][di].loc[:, ("maindata", k)].values
                        for trid_in in trids_in_tmp
                    ], axis=0)

            return

        xmod_out[trid_out][di] = {k: 0 for k in keys}
        for trid_in in trids_in_tmp:
            for k in keys:
                if k in xmod_in[trid_in][di]:
                    xmod_out[trid_out][di][k] += xmod_in[trid_in][di][k]

            for k in xmod_in[trid_in][di]:
                if k not in keys:
                    xmod_out[trid_out][di][k] = xmod_in[trid_in][di][k]