Calculate the smallest oriented coloring for a directed graph https://en.wikipedia.org/wiki/Oriented_coloring

oriented_coloring.py

from z3 import *
import matplotlib.pyplot as plt
import random, math

# pip install z3-solver matplotlib
# do not install z3, that's a different library

random.seed(0)

def find_ordered_coloring(edges, n):
    for k in range(n + 1):
        s = Solver()

        # make label variables
        labels = [Int(str(i)) for i in range(n)]

        # labels should be {0, ..., k - 1}
        for label in labels:
            s.add(0 <= label)
            s.add(label < k)

        # neighboring vertices should have different labels
        for i, j in edges:
            s.add(labels[i] != labels[j])

        # for each pair of label values
        for k1 in range(k):
            for k2 in range(k1 + 1, k):
                direction = Bool("d_%d_%d" % (k1, k2))

                # edges between vertices with labels k1 and k2 should
                # all go towards the same direction
                for i, j in edges:
                    direction1 = And(labels[i] == k1, labels[j] == k2)
                    direction2 = And(labels[i] == k2, labels[j] == k1)

                    s.add(Implies(direction1, direction))
                    s.add(Implies(direction2, Not(direction)))

        print(f"Checking k = {k}...")

        # check if a graph coloring with k colors exists
        if s.check() == sat:
            # return the result
            m = s.model()
            coloring = [m[labels[i]].as_long() for i in range(n)]
            return k, coloring

def make_adjacency_matrix(n, n_edges):
    edges = set()

    while len(edges) < n_edges:
        # create random edge
        i = random.randrange(n)
        j = random.randrange(n)

        # add edge if it does not exist yet
        # discard loops and reverse edges
        if i != j and (i, j) not in edges and (j, i) not in edges:
            edges.add((i, j))

    return edges

def plot_adjacency_matrix(edges, n, coloring):
    COLORS = [
        "red",
        "green",
        "blue",
        "yellow",
        "lightblue",
        "purple",
        "orange",
        "brown",
    ]

    # place vertices in a circle
    x = []
    y = []
    for i in range(n):
        angle = 2.0 * math.pi * i / n
        x.append(math.cos(angle))
        y.append(math.sin(angle))

    # draw edges
    for i, j in edges:
        ax, ay = x[i], y[i]
        bx, by = x[j], y[j]

        # move arrow head back slightly so they don't overlap
        dx = bx - ax
        dy = by - ay

        d = (dx * dx + dy * dy)**0.5

        s = 0.03

        ax += s * dx / d
        ay += s * dy / d

        plt.annotate("", xy=(ax, ay), xytext=(bx, by), arrowprops=dict(arrowstyle="->"))

    # draw vertices
    plt.scatter(x, y, c=[COLORS[c % len(COLORS)] for c in coloring], zorder=3)
    plt.show()

def make_circle(n):
    edges = []
    i = n - 1
    for j in range(n):
        edges.append((i, j))
        i = j
    return edges

def test():
    # ordered chromatic number for circle of size 5 should be 5
    assert(find_ordered_coloring(make_circle(5), 5)[0] == 5)

def main():
    # make random graph with 10 vertices and 15 edges
    n = 10
    edges = make_adjacency_matrix(n, n_edges=15)

    result = find_ordered_coloring(edges, n)

    if result:
        k, coloring = result
        print("Found oriented coloring with", k, "colors")
    else:
        coloring = [0] * n
        print("Oriented coloring does not exist")

    plot_adjacency_matrix(edges, n, coloring)

if __name__ == "__main__":
    main()