Entanglement Scheduling and TICK Slices

This example shows how an explicit scheduler controls preparation, entanglement, and measurement slices, and how those slices are serialized with TICK commands.

=== Entanglement Scheduling Demo ===

1. Creating graph state...
   Nodes: [0, 1, 2, 3]
   Input: [0]
   Output: [3]
   Edges: [(0, 1), (1, 2), (2, 3)]

2. Creating scheduler and solving schedule...
   Scheduling successful!
   Number of time slices: 5
   Preparation times: {1: 0, 2: 2, 3: 3}
   Measurement times: {0: 1, 1: 3, 2: 4}

3. Entanglement times (auto-scheduled)...
   Entanglement times: {(0, 1): 0, (1, 2): 2, (2, 3): 3}

4. Detailed timeline (Prep, Entangle, Measure):
   Time 0:
     Prepare: [1]
     Entangle: (0,1)
   Time 1:
     Measure: [0]
   Time 2:
     Prepare: [2]
     Entangle: (1,2)
   Time 3:
     Prepare: [3]
     Entangle: (2,3)
     Measure: [1]
   Time 4:
     Measure: [2]

5. Compiling pattern with scheduler-driven TICK commands...
   Pattern has 16 commands
   Maximum space usage: 2 qubits

   Pattern commands:
N: node=1
E: nodes=(0, 1)
TICK
M: node=0, plane=Plane.XY, angle=0.0
TICK
N: node=2
E: nodes=(1, 2)
TICK
N: node=3
E: nodes=(2, 3)
M: node=1, plane=Plane.XY, angle=0.0
TICK
M: node=2, plane=Plane.XY, angle=0.0
TICK
X: node=3
Z: node=3
0 more commands truncated. Change lim argument of print_pattern() to show more

6. Manual entanglement scheduling example...
   Manual schedule is valid: True
   Number of time slices: 4
   Pattern has 15 commands
   Maximum space usage: 2 qubits

Demo completed successfully!

from graphqomb.common import Plane, PlannerMeasBasis
from graphqomb.graphstate import GraphState
from graphqomb.pattern import print_pattern
from graphqomb.qompiler import qompile
from graphqomb.schedule_solver import ScheduleConfig, Strategy
from graphqomb.scheduler import Scheduler

# Create a simple graph state
print("=== Entanglement Scheduling Demo ===\n")
print("1. Creating graph state...")
node_labels = ["input", "middle1", "middle2", "output"]
edges = [("input", "middle1"), ("middle1", "middle2"), ("middle2", "output")]
inputs = ["input"]
outputs = ["output"]
meas_bases = {
    "input": PlannerMeasBasis(Plane.XY, 0.0),
    "middle1": PlannerMeasBasis(Plane.XY, 0.0),
    "middle2": PlannerMeasBasis(Plane.XY, 0.0),
}

graph, node_map = GraphState.from_graph(
    nodes=node_labels,
    edges=edges,
    inputs=inputs,
    outputs=outputs,
    meas_bases=meas_bases,
)

node0 = node_map["input"]
node1 = node_map["middle1"]
node2 = node_map["middle2"]
node3 = node_map["output"]

print(f"   Nodes: {list(graph.physical_nodes)}")
print(f"   Input: {list(graph.input_node_indices.keys())}")
print(f"   Output: {list(graph.output_node_indices.keys())}")
print(f"   Edges: {list(graph.physical_edges)}")

# Define flow
flow = {node0: {node1}, node1: {node2}, node2: {node3}}

# Create scheduler
print("\n2. Creating scheduler and solving schedule...")
scheduler = Scheduler(graph, flow)
config = ScheduleConfig(strategy=Strategy.MINIMIZE_SPACE)
success = scheduler.solve_schedule(config)

if success:
    print("   Scheduling successful!")
    print(f"   Number of time slices: {scheduler.num_slices()}")

    # Show preparation times
    prep_times = {k: v for k, v in scheduler.prepare_time.items() if v is not None}
    print(f"   Preparation times: {prep_times}")

    # Show measurement times
    meas_times = {k: v for k, v in scheduler.measure_time.items() if v is not None}
    print(f"   Measurement times: {meas_times}")

    # Show entanglement times (auto-scheduled by solve_schedule)
    print("\n3. Entanglement times (auto-scheduled)...")
    ent_times = {edge: time for edge, time in scheduler.entangle_time.items() if time is not None}
    print(f"   Entanglement times: {ent_times}")

    # Show detailed timeline
    print("\n4. Detailed timeline (Prep, Entangle, Measure):")
    timeline = scheduler.timeline
    for time_idx, (prep_nodes, ent_edges, meas_nodes) in enumerate(timeline):
        print(f"   Time {time_idx}:")
        if prep_nodes:
            print(f"     Prepare: {sorted(prep_nodes)}")
        if ent_edges:
            edges_str = [f"({min(e)},{max(e)})" for e in ent_edges]
            print(f"     Entangle: {', '.join(edges_str)}")
        if meas_nodes:
            print(f"     Measure: {sorted(meas_nodes)}")

    # Compile pattern (TICK commands are inserted automatically per time slice)
    print("\n5. Compiling pattern with scheduler-driven TICK commands...")
    pattern = qompile(graph, flow, scheduler=scheduler)
    print(f"   Pattern has {len(pattern.commands)} commands")
    print(f"   Maximum space usage: {pattern.max_space} qubits")

    print("\n   Pattern commands:")
    print_pattern(pattern, lim=30)

    # Manual entanglement scheduling example
    print("\n6. Manual entanglement scheduling example...")
    scheduler2 = Scheduler(graph, flow)
    scheduler2.manual_schedule(
        prepare_time={node1: 0, node2: 1, node3: 2},
        measure_time={node0: 1, node1: 2, node2: 3},
        entangle_time={
            (node0, node1): 0,
            (node1, node2): 1,
            (node2, node3): 2,
        },
    )

    scheduler2.validate_schedule()  # Will raise ValueError if invalid
    print("   Manual schedule is valid: True")
    print(f"   Number of time slices: {scheduler2.num_slices()}")

    pattern_manual = qompile(graph, flow, scheduler=scheduler2)
    print(f"   Pattern has {len(pattern_manual.commands)} commands")
    print(f"   Maximum space usage: {pattern_manual.max_space} qubits")

print("\nDemo completed successfully!")