Theory III: Computational Graphs

The previous notebooks have illustrated that if you have a set of pure functions with explicit interfaces, you can reason about the dependencies between them.

Such a set of functions thus defines a computational graph: nodes are the functions and the directed edges connecting them are the outputs and inputs of the functions.

We’ll re-use the functions from before and draw the computational graph:

Show code cell source

Hide code cell source

from math import sin, exp
import inspect


def Model1(a,b):
    
    # Complicated *pure* internal calculation
    x = a**2
    y = sin(b)
    z = exp(x * y)
    output = z / (1 + x)

    #--------- Outputs
    return output

    
def Model2(model1_output, c):
    
    # something else complicated and *pure*
    iteration_var = 0.531
    for i in range(1, 4):
        iteration_var  = c * iteration_var + model1_output

    #--------- Outputs
    return iteration_var


def Model3(model2_output, e, c, d):
    
    # something else also complicated and *pure*
    tolerance = 1e-6
    error = float('inf')
    x = model2_output

    while error > tolerance:
        old_x = x

        # Some nonlinear update
        x = (old_x - e) / (1 + abs(old_x - e))
        x = x + c * sin(x) + d
        error = abs(x - old_x)

    #--------- Outputs
    return x

def Model4(model1_output, f):
    return model1_output + f

def interface(f):
    return {f.__name__ : {"inputs" : list(inspect.signature(f).parameters.keys()), "outputs" : [f.__name__.lower() + "_output"]}}



import cottax.assemblies.jax as cx
from cottax.assemblies.jax import flat
from cottax.consumers.visualization.xdsm import PlainFormat
from cottax.consumers.visualization import draw

def to_graph(*models):
    return cx.absolute.declare(*(flat.function(m, m.__name__.lower() + "_output") for m in models)).graph

_ = draw(to_graph(Model1, Model2, Model3, Model4), formatter = PlainFormat(), boundary = True)

If an order exists in which the functions can be executed one-by-one, the computational graph is a directed acyclic graph (DAG). This notebook only considers these graphs; for cyclic graphs, see the next notebook.

Standard algorithms exist to determine a valid ordering of a DAG. To determine this order, we’ll use the networkx library. Below is some code to transfer our models to the networkx graph format.

Show code cell source

Hide code cell source

import networkx as nx


def models_to_graph(models):
    graph = nx.DiGraph()

    # Map each output to the model that produces it

    models = [interface(i) for i in models]
    
    output_to_model = {}

    for model_dict in models:
        for model_name, config in model_dict.items():
            graph.add_node(model_name)

            for output in config.get("outputs", []):
                output_to_model[output] = model_name

    # Create model -> model edges based on inputs
    for model_dict in models:
        for model_name, config in model_dict.items():

            for input_name in config.get("inputs", []):
                producer = output_to_model.get(input_name)

                if producer is not None:
                    # Several variables can join the same pair of models, so the edge
                    # keeps all of them rather than only the last one written.
                    if graph.has_edge(producer, model_name):
                        graph[producer][model_name]["inputs"].append(input_name)
                    else:
                        graph.add_edge(producer, model_name, inputs=[input_name])

    return graph

Then, we can compute a valid ordering [1]

nx_digraph = models_to_graph([Model4, Model3, Model2, Model1])

list(nx.topological_sort(nx_digraph))

['Model1', 'Model4', 'Model2', 'Model3']

The basic evaluation that is valid for any combination of models that form a DAG is then as follows (we also add a check that the input data structure matches what we’d expect):

Show code cell source

Hide code cell source

def external_inputs(models : set):
    # Get what base inputs are required and what outputs are produced by the models
    total_inputs = set()
    model_outputs= set()

    for model in models:
        interface_i = interface(model)[model.__name__]

        total_inputs.update(interface_i["inputs"])
        model_outputs.update(interface_i["outputs"])

    # These are the inputs that have to be defined externally: no model produces them!
    external_inputs = total_inputs - model_outputs

    return external_inputs

def TotalModel(data_structure, models : set):
    ds_start = data_structure.copy() # Keeping it pure: the original data structure is not modified, but a copy is used instead.

    # Let's also check that the provided data structure has all the required external inputs for the given models:
    if not set(external_inputs(models)) == set(ds_start.keys()):
        raise ValueError(f"The provided data structure keys {set(ds_start.keys())}" + 
                         f"do not match the required external inputs {set(external_inputs(models))} for the given models.")

    # Create the graph and order it:
    nx_digraph = models_to_graph(list(models))
    str_to_model = {model.__name__: model for model in models}
    models_in_order = list(nx.topological_sort(nx_digraph))

    # Rest the same as before:
    for model_name in models_in_order:
        model = str_to_model[model_name]
        # Get the interface for the current model
        interface_i = interface(model)[model.__name__]
        
        # Prepare the inputs for the current model based on its interface
        inputs = {input_name: ds_start[input_name] for input_name in interface_i["inputs"]}
        
        # Call the model with the prepared inputs
        output = model(**inputs)
        
        # Update the data structure with the output of the current model
        ds_start.update({interface_i["outputs"][0]: output})
    return ds_start

