import itertools
[docs]
def add_precursors_successors(
self, all_transforms, transforms_ids,
mapper, pipe_links, pipe_subend, transf, simu, mode="adjoint",
fetch_precursors=True):
"""Register a transform's sub-simulation and link it to its immediate neighbours.
Called once per ``(transform, sub-simulation)`` pair, twice (with
``fetch_precursors`` True then False) from
:func:`..fwd_pipe.fwd_adj_pipe`. Appends
``(simu, transf, "adjoint" if fetch_precursors else "forward")`` to
``transforms_ids`` and, unless already present in ``pipe_links``,
initializes its (empty) list of links — skipping dead-end transforms
that have no successors/precursors and are not a pipe end/start
(unless ``self.force_full_operator`` is set).
Then, for each input trid (if ``fetch_precursors``) or output trid
(otherwise) of this sub-simulation, walks its precursors/successors
and, for each of their matching sub-simulations, appends the
corresponding node id to ``pipe_links[transf_id]``. If the trid is
flagged with ``break_{fwd,adj}_onlyinit_pipe`` on either side, the
link is instead redirected to a same-transform self-loop and the
neighbour is recorded in ``pipe_subend`` (used later to seed the
virtual end-of-pipe node), so that only-init propagation stops there
instead of walking further into the neighbour.
If the transform has no precursors/successors at all in the relevant
direction, appends a self-loop link so the transform still appears
in its own pipe.
Args:
self: The parent object (obs operator), exposing
``force_full_operator``.
all_transforms: Namespace holding all registered transform instances,
used to check ``end_pipe``/``start_pipe`` flags.
transforms_ids: List of every registered ``(ddi, transform,
direction)`` node id, mutated in place.
mapper: Dictionary mapping each transform id to its inputs/outputs/
precursors/successors/subsimus metadata.
pipe_links: Mapping of each node id to the list of its immediate
neighbour node ids, mutated in place.
pipe_subend: List of neighbour node ids reached through a
``break_*_onlyinit_pipe`` boundary, mutated in place; used to
seed the virtual end-of-pipe node.
transf: Id of the transform being processed.
simu: Sub-simulation date/index being processed.
mode: Either ``"forward"`` or ``"adjoint"``, the direction of the
pipe currently being built (distinct from ``fetch_precursors``,
which selects which side of the transform to walk).
fetch_precursors: If True, walk the transform's inputs/precursors
(building the adjoint-direction node); if False, walk its
outputs/successors (building the forward-direction node).
Returns:
None. ``transforms_ids``, ``pipe_links`` and ``pipe_subend`` are
mutated in place.
"""
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
