Add causality and K-best example

causality-and-k-best/check_causality.py

@@ -0,0 +1,114 @@
+import json
+from collections import Counter
+
+include("butadienNetwork.py")
+
+with open("gurobi_flows.json") as f:
+    gurobi_flows = json.load(f)
+
+
+def canonical_smiles(smiles):
+    return str(Graph.fromSMILES(smiles).smiles)
+
+
+def vertex_key(v):
+    return str(v.graph.smiles)
+
+
+vertices_by_smiles = {}
+for v in dg.vertices:
+    vertices_by_smiles.setdefault(vertex_key(v), []).append(v)
+
+
+def current_vertex(smiles):
+    matches = vertices_by_smiles.get(canonical_smiles(smiles), [])
+    if len(matches) != 1:
+        raise RuntimeError(
+            f"{smiles}: expected one MØD vertex, found {len(matches)}"
+        )
+    return matches[0]
+
+
+def current_reaction(educts, products):
+    wanted_educts = Counter(canonical_smiles(x) for x in educts)
+    wanted_products = Counter(canonical_smiles(x) for x in products)
+
+    matches = []
+    for e in dg.edges:
+        current_educts = Counter(vertex_key(v) for v in e.sources)
+        current_products = Counter(vertex_key(v) for v in e.targets)
+        if current_educts == wanted_educts and current_products == wanted_products:
+            matches.append(e)
+
+    if len(matches) != 1:
+        raise RuntimeError(
+            f"{educts} -> {products}: expected one MØD edge, found {len(matches)}"
+        )
+    return matches[0]
+
+
+def check_causality(gurobi_flow, rank):
+    source_flow = {
+        current_vertex(smiles): value
+        for smiles, value in gurobi_flow["inputs"].items()
+    }
+    output_flow = {
+        current_vertex(smiles): value
+        for smiles, value in gurobi_flow["outputs"].items()
+    }
+
+    reaction_flow = {}
+    for reaction in gurobi_flow["reactions"]:
+        edge = current_reaction(reaction["educts"], reaction["products"])
+        reaction_flow[edge] = reaction_flow.get(edge, 0) + reaction["flow"]
+
+    flow = hyperflow.Model(dg)
+
+    for v in source_flow:
+        flow.addSource(v.graph)
+    for v in output_flow:
+        flow.addSink(v.graph)
+
+    for e in dg.edges:
+        flow.addConstraint(edgeFlow[e] == reaction_flow.get(e, 0))
+
+    for v in dg.vertices:
+        flow.addConstraint(inFlow[v] == source_flow.get(v, 0))
+        flow.addConstraint(outFlow[v] == output_flow.get(v, 0))
+
+    flow.findSolutions(maxNumSolutions=1)
+
+    try:
+        solution = next(iter(flow.solutions))
+    except StopIteration:
+        print(f"solution {rank}: not a valid MØD hyperflow")
+        return "INVALID"
+
+    query = causality.RealisabilityQuery(flow.dg)
+    dag = query.findDAG(solution)
+
+    if not dag:
+        print(f"solution {rank}: causally realisable: NO")
+        return "NO"
+
+    print(f"solution {rank}: causally realisable: YES")
+
+    printer = DGPrinter()
+    printer.graphvizPrefix = 'layout = "dot";'
+
+    def vertex_visible(v):
+        return solution.eval(vertexFlow[v]) != 0
+
+    printer.pushVertexVisible(vertex_visible)
+    dg.print(printer=printer, data=dag.getPrintData(False))
+
+    return "YES"
+
+
+results = []
+for rank, gurobi_flow in enumerate(gurobi_flows, start=1):
+    results.append(check_causality(gurobi_flow, rank))
+
+print("\n=== causality summary ===")
+for rank, result in enumerate(results, start=1):
+    print(f"solution {rank}: {result}")