Source code for pycif.plugins.obsoperators.standard.transforms.utils.dask.individual_transform

import numpy as np
import pandas as pd
from logging import debug, warning, info
from typing import List, Tuple
import threading
import os
from .......utils import path
from .......utils.check import batched_logging
from .. import aggregate_inout, deaggregate_inout




def do_individual_transform(*args, **kwargs):
    """Runs a single transform, batching its log output.

    Thin wrapper around _do_individual_transform: dask worker threads run
    many transforms concurrently, and without batching their debug/info
    calls interleave line-by-line with other threads' output. See
    pycif.utils.check.batched_logging.
    """
    with batched_logging():
        return _do_individual_transform(*args, **kwargs)



def _do_individual_transform(
    name, dask_datastore, aggregation_info,
    transform_pipe,
    datei_ref, rundir, workdir,
    controlvect, obsvect,
    mode="fwd", onlyinit=False,
    approx_operator=None,
    ignore_exceptions=False,
    init_inputs=None,
    check_transforms=False,
    do_simu=True,
    save_debug=False,
    ref_fwd_dir="",
    **kwargs
):
    """Run one transform for one sub-simulation as a Dask task.

    This is the unit of work scheduled by Dask (see
    :func:`..init_dask.add_transform`): given a transform id ``name =
    (ddi, transform, direction)``, gathers this transform's inputs and
    outputs from ``dask_datastore`` (the results of the precursor/
    successor tasks it depends on), reshapes them into the per-trid,
    per-date, per-neighbour layout expected by the transform, aggregates
    them (:func:`..aggregate_inout`), applies the transform's
    ``forward``/``adjoint`` method in a dedicated run sub-directory, then
    spreads the results back out per neighbour
    (:func:`..deaggregate_inout`) for downstream tasks to consume.

    If ``direction`` does not match the overall ``mode`` (e.g. an
    adjoint-direction node while running in forward mode), the transform
    is only initialized (``transform_onlyinit=True``) rather than
    actually executed, and its local mode is flipped accordingly.

    If ``ignore_exceptions`` is set, exceptions raised by the transform
    are swallowed unless the transform was expected to produce outputs
    required by ``init_inputs``, in which case they are re-raised.

    Args:
        name: ``(ddi, transform, direction)`` tuple identifying the task.
        dask_datastore: Dictionary of already-computed results for other
            ``(ddi, transform, direction)`` tasks, keyed the same way,
            each holding ``"inputs"``/``"outputs"`` datastores.
        aggregation_info: Unused here, kept for interface consistency with
            the Dask graph construction in :func:`..init_dask.add_transform`.
        transform_pipe: Object exposing the transform instances (by
            attribute name) and their shared ``mapper``.
        datei_ref: Reference start date of the run, used to special-case
            the very first sub-simulation with respect to
            ``approx_operator`` overlap handling.
        rundir: Base run directory under which a per-``ddi`` sub-directory
            is created for this transform's execution.
        workdir: Working directory passed through to the transform.
        controlvect: Control vector object passed through to the transform.
        obsvect: Observation vector object passed through to the transform.
        mode: Overall pipe execution mode: ``"fwd"``, ``"tl"`` or ``"adj"``.
        onlyinit: Whether to only initialize (not execute) the transform.
        approx_operator: Optional dict with ``"datei"``, ``"datef"`` and
            ``"overlap"`` describing the period over which an approximate
            operator should be used instead of the full one.
        ignore_exceptions: Whether to swallow exceptions raised by the
            transform unless required outputs are missing.
        init_inputs: Object describing which components/parameters are
            required at initialization, used to decide whether a
            swallowed exception must be re-raised.
        check_transforms: Whether to keep debug information needed by the
            adjoint/tangent-linear consistency test.
        do_simu: Whether to actually run the underlying model simulation.
        save_debug: Whether to dump debug datasets for this transform.
        ref_fwd_dir: Path to the reference forward run directory, used by
            adjoint/tangent-linear transforms that need forward outputs.
        **kwargs: Additional keyword arguments, unused here.

    Returns:
        dict: ``{"main": {"inputs": tmp_inputs, "outputs": tmp_outputs},
        "meta": ...}`` holding the transform's inputs/outputs re-nested
        per neighbour transform, ready to be consumed by downstream Dask
        tasks, plus any metadata attached to the datastore.
    """

    # print("AAAAAAAAA", name)
    # print(dask_datastore)

    # Get ID information
    ddi, transform, direction = name
    transf = getattr(transform_pipe, transform)
    mapper = transform_pipe.mapper
    transf_mapper = mapper[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"

    # Get dask thread ID
    tid = threading.get_ident()          # numeric thread ID
    name = threading.current_thread().name  # e.g. "ThreadPoolExecutor-0_3"

    # Some logging
    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'])),
        "Dask thread info:",
        f"  - thread ID: {tid}",
        f"  - thread name: {name}"
    ]
    debug('\n'.join(msg))

    # Fetch data
    fetch_inputs = {
        d: dask_datastore[d]["outputs"] if d[2] == "forward"
        else dask_datastore[d]["inputs"]
        for d in dask_datastore
        if d[2] == "forward" or (d[0] == ddi and d[1] == transform)}
    fetch_outputs = {
        d: dask_datastore[d]["inputs"] if d[2] == "adjoint"
        else dask_datastore[d]["outputs"]
        for d in dask_datastore
        if d[2] == "adjoint" or (d[0] == ddi and d[1] == transform)}

    # Rearranging dictionaries for inputs
    input_trids = list({key for subdict in fetch_inputs.values() for key in subdict})
    tmp_inputs = {k: {} for k in transf_mapper['inputs']}
    for trid_in in input_trids:
        if trid_in not in transf_mapper['subsimus'][ddi]['inputs']:
            continue
        try:
            tmp_inputs[trid_in] = {
                k: {} for k in transf_mapper['subsimus'][ddi]['inputs'][trid_in]}
        except:
            # print('qqqqq')
            print(__file__)
            import code
            code.interact(local=dict(locals(), **globals()))
        input_di = list({
            key for subdict in fetch_inputs.values()
            for key in subdict.get(trid_in, {})
        })
        for di_in in input_di:
            for d in fetch_inputs:
                if trid_in not in fetch_inputs[d]:
                    continue

                if di_in not in fetch_inputs[d][trid_in]:
                    continue

                if transform not in fetch_inputs[d][trid_in][di_in]:
                    continue

                if ddi not in fetch_inputs[d][trid_in][di_in][transform]:
                    continue

                if d[1] not in tmp_inputs[trid_in][di_in]:
                    tmp_inputs[trid_in][di_in][d[1]] = {}

                try:
                    tmp_inputs[trid_in][di_in][d[1]][d[0]] = \
                        fetch_inputs[d][trid_in][di_in][transform][ddi]
                except:
                    print('AAAAAAAAAAAAAA')
                    print(__file__)
                    import code
                    code.interact(local=dict(locals(), **globals()))
    # Rearranging dictionaries for outputs
    output_trids = list({key for subdict in fetch_outputs.values() for key in subdict})
    tmp_outputs = {k: {} for k in transf_mapper['outputs']}
    for trid_out in output_trids:
        if trid_out not in transf_mapper['subsimus'][ddi]['outputs']:
            continue

        try:
            tmp_outputs[trid_out] = {
                k: {} for k in transf_mapper['subsimus'][ddi]['outputs'][trid_out]}
        except:
            print('pppppp')
            print(__file__)
            import code
            code.interact(local=dict(locals(), **globals()))

        output_di = list({
            key for subdict in fetch_outputs.values()
            for key in subdict.get(trid_out, {})
        })
        for di_out in output_di:
            tmp_outputs[trid_out][di_out] = tmp_outputs[trid_out].get(di_out, {})
            for d in fetch_outputs:
                if trid_out not in fetch_outputs[d]:
                    continue

                if di_out not in fetch_outputs[d][trid_out]:
                    continue

                if transform not in fetch_outputs[d][trid_out][di_out]:
                    continue

                if ddi not in fetch_outputs[d][trid_out][di_out][transform]:
                    continue

                if d[1] not in tmp_outputs[trid_out][di_out]:
                    tmp_outputs[trid_out][di_out][d[1]] = {}

                try:
                    tmp_outputs[trid_out][di_out][d[1]][d[0]] = \
                        fetch_outputs[d][trid_out][di_out][transform][ddi]
                except:
                    print('BBBBBBBBBB')
                    print(__file__)
                    import code
                    code.interact(local=dict(locals(), **globals()))

    # print(__file__)
    # import code
    # code.interact(local=dict(locals(), **globals()))

    # Aggregate data
    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,
        "outputs": tmp_outputs
    }

    # print('YYYYYYYYYYYY')
    # print(fetch_inputs)
    # print(fetch_outputs)
    # print('TTTTTTTTTTTT')
    # print(tmp_datastore)

    # 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]

    # return {
    #     "main": {
    #         "inputs": f"main_inputs_{name}",
    #         "outputs": f"main_outputs_{name}"
    #     },
    #     "meta": f"meta_{name}"
    # }

    # Do the transform
    try:
        approx = False
        overlap = False
        if approx_operator is not None:
            approx_di = approx_operator['datei']
            approx_df = approx_operator['datef']
            approx_overlap = 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 == datei_ref:
                overlap = False

        transf = getattr(transform_pipe, transform)
        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 init_inputs is not None:
                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

    # print('HHHHHHHHHHHH')
    # print(tmp_datastore)
    # 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
    )

    # import logging
    # print('IIIIIIIIIIII')
    # print(
    #     logging.getLogger().level   # 10=DEBUG, 20=INFO, 30=WARNING
    # )

    # print('JJJJJJJJJJJJJ')
    # print(tmp_inputs)
    # print(tmp_outputs)
    # print('KKKKKKKKKKKK')

    return {
        "main": {"inputs": tmp_inputs, "outputs": tmp_outputs},
        "meta": tmp_datastore.get("metadata", {})
    }




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 ``(component, parameter)`` pairs by component and produces one
    line per component, listing its parameters (or nothing after the
    colon when the only parameter is the empty string).

    Args:
        trid_list: List of ``(component, parameter)`` tuples.

    Returns:
        str: The formatted, newline-joined description.
    """
    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)