From 37b8d1f5df1b9d9677ec11d4d743470e93673766 Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Tue, 7 May 2024 17:40:25 +0100
Subject: [PATCH 1/9] handle corrections on uncorrected qubits

---
 .../pytket_mbqc_py/graph_circuit.py           |  41 ++--
 .../pytket_mbqc_py/qubit_manager.py           |  14 ++
 pytket-mbqc-py/tests/test_graph_circuit.py    | 179 ++++++++++++++++++
 3 files changed, 218 insertions(+), 16 deletions(-)

diff --git a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
index c071d92a..4e638d12 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
@@ -240,8 +240,8 @@ def add_edge(self, vertex_one: int, vertex_two: int) -> None:
         ):
             raise Exception(
                 "This does not define a valid flow. "
-                f"In particular {vertex_two} is the flow of {self.flow_graph.predecessors(vertex_two)}, "
-                f"some of which are measured before {vertex_one}."
+                f"In particular {vertex_two} is the flow of {list(self.flow_graph.predecessors(vertex_two))}, "
+                f"some of which are measured after {vertex_one}."
             )
 
         # If this is the first future of vertex_one then it is taken to be its flow.
@@ -281,12 +281,18 @@ def _get_z_correction_expression(self, vertex: int) -> BitLogicExp:
         :return: Logical expression calculating the parity
             of the neighbouring x correction registers.
         """
+
+        neighbour_reg_list = [
+            self.vertex_x_corr_reg[neighbour][0]
+            for neighbour in self.entanglement_graph.neighbors(n=vertex)
+        ]
+
+        # This happens of this vertex has no neighbours.
+        if len(neighbour_reg_list) == 0:
+            return None
+
         return reduce(
-            lambda a, b: a ^ b,
-            [
-                self.vertex_x_corr_reg[neighbour][0]
-                for neighbour in self.entanglement_graph.neighbors(n=vertex)
-            ],
+            lambda a, b: a ^ b, neighbour_reg_list
         )
 
     def _apply_z_correction(self, vertex: int) -> None:
@@ -296,10 +302,12 @@ def _apply_z_correction(self, vertex: int) -> None:
 
         :param vertex: Vertex to be corrected.
         """
-        self.Z(
-            self.vertex_qubit[vertex],
-            condition=self._get_z_correction_expression(vertex=vertex),
-        )
+        condition = self._get_z_correction_expression(vertex=vertex)
+        if condition is not None:
+            self.Z(
+                self.vertex_qubit[vertex],
+                condition=condition,
+            )
 
     def _apply_classical_z_correction(self, vertex: int) -> None:
         """Apply Z correction on measurement result. This correction is calculated
@@ -308,11 +316,12 @@ def _apply_classical_z_correction(self, vertex: int) -> None:
 
         :param vertex: Vertex to be corrected.
         """
-        self.add_classicalexpbox_bit(
-            expression=self.qubit_meas_reg[self.vertex_qubit[vertex]][0]
-            ^ self._get_z_correction_expression(vertex=vertex),
-            target=[self.qubit_meas_reg[self.vertex_qubit[vertex]][0]],
-        )
+        condition = self._get_z_correction_expression(vertex=vertex)
+        if condition is not None:
+            self.add_classicalexpbox_bit(
+                expression=self.qubit_meas_reg[self.vertex_qubit[vertex]][0] ^ condition,
+                target=[self.qubit_meas_reg[self.vertex_qubit[vertex]][0]],
+            )
 
     def corrected_measure(self, vertex: int, t_multiple: int = 0) -> None:
         """Perform a measurement, applying the appropriate corrections.
diff --git a/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py b/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py
index 8e904372..d2b06ec4 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py
@@ -77,6 +77,13 @@ def physical_qubits_used(self) -> List[Qubit]:
             for qubit, initialised in self.qubit_initialised.items()
             if initialised
         ]
+    
+    @property
+    def initialised_qubits(self) -> List[Qubit]:
+        return [
+            qubit for qubit, initialised in self.qubit_initialised.items()
+            if initialised
+        ]
 
     def managed_measure(self, qubit: Qubit) -> None:
         """Measure the given qubit, storing the result in the
@@ -85,6 +92,13 @@ def managed_measure(self, qubit: Qubit) -> None:
 
         :param qubit: The qubit to be measured.
         :type qubit: Qubit