ds_start = {"a": 1.0, "b": 0.5, "c": 2.0, "d": 0.1, "e": 0.2, "f" : 3.0}

TotalModel(ds_start, {Model1, Model2, Model3, Model4})

{'a': 1.0,
 'b': 0.5,
 'c': 2.0,
 'd': 0.1,
 'e': 0.2,
 'f': 3.0,
 'model1_output': 0.8075731482210419,
 'model4_output': 3.8075731482210418,
 'model2_output': 9.901012037547295,
 'model3_output': 1.9113901980206618}

and it will error if we put something there that is not an external input for the given models:

ds_wrong = {"a": 1.0, "b": 0.5, "c": 2.0, "d": 0.1, "e": 0.2, "f" : 3.0, "model1_output" : 4.0}

try:
    TotalModel(ds_wrong, {Model1, Model2, Model3, Model4})
except ValueError as ve:
    print(f"ValueError: {ve}")

ValueError: The provided data structure keys {'model1_output', 'd', 'a', 'c', 'f', 'b', 'e'}do not match the required external inputs {'f', 'b', 'd', 'e', 'a', 'c'} for the given models.

Great! Let’s use an alternative for Model4 that uses Model3 output instead of Model1 output[3]

def Model4(a, model3_output):
    return a + model3_output



nx_digraph = models_to_graph([Model4, Model3, Model2, Model1])
print(list(nx.topological_sort(nx_digraph)))

['Model1', 'Model2', 'Model3', 'Model4']

The topological order has immediately changed to reflect this new dependency!

We can also immediately run it ( after removing the now-useless f input):

# remove `f` now!
ds_start = {"a": 1.0, "b": 0.5, "c": 2.0, "d": 0.1, "e": 0.2} 
TotalModel(ds_start, {Model1, Model2, Model3, Model4})

{'a': 1.0,
 'b': 0.5,
 'c': 2.0,
 'd': 0.1,
 'e': 0.2,
 'model1_output': 0.8075731482210419,
 'model2_output': 9.901012037547295,
 'model3_output': 1.9113901980206618,
 'model4_output': 2.911390198020662}

Alternative models

We can now solve the alternative problem. Let’s define a way to switch between two models:

def make_model4(model4_i : bool):
    if model4_i:
        def Model4(a, model3_output):
            return a + model3_output
    else:
        def Model4(model1_output, f):
            return model1_output + f
    return Model4

def make_model5(model5_i : bool):
    if model5_i:
        def Model5(model4_output, f):
            return model4_output + f
    else:
        def Model5(model2_output, model3_output, f):
            return model2_output + model3_output + f
    return Model5

We can immediately create the graph and see the various possibilities:

def create_graph(model4_i : bool, model5_i : bool):
    Model4 = make_model4(model4_i)
    Model5 = make_model5(model5_i)
    nx_digraph = models_to_graph([Model4, Model3, Model2, Model1, Model5])
    print("Topological ordering: ", list(nx.topological_sort(nx_digraph)))
    return draw(to_graph(Model1, Model2, Model3, Model4, Model5), formatter = PlainFormat(), boundary = True)

_ = create_graph(True, True)

Topological ordering:  ['Model1', 'Model2', 'Model3', 'Model4', 'Model5']

_ = create_graph(False, True)

Topological ordering:  ['Model1', 'Model4', 'Model2', 'Model5', 'Model3']

_ = create_graph(False, False)

Topological ordering:  ['Model1', 'Model4', 'Model2', 'Model3', 'Model5']

These graphs are rather cluttered and difficult to analyse by hand; this is precisely the point! An algorithm can analyse these very fast and create a valid ordering and do the analysing for you, as long as you have the pure functions and explicit data interfaces![4]

Summary

Explicit data interfaces and pure functions define a computational graph, which standard algorithms can use to define a valid ordering between the models.

Thus, given a collection of models, you don’t need to hand-order them, this can be done automatically!

In the next notebook, we’ll look at collections of models that cannot be ordered because of a dependency cycle.

---

[1]

In fact there are multiple valid ones: Model2 and Model4 can be executed in any order after Model1, and Model3 can be executed after both of them. This ordering is not unique, but it is valid.

[3]

The choice that the model’s output is the name is not very nice here; but we’ll keep that for brevity. In a more realistic setting you could make the input and output anything you want.

[4]

A reasonable question is if this isn’t much slower in execution than hand-ordering: the answer is no. A jax compilation traces through this graph once and creates a compiled function which is literally this graph. After compilation there’s no overhead at all! All overhead is in the graph construction and tracing, but since these are not large graphs, this is a very small cost. (a 1600 model graph takes under a few ms).