diff --git a/graphqomb/feedforward.py b/graphqomb/feedforward.py index fdb0102e..ce4ccd65 100644 --- a/graphqomb/feedforward.py +++ b/graphqomb/feedforward.py @@ -19,7 +19,7 @@ import typing_extensions -from graphqomb.common import Plane +from graphqomb.common import Axis, Plane, determine_pauli_axis from graphqomb.graphstate import BaseGraphState, odd_neighbors if sys.version_info >= (3, 10): @@ -277,3 +277,61 @@ def propagate_correction_map( # noqa: C901, PLR0912 new_zflow[parent] ^= {child_z} return new_xflow, new_zflow + + +def pauli_simplification( # noqa: C901, PLR0912 + graph: BaseGraphState, + xflow: Mapping[int, AbstractSet[int]], + zflow: Mapping[int, AbstractSet[int]] | None = None, +) -> tuple[dict[int, set[int]], dict[int, set[int]]]: + r"""Simplify the correction maps by removing redundant Pauli corrections. + + Parameters + ---------- + graph : `BaseGraphState` + Underlying graph state. + xflow : `collections.abc.Mapping`\[`int`, `collections.abc.Set`\[`int`\]\] + Correction map for X. + zflow : `collections.abc.Mapping`\[`int`, `collections.abc.Set`\[`int`\]\] | `None` + Correction map for Z. If `None`, it is generated from xflow by odd neighbors. + + Returns + ------- + `tuple`\[`dict`\[`int`, `set`\[`int`\]\], `dict`\[`int`, `set`\[`int`\]\]] + Updated correction maps for X and Z after simplification. + """ + if zflow is None: + zflow = {node: odd_neighbors(xflow[node], graph) - {node} for node in xflow} + + new_xflow = {k: set(vs) for k, vs in xflow.items()} + new_zflow = {k: set(vs) for k, vs in zflow.items()} + + inv_xflow: dict[int, set[int]] = {} + inv_zflow: dict[int, set[int]] = {} + for k, vs in xflow.items(): + for v in vs: + inv_xflow.setdefault(v, set()).add(k) + for k, vs in zflow.items(): + for v in vs: + inv_zflow.setdefault(v, set()).add(k) + + for node in graph.physical_nodes - graph.output_node_indices.keys(): + meas_basis = graph.meas_bases.get(node) + if meas_basis is None: + continue + meas_axis = determine_pauli_axis(meas_basis) + if meas_axis is None: + continue + + if meas_axis == Axis.X: + for parent in inv_xflow.get(node, set()): + new_xflow[parent] -= {node} + elif meas_axis == Axis.Z: + for parent in inv_zflow.get(node, set()): + new_zflow[parent] -= {node} + elif meas_axis == Axis.Y: + for parent in inv_xflow.get(node, set()) & inv_zflow.get(node, set()): + new_xflow[parent] -= {node} + new_zflow[parent] -= {node} + + return new_xflow, new_zflow diff --git a/graphqomb/greedy_scheduler.py b/graphqomb/greedy_scheduler.py new file mode 100644 index 00000000..714cde0f --- /dev/null +++ b/graphqomb/greedy_scheduler.py @@ -0,0 +1,550 @@ +"""Greedy heuristic scheduler for fast MBQC pattern scheduling. + +This module provides fast greedy scheduling algorithms as an alternative to +CP-SAT based optimization. The greedy algorithms provide approximate solutions +with speedup compared to CP-SAT, making them suitable for large-scale +graphs or when optimality is not critical. + +This module provides: + +- `greedy_minimize_time`: Fast greedy scheduler optimizing for minimal execution time +- `greedy_minimize_space`: Fast greedy scheduler optimizing for minimal qubit usage +""" + +from __future__ import annotations + +from graphlib import CycleError, TopologicalSorter +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Mapping + from collections.abc import Set as AbstractSet + + from graphqomb.graphstate import BaseGraphState + + +def greedy_minimize_time( + graph: BaseGraphState, + dag: Mapping[int, AbstractSet[int]], + max_qubit_count: int | None = None, +) -> tuple[dict[int, int], dict[int, int]]: + r"""Fast greedy scheduler optimizing for minimal execution time (makespan). + + This algorithm uses different strategies based on max_qubit_count: + - Without qubit limit: Prepare all nodes at time=0, measure in ASAP order + - With qubit limit: Use slice-by-slice scheduling with slack-filling + + Parameters + ---------- + graph : `BaseGraphState` + The graph state to schedule + dag : `collections.abc.Mapping`\[`int`, `collections.abc.Set`\[`int`\]\] + The directed acyclic graph representing measurement dependencies + max_qubit_count : `int` | `None`, optional + Maximum allowed number of active qubits. If None, no limit is enforced. + + Returns + ------- + `tuple`\[`dict`\[`int`, `int`\], `dict`\[`int`, `int`\]\] + A tuple of (prepare_time, measure_time) dictionaries + """ + unmeasured = graph.physical_nodes - graph.output_node_indices.keys() + input_nodes = set(graph.input_node_indices.keys()) + output_nodes = set(graph.output_node_indices.keys()) + + # Build inverse DAG: for each node, track which nodes must be measured before it + inv_dag: dict[int, set[int]] = {node: set() for node in graph.physical_nodes} + for parent, children in dag.items(): + for child in children: + inv_dag[child].add(parent) + + # Cache neighbors to avoid repeated set constructions in tight loops + neighbors_map = {node: graph.neighbors(node) for node in graph.physical_nodes} + + if max_qubit_count is None: + # Optimal strategy: prepare all nodes at time=0, measure in ASAP order + return _greedy_minimize_time_unlimited( + graph, inv_dag, neighbors_map, input_nodes, output_nodes + ) + + # With qubit limit: use slice-by-slice scheduling with slack-filling + return _greedy_minimize_time_limited( + graph, + dag, + inv_dag, + neighbors_map, + unmeasured, + input_nodes, + output_nodes, + max_qubit_count, + ) + + +def _greedy_minimize_time_unlimited( + graph: BaseGraphState, + inv_dag: Mapping[int, AbstractSet[int]], + neighbors_map: Mapping[int, AbstractSet[int]], + input_nodes: AbstractSet[int], + output_nodes: AbstractSet[int], +) -> tuple[dict[int, int], dict[int, int]]: + measure_time: dict[int, int] = {} + + # 1. Compute ASAP measurement times using topological order + # Neighbor constraint: assume all non-input nodes can be prepared at time 0 + # TopologicalSorter expects {node: dependencies}, which is inv_dag + try: + topo_order = list(TopologicalSorter(inv_dag).static_order()) + except CycleError as exc: + msg = "No nodes can be measured; possible cyclic dependency or incomplete preparation." + raise RuntimeError(msg) from exc + + for node in topo_order: + if node in output_nodes: + continue + # DAG constraint: must measure after all parents + parent_times = [measure_time[p] for p in inv_dag[node] if p in measure_time] + # Neighbor constraint: all neighbors must be prepared before measurement + # Input nodes are prepared at time -1, others at time 0 (earliest possible) + neighbor_prep_times = [ + -1 if n in input_nodes else 0 for n in neighbors_map[node] + ] + # Measure at the next time slot after all parents are measured + # AND after all neighbors are prepared + measure_time[node] = max( + max(parent_times, default=-1) + 1, + max(neighbor_prep_times, default=-1) + 1, + ) + + # 2. Compute ALAP preparation times (replace time=0 with latest possible) + prepare_time = alap_prepare_times(graph, measure_time) + + return prepare_time, measure_time + + +def _greedy_minimize_time_limited( # noqa: C901, PLR0912, PLR0913, PLR0917 + graph: BaseGraphState, + dag: Mapping[int, AbstractSet[int]], + inv_dag: Mapping[int, AbstractSet[int]], + neighbors_map: Mapping[int, AbstractSet[int]], + unmeasured: AbstractSet[int], + input_nodes: AbstractSet[int], + output_nodes: AbstractSet[int], + max_qubit_count: int, +) -> tuple[dict[int, int], dict[int, int]]: + prepare_time: dict[int, int] = {} + measure_time: dict[int, int] = {} + + # Make mutable copies + inv_dag_mut: dict[int, set[int]] = { + node: set(parents) for node, parents in inv_dag.items() + } + unmeasured_mut: set[int] = set(unmeasured) + + prepared: set[int] = set(input_nodes) + alive: set[int] = set(input_nodes) + + if len(alive) > max_qubit_count: + msg = "Initial number of active qubits exceeds max_qubit_count." + raise RuntimeError(msg) + + # Compute criticality for prioritizing preparations + criticality = _compute_criticality(dag, output_nodes) + + current_time = 0 + + while unmeasured_mut: + # Phase 1: Measure all ready nodes + # A node is ready if: + # - DAG dependencies are resolved (inv_dag_mut[node] is empty) + # - All neighbors are prepared + # - The node itself is prepared (if not an input node) + ready_to_measure: set[int] = set() + for node in unmeasured_mut: + if inv_dag_mut[node]: + continue # DAG dependencies not resolved + if not neighbors_map[node] <= prepared: + continue # Neighbors not prepared + if node not in input_nodes and node not in prepared: + continue # Self not prepared + ready_to_measure.add(node) + + for node in ready_to_measure: + measure_time[node] = current_time + unmeasured_mut.remove(node) + alive.discard(node) + + # Update DAG dependencies + for child in dag.get(node, ()): + inv_dag_mut[child].discard(node) + + # Phase 2: Prepare nodes using free capacity (slack-filling) + free_capacity = max_qubit_count - len(alive) + + if free_capacity > 0: + # Get unprepared nodes with their priority scores + unprepared = graph.physical_nodes - prepared + if unprepared: + prep_candidates = _get_prep_candidates_with_priority( + unprepared, + inv_dag_mut, + neighbors_map, + prepared, + unmeasured_mut, + output_nodes, + criticality, + ) + # Prepare top candidates within free capacity + for candidate, _score in prep_candidates[:free_capacity]: + prepare_time[candidate] = current_time + prepared.add(candidate) + alive.add(candidate) + + # Check if we made progress + if not ready_to_measure and free_capacity == 0 and unmeasured_mut: + # No measurements and no room to prepare - stuck + msg = ( + "Cannot schedule more measurements without exceeding max qubit count. " + "Please increase max_qubit_count." + ) + raise RuntimeError(msg) + + current_time += 1 + + # Safety check for infinite loops + if current_time > len(graph.physical_nodes) * 2: + msg = "Scheduling did not converge; possible cyclic dependency." + raise RuntimeError(msg) + + # Apply ALAP post-processing to minimize active volume + prepare_time = alap_prepare_times(graph, measure_time) + + return prepare_time, measure_time + + +def _compute_criticality( + dag: Mapping[int, AbstractSet[int]], + output_nodes: AbstractSet[int], +) -> dict[int, int]: + # Compute criticality (remaining DAG depth) for each node. + # Nodes with higher criticality should be prioritized for unblocking. + criticality: dict[int, int] = {} + + # TopologicalSorter(dag) returns nodes with no "dependencies" first. + # Since dag is {parent: children}, nodes with empty children come first (leaves). + # This is the correct order for computing criticality (leaves before roots). + try: + topo_order = list(TopologicalSorter(dag).static_order()) + except CycleError: + return {} + + for node in topo_order: + children_crits = [criticality.get(c, 0) for c in dag.get(node, ())] + criticality[node] = 1 + max(children_crits, default=0) + + # Output nodes have criticality 0 (they don't need to be measured) + for node in output_nodes: + criticality[node] = 0 + + return criticality + + +def _get_prep_candidates_with_priority( # noqa: PLR0913, PLR0917 + unprepared: AbstractSet[int], + inv_dag: Mapping[int, AbstractSet[int]], + neighbors_map: Mapping[int, AbstractSet[int]], + prepared: AbstractSet[int], + unmeasured: AbstractSet[int], + output_nodes: AbstractSet[int], + criticality: Mapping[int, int], +) -> list[tuple[int, float]]: + # Get preparation candidates sorted by priority score. + # Priority is based on how much preparing a node helps unblock measurements. + # Find nodes that are DAG-ready but blocked by missing neighbors + dag_ready_blocked: set[int] = set() + missing_map: dict[int, set[int]] = {} + + for node in unmeasured: + if inv_dag[node]: + continue # Not DAG-ready + missing = set(neighbors_map[node]) - set(prepared) + # Also check if the node itself needs preparation + if node not in prepared: + missing.add(node) + if missing: + dag_ready_blocked.add(node) + missing_map[node] = missing + + # Score each unprepared node + scores: list[tuple[int, float]] = [] + for candidate in unprepared: + score = 0.0 + for blocked_node in dag_ready_blocked: + if candidate in missing_map[blocked_node]: + crit = criticality.get(blocked_node, 1) + score += crit / len(missing_map[blocked_node]) + + # Apply penalty for output nodes (they stay alive forever) + if candidate in output_nodes: + score *= 0.5 + + scores.append((candidate, score)) + + # Sort by score descending (higher score = higher priority) + scores.sort(key=lambda x: -x[1]) + + return scores + + +def _determine_measure_nodes( + neighbors_map: Mapping[int, AbstractSet[int]], + measure_candidates: AbstractSet[int], + prepared: AbstractSet[int], + alive: AbstractSet[int], + max_qubit_count: int, +) -> tuple[set[int], set[int]]: + r"""Determine which nodes to measure without exceeding max qubit count. + + Parameters + ---------- + neighbors_map : `collections.abc.Mapping`\[`int`, `collections.abc.Set`\[`int`\]\] + Mapping from node to its neighbors. + measure_candidates : `collections.abc.Set`\[`int`\] + The candidate nodes available for measurement. + prepared : `collections.abc.Set`\[`int`\] + The set of currently prepared nodes. + alive : `collections.abc.Set`\[`int`\] + The set of currently active (prepared but not yet measured) nodes. + max_qubit_count : `int` + The maximum allowed number of active qubits. + + Returns + ------- + `tuple`\[`set`\[`int`\], `set`\[`int`\]\] + A tuple of (to_measure, to_prepare) sets indicating which nodes to measure and prepare. + + Raises + ------ + RuntimeError + If no nodes can be measured without exceeding the max qubit count. + """ + to_measure: set[int] = set() + to_prepare: set[int] = set() + + for node in measure_candidates: + # Neighbors that still need to be prepared for this node + new_neighbors = neighbors_map[node] - prepared + additional_to_prepare = new_neighbors - to_prepare + + # Projected number of active qubits after preparing these neighbors + projected_active = len(alive) + len(to_prepare) + len(additional_to_prepare) + + if projected_active <= max_qubit_count: + to_measure.add(node) + to_prepare |= new_neighbors + + if not to_measure: + msg = "Cannot schedule more measurements without exceeding max qubit count. Please increase max_qubit_count." + raise RuntimeError(msg) + + return to_measure, to_prepare + + +def greedy_minimize_space( # noqa: C901, PLR0914 + graph: BaseGraphState, + dag: Mapping[int, AbstractSet[int]], +) -> tuple[dict[int, int], dict[int, int]]: + r"""Fast greedy scheduler optimizing for minimal qubit usage (space). + + This algorithm uses a greedy approach to minimize the number of active + qubits at each time step: + 1. At each time step, select the next node to measure that minimizes the + projected number of alive qubits after any required preparations. + 2. Prepare neighbors of the measured node just before measurement. + + Parameters + ---------- + graph : `BaseGraphState` + The graph state to schedule + dag : `collections.abc.Mapping`\[`int`, `collections.abc.Set`\[`int`\]\] + The directed acyclic graph representing measurement dependencies + + Returns + ------- + `tuple`\[`dict`\[`int`, `int`\], `dict`\[`int`, `int`\]\] + A tuple of (prepare_time, measure_time) dictionaries + + Raises + ------ + RuntimeError + If no nodes can be measured at a given time step, indicating a possible + cyclic dependency or incomplete preparation. + """ + prepare_time: dict[int, int] = {} + measure_time: dict[int, int] = {} + + unmeasured = graph.physical_nodes - graph.output_node_indices.keys() + + try: + topo_order = list(TopologicalSorter(dag).static_order()) + except CycleError as exc: + msg = "No nodes can be measured; possible cyclic dependency or incomplete preparation." + raise RuntimeError(msg) from exc + topo_order.reverse() # from parents to children + topo_rank = {node: i for i, node in enumerate(topo_order)} + + # Build inverse DAG: for each node, track which nodes must be measured before it + inv_dag: dict[int, set[int]] = {node: set() for node in graph.physical_nodes} + for parent, children in dag.items(): + for child in children: + inv_dag[child].add(parent) + + prepared: set[int] = set(graph.input_node_indices.keys()) + alive: set[int] = set(graph.input_node_indices.keys()) + current_time = 0 + + # Cache neighbors once as the graph is static during scheduling + neighbors_map = {node: graph.neighbors(node) for node in graph.physical_nodes} + + measure_candidates: set[int] = {node for node in unmeasured if not inv_dag[node]} + + while unmeasured: + if not measure_candidates: + msg = "No nodes can be measured; possible cyclic dependency or incomplete preparation." + raise RuntimeError(msg) + + # calculate costs and pick the best node to measure + default_rank = len(topo_rank) + candidates = iter(measure_candidates) + best_node = next(candidates) + best_cost = _calc_activate_cost(best_node, neighbors_map, prepared, alive) + best_rank = topo_rank.get(best_node, default_rank) + for node in candidates: + cost = _calc_activate_cost(node, neighbors_map, prepared, alive) + rank = topo_rank.get(node, default_rank) + if cost < best_cost or (cost == best_cost and rank < best_rank): + best_cost = cost + best_rank = rank + best_node = node + + # Prepare neighbors at current_time + new_neighbors = neighbors_map[best_node] - prepared + needs_prep = bool(new_neighbors) + if needs_prep: + for neighbor in new_neighbors: + prepare_time[neighbor] = current_time + prepared.update(new_neighbors) + alive.update(new_neighbors) + + # Measure at current_time if no prep needed, otherwise at current_time + 1 + meas_time = current_time + 1 if needs_prep else current_time + measure_time[best_node] = meas_time + unmeasured.remove(best_node) + alive.remove(best_node) + + measure_candidates.remove(best_node) + + # Remove measured node from dependencies of all its children in the DAG + for child in dag.get(best_node, ()): + inv_dag[child].remove(best_node) + if not inv_dag[child] and child in unmeasured: + measure_candidates.add(child) + + current_time = meas_time + 1 + + return prepare_time, measure_time + + +def _calc_activate_cost( + node: int, + neighbors_map: Mapping[int, AbstractSet[int]], + prepared: AbstractSet[int], + alive: AbstractSet[int], +) -> int: + r"""Calculate the projected number of alive qubits if measuring this node next. + + If neighbors must be prepared, they become alive at the current time slice + while the node itself remains alive until the next slice. If no preparation + is needed, the node is measured in the current slice and removed. + + Parameters + ---------- + node : `int` + The node to evaluate. + neighbors_map : `collections.abc.Mapping`\[`int`, `collections.abc.Set`\[`int`\]\] + Cached neighbor sets for graph nodes. + prepared : `collections.abc.Set`\[`int`\] + The set of currently prepared nodes. + alive : `collections.abc.Set`\[`int`\] + The set of currently active (prepared but not yet measured) nodes. + + Returns + ------- + `int` + The activation cost for the node. + """ + new_neighbors = neighbors_map[node] - prepared + if new_neighbors: + return len(alive) + len(new_neighbors) + # No preparation needed -> node is measured in the current slice, so alive decreases by 1. + return max(len(alive) - 1, 0) + + +def alap_prepare_times( + graph: BaseGraphState, + measure_time: Mapping[int, int], +) -> dict[int, int]: + r"""Recompute preparation times using ALAP (As Late As Possible) strategy. + + Given fixed measurement times, this computes the latest possible preparation + time for each node while respecting the constraint that all neighbors must + be prepared before a node is measured. + + This post-processing reduces active volume (sum of qubit lifetimes) without + changing the measurement schedule or depth. + + Parameters + ---------- + graph : `BaseGraphState` + The graph state + measure_time : `collections.abc.Mapping`\[`int`, `int`\] + Fixed measurement times for non-output nodes + + Returns + ------- + `dict`\[`int`, `int`\] + ALAP preparation times for non-input nodes + """ + input_nodes = set(graph.input_node_indices.keys()) + + # deadline[v] = latest time v can be prepared + deadline: dict[int, int] = {} + + # For each measured node u, all its neighbors must be prepared before meas(u) + for u, meas_u in measure_time.items(): + for neighbor in graph.neighbors(u): + if neighbor in input_nodes: + continue # Input nodes don't need prep + if neighbor not in deadline: + deadline[neighbor] = meas_u - 1 + else: + deadline[neighbor] = min(deadline[neighbor], meas_u - 1) + + # For measured nodes, they must be prepared before their own measurement + for v, meas_v in measure_time.items(): + if v in input_nodes: + continue # Input nodes don't need prep + if v not in deadline: + deadline[v] = meas_v - 1 + else: + deadline[v] = min(deadline[v], meas_v - 1) + + # Handle nodes with no deadline yet (output nodes with no measured neighbors) + # These should be prepared at the latest possible time: max(measure_time) - 1 + # or 0 if there are no measurements + makespan = max(measure_time.values(), default=0) + for v in graph.physical_nodes - input_nodes: + if v not in deadline: + # No constraint from neighbors, prep as late as possible + deadline[v] = max(makespan - 1, 0) + + return deadline diff --git a/graphqomb/pattern.py b/graphqomb/pattern.py index b9718da4..c17307bb 100644 --- a/graphqomb/pattern.py +++ b/graphqomb/pattern.py @@ -101,6 +101,76 @@ def depth(self) -> int: """ return sum(1 for cmd in self.commands if isinstance(cmd, TICK)) + @property + def active_volume(self) -> int: + """Calculate tha active volume, summation of space for each timeslice. + + Returns + ------- + `int` + Active volume of the pattern + """ + return sum(self.space) + + @property + def volume(self) -> int: + """Calculate the volume, defined as max_space * depth. + + Returns + ------- + `int` + Volume of the pattern + """ + return self.max_space * self.depth + + @property + def idle_times(self) -> dict[int, int]: + r"""Calculate the idle times for each qubit in the pattern. + + Returns + ------- + `dict`\[`int`, `int`\] + A dictionary mapping each qubit index to its idle time. + """ + idle_times: dict[int, int] = {} + prepared_time: dict[int, int] = dict.fromkeys(self.input_node_indices, -1) + + current_time = 0 + for cmd in self.commands: + if isinstance(cmd, TICK): + current_time += 1 + elif isinstance(cmd, N): + prepared_time[cmd.node] = current_time + elif isinstance(cmd, M): + idle_times[cmd.node] = current_time - prepared_time[cmd.node] + + for output_node in self.output_node_indices: + if output_node in prepared_time: + idle_times[output_node] = current_time - prepared_time[output_node] + + return idle_times + + @property + def throughput(self) -> float: + """Calculate the number of measurements per TICK in the pattern. + + Returns + ------- + `float` + Number of measurements per TICK + + Raises + ------ + ValueError + If the pattern has zero depth (no TICK commands) + """ + num_measurements = sum(1 for cmd in self.commands if isinstance(cmd, M)) + num_ticks = self.depth + if num_ticks == 0: + msg = "Cannot calculate throughput for a pattern with zero depth (no TICK commands)." + raise ValueError(msg) + return num_measurements / num_ticks + def is_runnable(pattern: Pattern) -> None: """Check if the pattern is runnable. diff --git a/graphqomb/schedule_solver.py b/graphqomb/schedule_solver.py index 90cbeb05..6c495079 100644 --- a/graphqomb/schedule_solver.py +++ b/graphqomb/schedule_solver.py @@ -37,6 +37,7 @@ class ScheduleConfig: strategy: Strategy max_qubit_count: int | None = None max_time: int | None = None + use_greedy: bool = False @dataclass @@ -123,15 +124,65 @@ def _compute_alive_nodes_at_time( ctx.model.Add(q <= t).OnlyEnforceIf(a_meas) ctx.model.Add(q > t).OnlyEnforceIf(a_meas.Not()) + # alive <=> (a_pre AND NOT a_meas) + # A node is alive at time t if it has been prepared (prep <= t) + # and has not yet been measured (meas > t) alive = ctx.model.NewBoolVar(f"alive_{node}_{t}") + # Forward: alive => (a_pre AND NOT a_meas) ctx.model.AddImplication(alive, a_pre) ctx.model.AddImplication(alive, a_meas.Not()) - ctx.model.Add(a_pre - a_meas <= alive) + # Backward: (a_pre AND NOT a_meas) => alive + # Equivalent to: (NOT a_pre OR a_meas OR alive) + ctx.model.AddBoolOr([a_pre.Not(), a_meas, alive]) alive_at_t.append(alive) return alive_at_t +def _add_lifetime_cumulative_constraint( + ctx: _ModelContext, + node2prep: Mapping[int, cp_model.IntVar], + node2meas: Mapping[int, cp_model.IntVar], + *, + max_time: int, + capacity: int | cp_model.IntVar, + name_prefix: str, +) -> None: + """Add a cumulative constraint that limits the number of alive qubits. + + A qubit is considered alive at time t if it has been prepared (prep <= t) + and has not yet been measured (meas > t). This corresponds to an interval + [prep, meas) in the CP-SAT model. Input nodes are treated as prepared at + time 0, and output nodes are treated as alive until max_time. + """ + intervals: list[cp_model.IntervalVar] = [] + demands: list[int] = [] + + for node in ctx.graph.physical_nodes: + if node in node2prep and node in node2meas: + ctx.model.Add(node2prep[node] < node2meas[node]) + + start = ( + ctx.model.NewConstant(0) + if node in ctx.graph.input_node_indices + else node2prep[node] + ) + end = ( + ctx.model.NewConstant(max_time) + if node in ctx.graph.output_node_indices + else node2meas[node] + ) + + duration = ctx.model.NewIntVar(0, max_time, f"{name_prefix}_dur_{node}") + ctx.model.Add(end - start == duration) + intervals.append( + ctx.model.NewIntervalVar(start, duration, end, f"{name_prefix}_{node}") + ) + demands.append(1) + + ctx.model.AddCumulative(intervals, demands, capacity) + + def _set_minimize_space_objective( ctx: _ModelContext, node2prep: Mapping[int, cp_model.IntVar], @@ -140,9 +191,14 @@ def _set_minimize_space_objective( ) -> None: """Set objective to minimize the maximum number of qubits used at any time.""" max_space = ctx.model.NewIntVar(0, len(ctx.graph.physical_nodes), "max_space") - for t in range(max_time): - alive_at_t = _compute_alive_nodes_at_time(ctx, node2prep, node2meas, t) - ctx.model.Add(max_space >= sum(alive_at_t)) + _add_lifetime_cumulative_constraint( + ctx, + node2prep, + node2meas, + max_time=max_time, + capacity=max_space, + name_prefix="alive_interval", + ) ctx.model.Minimize(max_space) @@ -156,9 +212,14 @@ def _set_minimize_time_objective( """Set objective to minimize the total execution time.""" # Add space constraint if max_qubit_count is specified if max_qubit_count is not None: - for t in range(max_time): - alive_at_t = _compute_alive_nodes_at_time(ctx, node2prep, node2meas, t) - ctx.model.Add(sum(alive_at_t) <= max_qubit_count) + _add_lifetime_cumulative_constraint( + ctx, + node2prep, + node2meas, + max_time=max_time, + capacity=max_qubit_count, + name_prefix="alive_interval", + ) # Time objective: minimize makespan meas_vars = list(node2meas.values()) diff --git a/graphqomb/scheduler.py b/graphqomb/scheduler.py index 889f045d..b5d16fb5 100644 --- a/graphqomb/scheduler.py +++ b/graphqomb/scheduler.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, NamedTuple from graphqomb.feedforward import dag_from_flow +from graphqomb.greedy_scheduler import greedy_minimize_space, greedy_minimize_time from graphqomb.schedule_solver import ScheduleConfig, Strategy, solve_schedule if TYPE_CHECKING: @@ -484,14 +485,15 @@ def solve_schedule( config: ScheduleConfig | None = None, timeout: int = 60, ) -> bool: - r"""Compute the schedule using the constraint programming solver. + r"""Compute the schedule using constraint programming or greedy heuristics. Parameters ---------- config : `ScheduleConfig` | `None`, optional - The scheduling configuration. If None, defaults to MINIMIZE_SPACE strategy. + The scheduling configuration. If None, defaults to MINIMIZE_TIME strategy. timeout : `int`, optional - Maximum solve time in seconds, by default 60 + Maximum solve time in seconds for CP-SAT solver, by default 60. + Ignored when use_greedy=True. Returns ------- @@ -506,7 +508,17 @@ def solve_schedule( if config is None: config = ScheduleConfig(Strategy.MINIMIZE_TIME) - result = solve_schedule(self.graph, self.dag, config, timeout) + result: tuple[dict[int, int], dict[int, int]] | None + if config.use_greedy: + # Use fast greedy heuristics + if config.strategy == Strategy.MINIMIZE_TIME: + result = greedy_minimize_time(self.graph, self.dag, max_qubit_count=config.max_qubit_count) + else: # Strategy.MINIMIZE_SPACE + result = greedy_minimize_space(self.graph, self.dag) + else: + # Use CP-SAT solver for optimal solution + result = solve_schedule(self.graph, self.dag, config, timeout) + if result is None: return False diff --git a/tests/test_greedy_scheduler.py b/tests/test_greedy_scheduler.py new file mode 100644 index 00000000..af294e27 --- /dev/null +++ b/tests/test_greedy_scheduler.py @@ -0,0 +1,575 @@ +"""Test greedy scheduling algorithms.""" + +import pytest + +from graphqomb.graphstate import GraphState +from graphqomb.greedy_scheduler import ( + greedy_minimize_space, + greedy_minimize_time, +) +from graphqomb.schedule_solver import ScheduleConfig, Strategy +from graphqomb.scheduler import Scheduler + + +def test_greedy_minimize_time_simple() -> None: + """Test greedy_minimize_time on a simple graph.""" + # Create a simple 3-node chain graph + graph = GraphState() + node0 = graph.add_physical_node() + node1 = graph.add_physical_node() + node2 = graph.add_physical_node() + graph.add_physical_edge(node0, node1) + graph.add_physical_edge(node1, node2) + qindex = 0 + graph.register_input(node0, qindex) + graph.register_output(node2, qindex) + + flow = {node0: {node1}, node1: {node2}} + scheduler = Scheduler(graph, flow) + + # Run greedy scheduler + prepare_time, measure_time = greedy_minimize_time(graph, scheduler.dag) + + # Check that all non-input nodes have preparation times + assert node1 in prepare_time + assert node0 not in prepare_time # Input node should not be prepared + + # Check that all non-output nodes have measurement times + assert node0 in measure_time + assert node1 in measure_time + assert node2 not in measure_time # Output node should not be measured + + # Verify DAG constraints: node0 measured before node1 + assert measure_time[node0] < measure_time[node1] + + +def test_greedy_minimize_space_simple() -> None: + """Test greedy_minimize_space on a simple graph.""" + # Create a simple 3-node chain graph + graph = GraphState() + node0 = graph.add_physical_node() + node1 = graph.add_physical_node() + node2 = graph.add_physical_node() + graph.add_physical_edge(node0, node1) + graph.add_physical_edge(node1, node2) + qindex = 0 + graph.register_input(node0, qindex) + graph.register_output(node2, qindex) + + flow = {node0: {node1}, node1: {node2}} + scheduler = Scheduler(graph, flow) + + # Run greedy scheduler + prepare_time, measure_time = greedy_minimize_space(graph, scheduler.dag) + + # Check that all non-input nodes have preparation times + assert node1 in prepare_time + assert node0 not in prepare_time # Input node should not be prepared + + # Check that all non-output nodes have measurement times + assert node0 in measure_time + assert node1 in measure_time + assert node2 not in measure_time # Output node should not be measured + + # Verify DAG constraints + assert measure_time[node0] < measure_time[node1] + + +def _compute_max_alive_qubits( + graph: GraphState, + prepare_time: dict[int, int], + measure_time: dict[int, int], +) -> int: + """Compute the maximum number of alive qubits over time. + + A node is considered alive at time t if: + - It is an input node and t >= -1 and t < measurement time (if any), or + - It has a preparation time p and t >= p and t < measurement time (if any). + + Returns + ------- + int + The maximum number of alive qubits at any time step. + """ + # Determine time range to check + max_t = max(set(prepare_time.values()) | set(measure_time.values()), default=0) + + max_alive = len(graph.input_node_indices) # At least inputs are alive at t = -1 + for t in range(max_t + 1): + alive_nodes: set[int] = set() + for node in graph.physical_nodes: + # Determine preparation time + prep_t = -1 if node in graph.input_node_indices else prepare_time.get(node) + + if prep_t is None or t < prep_t: + continue + + # Determine measurement time (None for outputs or unscheduled) + meas_t = measure_time.get(node) + + if meas_t is None or t < meas_t: + alive_nodes.add(node) + + max_alive = max(max_alive, len(alive_nodes)) + + return max_alive + + +def test_greedy_minimize_time_with_max_qubit_count_respects_limit() -> None: + """Verify that greedy_minimize_time respects max_qubit_count.""" + graph = GraphState() + # chain graph: 0-1-2-3 + n0 = graph.add_physical_node() + n1 = graph.add_physical_node() + n2 = graph.add_physical_node() + n3 = graph.add_physical_node() + graph.add_physical_edge(n0, n1) + graph.add_physical_edge(n1, n2) + graph.add_physical_edge(n2, n3) + + qindex = 0 + graph.register_input(n0, qindex) + graph.register_output(n3, qindex) + + flow = {n0: {n1}, n1: {n2}, n2: {n3}} + scheduler = Scheduler(graph, flow) + + # Set max_qubit_count to 2 (a feasible value for this graph) + prepare_time, measure_time = greedy_minimize_time(graph, scheduler.dag, max_qubit_count=2) + + # Check basic properties + assert n1 in prepare_time + assert n0 not in prepare_time + assert n0 in measure_time + assert n2 in measure_time + assert n3 not in measure_time + + # Verify that the number of alive qubits never exceeds the limit + max_alive = _compute_max_alive_qubits(graph, prepare_time, measure_time) + assert max_alive <= 2 + + +def test_greedy_minimize_time_with_too_small_max_qubit_count_raises() -> None: + """Verify that greedy_minimize_time raises RuntimeError when max_qubit_count is too small.""" + graph = GraphState() + # chain graph: 0-1-2 (at least 2 qubits are needed) + n0 = graph.add_physical_node() + n1 = graph.add_physical_node() + n2 = graph.add_physical_node() + graph.add_physical_edge(n0, n1) + graph.add_physical_edge(n1, n2) + + qindex = 0 + graph.register_input(n0, qindex) + graph.register_output(n2, qindex) + + flow = {n0: {n1}, n1: {n2}} + scheduler = Scheduler(graph, flow) + + # max_qubit_count=1 is not feasible, so expect RuntimeError + with pytest.raises(RuntimeError, match="max_qubit_count"): + greedy_minimize_time(graph, scheduler.dag, max_qubit_count=1) + + +def test_greedy_scheduler_via_solve_schedule() -> None: + """Test greedy scheduler through Scheduler.solve_schedule with use_greedy=True.""" + # Create a simple graph + graph = GraphState() + node0 = graph.add_physical_node() + node1 = graph.add_physical_node() + node2 = graph.add_physical_node() + graph.add_physical_edge(node0, node1) + graph.add_physical_edge(node1, node2) + qindex = 0 + graph.register_input(node0, qindex) + graph.register_output(node2, qindex) + + flow = {node0: {node1}, node1: {node2}} + scheduler = Scheduler(graph, flow) + + # Test with greedy MINIMIZE_TIME + config = ScheduleConfig(strategy=Strategy.MINIMIZE_TIME, use_greedy=True) + success = scheduler.solve_schedule(config) + assert success + + # Verify schedule is valid + scheduler.validate_schedule() + + # Test with greedy MINIMIZE_SPACE + scheduler2 = Scheduler(graph, flow) + config = ScheduleConfig(strategy=Strategy.MINIMIZE_SPACE, use_greedy=True) + success = scheduler2.solve_schedule(config) + assert success + + # Verify schedule is valid + scheduler2.validate_schedule() + + +def test_greedy_vs_cpsat_correctness() -> None: + """Test that greedy scheduler produces valid schedules compared to CP-SAT.""" + # Create a slightly larger graph + graph = GraphState() + nodes = [graph.add_physical_node() for _ in range(5)] + + # Create a chain + for i in range(4): + graph.add_physical_edge(nodes[i], nodes[i + 1]) + + qindex = 0 + graph.register_input(nodes[0], qindex) + graph.register_output(nodes[4], qindex) + + flow = {nodes[i]: {nodes[i + 1]} for i in range(4)} + + # Test greedy scheduler + scheduler_greedy = Scheduler(graph, flow) + config = ScheduleConfig(strategy=Strategy.MINIMIZE_TIME, use_greedy=True) + success_greedy = scheduler_greedy.solve_schedule(config) + assert success_greedy + + # Verify greedy schedule is valid + scheduler_greedy.validate_schedule() + + # Test CP-SAT scheduler + scheduler_cpsat = Scheduler(graph, flow) + config = ScheduleConfig(strategy=Strategy.MINIMIZE_TIME, use_greedy=False) + success_cpsat = scheduler_cpsat.solve_schedule(config, timeout=10) + assert success_cpsat + + # Verify CP-SAT schedule is valid + scheduler_cpsat.validate_schedule() + + # Both should produce valid schedules + # Note: Greedy may not be optimal, so we don't compare quality here + + +def test_greedy_scheduler_larger_graph() -> None: + """Test greedy scheduler on a larger graph to ensure scalability.""" + # Create a larger graph with branching structure + graph = GraphState() + num_layers = 4 + nodes_per_layer = 3 + + # Build layered graph + all_nodes: list[list[int]] = [] + for layer in range(num_layers): + layer_nodes = [graph.add_physical_node() for _ in range(nodes_per_layer)] + all_nodes.append(layer_nodes) + + # Connect to previous layer (if not first layer) + if layer > 0: + for i, node in enumerate(layer_nodes): + # Connect to corresponding node in previous layer + prev_node = all_nodes[layer - 1][i] + graph.add_physical_edge(prev_node, node) + + # Register inputs (first layer) and outputs (last layer) + for i, node in enumerate(all_nodes[0]): + graph.register_input(node, i) + for i, node in enumerate(all_nodes[-1]): + graph.register_output(node, i) + + # Build flow (simple forward flow) + flow: dict[int, set[int]] = {} + for layer in range(num_layers - 1): + for i, node in enumerate(all_nodes[layer]): + if node not in graph.output_node_indices: + flow[node] = {all_nodes[layer + 1][i]} + + # Test greedy scheduler + scheduler = Scheduler(graph, flow) + config = ScheduleConfig(strategy=Strategy.MINIMIZE_TIME, use_greedy=True) + success = scheduler.solve_schedule(config) + assert success + + # Validate the schedule + scheduler.validate_schedule() + + # Check that we got reasonable results + assert scheduler.num_slices() > 0 + assert scheduler.num_slices() <= num_layers * 2 # Reasonable upper bound + + +@pytest.mark.parametrize("strategy", [Strategy.MINIMIZE_TIME, Strategy.MINIMIZE_SPACE]) +def test_greedy_scheduler_both_strategies(strategy: Strategy) -> None: + """Test greedy scheduler with both optimization strategies.""" + # Create a graph + graph = GraphState() + node0 = graph.add_physical_node() + node1 = graph.add_physical_node() + node2 = graph.add_physical_node() + node3 = graph.add_physical_node() + graph.add_physical_edge(node0, node1) + graph.add_physical_edge(node1, node2) + graph.add_physical_edge(node2, node3) + qindex = 0 + graph.register_input(node0, qindex) + graph.register_output(node3, qindex) + + flow = {node0: {node1}, node1: {node2}, node2: {node3}} + scheduler = Scheduler(graph, flow) + + # Test with specified strategy + config = ScheduleConfig(strategy=strategy, use_greedy=True) + success = scheduler.solve_schedule(config) + assert success + + # Validate schedule + scheduler.validate_schedule() + + +def test_greedy_minimize_space_wrapper() -> None: + """Test the greedy_minimize_space wrapper function.""" + # Create a simple graph + graph = GraphState() + node0 = graph.add_physical_node() + node1 = graph.add_physical_node() + node2 = graph.add_physical_node() + graph.add_physical_edge(node0, node1) + graph.add_physical_edge(node1, node2) + qindex = 0 + graph.register_input(node0, qindex) + graph.register_output(node2, qindex) + + flow = {node0: {node1}, node1: {node2}} + scheduler = Scheduler(graph, flow) + + # Test MINIMIZE_TIME + result = greedy_minimize_time(graph, scheduler.dag) + assert result is not None + prepare_time, measure_time = result + assert len(prepare_time) > 0 + assert len(measure_time) > 0 + + # Test MINIMIZE_SPACE + result = greedy_minimize_space(graph, scheduler.dag) + assert result is not None + prepare_time, measure_time = result + assert len(prepare_time) > 0 + assert len(measure_time) > 0 + + +def test_greedy_scheduler_dag_constraints() -> None: + """Test that greedy scheduler respects DAG constraints.""" + # Create a graph with more complex dependencies + graph = GraphState() + nodes = [graph.add_physical_node() for _ in range(6)] + + # Create edges forming a DAG structure + # 0 -> 2 -> 4 + # | + # 1 -> 3 -> 5 + graph.add_physical_edge(nodes[0], nodes[2]) + graph.add_physical_edge(nodes[2], nodes[4]) + graph.add_physical_edge(nodes[1], nodes[3]) + graph.add_physical_edge(nodes[3], nodes[5]) + graph.add_physical_edge(nodes[2], nodes[3]) + + graph.register_input(nodes[0], 0) + graph.register_input(nodes[1], 1) + graph.register_output(nodes[4], 0) + graph.register_output(nodes[5], 1) + + # Create flow with dependencies + flow = { + nodes[0]: {nodes[2]}, + nodes[1]: {nodes[3]}, + nodes[2]: {nodes[4]}, + nodes[3]: {nodes[5], nodes[1]}, # cyclic dependency to test DAG constraint handling + } + + scheduler = Scheduler(graph, flow) + config = ScheduleConfig(strategy=Strategy.MINIMIZE_TIME, use_greedy=True) + + # Note: This flow creates a cyclic DAG (nodes 3 and 4 have circular dependency) + # The greedy scheduler should raise RuntimeError for invalid flows + with pytest.raises(RuntimeError, match="No nodes can be measured"): + scheduler.solve_schedule(config) + + +def test_greedy_scheduler_edge_constraints() -> None: + """Test that greedy scheduler respects edge constraints (neighbor preparation).""" + # Create a simple graph + graph = GraphState() + node0 = graph.add_physical_node() + node1 = graph.add_physical_node() + node2 = graph.add_physical_node() + graph.add_physical_edge(node0, node1) + graph.add_physical_edge(node1, node2) + qindex = 0 + graph.register_input(node0, qindex) + graph.register_output(node2, qindex) + + flow = {node0: {node1}, node1: {node2}} + scheduler = Scheduler(graph, flow) + config = ScheduleConfig(strategy=Strategy.MINIMIZE_TIME, use_greedy=True) + success = scheduler.solve_schedule(config) + assert success + + # Validate edge constraints via validate_schedule + scheduler.validate_schedule() + + # Manually check: neighbors must be prepared before measurement + # node0 (input) is prepared at time -1, node1 prepared at some time + # node0 must be measured after node1 is prepared + # This is ensured by the auto-scheduled entanglement times + + # Check that entanglement times were auto-scheduled correctly + edge01 = (node0, node1) + edge12 = (node1, node2) + entangle01 = scheduler.entangle_time[edge01] + entangle12 = scheduler.entangle_time[edge12] + assert entangle01 is not None + assert entangle12 is not None + + # Entanglement must happen before measurement + meas0 = scheduler.measure_time[node0] + meas1 = scheduler.measure_time[node1] + assert meas0 is not None + assert meas1 is not None + assert entangle01 < meas0 + assert entangle12 < meas1 + + +def test_greedy_minimize_time_3x3_grid_optimal() -> None: + """Test that greedy_minimize_time achieves optimal depth on 3x3 grid. + + This is a regression test for the optimization that measures in ASAP order + based on DAG dependencies. With ALAP preparation, nodes are prepared as + late as possible, but depth should still be optimal. + Previously, the greedy algorithm produced depth=4 instead of optimal depth=3. + """ + # Create 3x3 grid graph + # Layout: + # 0 - 3 - 6 + # | | | + # 1 - 4 - 7 + # | | | + # 2 - 5 - 8 + # Inputs: 0, 1, 2 (left column) + # Outputs: 6, 7, 8 (right column) + graph = GraphState() + nodes = [graph.add_physical_node() for _ in range(9)] + + # Horizontal edges + for row in range(3): + for col in range(2): + graph.add_physical_edge(nodes[row + col * 3], nodes[row + (col + 1) * 3]) + + # Vertical edges + for row in range(2): + for col in range(3): + graph.add_physical_edge(nodes[row + col * 3], nodes[row + 1 + col * 3]) + + # Register inputs (left column) and outputs (right column) + for row in range(3): + graph.register_input(nodes[row], row) + graph.register_output(nodes[row + 6], row) + + # Flow: left to right + flow: dict[int, set[int]] = {} + for row in range(3): + flow[nodes[row]] = {nodes[row + 3]} # 0->3, 1->4, 2->5 + flow[nodes[row + 3]] = {nodes[row + 6]} # 3->6, 4->7, 5->8 + + scheduler = Scheduler(graph, flow) + + # Test greedy scheduler (no qubit limit) + prepare_time, measure_time = greedy_minimize_time(graph, scheduler.dag) + + # With ALAP, nodes are prepared as late as possible, not at time=0 + # Check that all non-input nodes have a prepare_time + for node in [3, 4, 5, 6, 7, 8]: + assert node in prepare_time, f"Node {node} should have a prepare_time" + + # Calculate depth + greedy_depth = max(measure_time.values()) + 1 + + # The optimal depth for a 3x3 grid is 3 (same as CP-SAT) + assert greedy_depth == 3, f"Expected depth=3, got depth={greedy_depth}" + + +def test_greedy_minimize_time_alap_preparation() -> None: + """Test that greedy_minimize_time uses ALAP preparation to minimize active volume.""" + graph = GraphState() + # Create a 4-node chain: 0-1-2-3 + n0 = graph.add_physical_node() + n1 = graph.add_physical_node() + n2 = graph.add_physical_node() + n3 = graph.add_physical_node() + graph.add_physical_edge(n0, n1) + graph.add_physical_edge(n1, n2) + graph.add_physical_edge(n2, n3) + + graph.register_input(n0, 0) + graph.register_output(n3, 0) + + flow = {n0: {n1}, n1: {n2}, n2: {n3}} + scheduler = Scheduler(graph, flow) + + prepare_time, measure_time = greedy_minimize_time(graph, scheduler.dag) + + # With ALAP, nodes should be prepared as late as possible + # n1 is neighbor of n0, so prep(n1) < meas(n0) + assert prepare_time[n1] == measure_time[n0] - 1 + # n2 is neighbor of n1, so prep(n2) < meas(n1) + assert prepare_time[n2] == measure_time[n1] - 1 + # n3 (output) is neighbor of n2, so prep(n3) < meas(n2) + assert prepare_time[n3] == measure_time[n2] - 1 + + # Input node should not have prepare_time + assert n0 not in prepare_time + + +def test_alap_reduces_active_volume() -> None: + """Test that ALAP preparation reduces active volume compared to ASAP.""" + graph = GraphState() + # Create a chain graph: 0-1-2-3 + n0 = graph.add_physical_node() + n1 = graph.add_physical_node() + n2 = graph.add_physical_node() + n3 = graph.add_physical_node() + graph.add_physical_edge(n0, n1) + graph.add_physical_edge(n1, n2) + graph.add_physical_edge(n2, n3) + graph.register_input(n0, 0) + graph.register_output(n3, 0) + + flow = {n0: {n1}, n1: {n2}, n2: {n3}} + scheduler = Scheduler(graph, flow) + + prepare_time, measure_time = greedy_minimize_time(graph, scheduler.dag) + + # With ALAP: n3 (output) should be prepared as late as possible + # n3 is neighbor of n2, so prep(n3) < meas(n2) + # This should be later than time=0 + assert prepare_time[n3] == measure_time[n2] - 1 + assert prepare_time[n3] > 0 # ALAP should delay preparation + + +def test_alap_preserves_depth() -> None: + """Test that ALAP does not increase depth.""" + # Create a 3x3 grid + graph = GraphState() + nodes = [graph.add_physical_node() for _ in range(9)] + + # Horizontal and vertical edges + for row in range(3): + for col in range(2): + graph.add_physical_edge(nodes[row + col * 3], nodes[row + (col + 1) * 3]) + for row in range(2): + for col in range(3): + graph.add_physical_edge(nodes[row + col * 3], nodes[row + 1 + col * 3]) + + for row in range(3): + graph.register_input(nodes[row], row) + graph.register_output(nodes[row + 6], row) + + flow: dict[int, set[int]] = {nodes[row]: {nodes[row + 3]} for row in range(3)} + flow.update({nodes[row + 3]: {nodes[row + 6]} for row in range(3)}) + + scheduler = Scheduler(graph, flow) + prepare_time, measure_time = greedy_minimize_time(graph, scheduler.dag) + + # Depth should still be optimal (3) + assert max(measure_time.values()) + 1 == 3