+
+        :raises Exception: Raised if the qubit to be measured does
+            not belong to the circuit.
         """
+        if not self.qubit_initialised[qubit]:
+            raise Exception(
+                f"The qubit {qubit} has not been initialised."
+            )
         self.qubit_list.insert(0, qubit)
         self.Measure(qubit=qubit, bit=self.qubit_meas_reg[qubit][0])
diff --git a/pytket-mbqc-py/tests/test_graph_circuit.py b/pytket-mbqc-py/tests/test_graph_circuit.py
index 80cd267a..bd2fef76 100644
--- a/pytket-mbqc-py/tests/test_graph_circuit.py
+++ b/pytket-mbqc-py/tests/test_graph_circuit.py
@@ -314,3 +314,182 @@ def test_cnot_early_measure(input_state, output_state):
     )
 
     assert result.get_counts(output_reg)[output_state] == n_shots
+
+def test_2q_t_gate_example():
+
+    api_offline = QuantinuumAPIOffline()
+    backend = QuantinuumBackend(device_name="H1-1LE", api_handler = api_offline)
+
+    graph_circuit = GraphCircuit(n_physical_qubits=6)
+
+    _, input_vertex_0 = graph_circuit.add_input_vertex()
+
+    # H[0]S[0]
+    graph_vertex_0_0 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(input_vertex_0, graph_vertex_0_0)
+    graph_circuit.corrected_measure(input_vertex_0, t_multiple=2)
+
+    # H[0]
+    graph_vertex_0_1 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_0_0, graph_vertex_0_1)
+    graph_circuit.corrected_measure(graph_vertex_0_0, t_multiple=0)
+
+    _, input_vertex_1 = graph_circuit.add_input_vertex()
+
+    # CZ[0,1]H[1]H[0]
+    graph_vertex_0_0 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_0_1, graph_vertex_0_0)
+    graph_circuit.corrected_measure(graph_vertex_0_1, t_multiple=0)
+
+    graph_vertex_1_0 = graph_circuit.add_graph_vertex()
+
+    graph_vertex_0_1 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_0_0, graph_vertex_0_1)
+
+    graph_circuit.add_edge(input_vertex_1, graph_vertex_1_0)
+    graph_circuit.add_edge(graph_vertex_0_0, graph_vertex_1_0)
+
+    graph_circuit.corrected_measure(input_vertex_1, t_multiple=0)
+
+    graph_vertex_1_1 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_1_0, graph_vertex_1_1)
+
+    # H[0]
+    graph_circuit.corrected_measure(graph_vertex_0_0, t_multiple=0)
+
+    # H[1]
+    graph_circuit.corrected_measure(graph_vertex_1_0, t_multiple=0)
+
+    # H[0]Z[0]
+    graph_vertex_0_0 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_0_1, graph_vertex_0_0)
+    graph_circuit.corrected_measure(graph_vertex_0_1, t_multiple=4)
+
+    # H[0]
+    graph_vertex_0_1 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_0_0, graph_vertex_0_1)
+
+    # CZ[0,1]H[0]H[1]
+    graph_vertex_0_2 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_0_1, graph_vertex_0_2)
+
+    graph_vertex_1_0 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_1_1, graph_vertex_1_0)
+    graph_circuit.corrected_measure(graph_vertex_1_1, t_multiple=0)
+
+    graph_circuit.corrected_measure(graph_vertex_0_0, t_multiple=0)
+    graph_circuit.corrected_measure(graph_vertex_0_1, t_multiple=0)
+
+    graph_vertex_0_1 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_0_2, graph_vertex_0_1)
+
+    graph_vertex_1_1 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_1_0, graph_vertex_1_1)
+
+    graph_circuit.add_edge(graph_vertex_0_2, graph_vertex_1_0)
+
+    # H[0]
+    graph_circuit.corrected_measure(graph_vertex_0_2, t_multiple=0)
+
+    # H[1]
+    graph_circuit.corrected_measure(graph_vertex_1_0, t_multiple=0)
+
+    # H[0]S[0]
+    graph_vertex_0_0 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_0_1, graph_vertex_0_0)
+    graph_circuit.corrected_measure(graph_vertex_0_1, t_multiple=2)
+
+    # H[0]
+    graph_vertex_0_1 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_0_0, graph_vertex_0_1)
+    graph_circuit.corrected_measure(graph_vertex_0_0, t_multiple=0)
+
+    outputs = graph_circuit.get_outputs()
+    out_meas_reg = graph_circuit.add_c_register(name='output measure', size=len(outputs))
+    for qubit, bit in zip(outputs.values(), out_meas_reg):
+        graph_circuit.Measure(qubit=qubit, bit=bit)
+
+    copmiled_graph_circuit = backend.get_compiled_circuit(circuit=graph_circuit)
+    n_shots = 1000
+    result = backend.run_circuit(circuit=copmiled_graph_circuit, n_shots=n_shots)
+    assert result.get_counts(cbits=out_meas_reg)[(1,0)] == n_shots
+
+def test_1q_t_gate_example():
+
+    ################################
+    # The following compiles to I
+
+    api_offline = QuantinuumAPIOffline()
+    backend = QuantinuumBackend(device_name="H1-1LE", api_handler = api_offline)
+
+    graph_circuit = GraphCircuit(n_physical_qubits=2)
+
+    _, input_vertex_0 = graph_circuit.add_input_vertex()
+
+    # H[0]
+    graph_vertex_1 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(input_vertex_0, graph_vertex_1)
+    graph_circuit.corrected_measure(input_vertex_0, t_multiple=0)
+
+    # H[0]T[0]
+    graph_vertex_0 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_1, graph_vertex_0)
+    graph_circuit.corrected_measure(graph_vertex_1, t_multiple=1)
+
+    # H[0]
+    graph_vertex_1 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_0, graph_vertex_1)
+    graph_circuit.corrected_measure(graph_vertex_0, t_multiple=0)
+
+    # H[0]T[0]S[0]Z[0]
+    graph_vertex_0 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_1, graph_vertex_0)
+    graph_circuit.corrected_measure(graph_vertex_1, t_multiple=7)
+
+    outputs = graph_circuit.get_outputs()
+    out_meas_reg = graph_circuit.add_c_register(name='output measure', size=len(outputs))
+    for qubit, bit in zip(outputs.values(), out_meas_reg):
+        graph_circuit.Measure(qubit=qubit, bit=bit)
+
+    copmiled_graph_circuit = backend.get_compiled_circuit(circuit=graph_circuit)
+
+    n_shots=1000
+    result = backend.run_circuit(circuit=copmiled_graph_circuit, n_shots=n_shots)
+    assert result.get_counts(cbits=out_meas_reg)[(0,)] == n_shots
+
+    ################################
+    # The following compiles to X
+
+    graph_circuit = GraphCircuit(n_physical_qubits=2)
+
+    _, input_vertex_0 = graph_circuit.add_input_vertex()
+
+    # H[0]
+    graph_vertex_1 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(input_vertex_0, graph_vertex_1)
+    graph_circuit.corrected_measure(input_vertex_0, t_multiple=0)
+
+    # H[0]T[0]
+    graph_vertex_0 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_1, graph_vertex_0)
+    graph_circuit.corrected_measure(graph_vertex_1, t_multiple=1)
+
+    # H[0]
+    graph_vertex_1 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_0, graph_vertex_1)
+    graph_circuit.corrected_measure(graph_vertex_0, t_multiple=0)
+
+    # H[0]T[0]S[0]Z[0]
+    graph_vertex_0 = graph_circuit.add_graph_vertex()
+    graph_circuit.add_edge(graph_vertex_1, graph_vertex_0)
+    graph_circuit.corrected_measure(graph_vertex_1, t_multiple=3)
+
+    outputs = graph_circuit.get_outputs()
+    out_meas_reg = graph_circuit.add_c_register(name='output measure', size=len(outputs))
+    for qubit, bit in zip(outputs.values(), out_meas_reg):
+        graph_circuit.Measure(qubit=qubit, bit=bit)
+
+    copmiled_graph_circuit = backend.get_compiled_circuit(circuit=graph_circuit)
+
+    result = backend.run_circuit(circuit=copmiled_graph_circuit, n_shots=n_shots)
+    assert result.get_counts(cbits=out_meas_reg)[(1,)] == n_shots 

From 20a61df490988be2c1d0d237a52738c1937c41c3 Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Wed, 8 May 2024 10:37:57 +0100
Subject: [PATCH 2/9] add random register manager tests

---
 pytket-mbqc-py/tests/test_random_register_manager.py | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 pytket-mbqc-py/tests/test_random_register_manager.py

diff --git a/pytket-mbqc-py/tests/test_random_register_manager.py b/pytket-mbqc-py/tests/test_random_register_manager.py
new file mode 100644
index 00000000..e69de29b

From 14e5ad1b844ddefe14655893569dbf0573566a7c Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Wed, 8 May 2024 11:24:29 +0100
Subject: [PATCH 3/9] Add random register manager

---
 pytket-mbqc-py/pytket_mbqc_py/__init__.py     |  1 +
 .../pytket_mbqc_py/graph_circuit.py           |  4 +-
 .../pytket_mbqc_py/random_register_manager.py | 57 +++++++++++++++++++
 .../tests/test_random_register_manager.py     | 45 +++++++++++++++
 4 files changed, 105 insertions(+), 2 deletions(-)
 create mode 100644 pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py

diff --git a/pytket-mbqc-py/pytket_mbqc_py/__init__.py b/pytket-mbqc-py/pytket_mbqc_py/__init__.py
index 63b01d73..60022e55 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/__init__.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/__init__.py
@@ -5,3 +5,4 @@
 from .cnot_block import CNOTBlocksGraphCircuit
 from .graph_circuit import GraphCircuit
 from .qubit_manager import QubitManager
+from .random_register_manager import RandomRegisterManager
diff --git a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
index 4e638d12..c4db9822 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
@@ -12,10 +12,10 @@
 from pytket.circuit.logic_exp import BitLogicExp
 from pytket.unit_id import BitRegister
 
-from pytket_mbqc_py.qubit_manager import QubitManager
+from pytket_mbqc_py.random_register_manager import RandomRegisterManager
 
 
-class GraphCircuit(QubitManager):
+class GraphCircuit(RandomRegisterManager):
     """Class for the automated construction of MBQC computations.
     In particular only graphs with valid flow can be constructed.
     Graph state construction and measurement corrections are added
