Source code for mrfmsim.shortcut

"""Shortcuts

The shortcut should work for both the Model and the Experiment classes.
"""

from mmodel.modifier import loop_input
from mmodel.utility import modelgraph_signature, modelgraph_returns
from mrfmsim import Node
from networkx.utils import nodes_equal
from mrfmsim.modifier import parse_fields, print_output, print_inputs
import networkx as nx


def print_shortcut(model, stdout_format_list, **pargs):
    """Shortcut to add print modifiers to nodes.

    :param model: mrfmsim experiment
    :param list stdout_format_list: list of format strings
        each format string should only contain one variable.

    If the variable is not in the inputs or outputs of the nodes,
    the variable is ignored. The printout prioritizes inputs in a
    node. If a parameter is a node input and output, the first
    occurrence is used (when it is the output of the node). The
    decision to make all output the same setting (pargs for print())
    is to simplify the implementation. If the user wants to have
    customized printout behavior, use the regular modifier for each
    node.
    """

    G = model.graph
    sig = modelgraph_signature(G)

    input_dict = {}
    output_dict = {}

    for format_string in stdout_format_list:
        field = parse_fields(format_string)[0]
        if field in sig.parameters:
            input_dict[field] = format_string
        else:
            output_dict[field] = format_string

    # input algorithm
    # find the first node that contains the input
    # use the default order of the nodes
    for input_, format_string in input_dict.items():
        for k in G.nodes:
            node_obj = G.get_node_object(k)
            if input_ in node_obj.inputs:
                node_mods = node_obj.modifiers
                G = G.edit_node(
                    k, modifiers=node_mods + [print_inputs(format_string, **pargs)]
                )
                break

    # output algorithm
    # find the output that outputs the parameter
    for output, format_string in output_dict.items():
        for k, v in nx.get_node_attributes(G, "output").items():
            if v == output:
                node_obj = G.get_node_object(k)
                node_mods = node_obj.modifiers
                G = G.edit_node(
                    k, modifiers=node_mods + [print_output(format_string, **pargs)]
                )
                break

    return model.edit(graph=G, returns=model.returns)


def loop_shortcut(model, parameter: str, name=None):
    """Shortcut to add a loop to a subgraph.

    The parameter needs to be in the graph signature. Otherwise,
    an exception is raised. For parameters that are not from the
    graph (generated by modifiers), use regular loops or change
    modifiers.

    :param model: executable model
    :param str parameter: loop parameter
    :param str name: name of the new model, defaults to old model name.
    :return: a new model with looped parameter
    """
    # check if the parameter is in the signature
    if parameter not in model.signature.parameters:
        raise Exception(f"Invalid shortcut: {repr(parameter)} is not a model input.")

    G = model.graph
    name = name or model.name
    modifiers = model.modifiers
    components = model.components

    loop_mod = loop_input(parameter)

    # If the parameter is in the signature but not in the graph.
    # This is due to the signature modifier on the model.
    # Use regular loop or change the modifiers.

    if parameter not in modelgraph_signature(G).parameters:
        raise Exception(
            f"{repr(parameter)} is not included in the graph."
            f" Use regular loop or change the modifier."
        )

    H = G.subgraph(inputs=[parameter])

    if nodes_equal(G.nodes, H.nodes):
        modifiers = [*modifiers, loop_mod]

    elif len(H.nodes()) == 1:
        node = list(H.nodes)[0]
        # if the looped node is only one node
        # add loop modifier to node attribute
        node_modifiers = H.nodes[node]["node_object"].modifiers
        node_modifiers = [*node_modifiers, loop_mod]
        G = G.edit_node(node, modifiers=node_modifiers)

    else:  # if there is more than one node
        sub_name = f"subnode_{parameter}"
        sub_des = (
            f"Submodel generated by loop_shortcut for parameter {repr(parameter)}."
        )
        H_returns = modelgraph_returns(H)  # empty list if the nodes don't have returns
        # if not H_returns:
        #     output = None
        # else:
        #     output = ", ".join(H_returns)
        output = ", ".join(H_returns)
        # create the model and substitute the subgraph
        looped_node = type(model)(
            f"submodel_{parameter}", H, model.handler, doc=sub_des
        )

        G = G.replace_subgraph(
            H, Node(sub_name, looped_node, output=output, modifiers=[loop_mod])
        )

    # reset returns to recreate the returns from the subgraphs
    return model.edit(name=name, graph=G, modifiers=modifiers, components=components)