Source code for kosmos.partitioning.partition

from dataclasses import dataclass, field

from qiskit import ClassicalRegister, QuantumCircuit, QuantumRegister

from kosmos.partitioning.circuit_converter import to_gate_list
from kosmos.topology.node import QuantumNode




@dataclass
class Partition:
    """Circuit partition with QPU assignment.

    Attributes:
        id (str): Identifier for the partition.
        circuit (QuantumCircuit): Quantum circuit assigned to this partition.
        logical_qubit_mapping (dict[int, int]): Mapping from logical to physical qubit indices.
        original_gate_indices (list[int]): Indices of gates from the original circuit that are
            in this partition.
        start_time (int | None): Scheduled start time for partition execution, if allocated.
        node (QuantumNode | None): The quantum node this partition is allocated to,
            if single-node partition.
        qubit_to_node_mapping (dict[int, QuantumNode]): Mapping from physical qubits to nodes
            for multi-node partitions.

    """

    id: str
    circuit: QuantumCircuit
    logical_qubit_mapping: dict[int, int]
    original_gate_indices: list[int] = None
    start_time: int | None = None
    node: QuantumNode | None = None
    qubit_to_node_mapping: dict[int, QuantumNode] = field(default_factory=dict)

    def __post_init__(self) -> None:
        """Set default value for original_gate_indices."""
        if self.original_gate_indices is None:
            self.original_gate_indices = []

    @property
    def min_gate_index(self) -> int:
        """Get earliest gate index from original circuit.

        Returns:
            int: Minimum gate index, or infinity if no gates.

        """
        return min(self.original_gate_indices) if self.original_gate_indices else float("inf")

    @property
    def physical_qubits_used(self) -> set[int]:
        """Get set of physical qubits used by this partition.

        Returns:
            set[int]: Qubit indices from the original circuit.

        """
        return set(self.logical_qubit_mapping.values())

    @property
    def is_multi_node(self) -> bool:
        """Check if partition is allocated acress multiple nodes.

        Returns:
            bool: True if partitions spans mutliple nodes.

        """
        return bool(self.qubit_to_node_mapping)

    @property
    def allocated_nodes(self) -> set[QuantumNode]:
        """Get all nodes this partition is allocated to.

        Returns:
            set[QuantumNode]: Set of nodes hosting qubits from this partition.

        """
        if self.is_multi_node:
            return set(self.qubit_to_node_mapping.values())
        if self.node is not None:
            return {self.node}
        return set()



    def get_node_for_qubit(self, physical_qubit: int) -> QuantumNode | None:
        """Get the node hosting a specific physical qubit.

        Args:
            physical_qubit (int): Physical qubit index.

        Returns:
            QuantumNode | None: Node hosting the qubit, or None if not allocated.

        """
        if self.is_multi_node:
            return self.qubit_to_node_mapping.get(physical_qubit)
        return self.node


    @property
    def end_time(self) -> int | None:
        """Calculate end time based on start time and circuit depth.

        Returns:
            int | None: Time the partition finishes, if a start time is given.

        """
        if self.start_time is None:
            return None
        return self.start_time + self.circuit.depth()





def create_subcircuit_from_gates(
    gate_list: list[dict],
) -> tuple[QuantumCircuit, dict[int, int]]:
    """Create a subcircuit from a list of gates.

    Args:
        gate_list (list[dict]): List of gate dictionaries.

    Returns:
        tuple[QuantumCircuit, dict[int, int]]: Subcircuit and logical to physical qubit mapping.

    """
    qubits_used = set()
    for gate in gate_list:
        qubits_used.update(gate["qubits"])

    logical_to_physical = dict(enumerate(sorted(qubits_used)))
    physical_to_logical = {phys: log for log, phys in logical_to_physical.items()}

    num_qubits = len(qubits_used)
    q = QuantumRegister(num_qubits, "q")
    c = ClassicalRegister(num_qubits, "c")
    subcircuit = QuantumCircuit(q, c)

    for gate in gate_list:
        logical_qubits = [physical_to_logical[q] for q in gate["qubits"]]
        gate_method = getattr(subcircuit, gate["name"])
        params = gate.get("params", [])
        gate_method(*params, *logical_qubits)

    return subcircuit, logical_to_physical





def partitions_from_assignments(
    circuit: QuantumCircuit,
    assignments: dict[int, int],
) -> list[Partition]:
    """Convert partitioning assignments to Partition objects.

    Args:
        circuit (QuantumCircuit): Original circuit.
        assignments (dict[int, int]): Gate index to partition ID mapping.

    Returns:
        list[Partition]: List of Partition objects with subcircuits.

    """
    gates = to_gate_list(circuit)

    qubit_to_partition: dict[int, int] = {}
    for gate_idx, partition_id in assignments.items():
        gate = gates[gate_idx]
        for qubit in gate["qubits"]:
            if qubit not in qubit_to_partition:
                qubit_to_partition[qubit] = partition_id

    complete_assignments = dict(assignments)
    for gate_idx, gate in enumerate(gates):
        if gate_idx not in complete_assignments:
            gate_qubits = gate["qubits"]

            if len(gate_qubits) == 1:
                qubit = gate_qubits[0]
                complete_assignments[gate_idx] = qubit_to_partition.get(qubit, 0)
            else:
                control_qubit = gate_qubits[0]
                complete_assignments[gate_idx] = qubit_to_partition.get(control_qubit, 0)

    partition_gates: dict[int, list[tuple[int, dict]]] = {}
    for gate_idx, partition_id in complete_assignments.items():
        if partition_id not in partition_gates:
            partition_gates[partition_id] = []
        partition_gates[partition_id].append((gate_idx, gates[gate_idx]))

    partitions = []
    for partition_id in sorted(partition_gates.keys()):
        gate_indices_and_gates = partition_gates[partition_id]
        gate_indices_and_gates.sort(key=lambda x: x[0])
        gate_indices = [g[0] for g in gate_indices_and_gates]
        gate_list = [g[1] for g in gate_indices_and_gates]

        subcircuit, logical_to_physical = create_subcircuit_from_gates(gate_list)

        partitions.append(
            Partition(
                id=f"partition_{partition_id}",
                circuit=subcircuit,
                logical_qubit_mapping=logical_to_physical,
                original_gate_indices=gate_indices,
            )
        )

    return partitions