diff --git a/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py b/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py
new file mode 100644
index 00000000..7687dcbb
--- /dev/null
+++ b/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py
@@ -0,0 +1,57 @@
+from pytket_mbqc_py.qubit_manager import QubitManager
+
+class RandomRegisterManager(QubitManager):
+
+    def __init__(self, n_physical_qubits):
+
+        super().__init__(
+            n_physical_qubits=n_physical_qubits
+        )
+
+    def generate_random_registers(self, n_registers, n_bits_per_reg = 3):
+
+        n_unused_qubits = len(self.qubit_list) - len(self.physical_qubits_used)
+        assert n_unused_qubits >= 0
+
+        if n_unused_qubits == 0:
+            raise Exception(
+                "There are no unused qubits "
+                + "which can be used to generate randomness."
+            )
+        
+        def get_random_bits(n_random_bits):
+
+            for _ in range(n_random_bits // n_unused_qubits):
+                 
+                qubit_list = [self.get_qubit() for _ in range(n_unused_qubits)]
+
+                for qubit in qubit_list:
+
+                    self.H(qubit=qubit)
+                    self.managed_measure(qubit=qubit)
+                    yield self.qubit_meas_reg[qubit][0]
+
+            qubit_list = [self.get_qubit() for _ in range(n_random_bits % n_unused_qubits)]
+
+            for qubit in qubit_list:
+
+                self.H(qubit=qubit)
+                self.managed_measure(qubit=qubit)
+                yield self.qubit_meas_reg[qubit][0]
+
+        for bit_index, bit in enumerate(get_random_bits(n_random_bits = n_bits_per_reg * n_registers)):
+
+            if bit_index % n_bits_per_reg == 0:
+                
+                reg = self.add_c_register(
+                    name=f'rand_{bit_index // n_bits_per_reg}',
+                    size=n_bits_per_reg,
+                )
+            
+            self.add_classicalexpbox_bit(
+                expression=reg[bit_index % n_bits_per_reg] | bit,
+                target=[reg[bit_index % n_bits_per_reg]],
+            )
+
+            if bit_index % n_bits_per_reg == n_bits_per_reg - 1:
+                yield reg
diff --git a/pytket-mbqc-py/tests/test_random_register_manager.py b/pytket-mbqc-py/tests/test_random_register_manager.py
index e69de29b..c265bae3 100644
--- a/pytket-mbqc-py/tests/test_random_register_manager.py
+++ b/pytket-mbqc-py/tests/test_random_register_manager.py
@@ -0,0 +1,45 @@
+from pytket_mbqc_py import RandomRegisterManager
+from pytket.circuit.display import render_circuit_jupyter
+from pytket.extensions.quantinuum import QuantinuumBackend, QuantinuumAPIOffline
+from itertools import product
+
+
+def test_random_register_manager():
+
+    # Here we test a total of 15 random bits generated on 2 qubits.
+    # Note that this mean one of the qubits must be used more than the other.
+    rand_reg_mngr = RandomRegisterManager(n_physical_qubits=2)
+    n_bits_per_reg=3
+    n_registers=5
+    reg_list = list(
+        rand_reg_mngr.generate_random_registers(
+            n_registers=n_registers,
+            n_bits_per_reg=n_bits_per_reg,
+        )
+    )
+
+    # Check that there are 3 registers, and that each has 3 bits.
+    assert len(reg_list) == n_registers
+    assert all(reg.size == n_bits_per_reg for reg in reg_list)
+
+    api_offline = QuantinuumAPIOffline()
+    backend = QuantinuumBackend(device_name="H1-1LE", api_handler = api_offline)
+    n_shots=100
+
+    compiled_circuit = backend.get_compiled_circuit(rand_reg_mngr)
+    result = backend.run_circuit(
+        circuit=compiled_circuit,
+        n_shots=n_shots,
+        seed=0,
+    )
+    result.get_counts(cbits=reg_list[0])
+
+    # Each bit string in each register should occur with equal probability.
+    # This may fail sometimes as we are performing a statistical check.
+    # With the correct seed it should pass every time.
+    for cbits in reg_list:
+        counts = result.get_counts(cbits=cbits)
+        assert all(
+            abs(counts[bit_string] - (n_shots / (2**n_bits_per_reg))) < (n_shots ** 0.5)
+            for bit_string in product([0,1], repeat=n_bits_per_reg)
+        )

From 130c9f471d6782bb930b8f70de27827c47007d6d Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Wed, 8 May 2024 14:55:43 +0100
Subject: [PATCH 4/9] Add documentation

---
 pytket-mbqc-py/pytket_mbqc_py/cnot_block.py   |  6 +-
 .../pytket_mbqc_py/graph_circuit.py           | 62 ++++++++++--
 .../pytket_mbqc_py/qubit_manager.py           | 10 +-
 .../pytket_mbqc_py/random_register_manager.py | 96 ++++++++++++++-----
 pytket-mbqc-py/source/index.rst               |  1 +
 pytket-mbqc-py/source/qubit_manager.rst       |  4 +-
 pytket-mbqc-py/tests/test_graph_circuit.py    | 61 ++++++++----
 .../tests/test_random_register_manager.py     | 21 ++--
 8 files changed, 195 insertions(+), 66 deletions(-)

diff --git a/pytket-mbqc-py/pytket_mbqc_py/cnot_block.py b/pytket-mbqc-py/pytket_mbqc_py/cnot_block.py
index 75effc85..9001376a 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/cnot_block.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/cnot_block.py
@@ -23,6 +23,7 @@ def __init__(
         n_physical_qubits: int,
         input_state: Tuple[int],
         n_layers: int,
+        n_registers: int,
     ) -> None:
         """Initialisation method.
 
@@ -56,7 +57,10 @@ def __init__(
             [[] for _ in range(n_rows)] for _ in range(n_layers)
         ]
 
-        super().__init__(n_physical_qubits=n_physical_qubits)
+        super().__init__(
+            n_physical_qubits=n_physical_qubits,
+            n_registers=n_registers,
+        )
 
         for layer in range(n_layers):
             # If this is the first layer then the control qubit of the first row needs
diff --git a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
index c4db9822..8556efd5 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
@@ -5,7 +5,7 @@
 """
 
 from functools import reduce
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Union
 
 import networkx as nx  # type:ignore
 from pytket import Qubit
@@ -20,7 +20,7 @@ class GraphCircuit(RandomRegisterManager):
     In particular only graphs with valid flow can be constructed.
     Graph state construction and measurement corrections are added
     automatically. A child of
-    :py:class:`~pytket_mbqc_py.qubit_manager.QubitManager`.
+    :py:class:`~pytket_mbqc_py.qubit_manager.RandomRegisterManager`.
 
     :ivar entanglement_graph: Graph detailing the graph state entanglement.
     :ivar flow_graph: Graph describing the flow dependencies of the graph state.
@@ -28,6 +28,8 @@ class GraphCircuit(RandomRegisterManager):
     :ivar vertex_measured: List indicating if vertex has been measured.
     :ivar vertex_x_corr_reg: List mapping vertex index to
         the classical register where the required X correction is stored.
+    :ivar vertex_random_reg: List mapping vertex to the corresponding
+        register of random bits.
     """
 
     entanglement_graph: nx.Graph
@@ -37,12 +39,14 @@ class GraphCircuit(RandomRegisterManager):
     def __init__(
         self,
         n_physical_qubits: int,
+        n_registers: int,
     ) -> None:
         """Initialisation method. Creates tools to track
         the graph state structure and the measurement corrections.
 
         :param n_physical_qubits: The number of physical qubits available.
-        :type n_physical_qubits: int
+        :n_registers: The number of random registers to generate. Defaults
+            to 100 random registers.
         """
         super().__init__(n_physical_qubits=n_physical_qubits)
 
@@ -60,6 +64,11 @@ def __init__(
         # vertex has been measured.
         self.vertex_x_corr_reg: List[BitRegister] = []
 
+        self.vertex_random_reg = list(
+            self.generate_random_registers(n_registers=n_registers)
+        )
+        self.add_barrier(units=self.qubits)
+
     def get_outputs(self) -> Dict[int, Qubit]:
         """Return the output qubits. Output qubits are those that
         are unmeasured, and which do not have a flow. This should
@@ -101,6 +110,27 @@ def get_outputs(self) -> Dict[int, Qubit]:
 
         # We need to correct all unmeasured qubits.
         for vertex in output_qubits.keys():
+            self.T(
+                self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][0]
+            )
+            self.S(
+                self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][0]
+            )
+            self.Z(
+                self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][0]
+            )
+
+            self.S(
+                self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][1]
+            )
+            self.Z(
+                self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][1]
+            )
+
+            self.Z(
+                self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][2]
+            )
+
             self._apply_x_correction(vertex=vertex)
             self._apply_z_correction(vertex=vertex)
 
@@ -140,6 +170,10 @@ def add_input_vertex(self) -> Tuple[Qubit, int]:
         qubit = self.get_qubit()
         index = self._add_vertex(qubit=qubit)
 
+        # In the case of input qubits, the initialisation is not random.
+        # As such the measurement angle does not need to be corrected.
+        self.add_c_setreg(0, self.vertex_random_reg[index])
+
         return (qubit, index)
 
     def add_graph_vertex(self) -> int:
@@ -151,6 +185,10 @@ def add_graph_vertex(self) -> int:
         self.H(qubit)
         index = self._add_vertex(qubit=qubit)
 
+        self.T(qubit, condition=self.vertex_random_reg[index][0])
+        self.S(qubit, condition=self.vertex_random_reg[index][1])
+        self.Z(qubit, condition=self.vertex_random_reg[index][2])
+
         return index
 
     @property
@@ -272,7 +310,7 @@ def _apply_x_correction(self, vertex: int) -> None:
             condition=self.vertex_x_corr_reg[vertex][0],
         )
 
-    def _get_z_correction_expression(self, vertex: int) -> BitLogicExp:
+    def _get_z_correction_expression(self, vertex: int) -> Union[None, BitLogicExp]:
         """Create logical expression by taking the parity of
         the X corrections that have to be applied to the neighbouring
         qubits.
@@ -291,9 +329,7 @@ def _get_z_correction_expression(self, vertex: int) -> BitLogicExp:
         if len(neighbour_reg_list) == 0:
             return None
 
-        return reduce(
-            lambda a, b: a ^ b, neighbour_reg_list
-        )
+        return reduce(lambda a, b: a ^ b, neighbour_reg_list)
 
     def _apply_z_correction(self, vertex: int) -> None:
         """Apply Z correction on qubit. This correction is calculated
@@ -319,7 +355,8 @@ def _apply_classical_z_correction(self, vertex: int) -> None:
         condition = self._get_z_correction_expression(vertex=vertex)
         if condition is not None:
             self.add_classicalexpbox_bit(
-                expression=self.qubit_meas_reg[self.vertex_qubit[vertex]][0] ^ condition,
+                expression=self.qubit_meas_reg[self.vertex_qubit[vertex]][0]
+                ^ condition,
                 target=[self.qubit_meas_reg[self.vertex_qubit[vertex]][0]],
             )
 
@@ -351,6 +388,15 @@ def corrected_measure(self, vertex: int, t_multiple: int = 0) -> None:
                 + f"are in the future of {vertex} and have already been measured."
             )
 
+        self.T(self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][0])
+        self.S(self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][0])
+        self.Z(self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][0])
+
+        self.S(self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][1])
+        self.Z(self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][1])
+
+        self.Z(self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][2])
+
         # This is actually optional for graph vertices
         # as the X correction commutes with the hadamard
         # basis measurement.
diff --git a/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py b/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py
index d2b06ec4..8d4bf472 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py
@@ -77,11 +77,13 @@ def physical_qubits_used(self) -> List[Qubit]:
             for qubit, initialised in self.qubit_initialised.items()
             if initialised
         ]
-    
+
     @property
     def initialised_qubits(self) -> List[Qubit]:
+        """Qubits which have been initialised."""
         return [
-            qubit for qubit, initialised in self.qubit_initialised.items()
+            qubit
+            for qubit, initialised in self.qubit_initialised.items()
             if initialised
         ]
 
