Added tests for substrate scheduler active volume functions and fixed…

… bug in tocks accounting

src/benchq/compilation/graph_states/circuit_compilers.py

@@ -53,8 +53,6 @@ def rbs_circuit_compiler(
         optimization: str,
         verbose: bool,
     ) -> GSCInfo:
-        if not isinstance(logical_architecture_name, str):
-            breakpoint()
         (
             compiled_graph_data,
             _,

src/benchq/compilation/graph_states/substrate_scheduler/substrate_scheduler.jl

@@ -74,7 +74,8 @@ function time_optimal_two_row_scheduler(asg, pauli_tracker, num_logical_qubits,
         nodes_to_satisfy_this_layer = setdiff(get_neighborhood(pauli_tracker.layering[layer_num], asg), satisfied_nodes)
         union!(satisfied_nodes, nodes_to_satisfy_this_layer)
         bars = []
-        for node in setdiff(nodes_to_satisfy_this_layer, nodes_satisfied_by_initialization)
+        subnodes_this_layer = setdiff(nodes_to_satisfy_this_layer, nodes_satisfied_by_initialization)
+        for node in subnodes_this_layer
             bar = [node_to_patch[node], node_to_patch[node]]
             for neighbor in asg.edge_data[node]
                 if node_to_patch[neighbor] < bar[1] && node in curr_physical_nodes
@@ -93,7 +94,7 @@ function time_optimal_two_row_scheduler(asg, pauli_tracker, num_logical_qubits,
             end
         end
         sort!(bars, by=x -> x[2])
-
+        
         # layer the bars
         while length(bars) > 0
             bar = bars[1]
@@ -254,7 +255,6 @@ function time_optimal_active_volume_scheduler(asg, pauli_tracker, num_logical_qu
     # Phase 2: For each layer of nodes in the remaining graph, group the multi-qubit stabilizer measurements 
     # into co-measurable sets to determine the number of tocks required to prepare the graph state of that layer.
     for layer_num in VerboseIterator(1:length(pauli_tracker.layering), verbose, "Scheduling Clifford operations...")
-
         # Get the subnodes this layer
         nodes_to_satisfy_this_layer = setdiff(get_neighborhood(pauli_tracker.layering[layer_num], asg), satisfied_nodes)
         union!(satisfied_nodes, nodes_to_satisfy_this_layer)
@@ -270,18 +270,13 @@ function time_optimal_active_volume_scheduler(asg, pauli_tracker, num_logical_qu
             end
         end
 
-        # Account for the number of tocks needed for graph state creation
-        # Construct Graphs.jl SimpleGraph from the subgraph of the AlgorithmSpecificGraph
-        layer_subgraph = construct_subgraph_from_asg(asg, subnodes_this_layer)
 
         # Compute number of tocks for graph state creation using greedy graph coloring on extension graph
-        num_tocks_for_graph_creation[layer_num] += compute_active_volume_tocks(layer_subgraph)
+        num_tocks_for_graph_creation[layer_num] += compute_active_volume_tocks_to_prepare_subgraph(asg.edge_data, subnodes_this_layer)
 
-        # TODO: determine what this does
         if layer_num < length(pauli_tracker.layering)
             union!(measured_nodes, pauli_tracker.layering[layer_num])
             for measured_node in pauli_tracker.layering[layer_num]
-                # patches_to_node[node_to_patch[measured_node]] = -1
                 measured_node = -1
             end
             added_nodes = pauli_tracker.layering[layer_num+1]
@@ -386,76 +381,98 @@ function active_volume_greedy_group_neighborhoods(asg, subnodes_this_layer)
     return num_tocks
 end
 
+"""
+Constructs a subgraph from the edge data of the AlgorithmSpecificGraph that contains
+only the nodes in the input set. The subgraph is constructed by adding edges between
+nodes that are neighbors in the original graph.
 
-function construct_subgraph_from_asg(asg::AlgorithmSpecificGraph, nodes::Set{UInt32})
+Args:
+    edge_data::List{List{Int}}          Edge data describing a graph
+    nodes::Set{Int}                     Set of nodes to include in the subgraph
+"""
+function construct_subgraph_from_edge_data(edge_data::Vector{Set{UInt32}}, nodes::Set{UInt32})
     # Create a SimpleGraph with the same number of nodes as the input set
     n = length(nodes)
     g = SimpleGraph(n)
 
     # Create a mapping from original node indices to new subgraph indices
     node_map = Dict{UInt32, Int64}(node => i for (i, node) in enumerate(nodes))
-
     for original_node in nodes
         i = node_map[original_node]  # Get the index in the subgraph
-        for neighbor in asg.edge_data[Int(original_node)]
+        for neighbor in edge_data[Int(original_node)]
             if haskey(node_map, neighbor)
                 j = node_map[neighbor]  # Get the index of the neighbor in the subgraph
                 if i != j
                     Graphs.add_edge!(g, i, j)  # Add edge between the two nodes in the subgraph
                 end
+            else
+                # Add a new entry to the node map for the neighbor
+                node_map[neighbor] = length(node_map) + 1
+                Graphs.add_vertex!(g)
+                Graphs.add_edge!(g, i, node_map[neighbor])
             end
         end
     end
-
     return g
 end
 
 
-# Function that generates the extension graph for a given graph such that
-# the output graph has edges added to the original graph between any two nodes 
-# that are distance 2 apart in the original graph.
-# This is used to group nodes that are 2-independent in the original graph.
-function generate_extension_graph(graph::SimpleGraph)
-    # Create a new graph with the same number of nodes as the input graph
-    n = nv(graph)
-    extension_graph = SimpleGraph(n)
+# Function that generates the extension graph for a given set of edges and
+# subnodes such that the output graph has edges added to the original graph 
+# between any two nodes that are distance 2 apart in the original edge data.
+# This is helpful for grouping nodes that are 2-independent in the original graph.
+function generate_extension_graph(edge_data::Vector{Set{UInt32}}, nodes::Set{UInt32})
+    # Initialize an empty graph with Int64 type for vertices
+    extension_graph = SimpleGraph{Int64}(length(nodes))
 
-    # Iterate over all nodes in the graph
-    for node in 1:n
-        # Iterate over the neighbors of the current node
-        for neighbor in neighbors(graph, node)
+    # Create a mapping from node indices to vertex indices
+    node_to_vertex_map = Dict{UInt32, Int}()
+    for (i, node) in enumerate(nodes)
+        node_to_vertex_map[node] = i
+    end
+
+    # Iterate over nodes and add edges
+    for node in nodes
+        node_index = node_to_vertex_map[node]
+        for neighbor in edge_data[node]
+            # Add edge between the current node and neighbor
+            # if that neighbor is in nodes
+            if haskey(node_to_vertex_map, neighbor) && neighbor != node
+                Graphs.add_edge!(extension_graph, node_index, node_to_vertex_map[neighbor])
+            end
+
             # Iterate over the neighbors of the neighbor node
-            for neighbor_node in neighbors(graph, neighbor)
-                # Check if the neighbor node is different from the current node
-                # and if it is not already connected to the current node
-                if neighbor_node != node && !(has_edge(extension_graph, node, neighbor_node))
-                    # Add an edge between the nodes in the extension graph
-                    Graphs.add_edge!(extension_graph, node, neighbor_node)
+            for nn_neighbor in edge_data[neighbor]
+                # Skip to avoid creating self loops in the graph
+                if nn_neighbor == node
+                    continue
+                end
+
+                if haskey(node_to_vertex_map, node) && haskey(node_to_vertex_map, nn_neighbor)
+                    nn_neighbor_index = node_to_vertex_map[nn_neighbor]
+                    Graphs.add_edge!(extension_graph, node_index, nn_neighbor_index)
                 end
             end
         end
     end
-
-    return extension_graph
+    return extension_graph, node_to_vertex_map
 end
 
 # Function that generates the extension graph and solves graph coloring using 
 # a specified algorithm from the Graphs.jl package. 
 # Note: this function doesn't keep track of the groups as would be
 # needed by a compiler, but only accounts for the number of tocks that would 
 # result from a lattice surgery compilation.
-function compute_active_volume_tocks(graph::Graph, coloring_algorithm=Graphs.degree_greedy_color)
+function compute_active_volume_tocks_to_prepare_subgraph(edge_data, subnodes, coloring_algorithm=Graphs.degree_greedy_color)
     # Check if the graph is empty
-    if nv(graph) == 0
+    if length(subnodes) == 0
         return 0
     end
 
     # Generate extension graph
-    extension_graph = generate_extension_graph(graph)
-
+    extension_graph, _ = generate_extension_graph(edge_data, subnodes)
     # Solve graph coloring using the specified algorithm
     coloring = coloring_algorithm(extension_graph)
-
     # Return the number of colors used
     return coloring.num_colors
 end

tests/benchq/compilation/test_ruby_slippers.py

@@ -398,7 +398,11 @@ def test_rbs_with_active_volume_gives_fewer_graph_creation_cycles_than_two_row()
 
     # given
     circuit = Circuit(
-        [H(0), *[CNOT(j, i) for i in range(1, 300) for j in range(2, 300)]]
+        [
+            H(0),
+            T(0),
+            *[CNOT(j, i) for i in range(1, 300) for j in range(2, 300)],
+        ]
     )
     optimization = "Time"
 
@@ -422,6 +426,52 @@ def test_rbs_with_active_volume_gives_fewer_graph_creation_cycles_than_two_row()
     )
 
 
+def test_active_volume_has_fewer_tocks_than_two_row():
+
+    # given
+    circuit = Circuit(
+        [
+            H(0),
+            H(1),
+            H(2),
+            H(3),
+            H(4),
+            T(0),
+            T(1),
+            T(2),
+            T(3),
+            T(4),
+            CNOT(0, 1),
+            CNOT(1, 2),
+            CNOT(2, 3),
+            CNOT(3, 4),
+            T(1),
+            T(2),
+            T(3),
+        ]
+    )
+    optimization = "Time"
+
+    compiled_data_two_row, _ = jl.run_ruby_slippers(
+        circuit,
+        verbose=True,
+        logical_architecture_name="two_row_bus",
+        optimization=optimization,
+    )
+
+    compiled_data_active_volume, _ = jl.run_ruby_slippers(
+        circuit,
+        verbose=False,
+        logical_architecture_name="active_volume",
+        optimization=optimization,
+    )
+
+    # then
+    assert sum(compiled_data_active_volume["graph_creation_tocks_per_layer"]) < sum(
+        compiled_data_two_row["graph_creation_tocks_per_layer"]
+    )
+
+
 ########################################################################################
 # Everything below here is testing utils
 ########################################################################################

tests/benchq/compilation/test_substrate_scheduler_components.py

@@ -0,0 +1,81 @@
+################################################################################
+# © Copyright 2022-2023 Zapata Computing Inc.
+################################################################################
+import os
+import pathlib
+
+import networkx as nx
+import numpy as np
+import pytest
+import stim
+from numba import njit
+from orquestra.integrations.qiskit.conversions import import_from_qiskit
+from orquestra.quantum.circuits import CNOT, CZ, Circuit, H, S, T, X
+from qiskit import QuantumCircuit
+
+from benchq.compilation.circuits import (
+    compile_to_native_gates,
+    pyliqtr_transpile_to_clifford_t,
+)
+from benchq.compilation.graph_states import jl
+
+
+def to_julia_set(python_set):
+    """Convert a Python set of integers to a Julia Set{UInt32}."""
+    return jl.Set[jl.UInt32]([jl.UInt32(e) for e in python_set])
+
+
+def to_julia_vector_of_sets(python_list_of_sets):
+    """Convert a Python list of sets of integers to a Julia Vector{Set{UInt32}}."""
+    return jl.Vector[jl.Set[jl.UInt32]]([to_julia_set(s) for s in python_list_of_sets])
+
+
+def test_generate_extension_graph():
+    # Define the edge_data and nodes in Python
+    edge_data = [{3}, {3}, {1, 2, 4}, {3, 5}, {4}]
+    nodes = {1, 2, 5}
+
+    # Define the expected outputs
+    expected_max_adj_list_size = 1
+    expected_min_adj_list_size = 0
+
+    # Convert Python sets and lists to Julia-compatible formats using utility functions
+    edge_data_julia_vector = to_julia_vector_of_sets(edge_data)
+    nodes_julia = to_julia_set(nodes)
+
+    # Call the Julia function
+    output, node_to_vertex_map = jl.generate_extension_graph(
+        edge_data_julia_vector, nodes_julia
+    )
+    output_adj_list = output.fadjlist
+
+    # then
+    assert (
+        max([len(neighbors) for neighbors in output_adj_list])
+        == expected_max_adj_list_size
+    )
+    assert (
+        min([len(neighbors) for neighbors in output_adj_list])
+        == expected_min_adj_list_size
+    )
+
+
+def test_compute_active_volume_tocks():
+    # Define the edge_data and nodes in Python
+    edge_data = [{3}, {3}, {1, 2, 4}, {3, 5}, {4}]
+    nodes = {1, 2, 5}
+
+    # Define the expected outputs
+    expected_active_volume_tocks = 2
+
+    # Convert Python sets and lists to Julia-compatible formats using utility functions
+    edge_data_julia_vector = to_julia_vector_of_sets(edge_data)
+    nodes_julia = to_julia_set(nodes)
+
+    # Call the Julia function
+    output = jl.compute_active_volume_tocks_to_prepare_subgraph(
+        edge_data_julia_vector, nodes_julia
+    )
+
+    # then
+    assert output == expected_active_volume_tocks