Source code for kosmos.dqc_scheduling.space_time_matrix

from kosmos.partitioning.partition import Partition
from kosmos.topology.node import QuantumNode

TWO_QUBIT_GATE_SIZE = 2




class SpaceTimeMatrix:
    """Matrix tracking QPU assignments for each qubit at each time step.

    Attributes:
        num_qubits (int): Number of qubits tracked in the matrix (rows).
        num_timesteps (int): Number of gate-timesteps in the simulation (columns).
        matrix (list[list[QuantumNode | None]]): Matrix that stores which Qubit is assigned to what
            node in the simulation.
        partition_info (dict[str, Partition]): Mapping from Partition objects, storing circuit
            segments and metadata.

    """

    def __init__(self, num_qubits: int, num_timesteps: int) -> None:
        """Initialize empty matrix.

        Args:
            num_qubits (int): Number of qubits tracked in the matrix (rows).
            num_timesteps (int): Number of gate-timesteps in the simulation (columns).

        """
        self.num_qubits: int = num_qubits
        self.num_timesteps: int = num_timesteps
        self.matrix: list[list[QuantumNode | None]] = [
            [None] * num_timesteps for _ in range(num_qubits)
        ]
        self.partition_info: dict[str, Partition] = {}



    def get_node_at(self, qubit: int, time: int) -> QuantumNode | None:
        """Get node owning a qubit at a specific time.

        Args:
            qubit (int): Qubit index.
            time (time): Time step.

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

        """
        if 0 <= qubit < self.num_qubits and 0 <= time < self.num_timesteps:
            return self.matrix[qubit][time]
        return None