@@ -97,8 +99,6 @@ def managed_measure(self, qubit: Qubit) -> None:
             not belong to the circuit.
         """
         if not self.qubit_initialised[qubit]:
-            raise Exception(
-                f"The qubit {qubit} has not been initialised."
-            )
+            raise Exception(f"The qubit {qubit} has not been initialised.")
         self.qubit_list.insert(0, qubit)
         self.Measure(qubit=qubit, bit=self.qubit_meas_reg[qubit][0])
diff --git a/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py b/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py
index 7687dcbb..213866a4 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py
@@ -1,57 +1,107 @@
-from pytket_mbqc_py.qubit_manager import QubitManager
+"""
+Tools for managing random bits. This include generating random bits
+to dedicated registers.
+"""
 
-class RandomRegisterManager(QubitManager):
+from collections.abc import Iterator
 
-    def __init__(self, n_physical_qubits):
+from pytket.unit_id import Bit, BitRegister
 
-        super().__init__(
-            n_physical_qubits=n_physical_qubits
-        )
+from pytket_mbqc_py.qubit_manager import QubitManager
 
-    def generate_random_registers(self, n_registers, n_bits_per_reg = 3):
 
-        n_unused_qubits = len(self.qubit_list) - len(self.physical_qubits_used)
-        assert n_unused_qubits >= 0
+class RandomRegisterManager(QubitManager):
+    """Class for generating random bits, and managing dedicated registers
+    where they are stored. Child of
+    :py:class:`~pytket_mbqc_py.qubit_manager.QubitManager`.
+    """
 
-        if n_unused_qubits == 0:
+    def generate_random_registers(
+        self,
+        n_registers: int,
+        n_bits_per_reg: int = 3,
+        max_n_randomness_qubits: int = 2,
+    ) -> Iterator[BitRegister]:
+        """Generate registers containing random bits. This is achieved
+        by initialising hadamard basis plus states and measuring them.
+
+        :param n_registers: The number of registers to generate.
+        :param n_bits_per_reg: The number of bits in each register,
+            defaults to 3.
+        :param max_n_randomness_qubits: The maximum number of qubits to use
+            to generate randomness, defaults to 2.
+        :raises Exception: Raised if there are no qubits left to use to
+            generate randomness.
+        :yield: Registers containing random bits.
+        """
+
+        if len(self.qubit_list) == 0:
             raise Exception(
                 "There are no unused qubits "
                 + "which can be used to generate randomness."
             )
-        
-        def get_random_bits(n_random_bits):
 
-            for _ in range(n_random_bits // n_unused_qubits):
-                 
-                qubit_list = [self.get_qubit() for _ in range(n_unused_qubits)]
+        # The number of qubits used is the smaller of the maximum number set
+        # by the user, or the number which is available.
+        n_randomness_qubits = min(len(self.qubit_list), max_n_randomness_qubits)
 
-                for qubit in qubit_list:
+        def get_random_bits(n_random_bits: int) -> Iterator[Bit]:
+            """An iterator over bits populated with random values.
+            These are created by initialising hadamard plus states and
+            measuring the state in the computation basis. At most
+            `n_randomness_qubits` qubit are created in the plus state, then
+            measured and reset as appropriate until the resulted number of bits
+            have been created.
 
