Theory II: Towards functional programming

As seen in the previous notebook, that particular style of imperative programming is problematic for a large, flexible, model framework. Let’s try to devise a better way to do this.

Hidden Dependencies

The main reason why it’s difficult to reason about which model comes before which other model (and whether there are cycles) is because of hidden dependencies. It isn’t obvious from the collection of models which model depends on which other model and how.

In the previous examples, these arise because of mutable state. The datastructure can be modified very deep within a complex function and you’d need to know that before knowing if other models depend on it (if such a model depends on the value modified deep within that function).

These functions have side effects: they don’t just return a value or output, they do things to their inputs, global state, or file-system.

If every function is pure (no side effects and thus the only thing they do is return a value), you’d have much less problems with hidden dependencies: you’d just have to look at its inputs and outputs rather than the function body.

Converting the previous examples is easy: instead of modifying the dictionary, we return the outputs only:

from math import sin, exp


def Model1(data_structure):
    #--------- Inputs
    a = data_structure['a']
    b = data_structure['b']

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

    #--------- Outputs
    return {"model1_output" : output}
    

def Model2(data_structure):

    #--------- Inputs
    model1_output = data_structure['model1_output']
    c              = data_structure['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 {"model2_output" : iteration_var}


def Model3(data_structure):
    #--------- Inputs
    model2_output = data_structure['model2_output']
    e              = data_structure['e']
    c              = data_structure['c']
    d              = data_structure['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 {"model3_output" : x}

These are now pure: we only have to look at the return statement to figure out what its effects are [1]. The entire complicated function body is not relevant anymore for reasoning about dependencies!

Now, we can somewhat reason about it, but the inputs are still fully implicit and the outputs are only visible in the dictionary: you have to look at the internals and figure it out from that. Instead, we can make it entirely explicit by using the model’s name to define the output and the function signature to define the inputs [2]

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 interface(f):
    return {f.__name__ : {"inputs" : list(inspect.signature(f).parameters.keys()), "outputs" : [f.__name__.lower() + "_output"]}}

print(interface(Model1))
{'Model1': {'inputs': ['a', 'b'], 'outputs': ['model1_output']}}

We’ve now made all the data dependencies explicit through the model interface. Because the models themselves are pure, we know that we cannot introduce any hidden additional dependences through side effects or mutable state!

You now don’t need to look at the internals of a function to determine where it should be ordered, you can just use the above interface function.

Let’s define the other models in the same way, and the total model:

Hide code cell source

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
print(interface(Model1))
print(interface(Model2))
print(interface(Model3))
{'Model1': {'inputs': ['a', 'b'], 'outputs': ['model1_output']}}
{'Model2': {'inputs': ['model1_output', 'c'], 'outputs': ['model2_output']}}
{'Model3': {'inputs': ['model2_output', 'e', 'c', 'd'], 'outputs': ['model3_output']}}

We can read this inputs outputs and reason that we (logically) can just do Model1, Model2, Model3 in that order.

But since we know the inputs and outputs of all models from the interface, we can just write the total model without even manually passing data:

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

    for model in models_in_order:
        # 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

base_input_dict = {"a": 1.0, "b": 2, "c": 3, "d": 4, "e": 5}

TotalModel(base_input_dict, models_in_order = [Model1, Model2, Model3])
{'a': 1.0,
 'b': 2,
 'c': 3,
 'd': 4,
 'e': 5,
 'model1_output': 1.2412888640075004,
 'model2_output': 30.473755232097503,
 'model3_output': 1.0547835473689968}

Great! We can even add a Model4 without really doing anything but looking at the ordering:

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

print(interface(Model4))
base_input_dict_f = {**base_input_dict, "f": 6}
TotalModel(base_input_dict_f, models_in_order = [Model1, Model4, Model2, Model3])
{'Model4': {'inputs': ['model1_output', 'f'], 'outputs': ['model4_output']}}
{'a': 1.0,
 'b': 2,
 'c': 3,
 'd': 4,
 'e': 5,
 'f': 6,
 'model1_output': 1.2412888640075004,
 'model4_output': 7.2412888640075,
 'model2_output': 30.473755232097503,
 'model3_output': 1.0547835473689968}

With this implementation, we haven’t necessarily fixed all the issues: you can still order them wrongly (which might run if the datastructure is pre-filled with something), you still have to order all of them by hand, and you still need to manually keep track of the interfaces.

The big thing however is that because you have these interfaces you can prove and check various properties without running the models!

For example, we can find out which inputs are required to be external (i.e. not produced by any other model):

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

print("Model1, Model2, Model3:        ", external_inputs({Model1, Model2, Model3}))
print("Model1, Model2, Model3, Model4:", external_inputs({Model1, Model2, Model3, Model4}))
Model1, Model2, Model3:         {'a', 'e', 'd', 'b', 'c'}
Model1, Model2, Model3, Model4: {'e', 'd', 'b', 'f', 'a', 'c'}

Clearly, model 4 requires extra f, as before, but now we can immediately just get that from the models themselves (without actually running them or giving it a datastructure)! We can then just check if we pass a datastructure with too little or too much external inputs and optionally just refuse that (eliminating the ‘correctness by accident’ problem).

The real nice thing is that we can even check whether the model order is correct!

def check_model_ordering(models_in_order):

    external_inputs_set = external_inputs(set(models_in_order))
    # Create a set to keep track of available outputs
    available_outputs = set()
    
    for model in models_in_order:        
        interface_i = interface(model)[model.__name__]        

        is_external_input = lambda input_name: input_name in external_inputs_set
        is_available_output = lambda input_name: input_name in available_outputs

        all_inputs_available = all(is_external_input(input_name) or is_available_output(input_name) 
                                   for input_name in interface_i["inputs"])

        # Check if all required inputs are available
        if not all_inputs_available:
            input_not_available = [input_name for input_name in interface_i["inputs"] 
                                   if not (is_external_input(input_name) or is_available_output(input_name))]
            print(f"{model.__name__} cannot be executed because not all inputs are available: [{', '.join(input_not_available)}]")
            return False
        
        # Add the output of the current model to the available outputs
        available_outputs.add(interface_i["outputs"][0])
    
    return True


print("Model1 -> Model2 -> Model3 -> Model4:", check_model_ordering([Model1, Model2, Model3, Model4]))
print("Model2 -> Model1 -> Model3 -> Model4:", check_model_ordering([Model2, Model1, Model3, Model4]))  
Model1 -> Model2 -> Model3 -> Model4: True
Model2 cannot be executed because not all inputs are available: [model1_output]
Model2 -> Model1 -> Model3 -> Model4: False

In fact, you can automatically generate a correct ordering using a so-called topological sort (see next notebooks)!

Further possible checks include that your models don’t write to the same output (which would be a problem; which one should be first), and that your models don’t have cycles, or if they have, where they actually are.

Computational graphs

We can visualize these dependencies: if every model is a circle, and arrows pointing from model outputs to model inputs, yielding a graph:

Hide code cell source

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)
../_images/9781b5a04b12e55af714dc52be99523dbad84cdee9dcaff1fd3fa8f4097ec04e.png

Such computational graphs can be used (like in the previous) examples to do the static correctness checking and ordering of models using standard algorithms, which we will explore in the next notebook.

Summary

In this notebook, we show that using pure functions and explicit data interfaces enables you to reason about the correctness of your model ordering and check for cycles without actually running the models.

In the next notebook, we’ll further explore these computational graphs (including the cycles we haven’t discussed yet in this context), algorithms to analyse them, and use them to order and run a set of models automatically.