# Copyright (c) 2023 - 2026 Chair for Design Automation, TUM
# All rights reserved.
#
# SPDX-License-Identifier: MIT
#
# Licensed under the MIT License
from __future__ import annotations
import pathlib
from datetime import datetime
from typing import TYPE_CHECKING, Any, cast
from .cycles import get_ion_chains
from .plotting import plot_state
from .scheduling import (
create_cycles_for_moves,
create_gate_info_list,
create_move_list,
create_priority_queue,
find_movable_cycles,
get_partitioned_priority_queues,
preprocess,
rotate_free_cycles,
)
if TYPE_CHECKING:
from .graph import Graph
from .ion_types import Edge
import contextlib
import json
def _snapshot_state_for_json(graph: Graph, t: int) -> dict[str, Any]:
"""
Build one JSON frame:
{ "t": t,
"ions": [ { "id": "$q_0$", "edge": ["(0, 0)", "(0, 1)"] }, ... ],
"gates": [ { "id": "...", "type": "XX", "qubits": ["$q_0$","$q_1$"], "edge": ["(r,c)","(r2,c2)"], "duration": 3, "pz": "pz1" }, ... ]
}
"""
def _fmt_node(p: tuple[float, float]) -> str:
return f"({p[0]}, {p[1]})"
state = get_ion_chains(graph)
frame: dict[str, Any] = {
"t": t,
"ions": [
{"id": f"$q_{ion}$", "edge": [str(state[ion][0]), str(state[ion][1])]} for ion in sorted(state.keys())
],
}
# Include executed gates buffered by main() for the previous step
pending = graph.executed_gates_next
if pending:
gates_out: list[dict[str, Any]] = []
for i, g in enumerate(pending):
try:
gid = str(g.get("id", f"g_{i}"))
gtype = str(g.get("type", "OP")).upper()
qubits = [f"$q_{q}$" if not str(q).startswith("$q_") else q for q in g.get("qubits", [])]
edge = None
if g.get("edge"):
edge = g["edge"]
elif g.get("edge_idc"):
a, b = g["edge_idc"]
edge = [_fmt_node(a), _fmt_node(b)]
duration = int(g.get("duration", 1))
pz_name = g.get("pz")
gates_out.append({
"id": gid,
"type": gtype,
"qubits": qubits,
**({"edge": edge} if edge else {}),
"duration": duration,
**({"pz": pz_name} if pz_name else {}),
})
except Exception:
continue
if gates_out:
frame["gates"] = gates_out
with contextlib.suppress(Exception):
graph.executed_gates_next = []
return frame
[docs]
def find_pz_order(graph: Graph, gate_info_list: dict[str, list[int]]) -> list[str]:
# find next processing zone that will execute a gate
pz_order = []
for gate in graph.sequence:
qubits = graph.get_gate_qubits(gate)
if len(qubits) == 1:
ion = qubits[0]
for pz in graph.pzs:
if ion in gate_info_list[pz.name]:
pz_order.append(pz.name)
break
elif len(qubits) == 2:
ion1, ion2 = qubits
for pz in graph.pzs:
if ion1 in gate_info_list[pz.name] and ion2 in gate_info_list[pz.name]:
pz_order.append(pz.name)
break
return pz_order
[docs]
def shuttle(
graph: Graph,
priority_queue: dict[int, str],
timestep: int,
cycle_or_paths: str,
unique_folder: pathlib.Path,
) -> None:
gate_info_list = create_gate_info_list(graph)
print(f"Gate info list: {gate_info_list}")
pz_executing_gate_order = find_pz_order(graph, gate_info_list)
print(f"Next processing zone executing gate: {pz_executing_gate_order}")
# new: stop moves (ions that are already in the correct processing zone for a two-qubit gate)
graph.stop_moves = []
ion1_needed_in_pz = None
# "swap" ions in the same processing zone if only one is needed
for pz in graph.pzs:
ions_at_pz = graph[pz.edge_idc[0]][pz.edge_idc[1]]["ions"]
if len(ions_at_pz) == 2:
ion1, ion2 = ions_at_pz
if ion2 not in gate_info_list[pz.name]:
graph[pz.edge_idc[0]][pz.edge_idc[1]]["ions"].remove(ion2)
graph[pz.edge_idc[0]][pz.edge_idc[1]]["ions"].insert(0, ion2)
print(f"swapped ion2, now: {ions_at_pz}")
# find the next processing zone that will execute a gate on ion1
# in case it is needed elsewhere
# solution to bug where ion1 was needed later in the current pz
# but also needed elsewhere
# so was never rotated because other ions where rotating into that pz
# and swapped with ion1 (ion1 would never be rotated then)
for pz_name in pz_executing_gate_order:
if ion1 in gate_info_list[pz_name]:
ion1_needed_in_pz = pz_name
break
# TODO ion1 swap also necessary? # new TODO
if ion1_needed_in_pz is not None and ion1_needed_in_pz != pz.name:
# ion1 not in gate_info_list[pz.name]:
graph[pz.edge_idc[0]][pz.edge_idc[1]]["ions"].remove(ion1)
graph[pz.edge_idc[0]][pz.edge_idc[1]]["ions"].insert(0, ion1)
print(f"swapped back ion1, now {ions_at_pz}")
# new: this could maybe be used to time the gates
# (bug where two ions were randomly in the correct pz already while only
# a 1-qubit gate was executed on one of them and next gate was
# the 2-qubit gate on them -> 3rd ion was moved into them)
if ion1 in gate_info_list[pz.name] and ion2 in gate_info_list[pz.name]:
graph.stop_moves.append(ion1)
graph.stop_moves.append(ion2)
elif len(ions_at_pz) == 1:
ion1 = ions_at_pz[0]
if ion1 in gate_info_list[pz.name] and len(gate_info_list[pz.name]) == 1:
graph.stop_moves.append(ion1)
preprocess(graph, priority_queue)
# Update ion chains after preprocess
graph.state = get_ion_chains(graph)
# print(f"priority queue: {priority_queue}")
part_prio_queues = get_partitioned_priority_queues(priority_queue)
all_cycles: dict[int, list[Edge]] = {}
# Iterate over all processing zones
# create move list for each pz -> needed to get all cycles
# priority queue later picks the cycles to rotate
for pz in graph.pzs:
prio_queue = part_prio_queues[pz.name]
move_list = create_move_list(graph, prio_queue, pz)
# print(f"Priority queue for {pz.name}: {prio_queue}")
# print(f"Move list for {pz.name}: {move_list}")
cycles = create_cycles_for_moves(graph, move_list, cycle_or_paths, pz)
# add cycles to all_cycles
all_cycles = {**all_cycles, **cycles}
# print(f"All cycles: {all_cycles}")
# now general priority queue picks cycles to rotate
chains_to_rotate = find_movable_cycles(graph, all_cycles, priority_queue, cycle_or_paths)
# print(f"Chains to rotate: {chains_to_rotate}")
rotate_free_cycles(graph, all_cycles, chains_to_rotate)
# new: postprocess?
# -> ions can already move into processing zone if they pass a junction
# Update ion chains after rotate
graph.state = get_ion_chains(graph)
preprocess(graph, priority_queue)
labels = (f"timestep {timestep}", f"remaining sequence: {graph.sequence}")
plot_state(
graph,
labels,
plot_ions=True,
show_plot=graph.plot,
save_plot=graph.save,
plot_cycle=False,
plot_pzs=False,
filename=unique_folder / f"{graph.arch}_timestep_{timestep}.png",
)
[docs]
def main(graph: Graph, sequence: list[int] | None, cycle_or_paths: str, record_path: str | None = None) -> int:
"""Run the inside-layout shuttle scheduler.
Integer gate IDs in ``sequence`` require corresponding entries in
``graph.gate_info``. When requested, timeline snapshots are written to
``record_path``.
Args:
graph: Configured inside architecture and scheduling state.
sequence: Optional ordered stable gate IDs, replacing ``graph.sequence``.
cycle_or_paths: Conflict-resolution strategy used for movement.
record_path: Optional path for JSON timeline snapshots.
Returns:
The number of simulated timesteps.
"""
timestep = 0
if sequence is not None:
graph.sequence = list(sequence)
timeline = []
if record_path is not None:
with contextlib.suppress(Exception):
pathlib.Path(record_path).parent.mkdir(parents=True, exist_ok=True)
# initial state BEFORE any movement
timeline.append(_snapshot_state_for_json(graph, timestep))
try:
with pathlib.Path(record_path).open("w", encoding="utf-8") as f:
json.dump({"timeline": timeline}, f, separators=(",", ":"), ensure_ascii=False)
except Exception:
pass
max_timesteps = 1e6
graph.state = get_ion_chains(graph)
unique_folder = pathlib.Path("runs") / datetime.now().strftime("%Y%m%d_%H%M%S")
if graph.save is True:
unique_folder.mkdir(exist_ok=True, parents=True)
plot_state(
graph,
labels=("Initial state", None),
plot_ions=False,
show_plot=graph.plot,
save_plot=graph.save,
plot_cycle=False,
plot_pzs=True,
filename=unique_folder / f"{graph.arch}_timestep_{timestep}.png",
)
graph.in_process = []
graph.locked_gates = cast("dict[int, str]", {})
while timestep < max_timesteps:
print("locked_gates", graph.locked_gates)
print(f"\nStarting timestep {timestep}")
print(f"Remaining sequence: {graph.sequence}")
# priority queue is dict with ions as keys and pz as values
# (for 2-qubit gates pz may not match the pz of the individual ion)
priority_queue, next_gate_at_pz = create_priority_queue(graph)
# check if ions are already in processing zone ->
# important for 2-qubit gates
# -> leave ion in processing zone if needed in a 2-qubit gate
for i in range(min(len(graph.pzs), len(graph.sequence))):
# only continue if previous ion was processed
gate = graph.sequence[i]
qubits = graph.get_gate_qubits(gate)
if len(qubits) == 2:
ion1, ion2 = qubits
for pz in graph.pzs:
state1 = graph.state[ion1]
state2 = graph.state[ion2]
next_gate_qubits = (
() if next_gate_at_pz[pz.name] == () else graph.get_gate_qubits(next_gate_at_pz[pz.name])
)
# append ion to in_process if it is in the correct processing zone
if state1 == pz.edge_idc and ion1 in next_gate_qubits and ion2 in next_gate_qubits:
graph.in_process.append(ion1)
# print(f"Added ion {ion1} to in_process")
if state2 == pz.edge_idc and ion1 in next_gate_qubits and ion2 in next_gate_qubits:
graph.in_process.append(ion2)
# print(f"Added ion {ion2} to in_process")
# print('in process before shuttling:', graph.in_process)
# shuttle one timestep
shuttle(graph, priority_queue, timestep, cycle_or_paths, unique_folder)
# reset ions in process
graph.in_process = []
# Check the state of each ion in the sequence
graph.state = get_ion_chains(graph)
processed_gates: list[int] = []
processed_execs: list[dict[str, Any]] = [] # NEW: gates executed this step
previous_ion_processed = True
pzs = graph.pzs.copy()
# go through the first gates in the sequence (as many as pzs or sequence length)
# for now, gates are processed in order
# (can only be processed in parallel if previous gates are processed)
for i in range(min(len(graph.pzs), len(graph.sequence))):
# only continue if previous ion was processed
if not previous_ion_processed:
break
gate = graph.sequence[i]
qubits = graph.get_gate_qubits(gate)
ion_processed = False
# wenn auf weg zu pz in anderer pz -> wird processed?
# Problem nur für 2-qubit gate? -> TODO fix
for pz in pzs:
if len(qubits) == 1:
ion = qubits[0]
if graph.state[ion] == pz.edge_idc:
print(f"Ion {ion} at Processing Zone {pz.name}")
processed_gates.insert(0, gate)
processed_execs.insert(
0,
{ # NEW
"id": f"t{timestep}_{pz.name}",
"type": "OP",
"qubits": [ion],
"edge_idc": pz.edge_idc,
"pz": pz.name,
"duration": 1,
},
)
ion_processed = True
# remove the processing zone from the list
# (it can only process one ion)
pzs.remove(pz)
# graph.in_process.append(ion)
break
elif len(qubits) == 2:
ion1, ion2 = qubits
state1 = graph.state[ion1]
state2 = graph.state[ion2]
# The following is now done at the beginning of next timestep
# (otherwise would do it double
# -> would leave the ones from last time step in in_process
# -> would not move even though
# they are move out of pz by preprocessing)
# append ion to in_process if it is in the correct processing zone
# if state1 == pz.edge_idc and ion1 in next_gate_at_pz[pz.name]
# and ion2 in next_gate_at_pz[pz.name]:
# graph.in_process.append(ion1)
# if state2 == pz.edge_idc and ion1 in next_gate_at_pz[pz.name]
# and ion2 in next_gate_at_pz[pz.name]: # also 1 qubit gate?
# graph.in_process.append(ion2)
# if both ions are in the processing zone, process the gate
if state1 == pz.edge_idc and state2 == pz.edge_idc:
print(f"Ions {ion1} and {ion2} at Processing Zone {pz.name}")
processed_gates.insert(0, gate)
processed_execs.insert(
0,
{ # NEW
"id": f"t{timestep}_{pz.name}",
"type": "OP",
"qubits": [ion1, ion2],
"edge_idc": pz.edge_idc,
"pz": pz.name,
"duration": 1,
},
)
ion_processed = True
# remove the processing zone from the list
# (it can only process one gate)
pzs.remove(pz)
# remove the locked pz of the processed two-qubit gate
if gate in graph.locked_gates and graph.locked_gates[gate] == pz.name:
graph.locked_gates.pop(gate)
break
else:
msg = "Invalid gate format"
raise ValueError(msg)
previous_ion_processed = ion_processed
# Remove processed ions from the sequence
for gate in processed_gates:
graph.sequence.remove(gate)
# Make executed gates available to the next snapshot
try:
graph.executed_gates_next = processed_execs
except Exception:
graph.executed_gates_next = []
if len(graph.sequence) == 0:
print(f"All ions have arrived at their destination in {timestep} timesteps")
break
timestep += 1
# when a timestep completes:
if record_path is not None:
timeline.append(_snapshot_state_for_json(graph, timestep))
try:
with pathlib.Path(record_path).open("w", encoding="utf-8") as f:
json.dump({"timeline": timeline}, f, separators=(",", ":"), ensure_ascii=False)
except Exception:
pass
if record_path is not None:
with pathlib.Path(record_path).open("w", encoding="utf-8") as f:
json.dump({"timeline": timeline}, f, separators=(",", ":"), ensure_ascii=False)
return timestep