+            Note that the physical bits may be reused so should be used
+            before the next step of the iteration.
+
+            :param n_random_bits: The number of random bits to generate.
+            :yield: Bits containing random values.
+            """
+            # We repeatedly initialise and measure qubits in groups
+            # of size n_randomness_qubits. This allows randomness generation
+            # to be done in parallel where possible.
+            for _ in range(n_random_bits // n_randomness_qubits):
+                # Initialise all qubits.
+                qubit_list = [self.get_qubit() for _ in range(n_randomness_qubits)]
+
+                # For each qubit, initialise and measure.
+                for qubit in qubit_list:
                     self.H(qubit=qubit)
                     self.managed_measure(qubit=qubit)
                     yield self.qubit_meas_reg[qubit][0]
 
-            qubit_list = [self.get_qubit() for _ in range(n_random_bits % n_unused_qubits)]
+            # This may leave a number of random bits less than
+            # n_randomness_qubits to generate. We do this here.
+            qubit_list = [
+                self.get_qubit() for _ in range(n_random_bits % n_randomness_qubits)
+            ]
 
             for qubit in qubit_list:
-
                 self.H(qubit=qubit)
                 self.managed_measure(qubit=qubit)
                 yield self.qubit_meas_reg[qubit][0]
 
-        for bit_index, bit in enumerate(get_random_bits(n_random_bits = n_bits_per_reg * n_registers)):
-
+        # For each bit, copy it's value to a persistent
+        # register. This is done in groups of n_bits_per_reg.
+        for bit_index, bit in enumerate(
+            get_random_bits(n_random_bits=n_bits_per_reg * n_registers)
+        ):
+            # If this is the first bit to co[y to a new register, get
+            # the new register.
             if bit_index % n_bits_per_reg == 0:
-                
                 reg = self.add_c_register(
-                    name=f'rand_{bit_index // n_bits_per_reg}',
+                    name=f"rand_{bit_index // n_bits_per_reg}",
                     size=n_bits_per_reg,
                 )
-            
+
+            # Copy the bit to the persistent register.
             self.add_classicalexpbox_bit(
                 expression=reg[bit_index % n_bits_per_reg] | bit,
                 target=[reg[bit_index % n_bits_per_reg]],
             )
 
+            # If this is the last bit to fill the register, yield the register.
             if bit_index % n_bits_per_reg == n_bits_per_reg - 1:
                 yield reg
diff --git a/pytket-mbqc-py/source/index.rst b/pytket-mbqc-py/source/index.rst
index ee837f98..942df6a2 100644
--- a/pytket-mbqc-py/source/index.rst
+++ b/pytket-mbqc-py/source/index.rst
@@ -12,6 +12,7 @@ Welcome to pytket-mbqc's documentation!
 
    qubit_manager
    graph_circuit
+   random_register_manager
 
 
 
diff --git a/pytket-mbqc-py/source/qubit_manager.rst b/pytket-mbqc-py/source/qubit_manager.rst
index 3ac356aa..cca5cb0d 100644
--- a/pytket-mbqc-py/source/qubit_manager.rst
+++ b/pytket-mbqc-py/source/qubit_manager.rst
@@ -11,4 +11,6 @@ Qubit Manager
 
     .. automethod:: QubitManager.managed_measure
 
-    .. autoproperty:: QubitManager.physical_qubits_used
\ No newline at end of file
+    .. autoproperty:: QubitManager.physical_qubits_used
+
+    .. autoproperty:: QubitManager.initialised_qubits
\ No newline at end of file
diff --git a/pytket-mbqc-py/tests/test_graph_circuit.py b/pytket-mbqc-py/tests/test_graph_circuit.py
index bd2fef76..23e119a4 100644
--- a/pytket-mbqc-py/tests/test_graph_circuit.py
+++ b/pytket-mbqc-py/tests/test_graph_circuit.py
@@ -6,7 +6,7 @@
 
 
 def test_plus_state():
-    circuit = GraphCircuit(n_physical_qubits=2)
+    circuit = GraphCircuit(n_physical_qubits=2, n_registers=3)
 
     input_qubit, vertex_one = circuit.add_input_vertex()
     circuit.H(input_qubit)
@@ -36,7 +36,10 @@ def test_plus_state():
 
 
 def test_x_gate():
-    circuit = GraphCircuit(n_physical_qubits=3)
+    circuit = GraphCircuit(
+        n_physical_qubits=3,
+        n_registers=3,
+    )
 
     _, vertex_one = circuit.add_input_vertex()
 
@@ -74,7 +77,10 @@ def test_x_gate():
     [((0, 0), (0, 0)), ((0, 1), (0, 1)), ((1, 0), (1, 1)), ((1, 1), (1, 0))],
 )
 def test_cnot(input_state, output_state):
-    circuit = GraphCircuit(n_physical_qubits=5)
+    circuit = GraphCircuit(
+        n_physical_qubits=5,
+        n_registers=10,
+    )
 
     target_qubit, vertex_one = circuit.add_input_vertex()
     if input_state[1]:
@@ -149,6 +155,7 @@ def test_cnot_block(input_state, output_state, n_layers):
         n_physical_qubits=n_physical_qubits,
         input_state=input_state,
         n_layers=n_layers,
+        n_registers=40,
     )
 
     output_vertex_quibts = circuit.get_outputs()
@@ -203,7 +210,7 @@ def test_large_cnot_block():
 
 
 def test_3_q_ghz():
-    graph_circuit = GraphCircuit(n_physical_qubits=5)
+    graph_circuit = GraphCircuit(n_physical_qubits=5, n_registers=5)
 
     input_quibt, input_vertex = graph_circuit.add_input_vertex()
 
@@ -255,7 +262,10 @@ def test_3_q_ghz():
     [((0, 0), (0, 0)), ((0, 1), (0, 1)), ((1, 0), (1, 1)), ((1, 1), (1, 0))],
 )
 def test_cnot_early_measure(input_state, output_state):
-    circuit = GraphCircuit(n_physical_qubits=3)
+    circuit = GraphCircuit(
+        n_physical_qubits=3,
+        n_registers=10,
+    )
 
     target_qubit, vertex_one = circuit.add_input_vertex()
     if input_state[1]:
@@ -315,12 +325,15 @@ def test_cnot_early_measure(input_state, output_state):
 
     assert result.get_counts(output_reg)[output_state] == n_shots
 
-def test_2q_t_gate_example():
 
+def test_2q_t_gate_example():
     api_offline = QuantinuumAPIOffline()
-    backend = QuantinuumBackend(device_name="H1-1LE", api_handler = api_offline)
+    backend = QuantinuumBackend(device_name="H1-1LE", api_handler=api_offline)
 
-    graph_circuit = GraphCircuit(n_physical_qubits=6)
+    graph_circuit = GraphCircuit(
+        n_physical_qubits=6,
+        n_registers=20,
+    )
 
     _, input_vertex_0 = graph_circuit.add_input_vertex()
 
@@ -405,24 +418,29 @@ def test_2q_t_gate_example():
     graph_circuit.corrected_measure(graph_vertex_0_0, t_multiple=0)
 
     outputs = graph_circuit.get_outputs()
-    out_meas_reg = graph_circuit.add_c_register(name='output measure', size=len(outputs))
+    out_meas_reg = graph_circuit.add_c_register(
+        name="output measure", size=len(outputs)
+    )
     for qubit, bit in zip(outputs.values(), out_meas_reg):
         graph_circuit.Measure(qubit=qubit, bit=bit)
 
     copmiled_graph_circuit = backend.get_compiled_circuit(circuit=graph_circuit)
     n_shots = 1000
     result = backend.run_circuit(circuit=copmiled_graph_circuit, n_shots=n_shots)
-    assert result.get_counts(cbits=out_meas_reg)[(1,0)] == n_shots
+    assert result.get_counts(cbits=out_meas_reg)[(1, 0)] == n_shots
 
-def test_1q_t_gate_example():
 
+def test_1q_t_gate_example():
     ################################
     # The following compiles to I
 
     api_offline = QuantinuumAPIOffline()
-    backend = QuantinuumBackend(device_name="H1-1LE", api_handler = api_offline)
+    backend = QuantinuumBackend(device_name="H1-1LE", api_handler=api_offline)
 
-    graph_circuit = GraphCircuit(n_physical_qubits=2)
+    graph_circuit = GraphCircuit(
+        n_physical_qubits=2,
+        n_registers=5,
+    )
 
     _, input_vertex_0 = graph_circuit.add_input_vertex()
 
@@ -447,20 +465,25 @@ def test_1q_t_gate_example():
     graph_circuit.corrected_measure(graph_vertex_1, t_multiple=7)
 
     outputs = graph_circuit.get_outputs()
-    out_meas_reg = graph_circuit.add_c_register(name='output measure', size=len(outputs))
+    out_meas_reg = graph_circuit.add_c_register(
+        name="output measure", size=len(outputs)
+    )
     for qubit, bit in zip(outputs.values(), out_meas_reg):
         graph_circuit.Measure(qubit=qubit, bit=bit)
 
     copmiled_graph_circuit = backend.get_compiled_circuit(circuit=graph_circuit)
 
-    n_shots=1000
+    n_shots = 1000
     result = backend.run_circuit(circuit=copmiled_graph_circuit, n_shots=n_shots)
     assert result.get_counts(cbits=out_meas_reg)[(0,)] == n_shots
 
     ################################
     # The following compiles to X
 
-    graph_circuit = GraphCircuit(n_physical_qubits=2)
+    graph_circuit = GraphCircuit(
+        n_physical_qubits=2,
+        n_registers=5,
+    )
 
     _, input_vertex_0 = graph_circuit.add_input_vertex()
 
@@ -485,11 +508,13 @@ def test_1q_t_gate_example():
     graph_circuit.corrected_measure(graph_vertex_1, t_multiple=3)
 
     outputs = graph_circuit.get_outputs()
-    out_meas_reg = graph_circuit.add_c_register(name='output measure', size=len(outputs))
+    out_meas_reg = graph_circuit.add_c_register(
+        name="output measure", size=len(outputs)
+    )
     for qubit, bit in zip(outputs.values(), out_meas_reg):
         graph_circuit.Measure(qubit=qubit, bit=bit)
 
     copmiled_graph_circuit = backend.get_compiled_circuit(circuit=graph_circuit)
 
     result = backend.run_circuit(circuit=copmiled_graph_circuit, n_shots=n_shots)
-    assert result.get_counts(cbits=out_meas_reg)[(1,)] == n_shots 
+    assert result.get_counts(cbits=out_meas_reg)[(1,)] == n_shots
diff --git a/pytket-mbqc-py/tests/test_random_register_manager.py b/pytket-mbqc-py/tests/test_random_register_manager.py
index c265bae3..e0867b88 100644
--- a/pytket-mbqc-py/tests/test_random_register_manager.py
+++ b/pytket-mbqc-py/tests/test_random_register_manager.py
@@ -1,16 +1,17 @@
-from pytket_mbqc_py import RandomRegisterManager
-from pytket.circuit.display import render_circuit_jupyter
-from pytket.extensions.quantinuum import QuantinuumBackend, QuantinuumAPIOffline
 from itertools import product
 
+from pytket.circuit.display import render_circuit_jupyter
+from pytket.extensions.quantinuum import QuantinuumAPIOffline, QuantinuumBackend
+
+from pytket_mbqc_py import RandomRegisterManager
 
-def test_random_register_manager():
 
+def test_random_register_manager():
     # Here we test a total of 15 random bits generated on 2 qubits.
     # Note that this mean one of the qubits must be used more than the other.
     rand_reg_mngr = RandomRegisterManager(n_physical_qubits=2)
-    n_bits_per_reg=3
-    n_registers=5
+    n_bits_per_reg = 3
+    n_registers = 5
     reg_list = list(
         rand_reg_mngr.generate_random_registers(
             n_registers=n_registers,
@@ -23,8 +24,8 @@ def test_random_register_manager():
     assert all(reg.size == n_bits_per_reg for reg in reg_list)
 
     api_offline = QuantinuumAPIOffline()
-    backend = QuantinuumBackend(device_name="H1-1LE", api_handler = api_offline)
-    n_shots=100
+    backend = QuantinuumBackend(device_name="H1-1LE", api_handler=api_offline)
+    n_shots = 100
 
     compiled_circuit = backend.get_compiled_circuit(rand_reg_mngr)
     result = backend.run_circuit(
@@ -40,6 +41,6 @@ def test_random_register_manager():
     for cbits in reg_list:
         counts = result.get_counts(cbits=cbits)
         assert all(
-            abs(counts[bit_string] - (n_shots / (2**n_bits_per_reg))) < (n_shots ** 0.5)
-            for bit_string in product([0,1], repeat=n_bits_per_reg)
+            abs(counts[bit_string] - (n_shots / (2**n_bits_per_reg))) < (n_shots**0.5)
+            for bit_string in product([0, 1], repeat=n_bits_per_reg)
         )

From 77c7e14b8a5ec789b06a0f62c800baf44c90c111 Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Wed, 8 May 2024 15:25:29 +0100
Subject: [PATCH 5/9] Repair long running tests

---
 pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py | 8 ++++++++
 pytket-mbqc-py/tests/test_graph_circuit.py     | 1 +
 2 files changed, 9 insertions(+)

diff --git a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
index 8556efd5..43c3344a 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
@@ -144,6 +144,9 @@ def _add_vertex(self, qubit: Qubit) -> int:
 
         :param qubit: Qubit to be added.
         :return: The vertex in the graphs corresponding to this qubit
+
+        :raises Exception: Raised if an insufficient number of random
+            registers were initialised.
         """
         self.vertex_qubit.append(qubit)
         self.vertex_measured.append(False)
@@ -156,6 +159,11 @@ def _add_vertex(self, qubit: Qubit) -> int:
         self.vertex_x_corr_reg.append(x_corr_reg)
         self.add_c_register(register=x_corr_reg)
 
+        if index >= len(self.vertex_random_reg):
+            raise Exception(
+                "An insufficient number of random registers were initialised."
+            )
+
         return index
 
     def add_input_vertex(self) -> Tuple[Qubit, int]:
diff --git a/pytket-mbqc-py/tests/test_graph_circuit.py b/pytket-mbqc-py/tests/test_graph_circuit.py
index 23e119a4..2e3eda4c 100644
--- a/pytket-mbqc-py/tests/test_graph_circuit.py
+++ b/pytket-mbqc-py/tests/test_graph_circuit.py
@@ -188,6 +188,7 @@ def test_large_cnot_block():
         n_physical_qubits=n_physical_qubits,
         input_state=input_state,
         n_layers=n_layers,
+        n_registers=40,
     )
 
     output_vertex_quibts = circuit.get_outputs()

From c878663e0812960d6cdc8a6ca24df8d4931297d6 Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Thu, 9 May 2024 19:19:16 +0100
Subject: [PATCH 6/9] Simplify docs source

---
 pytket-mbqc-py/pytket_mbqc_py/cnot_block.py     |  8 +++++---
 pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py  |  7 +++----
 pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py  |  3 +--
 .../pytket_mbqc_py/random_register_manager.py   |  8 +++-----
 pytket-mbqc-py/source/graph_circuit.rst         | 17 +++--------------
 pytket-mbqc-py/source/qubit_manager.rst         | 15 +++------------
 .../source/random_register_manager.rst          |  7 +++++++
 7 files changed, 25 insertions(+), 40 deletions(-)
 create mode 100644 pytket-mbqc-py/source/random_register_manager.rst

