From 08bce98b6d28ab197161b1badc402563a0053b50 Mon Sep 17 00:00:00 2001
From: Daniel Tsvetkov <danieltcv@gmail.com>
Date: Sun, 29 Mar 2020 10:03:00 +0200
Subject: [PATCH] sequential measurements are memoized

---
 lib.py | 43 +++++++++++++++++++++++++++++++++++++------
 1 file changed, 37 insertions(+), 6 deletions(-)

diff --git a/lib.py b/lib.py
index 6decf7f..afbc93c 100644
--- a/lib.py
+++ b/lib.py
@@ -164,6 +164,7 @@ class State(Vector):
             super().__init__(m, *args, **kwargs)
         self.name = name
         self.measurement_result = None
+        self.last_basis = None
         # TODO: SHOULD WE NORMALIZE?
         # if not self._is_normalized():
         #     raise TypeError("Not a normalized state vector")
@@ -287,6 +288,7 @@ class State(Vector):
         Measures in the computational basis.
         Irreversable operation. Measuring again will result in the same result
         TODO: Should we memoize per basis?
+              Currently it's memoized if the basis is the same
               If it's measured twice, should it return the same state?
               What about if measured twice but in different bases?
               E.g. measure1 -> computation -> A
@@ -297,15 +299,15 @@ class State(Vector):
                    random weighted?
         :return: binary representation of the measured qubit (e.g. "011")
         """
-        # if self.measurement_result:
-        #     return self.measurement_result
         weights, m_ops = [], []
         if not basis:
+            basis = []
             for j in range(len(self)):
                 # Default is computational basis
                 e_j = self.get_computational_basis_vector(j)
                 m = MeasurementOperator.create_from_basis(e_j)
                 m_ops.append(m)
+                basis.append(State(e_j))
         else:
             assert len(basis) == len(self)
             for i, e_j in enumerate(basis):
@@ -314,6 +316,9 @@ class State(Vector):
                     assert e_j.is_orthogonal(basis[0])
                 m_ops.append(MeasurementOperator.create_from_basis(e_j))
 
+        if self.measurement_result and self.last_basis == basis:
+            return self.measurement_result
+
         for m_op in m_ops:
             prob = self.get_prob_from_measurement_op(m_op)
             weights += [prob]
@@ -327,13 +332,14 @@ class State(Vector):
             # TODO: This may be wrong
             # normalize the weights
             weights = list(np.array(weights) / sum(weights))
+            # assert np.isclose(sum(weights), 1.0)
 
         format_str = self.get_fmt_of_element()
         choices = empty_choices + [format_str.format(i) for i in
                                    range(len(weights))]
         weights = empty_weights + weights
         self.measurement_result = random.choices(choices, weights)[0]
-
+        self.last_basis = basis
         return self.measurement_result
 
     def measure_with_op(self, mo):
@@ -841,7 +847,8 @@ T = lambda phi: Gate([[1, 0],
 #
 # Decomposed CNOT :
 # reverse engineered from
-# https://quantumcomputing.stackexchange.com/questions/4252/how-to-derive-the-cnot-matrix-for-a-3-qbit-system-where-the-control-target-qbi
+# https://quantumcomputing.stackexchange.com/questions/4252/how-to-derive-the
+# -cnot-matrix-for-a-3-qbit-system-where-the-control-target-qbi
 #
 # CNOT(q1, I, q2):
 # |0><0| x I_2 x I_2 + |1><1| x I_2 x X
@@ -890,7 +897,6 @@ SWAP = Gate([
     CNOT.on(CNOT.on(CNOT.on(state, [x, y]), [y, x]), [x, y])
 )
 
-
 TOFF = Gate([
     [1, 0, 0, 0, 0, 0, 0, 0],
     [0, 1, 0, 0, 0, 0, 0, 0],
@@ -1040,7 +1046,8 @@ def test_partials():
     assert CNOT.on(CNOT.on(CNOT.on(s("|10>")), which_qbit=[1, 0])) == s("|01>")
 
     # Test SWAP via 3 successive CNOTs, the 0<->2 are swapped
-    assert CNOT.on(CNOT.on(CNOT.on(s("|100>"), [0, 2]), [2, 0]), [0, 2]) == s("|001>")
+    assert CNOT.on(CNOT.on(CNOT.on(s("|100>"), [0, 2]), [2, 0]), [0, 2]) == s(
+        "|001>")
 
     # apply on 0, 1 of 3qbit state
     assert CNOT.on(s("|000>"), which_qbit=[0, 1]) == s("|000>")
@@ -1085,6 +1092,28 @@ def test_partials():
     assert TOFF.on(s("|10100>"), which_qbit=[0, 2, 4]) == s("|10101>")
 
 
+def test_sequential_measurements():
+    st = H.on(s("|0>"))
+    meas1 = st.measure()
+    assert st.measurement_result is not None
+    assert st.last_basis == computational_basis
+
+    # Measuring more times should return the same result
+    for i in range(10):
+        meas2 = st.measure()
+        assert meas2 == meas1
+
+    # Measuring in different basis might not yield the same result
+    meas3 = st.measure(basis=plus_minus_basis)
+    assert st.measurement_result is not None
+    assert st.last_basis == plus_minus_basis
+
+    # But measuring more times should still return the same result
+    for i in range(10):
+        meas4 = st.measure(basis=plus_minus_basis)
+        assert meas4 == meas3
+
+
 def test():
     # Test properties of Hilbert vector space
     # The four postulates of Quantum Mechanics
@@ -1189,6 +1218,8 @@ def test():
     # Test measurement operators
     test_measurement_ops()
 
+    test_sequential_measurements()
+
     # IV: Compositing | The state space of a composite physical system
     #                   is the tensor product of the state spaces
     #                   of the component physical systems.