diff --git a/src/qibo/models/circuit.py b/src/qibo/models/circuit.py index b909bb8e56..4aff6fd184 100644 --- a/src/qibo/models/circuit.py +++ b/src/qibo/models/circuit.py @@ -43,9 +43,6 @@ class _Queue(list): def __init__(self, nqubits): super().__init__(self) self.nqubits = nqubits - self.moments = [nqubits * [None]] - self.moment_index = nqubits * [0] - self.nmeasurements = 0 def to_fused(self): """Transform all gates in queue to :class:`qibo.gates.FusedGate`.""" @@ -88,24 +85,30 @@ def from_fused(self): queue.append(gate.gates[0]) return queue - def append(self, gate: gates.Gate): - super().append(gate) - if gate.qubits: - qubits = gate.qubits - else: # special gate acting on all qubits - qubits = tuple(range(self.nqubits)) - - if isinstance(gate, gates.M): - self.nmeasurements += 1 - - # calculate moment index for this gate - idx = max(self.moment_index[q] for q in qubits) - for q in qubits: - if idx >= len(self.moments): - # Add a moment - self.moments.append(len(self.moments[-1]) * [None]) - self.moments[idx][q] = gate - self.moment_index[q] = idx + 1 + @property + def nmeasurements(self): + return len(list(filter(lambda gate: isinstance(gate, gates.M), self))) + + @property + def moments(self): + moments = [self.nqubits * [None]] + moment_index = self.nqubits * [0] + for gate in self: + qubits = ( + gate.qubits + if not isinstance(gate, gates.CallbackGate) + else tuple(range(self.nqubits)) # special gate acting on all qubits + ) + + # calculate moment index for this gate + idx = max(moment_index[q] for q in qubits) + for q in qubits: + if idx >= len(moments): + # Add a moment + moments.append(len(moments[-1]) * [None]) + moments[idx][q] = gate + moment_index[q] = idx + 1 + return moments class Circuit: diff --git a/tests/test_models_circuit.py b/tests/test_models_circuit.py index 8d83b9d83e..5181f57a61 100644 --- a/tests/test_models_circuit.py +++ b/tests/test_models_circuit.py @@ -19,8 +19,10 @@ def test_parametrizedgates_class(): def test_queue_class(): + from qibo.callbacks import EntanglementEntropy from qibo.models.circuit import _Queue + entropy = EntanglementEntropy([0]) queue = _Queue(4) gatelist = [ gates.H(0), @@ -29,12 +31,14 @@ def test_queue_class(): gates.H(2), gates.CNOT(1, 2), gates.Y(3), + gates.CallbackGate(entropy), ] for g in gatelist: queue.append(g) assert queue.moments == [ [gatelist[0], gatelist[1], gatelist[3], gatelist[5]], [gatelist[2], gatelist[4], gatelist[4], None], + [gatelist[6] for _ in range(4)], ]