import pandas as pd
import numpy as np
import copy
from collections import Counter
from logging import info, debug
import itertools
from .precursors_successors import add_precursors_successors
from .dask.init_dask import init_dask
from ......utils.parallel import thread
[docs]
def fwd_adj_pipe(self, all_transforms, mapper, mode="forward"):
"""Compute the optimal execution order for one direction of the pipe.
Builds, for every ``(sub-simulation, transform)`` pair, the graph of
immediate precursors/successors (``pipe_links``, saved on ``self`` as
``pipe_links_fwd``/``pipe_links_adj`` for later memory cleaning), then
appends a virtual entry transform (``final_toobsvect`` in forward
mode, ``final_fromcontrol`` in adjoint mode) that gathers every
pipe-end node so the whole graph can be walked from a single root.
If ``self.use_dask`` is set, the raw links are stored
(``pipe_links_fwd_dask``/``pipe_links_adj_dask``) and, for the
adjoint call, the function returns early since Dask handles ordering
itself.
Otherwise, it iteratively expands the graph from the virtual root
(vectorized with pandas for speed) to produce, for each node, the
chronologically correct sequence of precursors, marking as
"dead branches" those adjoint/forward paths that do not eventually
reach a control-vector-relevant ``fromcontrol`` transform (unless
``self.force_full_operator`` is set). Dead-end nodes (with no
successors in the current direction) are then pruned iteratively.
The final ordered list of ``(ddi, transform, direction)`` tuples is
stored on ``self.period_order_fwd`` (mode ``"forward"``) or
``self.period_order_adj`` (mode ``"adjoint"``).
Args:
self: The obs operator, exposing ``use_dask``, ``force_full_operator``
and ``datef``, and receiving the computed ``pipe_links_*`` and
``period_order_*`` attributes.
all_transforms: Namespace holding all registered transform instances,
in pipe order.
mapper: Dictionary mapping each transform id to its ``subsimus``,
``inputs`` and ``outputs`` metadata.
mode: Either ``"forward"`` or ``"adjoint"``, selecting which
direction of the pipe to order.
"""
# Loop over transformations and periods and find precursors to be
# computed just before
debug("First find all precursors and periods for all transforms")
pipe_links = {}
transforms_ids = []
pipe_subend = []
# For each transform, loop over all its sub-simulations
# TODO: this should be accelerated somehow. Very slow for large cases...
list_transforms = all_transforms.attributes
if mode == "adjoint":
list_transforms = list_transforms[::-1]
for i, transf in enumerate(list_transforms):
list_subsimus = sorted(mapper[transf]["subsimus"].keys())
if mode == "adjoint":
list_subsimus = list_subsimus[::-1]
for simu in list_subsimus:
add_precursors_successors(
self, all_transforms, transforms_ids, mapper,
pipe_links, pipe_subend, transf, simu,
fetch_precursors=True, mode=mode
)
add_precursors_successors(
self, all_transforms, transforms_ids, mapper,
pipe_links, pipe_subend, transf, simu,
fetch_precursors=False, mode=mode
)
# Save pipe_links for later use to clean the memory
debug("Saving pipe_links for later use")
if mode == "forward":
self.pipe_links_fwd = copy.deepcopy(pipe_links)
else:
self.pipe_links_adj = copy.deepcopy(pipe_links)
debug("Add a virtual transform at the end of the pipe")
if mode == "forward":
# Add a virtual transformation of which the precursors are toobsvect
# To recursively walk the path leading to them
transforms_ids.append(("", "final_toobsvect", "forward"))
transforms_ids.append(("", "final_toobsvect", "adjoint"))
pipe_links[("", "final_toobsvect", "adjoint")] = [
transf_id for transf_id in pipe_links
if transf_id in pipe_subend
or (transf_id[2] == "adjoint"
and (getattr(all_transforms, transf_id[1]).end_pipe
or (self.force_full_operator
and pipe_links[(transf_id[0], transf_id[1], "forward")] == []
)
)
)
]
for transf_id in pipe_links:
if transf_id[1] not in mapper or transf_id[2] != "forward":
continue
if (getattr(all_transforms, transf_id[1]).end_pipe
or (self.force_full_operator
and pipe_links[transf_id] == [])
):
pipe_links[transf_id] = [("", "final_toobsvect", "forward")]
else:
# Add a virtual transformation of which the successors are fromcontrol
# To recursively walk the path leading to them
transforms_ids.append(("", "final_fromcontrol", "adjoint"))
transforms_ids.append(("", "final_fromcontrol", "forward"))
pipe_links[("", "final_fromcontrol", "forward")] = [
transf_id for transf_id in pipe_links
if transf_id in pipe_subend
or (transf_id[2] == "forward"
and (getattr(all_transforms, transf_id[1]).start_pipe
or (self.force_full_operator
and pipe_links[(transf_id[0], transf_id[1], "adjoint")] == []
)
)
)
]
for transf_id in pipe_links:
if transf_id[1] not in mapper or transf_id[2] != "adjoint":
continue
if (getattr(all_transforms, transf_id[1]).start_pipe
or (self.force_full_operator and pipe_links[transf_id] == [])
):
# If not force_full_operator, walk the path only if fromcontrol
# associated to a control vector variables
if getattr(all_transforms, transf_id[1]).plugin.name == "fromcontrol" \
and not self.force_full_operator:
out_mapper = mapper[transf_id[1]]["outputs"]
iscontrol = False
for trid in out_mapper:
out_tracer = out_mapper[trid]["tracer"]
if not getattr(out_tracer, "iscontrol", False):
continue
else:
iscontrol = True
if not iscontrol:
continue
pipe_links[transf_id] = [("", "final_fromcontrol", "adjoint")]
# Initialize the DASK tree
if self.use_dask:
if mode == "forward":
self.pipe_links_fwd_dask = copy.deepcopy(pipe_links)
else:
self.pipe_links_adj_dask = copy.deepcopy(pipe_links)
return
# Sort transforms in the pipe entry depending
# on the dates and overall transform order
ascending = mode == "forward"
entry_transforms = pipe_links[transforms_ids[-1]]
if mode == "forward":
transforms_dates = [t[0] for t in entry_transforms]
else:
transforms_dindex = [
sorted(list(mapper[t[1]]["subsimus"].keys())).index(t[0])
for t in entry_transforms
]
transforms_dates = [
sorted(list(mapper[t[1]]["subsimus"].keys()))[i + 1]
if i + 1 < len(mapper[t[1]]["subsimus"]) else self.datef
for t, i in zip(entry_transforms, transforms_dindex)
]
transform_order = pd.DataFrame(
{"date": transforms_dates,
"id": [all_transforms.attributes.index(t[1])
for t in entry_transforms],
"in_pipeend": [
getattr(all_transforms, transf_id[1]).end_pipe
for transf_id in entry_transforms
]}
).sort_values(["date", "in_pipeend", "id"], ascending=ascending).index
pipe_links[transforms_ids[-1]] = [
entry_transforms[i] for i in transform_order]
# Turn pipe_links to list of indexes to speed up
debug("Turn ids to indexes to speed up computations")
transforms_ids_df = pd.Series(
index=transforms_ids, data=range(len(transforms_ids))
)
transforms_ids_df.index = transforms_ids_df.index.map(str)
pipe_links_inds = [
[
# transforms_ids_df.index(precur_id)
transforms_ids_df.loc[str(precur_id)]
for precur_id in pipe_links.get(transf_id, [])]
for transf_id in transforms_ids
]
# Stop here if no transforms to run
if len(pipe_links_inds[-1]) == 0:
if mode == "forward":
self.period_order_fwd = []
else:
self.period_order_adj = []
return
# Turns to pandas
debug("Now compute the optimal order")
ref_ds = pd.DataFrame(
data={"precursors": 0},
index=range(len(transforms_ids)))
ref_ds = ref_ds.loc[ref_ds.index.repeat([len(p) for p in pipe_links_inds])]
ref_ds.loc[:, "precursors"] = list(itertools.chain(*pipe_links_inds))
# Branches to prune and initial state for iteration on pipelines
# Include forward dead ends to propagate info on dead branches
prune_branches = [False for t in pipe_links_inds[-1]]
starting_point = pipe_links_inds[-1]
if mode == "adjoint":
prune_branches = []
for transf_id in pipe_links[("", "final_fromcontrol", "forward")]:
if self.force_full_operator:
prune_branches.append(False)
continue
if getattr(all_transforms, transf_id[1]).plugin.name == "fromcontrol":
out_mapper = mapper[transf_id[1]]["outputs"]
iscontrol = False
for trid in out_mapper:
out_tracer = out_mapper[trid]["tracer"]
if not getattr(out_tracer, "iscontrol", False):
continue
else:
iscontrol = True
prune_branches.append(not iscontrol)
else:
prune_branches.append(False)
# Initiate pipe with tooobsvect transforms in backwards mode
# or fromcontrol in forward mode if adjoint
ds = pd.DataFrame(
data={
"precursors": starting_point,
"dead_branch": prune_branches
},
index=[len(pipe_links_inds) - 2] * len(starting_point))
tmp_ds = copy.deepcopy(ds)
niterations = 0
ref_len = np.inf
while ref_len != len(ds) or np.any(tmp_ds.index.values != ds.index.values):
niterations += 1
ref_len = len(ds)
tmp_ds = copy.deepcopy(ds)
precursors = ref_ds.loc[
ds["precursors"].loc[ds["precursors"].isin(ref_ds.index)]
].dropna()
target_index = pd.RangeIndex(len(ds)).repeat(
[len(pipe_links_inds[i]) + 1
for i in ds["precursors"].astype(int)]
)
orig_index = np.append([0], np.cumsum(
[len(pipe_links_inds[i]) + 1
for i in ds["precursors"].astype(int)]))
ds = ds.iloc[target_index]
ds.iloc[np.delete(np.arange(len(ds)), orig_index[:-1]),
0] = precursors.values
new_index = tmp_ds.iloc[target_index]["precursors"].values
new_index[orig_index[:-1]] = \
tmp_ds.index.values
ds.index = new_index
groups = ds.groupby("precursors").prod().astype(bool)
ds = ds.drop_duplicates(subset="precursors", keep="last")
# Propagate information about dead branches
ds.loc[:, "dead_branch"] = groups.loc[
ds.drop_duplicates(subset="precursors", keep="last")["precursors"]
].values
# Prune dead branches
pipe_indexes = ds["precursors"].values.astype(int)
dead_branches = ds["dead_branch"].values
dead_branches = {
transforms_ids[ind]: dead
for ind, dead in zip(pipe_indexes, dead_branches)
}
info(f"Optimal order computed in {niterations} iterations!")
# Get final pipe_list including forward and backward transformations
debug("Turning back to ids")
pipe_indexes = ds["precursors"].values.astype(int)
all_pipe_transforms = [transforms_ids[i] for i in pipe_indexes]
# Cleaning the pipe from dead-ends
debug("Filtering transform with no successor")
if self.force_full_operator:
final_index = all_pipe_transforms.index(
("", "final_toobsvect", "forward") if mode == "forward"
else ("", "final_fromcontrol", "adjoint")
)
all_pipe_transforms_filtered = all_pipe_transforms[:final_index]
else:
dir_prune = "forward" if mode == "adjoint" else "adjoint"
tmp_filtered = []
all_pipe_transforms_filtered = copy.deepcopy(all_pipe_transforms)
dead_end = []
dead_end_df = pd.Series(
index=all_pipe_transforms, data=False)
dead_end_df.index = dead_end_df.index.map(str)
while len(all_pipe_transforms_filtered) != len(tmp_filtered):
tmp_filtered = copy.deepcopy(all_pipe_transforms_filtered)
dead_end_set = set(dead_end)
all_pipe_transforms_filtered = [
t for i, t in enumerate(tmp_filtered)
if (
pipe_links.get(t, []) != []
and ((t[2] != mode) or
(not dead_branches.get((t[0], t[1], dir_prune), False)
and t[2] == mode))
)
and np.any([s not in dead_end_set for s in pipe_links.get(t, [])])
]
tmp_filtered_df = pd.Series(
index=tmp_filtered,
data=range(len(tmp_filtered))
)
tmp_filtered_df.index = tmp_filtered_df.index.map(str)
dead_end_tmp = tmp_filtered_df.loc[
tmp_filtered_df.index.difference(list(map(
str,
all_pipe_transforms_filtered
+ [('', 'final_toobsvect', 'forward'),
("", "final_fromcontrol", "adjoint")]))
)
]
dead_end.extend(
[tmp_filtered[i] for i in dead_end_tmp]
)
debug("Returning result")
if mode == "forward":
self.period_order_fwd = [
t for t in all_pipe_transforms_filtered
if (t[0], t[1], "forward") in all_pipe_transforms
]
else:
self.period_order_adj = [
t for t in all_pipe_transforms_filtered
if (t[0], t[1], "adjoint") in all_pipe_transforms
]
