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): _description_
tmp_outputs (dict): _description_
transform_mode (str): _description_
transform_onlyinit (bool): _description_
mapper (dict): _description_
check_transforms (bool, optional): _description_. Defaults to False.
Raises:
Exception: _description_
TypeError: _description_
Returns:
_type_: _description_
"""
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
):
# 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