diff --git a/pytket-mbqc-py/pytket_mbqc_py/cnot_block.py b/pytket-mbqc-py/pytket_mbqc_py/cnot_block.py
index 9001376a..a18bf638 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/cnot_block.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/cnot_block.py
@@ -30,15 +30,17 @@ def __init__(
         :param n_physical_qubits: The maximum number of physical qubits
             available. These qubits will be reused and so the total
             number of 'logical' qubits may be larger.
-        :type n_physical_qubits: int
         :param input_state: Integer tuple describing the input
             to the circuit. This will be a classical binary
             string and so the outcome is deterministic.
-        :type input_state: Tuple[int]
         :param n_layers: The number of layers of CNOT gates.
-        :type n_layers: int
+        :param n_registers: The number of classical registers to use
+            for random number generation.
         """
 
+        # TODO: n_registers should be calculated from the given
+        # input variables.
+
         self.input_state = input_state
         self.n_layers = n_layers
 
diff --git a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
index 43c3344a..c49ca23f 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
@@ -19,8 +19,7 @@ class GraphCircuit(RandomRegisterManager):
     """Class for the automated construction of MBQC computations.
     In particular only graphs with valid flow can be constructed.
     Graph state construction and measurement corrections are added
-    automatically. A child of
-    :py:class:`~pytket_mbqc_py.qubit_manager.RandomRegisterManager`.
+    automatically.
 
     :ivar entanglement_graph: Graph detailing the graph state entanglement.
     :ivar flow_graph: Graph describing the flow dependencies of the graph state.
@@ -45,7 +44,7 @@ def __init__(
         the graph state structure and the measurement corrections.
 
         :param n_physical_qubits: The number of physical qubits available.
-        :n_registers: The number of random registers to generate. Defaults
+        :param n_registers: The number of random registers to generate. Defaults
             to 100 random registers.
         """
         super().__init__(n_physical_qubits=n_physical_qubits)
@@ -321,7 +320,7 @@ def _apply_x_correction(self, vertex: int) -> None:
     def _get_z_correction_expression(self, vertex: int) -> Union[None, BitLogicExp]:
         """Create logical expression by taking the parity of
         the X corrections that have to be applied to the neighbouring
-        qubits.
+        qubits. If there are no neighbours then None will be returned.
 
         :param vertex: Vertex to be corrected.
         :return: Logical expression calculating the parity
diff --git a/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py b/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py
index 8d4bf472..05626e7b 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py
@@ -14,8 +14,7 @@ class QubitManager(Circuit):
     Manages a collection of qubits. In particular maintains
     a list of qubits which are not in use. This can be added
     to by measuring qubits (making them available again)
-    and drawn from by initialising qubits. This is a child of
-    :py:class:`~pytket.circuit.Circuit`.
+    and drawn from by initialising qubits.
 
     :ivar qubit_list: List of available quits.
     :ivar qubit_initialised: The qubits which have been added
diff --git a/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py b/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py
index 213866a4..d3bfec3c 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py
@@ -12,8 +12,7 @@
 
 class RandomRegisterManager(QubitManager):
     """Class for generating random bits, and managing dedicated registers
-    where they are stored. Child of
-    :py:class:`~pytket_mbqc_py.qubit_manager.QubitManager`.
+    where they are stored.
     """
 
     def generate_random_registers(
@@ -26,10 +25,9 @@ def generate_random_registers(
         by initialising hadamard basis plus states and measuring them.
 
         :param n_registers: The number of registers to generate.
-        :param n_bits_per_reg: The number of bits in each register,
-            defaults to 3.
+        :param n_bits_per_reg: The number of bits in each register.
         :param max_n_randomness_qubits: The maximum number of qubits to use
-            to generate randomness, defaults to 2.
+            to generate randomness.
         :raises Exception: Raised if there are no qubits left to use to
             generate randomness.
         :yield: Registers containing random bits.
diff --git a/pytket-mbqc-py/source/graph_circuit.rst b/pytket-mbqc-py/source/graph_circuit.rst
index 05d998b5..79b97a8d 100644
--- a/pytket-mbqc-py/source/graph_circuit.rst
+++ b/pytket-mbqc-py/source/graph_circuit.rst
@@ -2,17 +2,6 @@ Graph Circuit
 =============
 
 .. automodule:: pytket_mbqc_py.graph_circuit
-
-.. autoclass:: pytket_mbqc_py.graph_circuit.GraphCircuit
-
-    .. automethod:: GraphCircuit.__init__
-
-    .. automethod:: GraphCircuit.add_input_vertex
-    
-    .. automethod:: GraphCircuit.add_graph_vertex
-
-    .. automethod:: GraphCircuit.add_edge
-
-    .. automethod:: GraphCircuit.corrected_measure
-
-    .. automethod:: GraphCircuit.get_outputs
\ No newline at end of file
+    :members:
+    :show-inheritance:
+    :special-members:
\ No newline at end of file
diff --git a/pytket-mbqc-py/source/qubit_manager.rst b/pytket-mbqc-py/source/qubit_manager.rst
index cca5cb0d..0ac80d25 100644
--- a/pytket-mbqc-py/source/qubit_manager.rst
+++ b/pytket-mbqc-py/source/qubit_manager.rst
@@ -2,15 +2,6 @@ Qubit Manager
 =============
 
 .. automodule:: pytket_mbqc_py.qubit_manager
-
-.. autoclass:: pytket_mbqc_py.qubit_manager.QubitManager
-
-    .. automethod:: QubitManager.__init__
-
-    .. automethod:: QubitManager.get_qubit
-
-    .. automethod:: QubitManager.managed_measure
-
-    .. autoproperty:: QubitManager.physical_qubits_used
-
-    .. autoproperty:: QubitManager.initialised_qubits
\ No newline at end of file
+    :members:
+    :show-inheritance:
+    :special-members:
\ No newline at end of file
diff --git a/pytket-mbqc-py/source/random_register_manager.rst b/pytket-mbqc-py/source/random_register_manager.rst
new file mode 100644
index 00000000..4042191a
--- /dev/null
+++ b/pytket-mbqc-py/source/random_register_manager.rst
@@ -0,0 +1,7 @@
+Random Register Manager
+=======================
+
+.. automodule:: pytket_mbqc_py.random_register_manager
+    :members:
+    :show-inheritance:
+    :special-members:
\ No newline at end of file

From 9528995b46dac7854d2e57ad72a00e4fff6be904 Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Fri, 10 May 2024 11:21:56 +0100
Subject: [PATCH 7/9] Combine corrections where possible

---
 .../pytket_mbqc_py/graph_circuit.py           | 103 ++++++++++++++----
 .../pytket_mbqc_py/qubit_manager.py           |   1 -
 2 files changed, 81 insertions(+), 23 deletions(-)

diff --git a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
index c49ca23f..a36d74cf 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
@@ -130,7 +130,14 @@ def get_outputs(self) -> Dict[int, Qubit]:
                 self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][2]
             )
 
-            self._apply_x_correction(vertex=vertex)
+            # self._apply_x_correction(vertex=vertex)
+
+            # Apply X correction according to correction register.
+            self.X(
+                self.vertex_qubit[vertex],
+                condition=self.vertex_x_corr_reg[vertex][0],
+            )
+
             self._apply_z_correction(vertex=vertex)
 
         return output_qubits
@@ -305,17 +312,17 @@ def add_edge(self, vertex_one: int, vertex_two: int) -> None:
             v_of_edge=vertex_two,
         )
 
-    def _apply_x_correction(self, vertex: int) -> None:
-        """Apply X correction. This correction is drawn from
-        the x correction register which is altered
-        by the corrected measure method as appropriate.
+    # def _apply_x_correction(self, vertex: int) -> None:
+    #     """Apply X correction. This correction is drawn from
+    #     the x correction register which is altered
+    #     by the corrected measure method as appropriate.
 
-        :param vertex: The vertex to be corrected.
-        """
-        self.X(
-            self.vertex_qubit[vertex],
-            condition=self.vertex_x_corr_reg[vertex][0],
-        )
+    #     :param vertex: The vertex to be corrected.
+    #     """
+    #     self.X(
+    #         self.vertex_qubit[vertex],
+    #         condition=self.vertex_x_corr_reg[vertex][0],
+    #     )
 
     def _get_z_correction_expression(self, vertex: int) -> Union[None, BitLogicExp]:
         """Create logical expression by taking the parity of
@@ -374,7 +381,6 @@ def corrected_measure(self, vertex: int, t_multiple: int = 0) -> None:
         :param vertex: Vertex to be measured.
         :param t_multiple: The angle in which to measure, defaults to 0.
             This defines the rotated hadamard basis to measure in.
-        :type t_multiple: int, optional
         :raises Exception: Raised if this vertex has already been measured.
         :raises Exception: Raised if there are vertex in the past of this
             one which have not been measured.
@@ -395,19 +401,72 @@ def corrected_measure(self, vertex: int, t_multiple: int = 0) -> None:
                 + f"are in the future of {vertex} and have already been measured."
             )
 
-        self.T(self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][0])
-        self.S(self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][0])
-        self.Z(self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][0])
+        # Apply X correction according to correction register.
+        self.X(
+            self.vertex_qubit[vertex],
+            condition=self.vertex_x_corr_reg[vertex][0],
+        )
 
-        self.S(self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][1])
-        self.Z(self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][1])
+        self.T(
+            self.vertex_qubit[vertex],
+            # Required to invert random T from initialisation.
+            condition=self.vertex_random_reg[vertex][0],
+        )
+        self.S(
+            self.vertex_qubit[vertex],
+            # Required to invert random T from initialisation.
+            # This additional term is required to account for the case where
+            # the correcting T is commuted through an X correction.
+            condition=(
+                self.vertex_random_reg[vertex][0]
+                & self.vertex_x_corr_reg[vertex][0]
+            ),
+        )
 
-        self.Z(self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][2])
+        self.S(
+            self.vertex_qubit[vertex],
+            # Required to invert random T from initialisation.
+            condition=self.vertex_random_reg[vertex][0],
+        )
 
-        # This is actually optional for graph vertices
-        # as the X correction commutes with the hadamard
-        # basis measurement.
-        self._apply_x_correction(vertex=vertex)
+        self.S(
+            self.vertex_qubit[vertex],
+            # Required to invert random S from initialisation.
+            condition=self.vertex_random_reg[vertex][1],
+        )
+
+        self.Z(
+            self.vertex_qubit[vertex],
+            condition=(
+                # Required to invert random T from initialisation.
+                self.vertex_random_reg[vertex][0]
+                # Required to invert random S from initialisation.
+                ^ self.vertex_random_reg[vertex][1]
+                # Required to invert random Z from initialisation.
+                ^ self.vertex_random_reg[vertex][2]
+                # Required to invert random S from initialisation.
+                # This additional term is required to account for the case where
+                # the correcting S is commuted through an X correction.
+                ^ (
+                    self.vertex_random_reg[vertex][1]
+                    & self.vertex_x_corr_reg[vertex][0]
+                )
+                # Required to invert random T from initialisation.
+                # This additional term is required to account for the case where
+                # the correcting S is commuted through an X correction.
+                ^ (
+                    self.vertex_random_reg[vertex][0]
+                    & self.vertex_x_corr_reg[vertex][0]
+                )
+                # Required to invert random T from initialisation.
+                # This additional term is required to account for the case where
+                # the correcting T is commuted through an X correction.
+                ^ (
+                    self.vertex_random_reg[vertex][0]
+                    & self.vertex_x_corr_reg[vertex][0]
+                )
+            ),
+        )
 
         inverse_t_multiple = 8 - t_multiple
         inverse_t_multiple = inverse_t_multiple % 8
diff --git a/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py b/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py
index 05626e7b..4ad92cce 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/qubit_manager.py
@@ -92,7 +92,6 @@ def managed_measure(self, qubit: Qubit) -> None:
         the list of available qubits.
 
         :param qubit: The qubit to be measured.
-        :type qubit: Qubit
 
         :raises Exception: Raised if the qubit to be measured does
             not belong to the circuit.

From 7d09af1b310710675f1d2f99164b3feb01916b0e Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Fri, 10 May 2024 11:27:05 +0100
Subject: [PATCH 8/9] Rename initialisation register list

---
 .../pytket_mbqc_py/graph_circuit.py           | 72 +++++++------------
 1 file changed, 26 insertions(+), 46 deletions(-)

diff --git a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
index a36d74cf..48c85e65 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
@@ -27,8 +27,10 @@ class GraphCircuit(RandomRegisterManager):
     :ivar vertex_measured: List indicating if vertex has been measured.
     :ivar vertex_x_corr_reg: List mapping vertex index to
         the classical register where the required X correction is stored.
-    :ivar vertex_random_reg: List mapping vertex to the corresponding
-        register of random bits.
+    :ivar vertex_init_reg: List mapping vertex to the register describing
+        the state it was initialised in. In particular this is a 3 bit register
+        with the 0th entry giving the T rotation, the 1st giving the
+        S rotation, and the 2nd giving the Z rotation.
     """
 
     entanglement_graph: nx.Graph
@@ -63,7 +65,7 @@ def __init__(
         # vertex has been measured.
         self.vertex_x_corr_reg: List[BitRegister] = []
 
-        self.vertex_random_reg = list(
+        self.vertex_init_reg = list(
             self.generate_random_registers(n_registers=n_registers)
         )
         self.add_barrier(units=self.qubits)
@@ -109,26 +111,14 @@ def get_outputs(self) -> Dict[int, Qubit]:
 
         # We need to correct all unmeasured qubits.
         for vertex in output_qubits.keys():
-            self.T(
-                self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][0]
-            )
-            self.S(
-                self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][0]
-            )
-            self.Z(
-                self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][0]
-            )
+            self.T(self.vertex_qubit[vertex], condition=self.vertex_init_reg[vertex][0])
+            self.S(self.vertex_qubit[vertex], condition=self.vertex_init_reg[vertex][0])
+            self.Z(self.vertex_qubit[vertex], condition=self.vertex_init_reg[vertex][0])
 
