import copy
import pandas as pd
from logging import debug
import numpy as np
import datetime
import xarray as xr
from ......utils.datastores.empty import init_empty
from ......utils.check.errclass import CifError, CifTypeError
print_all = False
transfor2print = [
# "time_interpolation_std_00006",
# "regrid_std_00008",
# "vertical_interpolation_std_00007",
# "loadfromoutputs_std_00036"
# 'satellites_std_00001',
# 'time_interpolation_std_00012'
# 'dump2inputs_std_00008'
]
# transfor2print = [
# 'dump2inputs_std_00030',
# 'time_interpolation_std_00007',
# 'loadfromoutputs_std_00043'
# ]
# transfor2print = [
# 'fromcontrol_std_00011',
# 'time_interpolation_std_00012',
# 'dump2inputs_std_00006'
# ]
# mode2print = "adj"
mode2print = "fwd"
open_debug = True
[docs]
def aggregate_inout(
transform, ddi,
tmp_inputs, tmp_outputs,
transform_mode, transform_onlyinit,
mapper,
check_transforms=False
):
"""Aggregate inputs and outputs from precursors and successors respectively.
There can be only one precursor per trid, with the possibility of having several sub-dates.
By construction, it is not possible to have several precursors for a given trid to avoid ambiguity.
This could become possible in the future, but it would require further complexity in identifying datastores
(i.e., including precursor and successor in dictionaries, at the cost of reduced readability).
Outputs can have several successors for a given trid. This implies that adjoint sensitivities must be properly
propagated backwards.
Args:
transform (str): The name of the transform currently being processed
ddi (datetime.datetime): The period being processed
tmp_inputs (dict): Per-trid, per-date datastore of raw inputs as
fetched from precursors (nested under the single precursor id
and its own sub-simulation date), as produced by
:func:`fetch_inoutputs.fetch_inputs_outputs`. Collapsed in
place to ``{trid: {date: <data>}}``.
tmp_outputs (dict): Per-trid, per-date datastore of raw outputs as
fetched from successors (nested under each successor id and
its own sub-simulation date). Collapsed in place to
``{trid: {date: <data>}}``, summing (or concatenating, for
sparse data) contributions from multiple successors.
transform_mode (str): Either ``"fwd"`` or ``"adj"``, the direction
currently being executed.
transform_onlyinit (bool): Whether the transform is only being
initialized (no actual data processing), used to decide
whether to compute the ``delta_adj_out`` debug diagnostic.
mapper (dict): Dictionary mapping each transform id to its
precursors/successors/outputs metadata.
check_transforms (bool, optional): Whether to keep track of the
original adjoint output (``adj_out_original``) so that
:func:`check_adjtltest.check_adjtltest` can later verify
consistency. Defaults to False.
Raises:
CifError: If more than one precursor is found for a given input
trid (should never happen by construction).
CifTypeError: If a transform mixes sparse and non-sparse outputs
for the same trid.
Returns:
tuple: ``(tmp_inputs, tmp_outputs)``, the same dictionaries passed
in, mutated and collapsed to one entry per ``(trid, date)``.
"""
if transform in transfor2print or print_all:
debug(f'aggreget_inout: origin_in: {transform} / {ddi} / {tmp_inputs}')
if not transform_onlyinit and transform_mode == mode2print and open_debug:
print("AAA")
print(__file__)
import code
code.interact(local=dict(locals(), **globals()))
# Get variables from mapper
transf_mapper = mapper[transform]
precursors = transf_mapper["precursors"]
successors = transf_mapper["successors"]
# Loop on inputs and check precursors
for trid in tmp_inputs:
trid_dict = {}
for di in tmp_inputs[trid]:
nprecursors = len(tmp_inputs[trid][di])
# If no data, continue
if nprecursors == 0:
trid_dict = {di: {}}
continue
# If several precursor raise error.
# Should never happen by construction...
if nprecursors > 1:
raise CifError(
f"Several possible precursors for {trid} for transform {transform}:\n"
f"{precursors[trid].keys()}\n"
"This should not happen. Please check your yaml with a core developer."
)
# Turn input to list if only one precursor to fit expected format
precursor = list(tmp_inputs[trid][di].keys())[0]
# subsimus_in = mapper[precursor]["subsimus"]
subsimus_in = tmp_inputs[trid][di][precursor]
trid_dict[di] = {}
for idate, precursor_di in enumerate(subsimus_in):
# Fill values from precursor
trid_dict[di] = tmp_inputs[trid][di][precursor].get(
precursor_di, {})
# Save input to compute adjoint delta in check_transform mode
if check_transforms and transform_mode == "adj" and not transform_onlyinit:
if "adj_out" in trid_dict[di]:
trid_dict[di]["adj_out_original"] = \
copy.deepcopy(trid_dict[di]["adj_out"])
tmp_inputs[trid] = trid_dict
if transform in transfor2print or print_all:
debug(f'aggreget_inout: target_in: {transform} / {ddi} / {tmp_inputs}')
debug(f'aggreget_inout: origin_out: {transform} / {ddi} / {tmp_outputs}')
if not transform_onlyinit and transform_mode == mode2print and open_debug:
print("BBB")
print(__file__)
import code
code.interact(local=dict(locals(), **globals()))
for trid in tmp_outputs:
for di in tmp_outputs[trid]:
if len(successors[trid]) == 0:
continue
# Check whether there is a need to aggregate several data
aggregated_out = None
to_aggregate = False
for isuccessor, successor in enumerate(successors[trid]):
if successor not in tmp_outputs[trid][di]:
continue
for idate, successor_di in enumerate(tmp_outputs[trid][di][successor]):
# Initialize if None
if aggregated_out is None:
aggregated_out = tmp_outputs[trid][di][successor][successor_di]
# Otherwise means that there is some aggregation to do
else:
to_aggregate = True
# Stop here if no aggregation is needed
if not to_aggregate:
tmp_outputs[trid][di] = aggregated_out if aggregated_out is not None else {}
continue
# Loop over successors and their sub-dates
# If aggregation
aggregated_out = []
is_sparse = None
aggreg_index = 0
for isuccessor, successor in enumerate(successors[trid]):
if successor not in tmp_outputs[trid][di]:
continue
# Check sparse
out_mapper = transf_mapper["outputs"][trid]
if out_mapper.get("sparse_data", False) or out_mapper.get("sampled", False):
if is_sparse is None:
is_sparse = True
elif is_sparse:
raise CifTypeError(
f"The transform {transform} has sparse and non sparse outputs. "
"I cannot deal with that yet..."
)
# Now loop over successor dates
subsimus_out = transf_mapper["subsimus"][ddi]["outputs"][trid]
subsimus_in = mapper[successor]["subsimus"]
for idate, successor_di in enumerate(subsimus_in):
if di not in subsimus_in[successor_di]['inputs'].get(trid, []):
continue
aggreg_index += 1
# Fetch data to aggregate and turn to empty datastore if sparse
if successor_di not in tmp_outputs[trid][di][successor]:
debug(
f"No value in {successor_di} for {transform} / {successor} / {di}"
)
dict_to_aggregate = tmp_outputs[trid][di][successor].get(
successor_di,
init_empty() if is_sparse else {}
)
if is_sparse:
if len(dict_to_aggregate) == 0:
dict_to_aggregate = init_empty()
# Aggregate data
aggregated_out.append(dict_to_aggregate)
# Deal differently sparse and non sparse data
if is_sparse:
if isinstance(aggregated_out[-1], pd.DataFrame):
aggregated_out[-1][
("metadata", "itransform")] = isuccessor
aggregated_out[-1][
("metadata", "idate")] = idate
if transform in transfor2print or print_all:
debug(
f"{transform} / {di} / {successor} / {successor_di} / {idate} / {len(aggregated_out[-1])}")
# Now aggregate sparse or sum arrays
if is_sparse:
tmp_outputs[trid][di] = pd.concat(aggregated_out)
# Stores idates and issuccessors in an unalterable place
if 'aggregation_index' not in transf_mapper:
transf_mapper['aggregation_index'] = {}
if trid not in transf_mapper['aggregation_index']:
transf_mapper['aggregation_index'][trid] = {}
if ddi not in transf_mapper['aggregation_index'][trid]:
transf_mapper['aggregation_index'][trid][ddi] = {}
if transform_mode == 'adj':
transf_mapper['aggregation_index'][trid][ddi][di] = {
"idate": copy.deepcopy(
tmp_outputs[trid][di][("metadata", "idate")].values
),
"itransform": copy.deepcopy(
tmp_outputs[trid][di][(
"metadata", "itransform")].values
)
}
else:
if np.all([d == {} for d in aggregated_out]):
tmp_outputs[trid][di] = {}
continue
to_sum = 0
other_outputs = {}
for a in aggregated_out:
# Update other_outputs by attributes other than adj_out
for attr in a:
if attr == "adj_out":
to_sum += a["adj_out"]
else:
other_outputs[attr] = a[attr]
tmp_outputs[trid][di] = {
"adj_out": to_sum,
"original_outputs": copy.deepcopy(aggregated_out),
**other_outputs
}
if transform in transfor2print or print_all:
debug(f'aggreget_inout: target_out: {transform} / {ddi} / {tmp_outputs}')
if 'aggregation_index' not in transf_mapper:
debug(f'aggregation_index not in transf_mapper for {transform}')
elif trid not in transf_mapper['aggregation_index']:
debug(f'{trid} not in aggregation_index for {transform}')
elif ddi not in transf_mapper['aggregation_index'][trid]:
debug(f'{ddi} not in aggregation_index for {transform} / {trid}')
else:
for trid in tmp_outputs:
for di in tmp_outputs[trid]:
if di not in transf_mapper['aggregation_index'][trid][ddi]:
debug(f'{di} not in transf_mapper for {transform} / {trid} / {ddi}')
continue
debug(
f"aggregation index: idate {transf_mapper['aggregation_index'][trid][ddi][di]['idate'].shape}")
debug(
f"aggregation index: itransform {transf_mapper['aggregation_index'][trid][ddi][di]['itransform'].shape}")
return tmp_inputs, tmp_outputs
[docs]
def deaggregate_inout(
transform, transform_mode,
transform_onlyinit,
ddi,
tmp_datastore,
mapper,
check_transforms=False
):
"""Spread a transform's collapsed inputs/outputs back out to precursors/successors.
This is the inverse of :func:`aggregate_inout`: it takes the
already-computed ``tmp_datastore["inputs"]``/``["outputs"]`` for one
sub-simulation (``ddi``) of ``transform``, and re-nests each
``{trid: {date: <data>}}`` entry back under the relevant precursor
(for inputs) or successor (for outputs) and their own sub-simulation
date, so that each neighbour transform finds its expected data
shape in the shared datastore.
For outputs backed by a :class:`pandas.DataFrame` (sparse data), the
per-successor/per-date split uses the ``aggregation_index`` recorded
by :func:`aggregate_inout` (falling back to the current
successor/date index if unavailable). For dict-based (non-sparse)
outputs, the ``original_outputs`` saved during aggregation is used to
restore each successor's individual ``adj_out`` contribution.
Args:
transform: The name of the transform currently being processed.
transform_mode: Either ``"fwd"`` or ``"adj"``, the direction
currently being executed.
transform_onlyinit: Whether the transform is only being
initialized (no actual data processing).
ddi: The sub-simulation index being processed.
tmp_datastore: This transform's ``{"inputs": ..., "outputs": ...}``
datastore for sub-simulation ``ddi``, holding the collapsed
``{trid: {date: <data>}}`` entries to spread back out.
mapper: Dictionary mapping each transform id to its precursors/
successors/subsimus metadata.
check_transforms: Whether to compute the ``delta_adj_out`` debug
diagnostic (difference between the adjoint output before and
after aggregation) for precursor inputs.
Returns:
tuple: ``(tmp_inputs, tmp_outputs)``, the input/output datastores
re-nested per precursor/successor and their own sub-simulation date.
"""
# Get variables from mapper
transf_mapper = mapper[transform]
precursors = transf_mapper["precursors"]
successors = transf_mapper["successors"]
# Deal with inputs
tmp_inputs = tmp_datastore["inputs"]
if transform in transfor2print or print_all:
debug(f'deaggreget_inout: origin_in: {transform} / {ddi} / {tmp_inputs}')
for trid in transf_mapper["subsimus"][ddi]["inputs"]:
trid_dict = {}
trid_inputs = tmp_inputs.get(trid, {})
# for di in transf_mapper["subsimus"][ddi]["inputs"][trid]:
for di in trid_inputs:
di_inputs = trid_inputs[di]
di_dict = {}
for precursor in precursors.get(trid, []):
precursor_dict = {}
subsimus_in = mapper[precursor]["subsimus"]
for precursor_di in subsimus_in:
# Skip if subsimu not in precursor subsimulations
if di not in subsimus_in[precursor_di]['outputs'].get(trid, []):
continue
precursor_dict[precursor_di] = di_inputs
# if type(di_inputs) == pd.core.frame.DataFrame:
# precursor_dict[precursor_di] = di_inputs
# else:
# precursor_dict[precursor_di] = di_inputs.get(
# precursor, di_inputs)
# Now compute delta_adj_out for check_transforms
if check_transforms:
if "adj_out_original" in di_inputs:
precursor_dict[precursor_di]["delta_adj_out"] = (
di_inputs["adj_out"] -
di_inputs["adj_out_original"]
)
del di_inputs["adj_out_original"]
di_dict[precursor] = precursor_dict
trid_dict[di] = di_dict
tmp_inputs[trid] = trid_dict
if transform in transfor2print or print_all:
debug(f'deaggreget_inout: target_in: {transform} / {ddi} / {tmp_inputs}')
# Deal with outputs
tmp_outputs = tmp_datastore["outputs"]
if transform in transfor2print or print_all:
debug(f'deaggreget_inout: origin_out: {transform} / {ddi} / {tmp_outputs}')
if not transform_onlyinit and transform_mode == mode2print and open_debug:
print("BBB")
print(__file__)
import code
code.interact(local=dict(locals(), **globals()))
for trid in transf_mapper["subsimus"][ddi]["outputs"]:
trid_dict = {}
trid_outputs = tmp_outputs.get(trid, {})
subsimus_out = transf_mapper["subsimus"][ddi]["outputs"][trid]
trid_successors = successors.get(trid, [])
aggreg_index = 0
for di in subsimus_out:
di_dict = {}
# Skip if no outputs for this sub-simulation
if di not in trid_outputs:
trid_dict[di] = di_dict
continue
di_outputs = trid_outputs[di]
for isuccessor, successor in enumerate(trid_successors):
successor_dict = {}
subsimus_in = mapper[successor]["subsimus"]
for idate, successor_di in enumerate(subsimus_in):
# Skip if subsimu not in successor subsimulations
if di not in subsimus_in[successor_di]['inputs'].get(trid, []):
continue
aggreg_index += 1
successor_di_dict = {}
# Dataframe processing
if type(di_outputs) == pd.core.frame.DataFrame:
# Use information from aggregation index if any
ref_idates = transf_mapper.get('aggregation_index', {}).get(trid, {}).get(
ddi, {}).get(di, {}).get("idate", idate * np.ones(len(di_outputs)))
ref_itransforms = transf_mapper.get('aggregation_index', {}).get(trid, {}).get(
ddi, {}).get(di, {}).get("itransform", isuccessor * np.ones(len(di_outputs)))
mask = (ref_itransforms == isuccessor) & (
ref_idates == idate)
successor_di_dict = di_outputs.loc[mask]
# Array processing
else:
if successor in di_outputs:
raise CifError("Does this ever happen?")
successor_di_dict = {
v: di_outputs[v] for v in di_outputs
if v != "original_outputs"
}
# Replacing original adj_out before aggregation
# If available
if "adj_out" in successor_di_dict \
and "original_outputs" in di_outputs:
original_outputs = di_outputs["original_outputs"][aggreg_index - 1]
if "adj_out" in original_outputs:
successor_di_dict["adj_out"] = copy.deepcopy(
original_outputs["adj_out"])
else:
successor_di_dict["adj_out"] = \
0. * successor_di_dict["adj_out"]
successor_dict[successor_di] = successor_di_dict
di_dict[successor] = successor_dict
trid_dict[di] = di_dict
tmp_outputs[trid] = trid_dict
if transform in transfor2print or print_all:
debug(f'deaggreget_inout: target_out: {transform} / {ddi} / {tmp_outputs}')
return tmp_inputs, tmp_outputs
