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
):

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