Source code for pycif.plugins.obsoperators.standard.transforms.utils.precursors_successors

import itertools




def add_precursors_successors(
        self, all_transforms, transforms_ids,
        mapper, pipe_links, pipe_subend, transf, simu, mode="adjoint",
        fetch_precursors=True):
    """Find the precursors/successors in backward and forward mode

    :param self:
    :param all_transforms:
    :param transforms_ids:
    :param mapper:
    :param pipe_links:
    :param transf:
    :param simu:
    :param mode:
    :param precursors:
    :param fetch_precursors:
    :return:
    """

    out_mode = "adjoint" if fetch_precursors else "forward"
    transf_id = (simu, transf, out_mode)

    if transf_id not in transforms_ids:
        transforms_ids.append(transf_id)

    # If the transform is not already in the pipe,
    # add it only if it has some successors (resp. precursors),
    # or if it is related to the observation vector (resp. control vector)
    # in adjoint (resp. forward) mode
    # Other transformations are dead-ends and don't need to be computed
    if transf_id not in pipe_links:
        all_successors = \
            [success for trid in mapper[transf]["successors"]
             for success in mapper[transf]["successors"][trid]]
        if not all_successors:
            if not getattr(all_transforms, transf).end_pipe \
                    and not self.force_full_operator:
                return

        all_precursors = \
            [precurs for trid in mapper[transf]["precursors"]
             for precurs in mapper[transf]["precursors"][trid]]
        if not all_precursors:
            if not getattr(all_transforms, transf).start_pipe \
                    and not self.force_full_operator:
                return

        pipe_links[transf_id] = []

    # Loop on all precursors and check whether they produce outputs
    # necessary for the present sub-simulation
    transf_subsimus = mapper[transf]["subsimus"][simu]
    if fetch_precursors:
        neighbours_id = "precursors"
        transform_inout = "inputs"
        neighbour_inout = "outputs"
    else:
        neighbours_id = "successors"
        transform_inout = "outputs"
        neighbour_inout = "inputs"

    transf_neighbours = mapper[transf][neighbours_id]
    for trid in transf_neighbours:
        if trid not in transf_subsimus[transform_inout]:
            continue

        # Loop on the input sub-simulation for that trid
        trid_subsimus = transf_subsimus[transform_inout][trid]
        if mode == "adjoint":
            trid_subsimus = list(trid_subsimus.keys())[::-1]
        trid_neighbours = transf_neighbours[trid]
        for neigh_simu, neighbour in itertools.product(
                trid_subsimus, trid_neighbours):
            neighbour_subsimus = mapper[neighbour]["subsimus"]
            is_in_neighbours = [
                tmp_simu for tmp_simu in neighbour_subsimus
                if neigh_simu in neighbour_subsimus[
                    tmp_simu][neighbour_inout].get(trid, [])
            ]
            if mode == "adjoint":
                is_in_neighbours = is_in_neighbours[::-1]
            for tmp_simu in is_in_neighbours:
                neigh_id = (tmp_simu, neighbour,
                            "adjoint" if fetch_precursors else "forward")
                if neigh_id in transforms_ids:
                    transforms_ids.remove(neigh_id)

                transforms_ids.append(neigh_id)

                # Do not add precursor in onlyinit mode if
                # no propagation of information
                break_id = f"break_{'fwd' if mode == 'forward' else 'adj'}" \
                    f"_onlyinit_pipe"
                break_transf = mapper[transf][transform_inout][trid].get(
                    break_id, False)
                break_neighbour = mapper[neighbour][neighbour_inout][trid].get(
                    break_id, False)
                if (
                        ((fetch_precursors and mode == "forward")
                         or (not fetch_precursors and mode == "adjoint"))
                        and (break_transf or break_neighbour)
                ):
                    if (transf_id[0], transf_id[1], mode) \
                            not in pipe_links[transf_id]:
                        pipe_links[transf_id].append(
                            (transf_id[0], transf_id[1], mode))
                    pipe_subend.append(neigh_id)
                    # setattr(
                    #     getattr(all_transforms, neigh_id[1]),
                    #     "end_pipe" if mode == "forward" else "start_pipe",
                    #     True
                    # )
                    continue

                # # Do not add precursor in onlyinit mode if
                # # no propagation of information
                # if fetch_precursors and mode == "forward":
                #     if mapper[transf][transform_inout][trid].get(
                #             "break_fwd_onlyinit_pipe", False):
                #         pipe_links[transf_id].append(
                #             (transf_id[0], transf_id[1], "forward"))
                #         getattr(all_transforms, neigh_id[1]).end_pipe = True
                #         continue
                #
                # if not fetch_precursors and mode == "adjoint":
                #     if mapper[transf][transform_inout][trid].get(
                #             "break_adj_onlyinit_pipe", False):
                #         pipe_links[transf_id].append(
                #             (transf_id[0], transf_id[1], "adjoint"))
                #         getattr(all_transforms, neigh_id[1]).start_pipe = True

                pipe_links[transf_id].append(neigh_id)

    # When no precursor, set precursor to same transform in forward
    if len(list(itertools.chain(*transf_neighbours.values()))) == 0:
        if fetch_precursors and mode == "forward":
            pipe_links[transf_id].append(
                (transf_id[0], transf_id[1], "forward")
            )
            return

        if not fetch_precursors and mode == "adjoint":
            pipe_links[transf_id].append(
                (transf_id[0], transf_id[1], "adjoint")
            )
            return