Ordering Equations with NetworkX

I recently came across an interesting problem. I needed to set values for a bunch of variables, call them $x_1, \ldots, x_n$, and each one had a known rule of the form:

$$ x_i = f(x_{i_1}, ..., x_{i_{k}}) $$

so that each variable needed be initialized as a function of some set of other variables (possibly empty). This type of problem is familiar to many economists, (e.g., setting initial values in the popular macroeconomics software Dynare), but likely occurs in many other scientific computing applications as well. 

As anyone who has tackled this situation knows, the tricky part is ordering the rules so that each $x_i$ is initialized after all the variables it depends on. For small problems, this can be done manually without too much trouble. As the size of the problem grows, however, the complexity of this task can easily get out of hand. For example, in my application, the dependency structure turned out to be:

Yikes. Doing this manually would be a nightmare. Let's automate it!

My approach was to create a directed graph, which is a collection of nodes combined with one-way arrows connecting those nodes. I constructed a graph with the $\{x_i\}$ as the nodes, and then added an arrow from each $x_i$ to every $x_j$ that depends on it. Once the graph is constructed, running a topological sort generates a list of the variables in order so that each node is listed after all of its dependencies -- exactly what we needed.

To demonstrate, I'll show a simple example:

a = 2.0 * c;
b = a + c;
c = 5.0;

In this case, the solution is to initialize $c$ first, then $a$, then $b$. 

To implement the directed graph, I used the excellent Python package NetworkX. I started with a dict of (var, rule) pairs. For each rule, I extracted all the variable names using Python's regular expressions (re) module. I then constructed the graph and ran the sort. Here's the code:

import networkx as nx
import re

"""Construct regular expressions pattern to search for variable names.
This will find all blocks of letters, numbers, or underscores
starting with a letter or underscore."""

pattern = re.compile('[a-zA-Z_]\w*')

# Initialize the graph
G = nx.DiGraph()

# Example set of rules
rules = {
    'a' : '2.0 * c',
    'b' : 'a + c',
    'c' : '5.0',
}

# Loop through the rules
for var, rule in rules.items():

    # Add variable to graph
    G.add_node(var)

    # Get list of dependencies
    dependencies = re.findall(pattern, rule)

    # Add edges
    for dep in dependencies:
        G.add_edge(dep, var)

# Get sorted variable list
sorted_vars = nx.topological_sort(G)

# Get sorted list of rules
sorted_rules = [(var, rules[var]) for var in sorted_vars]

# Print results
for var, rule in sorted_rules:
    print('{0} = {1};\n'.format(var, rule))

The output is:

c = 5.0;
a = 2.0 * c;
b = a + c;

just like we wanted.