refactor strategies

Author: coruscating

docs/index.rst

@@ -19,3 +19,15 @@ If you use GADD in your work, please cite:
       journal={arXiv preprint arXiv:2403.02294},
       year={2024}
     }
+
+
+.. toctree::
+   :hidden:
+   :maxdepth: 2
+   :caption: API Reference
+
+   api/gadd
+   api/sequences
+   api/utility_functions
+   api/group_operations
+   api/circuit_padding

gadd/__init__.py

@@ -1,7 +1,7 @@
 """GADD: Genetic Algorithm for Dynamical Decoupling optimization."""
 
 from .gadd import GADD, TrainingConfig, TrainingState, TrainingResult
-from .sequences import DDSequence, DDStrategy, StandardSequences
+from .strategies import DDSequence, DDStrategy, StandardSequences
 from .utility_functions import (
     UtilityFunction,
     SuccessProbability,

gadd/circuit_padding.py

@@ -16,7 +16,7 @@
     PadDynamicalDecoupling,
 )
 
-from .sequences import DDStrategy
+from .strategies import DDStrategy
 
 
 class DDPulse:

gadd/gadd.py

@@ -9,8 +9,6 @@
 import time
 import os
 
-import rustworkx as rx
-
 import numpy as np
 from numpy.random import BitGenerator, Generator, SeedSequence, default_rng
 
@@ -20,7 +18,7 @@
 from qiskit_ibm_runtime import Sampler
 import matplotlib.pyplot as plt
 
-from .sequences import DDStrategy, DDSequence, StandardSequences
+from .strategies import DDStrategy, DDSequence, StandardSequences, ColorAssignment
 from .group_operations import complete_sequence_to_identity
 from .circuit_padding import apply_dd_strategy as _apply_dd_strategy
 from .utility_functions import UtilityFunction, SuccessProbability
@@ -38,7 +36,7 @@ class TrainingConfig:
         mutation_probability (float): Initial probability of mutation.
         optimization_level (int): Qiskit transpilation optimization level.
         shots (int): Number of shots for quantum circuit execution.
-        num_colors (int): Number of distinct sequences per strategy (``C`` in the paper).
+        num_colors (int): Number of distinct sequences per strategy (``k`` in the paper).
         group_size (int): Size of the decoupling group (``|G|`` in the paper).
         mode (str): Mode for generating initial population.
         dynamic_mutation (bool): Whether to dynamically adjust mutation probability.
@@ -112,7 +110,7 @@ def __init__(
         self,
         backend: Optional[Backend] = None,
         utility_function: Optional[UtilityFunction] = None,
-        coloring: Optional[Dict] = None,
+        coloring: Optional[Union[Dict, ColorAssignment]] = None,
         seed: Optional[Union[int, SeedSequence, BitGenerator, Generator]] = None,
         config: Optional[TrainingConfig] = None,
     ):
