Source code for mmodel.shortcut

"""Shortcuts

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

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




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

    :param model: model instance
    :param list stdout_format_list: list of format strings.
        Each format string should only contain one variable.
    :param pargs: keyword arguments for the print function

    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 outputs use 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.signature.parameters:
                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 (signature modified by modifiers), use regular loops or change
    the modifiers.

    :param model: executable model
    :param str parameter: loop parameter
    :param str name: name of the new model; defaults to the old model name.
    :return: a new model with the looped parameter
    """

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

    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 a regular loop or change the modifiers.

    if parameter not in modelgraph_signature(G).parameters:
        raise Exception(
            f"{repr(parameter)} is not a graph model parameter,"
            f" check parameter input or use a regular loop."
        )

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

    # if the nodes are equal, add the loop modifier to the model
    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

        output = ", ".join(H_returns) or None
        # this is necessary when the subgraph has no returns

        # create the model and substitute the subgraph
        looped_node = type(model)(
            f"submodel_{parameter}", H, model.handler, doc=sub_des
        )

        # use the graph default node type
        node_cls = G.graph["node_type"]
        G = G.replace_subgraph(
            H, node_cls(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)