partial measurement

lib.py

@@ -264,17 +264,28 @@ class State(Vector):
         """If the inner (dot) product is zero, this vector is orthogonal to other"""
         return self.inner(other) == 0
 
-    def get_prob(self, j):
-        """pr(j) = |<e_j|self>|^2"""
-        # j_th basis vector
-        e_j = State([[1] if i == int(j) else [0] for i in range(len(self.m))])
-        return np.absolute((e_j.inner(self)).m.item(0)) ** 2
+    def get_prob_from_measurement_op(self, m_op):
+        assert type(m_op) is MeasurementOperator
+        return m_op.get_prob(self)
+
+    def get_computational_basis_vector(self, j):
+        return Vector([[1] if i == int(j) else [0] for i in range(len(self.m))])
+
+    def get_prob(self, j, basis=None):
+        """pr(j) = |<e_j|self>|^2
+        """
+        if basis is None:
+            basis = self.get_computational_basis_vector(j)
+        m = MeasurementOperator.create_from_basis(basis)
+        return m.get_prob(self)
 
     def get_fmt_of_element(self):
         return "{:0" + str(int(np.ceil(np.log2(len(self))))) + "b}"
 
-    def measure(self):
+    def measure(self, basis=None):
         """
+        m_ops: a set of MeasurementOperators.
+               e.g. for computational basis, the m_ops is [|0><0|, |1><1|]
         Measures in the computational basis.
         Irreversable operation. Measuring again will result in the same result
         TODO: Generalize the method so it takes a basis
@@ -289,12 +300,26 @@ class State(Vector):
         """
         if self.measurement_result:
             return self.measurement_result
-        weights = [self.get_prob(j) for j in range(len(self))]
+        weights = []
+        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)
+                basis.append(m)
+
+        for i in range(len(basis)):
+            m_op = basis[i]
+            prob = self.get_prob_from_measurement_op(m_op)
+            weights += [prob]
+
         empty_prob = 1.0 - sum(weights)
         format_str = self.get_fmt_of_element()
         choices = [REPR_EMPTY_SET] + [format_str.format(i) for i in range(len(weights))]
         weights = [empty_prob] + weights
         self.measurement_result = random.choices(choices, weights)[0]
+
         return self.measurement_result
 
     def measure_with_op(self, mo):
@@ -321,19 +346,16 @@ class State(Vector):
         # [0, 0, 1, 1, 0, 0, 1, 1]
         partial_measurement_of_qbit = [int(b[qubit_n - 1]) for b in bin_repr]
         # [0, 1, 4, 5]
-        indexes_for_p_0 = [i for i, index in enumerate(partial_measurement_of_qbit) if index == 0]
-        weights_0 = [self.get_prob(j) for j in indexes_for_p_0]
-        choices_0 = [format_str.format(i) for i in indexes_for_p_0]
-        measurement_result = random.choices(choices_0, weights_0)[0][qubit_n - 1]
-        # TODO: Verify if this is the correct collapse to lower dimension after partial measurement
-        #       https://www.youtube.com/watch?v=MG_9JWsrKtM&list=PL1826E60FD05B44E4&index=16
-        normalization_factor = np.sqrt(np.sum(np.abs([self.m[i][0] for i in indexes_for_p_0]) ** 2))
-        # TODO: This can be 0...
-        post_measurement_state = [
-            [self.m[i][0] / normalization_factor] for i in indexes_for_p_0
-        ]
-        self.m = s('|' + measurement_result + '>') * State(post_measurement_state)
-        return measurement_result
+        weights, choices = defaultdict(list), defaultdict(list)
+        for result in [1, 0]:
+            indexes_for_p_0 = [i for i, index in enumerate(partial_measurement_of_qbit) if index == result]
+            weights[result] = [self.get_prob(j) for j in indexes_for_p_0]
+            choices[result] = [format_str.format(i) for i in indexes_for_p_0]
+        weights_01 = [sum(weights[0]), sum(weights[1])]
+        measurement_result = random.choices([0, 1], weights_01)[0]
+        normalization_factor = np.sqrt(sum([np.abs(i) ** 2 for i in weights[measurement_result]]))
+        new_m = weights[measurement_result] / normalization_factor
+        return str(measurement_result), State(new_m.reshape((len(new_m), 1)))
 
     def pretty_print(self):
         format_str = self.get_fmt_of_element() + " | {}"
@@ -808,7 +830,7 @@ class MeasurementOperator(HermitianOperator):
     """Measurement operators are Hermitians: <ψ|M†_m M_m|ψ>"""
 
     @classmethod
-    def create_from_prob(cls, matrix: Matrix, *args, **kwargs):
+    def create_from_basis(cls, matrix: Matrix, *args, **kwargs):
         """returns |M†_m><M_m|"""
         return cls(matrix.conjugate_transpose().x(matrix).m, *args, **kwargs)
 
@@ -821,8 +843,8 @@ class MeasurementOperator(HermitianOperator):
         return state_ct.m.dot(self.m.dot(state.m)).item()
 
 
-m0 = MeasurementOperator.create_from_prob(_0)
-m1 = MeasurementOperator.create_from_prob(_1)
+m0 = MeasurementOperator.create_from_basis(_0)
+m1 = MeasurementOperator.create_from_basis(_1)
 
 
 def test_measurement_ops():
@@ -930,6 +952,20 @@ def test():
     assert np.isclose(_p.get_prob(0), 0.5)  # Probability for |+> in 0 is 0.5
     assert np.isclose(_p.get_prob(1), 0.5)  # Probability for |+> in 1 is 0.5
 
+    assert _0.measure() == '0'
+    assert _1.measure() == '1'
+
+    assert s("|10>").measure() == '10'
+    assert s("|10>").measure_partial(1) == ("1", _0)
+    assert s("|10>").measure_partial(2) == ("0", _1)
+
+    # Maximally entangled
+    result, pms = b_phi_p.measure_partial(1)
+    if result == "0":
+        assert pms == _0
+    elif result == "1":
+        assert pms == _1
+
     # Test measurement operators
     test_measurement_ops()
 
@@ -1223,9 +1259,9 @@ def test_quantum_processor():
 
 def test_light():
     # http://alienryderflex.com/polarizer/
-    hor_filter = MeasurementOperator.create_from_prob(Matrix(_0.m), name='h')
-    diag_filter = MeasurementOperator.create_from_prob(Matrix(_p.m), name='d')
-    ver_filter = MeasurementOperator.create_from_prob(Matrix(_1.m), name='v')
+    hor_filter = MeasurementOperator.create_from_basis(Matrix(_0.m), name='h')
+    diag_filter = MeasurementOperator.create_from_basis(Matrix(_p.m), name='d')
+    ver_filter = MeasurementOperator.create_from_basis(Matrix(_1.m), name='v')
 
     def random_light():
         random_pol = Vector([[np.random.uniform(0, 1)], [np.random.uniform(0, 1)]])
@@ -1261,12 +1297,12 @@ def test_light():
 
 if __name__ == "__main__":
     test()
-    test_quantum_processor()
-    test_light()
+    # test_quantum_processor()
+    # test_light()
 
     # Partial measurement - I want to measure just the first qubit...
     # TODO - makes sense???
     # m0_ = m0 * I
-    # post_measurement = b_phi_m.measure_with_op(m0_)
+    # s("|10>").measure_partial(1)
     # result_0 = post_measurement.measure()
     # test_measure_partial()