-            self.S(
-                self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][1]
-            )
-            self.Z(
-                self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][1]
-            )
+            self.S(self.vertex_qubit[vertex], condition=self.vertex_init_reg[vertex][1])
+            self.Z(self.vertex_qubit[vertex], condition=self.vertex_init_reg[vertex][1])
 
-            self.Z(
-                self.vertex_qubit[vertex], condition=self.vertex_random_reg[vertex][2]
-            )
+            self.Z(self.vertex_qubit[vertex], condition=self.vertex_init_reg[vertex][2])
 
             # self._apply_x_correction(vertex=vertex)
 
@@ -165,7 +155,7 @@ def _add_vertex(self, qubit: Qubit) -> int:
         self.vertex_x_corr_reg.append(x_corr_reg)
         self.add_c_register(register=x_corr_reg)
 
-        if index >= len(self.vertex_random_reg):
+        if index >= len(self.vertex_init_reg):
             raise Exception(
                 "An insufficient number of random registers were initialised."
             )
@@ -186,7 +176,7 @@ def add_input_vertex(self) -> Tuple[Qubit, int]:
 
         # In the case of input qubits, the initialisation is not random.
         # As such the measurement angle does not need to be corrected.
-        self.add_c_setreg(0, self.vertex_random_reg[index])
+        self.add_c_setreg(0, self.vertex_init_reg[index])
 
         return (qubit, index)
 
@@ -199,9 +189,9 @@ def add_graph_vertex(self) -> int:
         self.H(qubit)
         index = self._add_vertex(qubit=qubit)
 
-        self.T(qubit, condition=self.vertex_random_reg[index][0])
-        self.S(qubit, condition=self.vertex_random_reg[index][1])
-        self.Z(qubit, condition=self.vertex_random_reg[index][2])
+        self.T(qubit, condition=self.vertex_init_reg[index][0])
+        self.S(qubit, condition=self.vertex_init_reg[index][1])
+        self.Z(qubit, condition=self.vertex_init_reg[index][2])
 
         return index
 
@@ -410,7 +400,7 @@ def corrected_measure(self, vertex: int, t_multiple: int = 0) -> None:
         self.T(
             self.vertex_qubit[vertex],
             # Required to invert random T from initialisation.
-            condition=self.vertex_random_reg[vertex][0],
+            condition=self.vertex_init_reg[vertex][0],
         )
         self.S(
             self.vertex_qubit[vertex],
@@ -418,53 +408,43 @@ def corrected_measure(self, vertex: int, t_multiple: int = 0) -> None:
             # This additional term is required to account for the case where
             # the correcting T is commuted through an X correction.
             condition=(
-                self.vertex_random_reg[vertex][0]
-                & self.vertex_x_corr_reg[vertex][0]
+                self.vertex_init_reg[vertex][0] & self.vertex_x_corr_reg[vertex][0]
             ),
         )
 
         self.S(
             self.vertex_qubit[vertex],
             # Required to invert random T from initialisation.
-            condition=self.vertex_random_reg[vertex][0],
+            condition=self.vertex_init_reg[vertex][0],
         )
 
         self.S(
             self.vertex_qubit[vertex],
             # Required to invert random S from initialisation.
-            condition=self.vertex_random_reg[vertex][1],
+            condition=self.vertex_init_reg[vertex][1],
         )
 
         self.Z(
             self.vertex_qubit[vertex],
             condition=(
                 # Required to invert random T from initialisation.
-                self.vertex_random_reg[vertex][0]
+                self.vertex_init_reg[vertex][0]
                 # Required to invert random S from initialisation.
-                ^ self.vertex_random_reg[vertex][1]
+                ^ self.vertex_init_reg[vertex][1]
                 # Required to invert random Z from initialisation.
-                ^ self.vertex_random_reg[vertex][2]
+                ^ self.vertex_init_reg[vertex][2]
                 # Required to invert random S from initialisation.
                 # This additional term is required to account for the case where
                 # the correcting S is commuted through an X correction.
-                ^ (
-                    self.vertex_random_reg[vertex][1]
-                    & self.vertex_x_corr_reg[vertex][0]
-                )
+                ^ (self.vertex_init_reg[vertex][1] & self.vertex_x_corr_reg[vertex][0])
                 # Required to invert random T from initialisation.
                 # This additional term is required to account for the case where
                 # the correcting S is commuted through an X correction.
-                ^ (
-                    self.vertex_random_reg[vertex][0]
-                    & self.vertex_x_corr_reg[vertex][0]
-                )
+                ^ (self.vertex_init_reg[vertex][0] & self.vertex_x_corr_reg[vertex][0])
                 # Required to invert random T from initialisation.
                 # This additional term is required to account for the case where
                 # the correcting T is commuted through an X correction.
-                ^ (
-                    self.vertex_random_reg[vertex][0]
-                    & self.vertex_x_corr_reg[vertex][0]
-                )
+                ^ (self.vertex_init_reg[vertex][0] & self.vertex_x_corr_reg[vertex][0])
             ),
         )
 

