Source code for pycif.plugins.obsoperators.standard.transforms.do_transforms
import copy
import datetime
import glob
import os
import re
import tracemalloc
from logging import debug
from typing import List, Tuple
import numpy as np
import pandas as pd
from .....utils import path
from .utils import (
aggregate_inout,
check_adjtltest,
clean_memory,
deaggregate_inout,
dump_debug,
fetch_inputs_outputs,
submit_and_kill,
)
[docs]
def input_output_msg(trid_list: List[Tuple[str, str]]) -> str:
"""Format a list of tracer IDs as a human-readable multi-line string.
Groups parameters by component and produces output of the form::
- component1: param1, param2
- component2: param1
Args:
trid_list (list[tuple[str, str]]): list of ``(component, parameter)``
tracer-ID pairs to format.
Returns:
str: multi-line string, one component per line, suitable for logging.
"""
trid_dict = {}
for component, parameter in trid_list:
if component in trid_dict:
trid_dict[component].append(parameter)
else:
trid_dict[component] = [parameter]
msg = []
for component, param_list in trid_dict.items():
comp_msg = f" - {component}: "
# If param_list does not contains only an empty string
if len(param_list) > 1 or param_list[0]:
comp_msg = comp_msg + ", ".join(param_list)
msg.append(comp_msg)
return '\n'.join(msg)
[docs]
def do_transforms(
self,
transform_pipe,
mapper,
controlvect,
obsvect,
mode,
rundir,
workdir,
do_simu=True,
onlyinit=False,
check_transforms=False,
adj_test_threshold=10,
save_debug=False,
ignore_exceptions=False,
ref_fwd_dir="",
dump_metadata_only=False,
**kwargs
):
"""Execute the full ordered transform pipeline for a single operator call.
Iterates over all ``(date, transform, direction)`` entries in
``self.period_order_fwd`` (or ``self.period_order_adj`` in adjoint mode)
and applies each transform in sequence. Key responsibilities:
* **Datastore initialisation** — builds a nested
``transform_pipe.datastore[transform][date]`` dictionary on the first
call to track intermediate inputs/outputs shared between transforms.
* **Restart support** — reads ``rundir/finished_transforms.txt`` to skip
transforms already completed in a previous interrupted run when
``self.autorestart`` is enabled.
* **Input/output routing** — uses
:func:`~.utils.fetch_inoutputs.fetch_inputs_outputs` and
:func:`~.utils.aggreg_deaggreg_inout.aggregate_inout` to gather inputs
from precursor datastores, and
:func:`~.utils.aggreg_deaggreg_inout.deaggregate_inout` to redistribute
outputs to successor datastores.
* **Approximate operator** — when ``self.approx_operator`` is set
(parallel mode), transforms outside the segment window execute in
dry-run (*onlyinit*) mode only.
* **Memory monitoring** — tracks peak memory with :mod:`tracemalloc`
when ``self.monitor_memory`` is enabled.
* **Memory cleaning** — releases unused datastore entries after each
transform when ``self.clean_memory`` is enabled.
* **Autokill / restart** — kills the job and resubmits if the elapsed
wall-clock time exceeds ``self.autokill_time``.
* **Adjoint / TL test** — when ``check_transforms`` is ``True``, saves
copies of each transform's in/outputs and calls
:func:`~.utils.check_adjtltest.check_adjtltest` at the end of the
adjoint pass.
Args:
self (ObsOperator): the obs-operator plugin instance.
transform_pipe: the :class:`~pycif.utils.classes.transforms.Transform`
pipeline object populated by :func:`init_transform`.
mapper (dict): the pipeline mapper dictionary.
controlvect (ControlVect): control-vector object.
obsvect (ObsVect): observation-vector object.
mode (str): execution mode — one of ``'fwd'``, ``'tl'``, or
``'adj'``.
rundir (str): the run sub-directory for this operator call.
workdir (str): parent working directory.
do_simu (bool, optional): if ``False``, skip actual transform
execution (used internally for dry runs). Defaults to ``True``.
onlyinit (bool, optional): if ``True``, run all transforms in
initialisation / dry-run mode only. Defaults to ``False``.
check_transforms (bool, optional): if ``True``, validate each
transform's adjoint / TL identity. Defaults to ``False``.
adj_test_threshold (float, optional): relative tolerance used by the
adjoint test. Defaults to ``10``.
save_debug (bool, optional): if ``True``, dump the inputs and outputs
of each transform to *rundir* for post-run inspection.
Defaults to ``False``.
ignore_exceptions (bool, optional): if ``True``, non-fatal errors
inside individual transforms are swallowed and execution
continues. Defaults to ``False``.
ref_fwd_dir (str, optional): path to the reference forward run
directory, forwarded to each transform for adjoint input
location. Defaults to ``""``.
dump_metadata_only (bool, optional): passed through to
:func:`~.utils.dump_debug.dump_debug` when ``save_debug`` is
``True``. If ``True``, writes lightweight plain-text metadata
summaries (dimensions, min/max, NaN presence) instead of full
NetCDF/datastore files. Greatly reduces wall-time and disk
overhead while still allowing data-flow inspection.
Defaults to ``False``.
**kwargs: additional keyword arguments (ignored).
"""
# Initialize the datastore with all precursors/successors and
# sub-simulations if not already done
if not hasattr(transform_pipe, "datastore"):
transform_pipe.datastore = {}
for transform in transform_pipe.attributes:
transf_mapper = mapper[transform]
inputs = transf_mapper["inputs"]
outputs = transf_mapper["outputs"]
precursors = transf_mapper["precursors"]
successors = transf_mapper["successors"]
subsimus = transf_mapper["subsimus"]
transform_pipe.datastore[transform] = {
ddi: {
"inputs": {
trid: {di: {precursor: {}
for precursor in precursors.get(trid, [])}
for di in subsimus[ddi]["inputs"][trid]}
for trid in subsimus[ddi]["inputs"]},
"outputs": {
trid: {di: {successor: {}
for successor in successors.get(trid, [])}
for di in subsimus[ddi]["outputs"][trid]}
for trid in subsimus[ddi]["outputs"]}}
for ddi in subsimus
}
missingperiod = False
# Keep in memory what directories have been recently created
# It is used as an information to re-compute or not transforms
created_directories = {}
# Start tracemalloc process
if self.monitor_memory:
tracemalloc.start()
with open(f"{rundir}/memory_usage.txt", "w") as f:
pass
# Reload transforms already computed
finished_transforms = []
file_transforms = os.path.join(rundir, "finished_transforms.txt")
if os.path.isfile(file_transforms) and self.autorestart:
if os.path.getsize(file_transforms) > 0:
finished_transforms = pd.read_csv(
file_transforms, sep=";", header=None,
infer_datetime_format=True, parse_dates=[0]
)
finished_transforms = [
(row[1][0].to_pydatetime(), row[1][1], row[1][2])
for row in finished_transforms.iterrows()]
# Duplicate previous file if any
if os.path.isfile(file_transforms):
list_transform_files = glob.glob(f"{file_transforms}.*")
if not list_transform_files:
max_index = 0
else:
list_matches = [re.search(file_transforms + '\\.(\\d{3})', f)
for f in list_transform_files]
if [m for m in list_matches if m is not None] == []:
max_index = 0
else:
max_index = max([int(m.group(1)) for m in list_matches
if m is not None])
os.rename(file_transforms, f"{file_transforms}.{max_index + 1:03}")
# Initialize file from scratch
open(file_transforms, "w").close()
# Re-ordering the tranformations if not in fwd mode
period_order = self.period_order_fwd[:]
if mode == "adj":
period_order = self.period_order_adj[:]
# Loop over transforms
for ddi, transform, direction in period_order:
transf_mapper = mapper[transform]
transf = getattr(transform_pipe, transform)
# Adapt mode depending on forward or backward in period_order
transform_mode = mode
transform_onlyinit = onlyinit
if (direction == "adjoint" and mode in ["fwd", "tl"]) \
or (direction == "forward" and mode == "adj"):
transform_onlyinit = True
transform_mode = "fwd" if mode == "adj" else "adj"
msg = [
f"Doing transform {transform}: {transf.plugin.name} in "
f"{transform_mode} mode (onlyinit = {transform_onlyinit}), "
f"period {ddi}",
"From inputs:",
input_output_msg(list(transf_mapper['inputs'])),
"To outputs:",
input_output_msg(list(transf_mapper['outputs'])),
]
# Some logging
debug('\n'.join(msg))
# Fetch datastore from relevant transformation datastores
# TODO: explicitly build the pipeline and contraining from which
# transform to extract; also allow cleaning un-used datastores to
# avoid memory leaks
inputs = transf_mapper["inputs"]
outputs = transf_mapper["outputs"]
precursors = transf_mapper["precursors"]
successors = transf_mapper["successors"]
subsimus = transf_mapper["subsimus"][ddi]
tmp_datastore = transform_pipe.datastore[transform][ddi]
# Fetch inputs/outputs from other transforms
tmp_inputs, tmp_outputs = fetch_inputs_outputs(
transform, ddi, transform_pipe, tmp_datastore,
subsimus, successors, precursors, transf_mapper,
mapper)
# Simplify the input/output dictionary in the case that there is no
# multiple precursor/successor for each trid
tmp_inputs, tmp_outputs = aggregate_inout(
transform, ddi,
tmp_inputs, tmp_outputs,
transform_mode, transform_onlyinit,
mapper,
check_transforms=check_transforms
)
# Update tmp_datastore
tmp_datastore["inputs"] = tmp_inputs
tmp_datastore["outputs"] = tmp_outputs
# Create sub directory if needed
# If needs to be created, it means that the simulation was not
# already done, thus should not reload from here
runsubdir = os.path.join(rundir, ddi.strftime("%Y-%m-%d_%H-%M"))
_, created = path.init_dir(runsubdir)
if runsubdir not in created_directories:
created_directories[runsubdir] = created
created = created_directories[runsubdir]
# Do simu if not already done
do_simu = (
(ddi, transform, direction) not in finished_transforms
or not self.autorestart
or missingperiod
)
missingperiod = do_simu if not transform_onlyinit else missingperiod
# Save outputs if save_debug activated
if save_debug and not transform_onlyinit:
dump_debug(
transform, transf_mapper, tmp_datastore, runsubdir, ddi,
entry="inputs" if transform_mode in ["fwd", "tl"] else "outputs",
dump_metadata_only=dump_metadata_only
)
# Informs the transform about the reference directory if not already set
if transform_mode in ["adj"] and not transform_onlyinit:
if not hasattr(transf, "adj_refdir"):
transf.adj_refdir = self.adj_refdir
elif transform_mode in ["fwd", "tl"] and not transform_onlyinit:
if not hasattr(transf, "adj_refdir"):
transf.adj_refdir = rundir
# Do the transform
try:
approx = False
overlap = False
if hasattr(self, "approx_operator"):
approx_di = self.approx_operator.datei
approx_df = self.approx_operator.datef
approx_overlap = self.approx_operator.overlap
approx = not approx_di <= ddi < approx_df
overlap = \
approx_di <= ddi \
< approx_di + pd.to_timedelta(approx_overlap)
# Special case if ddi == datei
if ddi == self.datei:
overlap = False
apply_transform = transf.forward if transform_mode in ["fwd", "tl"] \
else transf.adjoint
apply_transform(
tmp_datastore,
controlvect,
obsvect,
transf_mapper,
ddi,
ddi,
transform_mode,
runsubdir,
workdir,
do_simu=do_simu,
onlyinit=transform_onlyinit,
save_debug=save_debug,
approx_transf=approx,
overlap=overlap,
ref_fwd_dir=ref_fwd_dir,
check_transforms=check_transforms
)
except Exception as e:
if not ignore_exceptions:
raise e
else:
# Raise the error if outputs where required in init_inputs
list_outputs = tmp_datastore["outputs"].keys()
list_outputs_components = set([c[0] for c in list_outputs])
if hasattr(self, "init_inputs"):
init_inputs = self.init_inputs
for cmp in init_inputs.components.attributes:
if cmp not in list_outputs_components:
continue
list_params = getattr(init_inputs.components, cmp)
if list_params == []:
raise e
list_outputs_params = \
set(c[1] for c in list_outputs if c[0] == cmp)
if np.any(np.isin(list_outputs_params, list_params)):
raise e
# Keeps the running directory in memory for later adjoint simulations
if transform_mode in ["fwd", "tl"] and not transform_onlyinit:
transf.adj_refdir = f"{runsubdir}/../"
# Save outputs if save_debug activated
if save_debug:
entry = "outputs" if transform_mode in ["fwd", "tl"] else "inputs"
dump_debug(
transform, transf_mapper, tmp_datastore, runsubdir, ddi,
entry=entry, transform_onlyinit=transform_onlyinit,
dump_metadata_only=dump_metadata_only
)
# Redistribute the datastore accounting for successor/precursors
# and inputs/outputs sub-simulations
# There is a hiccup for sparse data when debugging as transform IDs make groups
# in the dataframe, similarly to a dictionary, which is not the expected
# behaviour.
tmp_inputs, tmp_outputs = deaggregate_inout(
transform,
transform_mode,
transform_onlyinit,
ddi, tmp_datastore,
mapper,
check_transforms=check_transforms
)
# Save the datastore for later
transform_pipe.datastore[transform][ddi] = {
"inputs": tmp_inputs,
"outputs": tmp_outputs
}
# Check if individual transformations pass the test of the adjoint
if check_transforms and not transform_onlyinit:
if transform_mode in ["fwd", "tl"]:
if not hasattr(self, "data_tl"):
self.data_tl = {transform: {}}
if not transform in self.data_tl:
self.data_tl[transform] = {}
self.data_tl[transform][ddi] = copy.deepcopy(
transform_pipe.datastore[transform][ddi])
elif transform_mode == "adj":
if not hasattr(self, "data_adj"):
self.data_adj = {transform: {}}
if not transform in self.data_adj:
self.data_adj[transform] = {}
self.data_adj[transform][ddi] = copy.deepcopy(
transform_pipe.datastore[transform][ddi])
# Monitor memory usage
if self.monitor_memory:
current, peak = tracemalloc.get_traced_memory()
debug(f"Current memory usage is {current / 1024 ** 2}MB; "
f"Peak was {peak / 1024 ** 2}MB")
with open(f"{rundir}/memory_usage.txt", "a") as f:
f.write(f"{transform} / {ddi}, {direction} : "
f"{current / 1024 ** 2}MB / {peak / 1024 ** 2}MB\n")
# Clean the datastore if no successor/precursor use it anymore
if self.clean_memory:
transform_pipe = clean_memory(
self, transform_pipe, mapper, period_order,
ddi, transform, direction, mode, transform_onlyinit)
# Keep in memory that the transform has been properly processed
finished_transforms.append((ddi, transform, direction))
with open(file_transforms, "a") as f:
f.write(f"{ddi};{transform};{direction}\n")
# Kill the job a re-submit one if reaching the time limit
if hasattr(self, "autokill_time"):
start_time = self.reference_instances["reference_setup"].simu_start_time
elapsed_time = datetime.datetime.now() - start_time
if elapsed_time > pd.to_timedelta(self.autokill_time):
submit_and_kill(self)
# Clean datastore as not used anymore
del transform_pipe.datastore
# Check test of the adjoint for individual transforms
if not (mode == "adj" and check_transforms):
return
check_adjtltest(self, self.period_order_fwd, mapper, transform_pipe)