From 051ec97d804b47ae903101247c95bc6f5e7cc336 Mon Sep 17 00:00:00 2001
From: SeanvdMeer <18538762+minisean@users.noreply.github.com>
Date: Fri, 13 Feb 2026 16:52:58 +0100
Subject: [PATCH] Fixed initial state indexing mismatch

---
 .../library/repetition_code/circuit_components.py         | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/qce_circuit/library/repetition_code/circuit_components.py b/src/qce_circuit/library/repetition_code/circuit_components.py
index 2ebfa20..eab75af 100644
--- a/src/qce_circuit/library/repetition_code/circuit_components.py
+++ b/src/qce_circuit/library/repetition_code/circuit_components.py
@@ -531,7 +531,7 @@ def get_operations(self, initial_state: InitialStateContainer, **kwargs) -> List
         result.extend([
             initial_state.get_data_qubit_operation(
                 qubit_index=self.map_qubit_id_to_circuit_index(
-                    qubit_id=self.data_qubit_ids[initial_state_index],
+                    qubit_id=initial_state_index if isinstance(initial_state_index, IQubitID) else self.data_qubit_ids[initial_state_index],
                 ),
                 initial_state_index=initial_state_index,
                 **kwargs,
@@ -541,7 +541,7 @@ def get_operations(self, initial_state: InitialStateContainer, **kwargs) -> List
         result.extend([
             initial_state.get_data_qubit_operation(
                 qubit_index=self.map_qubit_id_to_circuit_index(
-                    qubit_id=self.ancilla_qubit_ids[initial_state_index],
+                    qubit_id=initial_state_index if isinstance(initial_state_index, IQubitID) else self.ancilla_qubit_ids[initial_state_index],
                 ),
                 initial_state_index=initial_state_index,
                 **kwargs,
@@ -821,7 +821,7 @@ def get_operations(self, initial_state: InitialStateContainer, **kwargs) -> List
         result.extend([
             initial_state.get_data_qubit_operation(
                 qubit_index=self.map_qubit_id_to_circuit_index(
-                    qubit_id=self.data_qubit_ids[initial_state_index],
+                    qubit_id=initial_state_index if isinstance(initial_state_index, IQubitID) else self.data_qubit_ids[initial_state_index],
                 ),
                 initial_state_index=initial_state_index,
                 **kwargs,
@@ -831,7 +831,7 @@ def get_operations(self, initial_state: InitialStateContainer, **kwargs) -> List
         result.extend([
             initial_state.get_ancilla_qubit_operation(
                 qubit_index=self.map_qubit_id_to_circuit_index(
-                    qubit_id=self.ancilla_qubit_ids[initial_state_index],
+                    qubit_id=initial_state_index if isinstance(initial_state_index, IQubitID) else self.ancilla_qubit_ids[initial_state_index],
                 ),
                 initial_state_index=initial_state_index,
                 **kwargs,