@@ -126,16 +124,21 @@ def __init__(
 
         # Set up coloring
         if coloring is None and backend is not None:
-            # Default greedy coloring of coupling map
-            if hasattr(backend, "coupling_map") and backend.coupling_map:
-                self._coloring = rx.graph_greedy_color(
-                    backend.coupling_map.graph.to_undirected()
-                )
-            else:
-                # Fallback: all qubits same color
-                self._coloring = {i: 0 for i in range(backend.num_qubits)}
+            # Default coloring from backend
+            self._coloring = ColorAssignment(backend=backend)
+        elif isinstance(coloring, dict):
+            # Convert dict to ColorAssignment
+            # Assume dict is qubit->color mapping
+            color_to_qubits = {}
+            for qubit, color in coloring.items():
+                if color not in color_to_qubits:
+                    color_to_qubits[color] = []
+                color_to_qubits[color].append(qubit)
+            self._coloring = ColorAssignment.from_manual_assignment(color_to_qubits)
+        elif isinstance(coloring, ColorAssignment):
+            self._coloring = coloring
         else:
-            self._coloring = coloring or {}
+            self._coloring = None
 
         # Decoupling group - matches paper's group G
         self._decoupling_group = ["Ip", "Im", "Xp", "Xm", "Yp", "Ym", "Zp", "Zm"]
@@ -166,11 +169,22 @@ def coloring(self):
 
     @coloring.setter
     def coloring(self, coloring):
-        if not isinstance(coloring, dict):
+        if isinstance(coloring, dict):
+            # Convert dict to ColorAssignment
+            color_to_qubits = {}
+            for qubit, color in coloring.items():
+                if color not in color_to_qubits:
+                    color_to_qubits[color] = []
+                color_to_qubits[color].append(qubit)
+            self._coloring = ColorAssignment.from_manual_assignment(color_to_qubits)
+        elif isinstance(coloring, ColorAssignment):
+            self._coloring = coloring
+        elif coloring is None:
+            self._coloring = None
+        else:
             raise TypeError(
-                "Coloring must be a dictionary keyed by qubit index with color values"
+                "Coloring must be a dictionary, ColorAssignment instance, or None"
             )
-        self._coloring = coloring
 
     def apply_dd(
         self,
@@ -197,13 +211,15 @@ def apply_dd(
         # Use provided backend or fall back to instance backend
         backend = backend or self._backend
 
-        # Get coloring for the backend
-        if backend and hasattr(backend, "coupling_map") and backend.coupling_map:
-            coloring = rx.graph_greedy_color(backend.coupling_map.graph.to_undirected())
+        # Get coloring for the circuit
+        if self._coloring is not None:
+            coloring_dict = self._coloring.to_dict()
+        elif backend:
+            color_assignment = ColorAssignment(backend=backend)
+            coloring_dict = color_assignment.to_dict()
         else:
-            coloring = self._coloring or {
-                i: 0 for i in range(target_circuit.num_qubits)
-            }
+            # Fallback: all qubits same color
+            coloring_dict = {i: 0 for i in range(target_circuit.num_qubits)}
 
         # Get instruction durations from backend if available
         instruction_durations = None
@@ -217,7 +233,7 @@ def apply_dd(
         return _apply_dd_strategy(
             target_circuit,
             strategy,
-            coloring,
+            coloring_dict,
             instruction_durations=instruction_durations,
             staggered=staggered,
         )

gadd/sequences.py renamed to gadd/strategies.py

@@ -1,8 +1,15 @@
+"""
+DD strategies, sequences, and coloring assignments.
+"""
+
 from dataclasses import dataclass
 from typing import List, Dict, Optional, Union
+
+import numpy as np
+import rustworkx as rx
+
 from qiskit.circuit.library import IGate, XGate, YGate, U1Gate, RZGate
 from qiskit import QuantumCircuit
-import numpy as np
 
 
 # Default gateset for DD pulses
@@ -211,34 +218,127 @@ def from_dict(cls, data: Dict[str, any]) -> "DDStrategy":
 
 
 class ColorAssignment:
-    """Assignment of device qubits to colors."""
+    """Assignment of device qubits to colors based on connectivity graph."""
 
-    def __init__(self, assignments: Dict[int, List[int]]):
+    def __init__(
+        self, graph: Optional[rx.PyGraph] = None, backend: Optional["Backend"] = None
+    ):
         """
         Initialize color assignment.
 
         Args:
-            assignments: Dictionary mapping colors to lists of qubit indices
+            graph: Connectivity graph where nodes are qubits and edges are connections.
+                If None and backend is provided, will extract from backend.
+            backend: Backend to extract connectivity from if graph not provided.
+
+        Raises:
+            ValueError: If neither graph nor backend is provided.
         """
-        self.assignments = assignments
-        self._validate()
+        if graph is None and backend is None:
+            raise ValueError("Must provide either graph or backend")
+
+        if graph is None:
+            # Extract from backend
+            if hasattr(backend, "coupling_map") and backend.coupling_map:
+                self.graph = backend.coupling_map.graph.to_undirected()
+            else:
+                # Fallback: create complete graph
+                n_qubits = backend.num_qubits if hasattr(backend, "num_qubits") else 1
+                self.graph = rx.PyGraph()
+                self.graph.add_nodes_from(range(n_qubits))
+        else:
+            self.graph = graph
+
+        # Perform graph coloring
+        self._color_map = rx.graph_greedy_color(self.graph)
+
+        # Create assignments dictionary (color -> list of qubits)
+        self.assignments = {}
+        for qubit, color in self._color_map.items():
+            if color not in self.assignments:
+                self.assignments[color] = []
+            self.assignments[color].append(qubit)
+
         # Create reverse mapping for efficiency
-        self._qubit_to_color = {}
+        self._qubit_to_color = self._color_map.copy()
+
+    @classmethod
+    def from_circuit(cls, circuit: QuantumCircuit) -> "ColorAssignment":
+        """
+        Create color assignment from circuit connectivity.
+
+        Args:
+            circuit: Quantum circuit to extract connectivity from.
+
+        Returns:
+            ColorAssignment based on circuit structure.
+        """
+        # Build connectivity graph from circuit
+        graph = rx.PyGraph()
+        qubits = set()
+        edges = set()
+
+        for instruction in circuit.data:
+            if instruction.operation.name in ["cx", "ecr", "cz"]:  # Two-qubit gates
+                qubits_involved = [q._index for q in instruction.qubits]
+                if len(qubits_involved) == 2:
+                    q1, q2 = qubits_involved
+                    qubits.add(q1)
+                    qubits.add(q2)
+                    edges.add((min(q1, q2), max(q1, q2)))
+            else:
+                # Track all qubits
+                for q in instruction.qubits:
+                    qubits.add(q._index)
+
+        # Add all qubits as nodes
+        graph.add_nodes_from(sorted(qubits))
+
+        # Add edges
+        for q1, q2 in edges:
+            graph.add_edge(q1, q2, None)
+
+        return cls(graph=graph)
+
+    @classmethod
+    def from_manual_assignment(
+        cls, assignments: Dict[int, List[int]]
+    ) -> "ColorAssignment":
+        """
+        Create from manual color assignments.
+
+        Args:
+            assignments: Dictionary mapping colors to lists of qubit indices.
+
+        Returns:
+            ColorAssignment with specified assignments.
+        """
+        # Create a graph where qubits with different colors are connected
+        graph = rx.PyGraph()
+        all_qubits = set()
+        for qubits in assignments.values():
+            all_qubits.update(qubits)
+
+        graph.add_nodes_from(sorted(all_qubits))
+
+        # Connect qubits with different colors
+        colors = list(assignments.keys())
+        for i in range(len(colors)):
+            for j in range(i + 1, len(colors)):
+                for q1 in assignments[colors[i]]:
+                    for q2 in assignments[colors[j]]:
+                        graph.add_edge(q1, q2, None)
+
+        # Create instance and override the computed coloring
+        instance = cls(graph=graph)
+        instance.assignments = assignments
+        instance._qubit_to_color = {}
         for color, qubits in assignments.items():
             for qubit in qubits:
-                self._qubit_to_color[qubit] = color
+                instance._qubit_to_color[qubit] = color
+        instance._color_map = instance._qubit_to_color.copy()
 
-    def _validate(self):
-        """Validate color assignments."""
-        # Check for overlapping qubit assignments
-        assigned = set()
-        for color, qubits in self.assignments.items():
-            if not isinstance(qubits, list):
-                raise TypeError(f"Qubits for color {color} must be a list")
-            overlap = assigned.intersection(qubits)
-            if overlap:
-                raise ValueError(f"Qubits {overlap} assigned to multiple colors")
-            assigned.update(qubits)
+        return instance
 
     def get_color(self, qubit: int) -> Optional[int]:
         """Get color assigned to a qubit."""
@@ -252,3 +352,24 @@ def get_qubits(self, color: int) -> List[int]:
     def n_colors(self) -> int:
         """Number of colors used in assignment."""
         return len(self.assignments)
+
+    def to_dict(self) -> Dict[int, int]:
+        """
+        Convert to qubit->color mapping dictionary.
+
+        Returns:
+            Dictionary mapping qubit indices to color values.
+        """
+        return self._color_map.copy()
+
+    def validate_coloring(self) -> bool:
+        """
+        Validate that the coloring is proper (no adjacent nodes have same color).
+
+        Returns:
+            True if coloring is valid, False otherwise.
+        """
+        for edge in self.graph.edge_list():
+            if self._color_map[edge[0]] == self._color_map[edge[1]]:
+                return False
+        return True

tests/test_circuit_padding.py

@@ -10,7 +10,7 @@
     get_instruction_duration,
     apply_dd_strategy,
 )
-from gadd.sequences import DDSequence, DDStrategy
+from gadd.strategies import DDSequence, DDStrategy
 
 
 class TestDDPulse(unittest.TestCase):