From 2be8f07e297c7917ecabd4910eb1bea1c4bde5a4 Mon Sep 17 00:00:00 2001
From: Dan Mills <52407433+daniel-mills-cqc@users.noreply.github.com>
Date: Fri, 10 May 2024 11:57:46 +0100
Subject: [PATCH 9/9] Add aditional documentation

---
 pytket-mbqc-py/pytket_mbqc_py/cnot_block.py   |  3 +-
 .../pytket_mbqc_py/graph_circuit.py           | 53 +++++++++++--------
 .../pytket_mbqc_py/random_register_manager.py | 10 ++--
 3 files changed, 39 insertions(+), 27 deletions(-)

diff --git a/pytket-mbqc-py/pytket_mbqc_py/cnot_block.py b/pytket-mbqc-py/pytket_mbqc_py/cnot_block.py
index a18bf638..dcf4cf50 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/cnot_block.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/cnot_block.py
@@ -35,7 +35,8 @@ def __init__(
             string and so the outcome is deterministic.
         :param n_layers: The number of layers of CNOT gates.
         :param n_registers: The number of classical registers to use
-            for random number generation.
+            for state preparation information. Note that there should
+            be at least one register per logical qubit.
         """
 
         # TODO: n_registers should be calculated from the given
diff --git a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
index 48c85e65..0817cf16 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/graph_circuit.py
@@ -46,8 +46,10 @@ def __init__(
         the graph state structure and the measurement corrections.
 
         :param n_physical_qubits: The number of physical qubits available.
-        :param n_registers: The number of random registers to generate. Defaults
-            to 100 random registers.
+        :param n_registers: The number of state initialisation registers
+            to generate. Each register describes the state that the logical
+            qubit is initialised in. Note that the number of such registers
+            should be at least the number of logical qubits.
         """
         super().__init__(n_physical_qubits=n_physical_qubits)
 
@@ -65,6 +67,11 @@ def __init__(
         # vertex has been measured.
         self.vertex_x_corr_reg: List[BitRegister] = []
 
+        # Generate one random register per vertex.
+        # When qubits are added they will be initialised in this
+        # random register. This in except for the case of input qubits
+        # which are initialised in the 0 state, and in which case this
+        # register is overwritten.
         self.vertex_init_reg = list(
             self.generate_random_registers(n_registers=n_registers)
         )
@@ -111,23 +118,29 @@ def get_outputs(self) -> Dict[int, Qubit]:
 
         # We need to correct all unmeasured qubits.
         for vertex in output_qubits.keys():
+            # Corrections are first applied to invert the initialisation
             self.T(self.vertex_qubit[vertex], condition=self.vertex_init_reg[vertex][0])
             self.S(self.vertex_qubit[vertex], condition=self.vertex_init_reg[vertex][0])
-            self.Z(self.vertex_qubit[vertex], condition=self.vertex_init_reg[vertex][0])
 
             self.S(self.vertex_qubit[vertex], condition=self.vertex_init_reg[vertex][1])
-            self.Z(self.vertex_qubit[vertex], condition=self.vertex_init_reg[vertex][1])
 
-            self.Z(self.vertex_qubit[vertex], condition=self.vertex_init_reg[vertex][2])
-
-            # self._apply_x_correction(vertex=vertex)
+            self.Z(
+                self.vertex_qubit[vertex],
+                condition=(
+                    self.vertex_init_reg[vertex][2]
+                    ^ self.vertex_init_reg[vertex][1]
+                    ^ self.vertex_init_reg[vertex][0]
+                ),
+            )
 
             # Apply X correction according to correction register.
+            # These are corrections resulting from the measurements
             self.X(
                 self.vertex_qubit[vertex],
                 condition=self.vertex_x_corr_reg[vertex][0],
             )
 
+            # Apply Z corrections resulting from measurements.
             self._apply_z_correction(vertex=vertex)
 
         return output_qubits
@@ -141,7 +154,7 @@ def _add_vertex(self, qubit: Qubit) -> int:
         :param qubit: Qubit to be added.
         :return: The vertex in the graphs corresponding to this qubit
 
-        :raises Exception: Raised if an insufficient number of random
+        :raises Exception: Raised if an insufficient number of initialisation
             registers were initialised.
         """
         self.vertex_qubit.append(qubit)
@@ -157,7 +170,8 @@ def _add_vertex(self, qubit: Qubit) -> int:
 
         if index >= len(self.vertex_init_reg):
             raise Exception(
-                "An insufficient number of random registers were initialised."
+                "An insufficient number of initialisation registers "
+                + "were initialised."
             )
 
         return index
@@ -175,7 +189,7 @@ def add_input_vertex(self) -> Tuple[Qubit, int]:
         index = self._add_vertex(qubit=qubit)
 
         # In the case of input qubits, the initialisation is not random.
-        # As such the measurement angle does not need to be corrected.
+        # As such the initialisation register should be set to 0.
         self.add_c_setreg(0, self.vertex_init_reg[index])
 
         return (qubit, index)
@@ -189,6 +203,8 @@ def add_graph_vertex(self) -> int:
         self.H(qubit)
         index = self._add_vertex(qubit=qubit)
 
+        # The graph state is randomly initialised based on the
+        # initialisation register.
         self.T(qubit, condition=self.vertex_init_reg[index][0])
         self.S(qubit, condition=self.vertex_init_reg[index][1])
         self.Z(qubit, condition=self.vertex_init_reg[index][2])
@@ -302,18 +318,6 @@ def add_edge(self, vertex_one: int, vertex_two: int) -> None:
             v_of_edge=vertex_two,
         )
 
-    # def _apply_x_correction(self, vertex: int) -> None:
-    #     """Apply X correction. This correction is drawn from
-    #     the x correction register which is altered
-    #     by the corrected measure method as appropriate.
-
-    #     :param vertex: The vertex to be corrected.
-    #     """
-    #     self.X(
-    #         self.vertex_qubit[vertex],
-    #         condition=self.vertex_x_corr_reg[vertex][0],
-    #     )
-
     def _get_z_correction_expression(self, vertex: int) -> Union[None, BitLogicExp]:
         """Create logical expression by taking the parity of
         the X corrections that have to be applied to the neighbouring
@@ -392,6 +396,7 @@ def corrected_measure(self, vertex: int, t_multiple: int = 0) -> None:
             )
 
         # Apply X correction according to correction register.
+        # This is to correct for measurement outcomes.
         self.X(
             self.vertex_qubit[vertex],
             condition=self.vertex_x_corr_reg[vertex][0],
@@ -448,6 +453,10 @@ def corrected_measure(self, vertex: int, t_multiple: int = 0) -> None:
             ),
         )
 
+        # Rotate measurement basis.
+        # TODO: These measurements should be combined with the above
+        # so that the measurement angles are hidden by the initialisation
+        # angles.
         inverse_t_multiple = 8 - t_multiple
         inverse_t_multiple = inverse_t_multiple % 8
         if inverse_t_multiple // 4:
diff --git a/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py b/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py
index d3bfec3c..8fc6a694 100644
--- a/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py
+++ b/pytket-mbqc-py/pytket_mbqc_py/random_register_manager.py
@@ -27,7 +27,9 @@ def generate_random_registers(
         :param n_registers: The number of registers to generate.
         :param n_bits_per_reg: The number of bits in each register.
         :param max_n_randomness_qubits: The maximum number of qubits to use
-            to generate randomness.
+            to generate randomness. If a number of qubits less than this
+            number are actually available then the number available will
+            be used.
         :raises Exception: Raised if there are no qubits left to use to
             generate randomness.
         :yield: Registers containing random bits.
@@ -48,8 +50,8 @@ def get_random_bits(n_random_bits: int) -> Iterator[Bit]:
             These are created by initialising hadamard plus states and
             measuring the state in the computation basis. At most
             `n_randomness_qubits` qubit are created in the plus state, then
-            measured and reset as appropriate until the resulted number of bits
-            have been created.
+            measured and reset as appropriate until the requested number
+            of bits have been created.
 
             Note that the physical bits may be reused so should be used
             before the next step of the iteration.
@@ -86,7 +88,7 @@ def get_random_bits(n_random_bits: int) -> Iterator[Bit]:
         for bit_index, bit in enumerate(
             get_random_bits(n_random_bits=n_bits_per_reg * n_registers)
         ):
-            # If this is the first bit to co[y to a new register, get
+            # If this is the first bit to copy to a new register, create
             # the new register.
             if bit_index % n_bits_per_reg == 0:
                 reg = self.add_c_register(