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"):
# 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
]
