diff --git a/lib.py b/lib.py
index a58ce2c..9a32782 100644
--- a/lib.py
+++ b/lib.py
@@ -153,13 +153,13 @@ class Vector(Matrix):
         return self.m.shape[1] == 1
 
 
-VectorOrList = Union[Vector, ListOrNdarray]
+VectorOrList = Union[Vector, Matrix, ListOrNdarray]
 
 
 class State(Vector):
     def __init__(self, m: VectorOrList = None, name: str = '', *args, **kwargs):
         """State vector representing quantum state"""
-        if type(m) is Vector:
+        if type(m) in [Vector, Matrix]:
             super().__init__(m.m, *args, **kwargs)
         else:
             super().__init__(m, *args, **kwargs)
@@ -300,17 +300,22 @@ class State(Vector):
         """
         if self.measurement_result:
             return self.measurement_result
-        weights = []
+        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)
-                basis.append(m)
+                m_ops.append(m)
+        else:
+            assert len(basis) == len(self)
+            for i, e_j in enumerate(basis):
+                # All should be orthogonal
+                if i != 0:
+                    assert e_j.is_orthogonal(basis[0])
+                m_ops.append(MeasurementOperator.create_from_basis(e_j))
 
-        for i in range(len(basis)):
-            m_op = basis[i]
+        for m_op in m_ops:
             prob = self.get_prob_from_measurement_op(m_op)
             weights += [prob]
 
@@ -959,6 +964,10 @@ def test():
     assert s("|10>").measure_partial(1) == ("1", _0)
     assert s("|10>").measure_partial(2) == ("0", _1)
 
+    # measure in arbitrary basis
+    _0.measure(basis=[_p, _m])
+    s("|00>").measure(basis=[b_phi_p, b_phi_m, b_psi_m, b_psi_p])
+
     # Maximally entangled
     result, pms = b_phi_p.measure_partial(1)
     if result == "0":
@@ -1295,14 +1304,25 @@ def test_light():
     experiment(3, random_lights, [ver_filter, diag_filter, hor_filter])
 
 
+def test_krisi_measurement():
+    beta = 2.6
+    b_0 = State(1 / np.sqrt(2) * (s("|100>") - s("|010>")))
+    b_1 = State(1 / np.sqrt(2) * (s("|011>") - s("|101>")))
+    meas = s("|000>").measure(basis=[
+        s("|000>"),
+        State(np.cos(beta) * b_0 + np.sin(beta) * s("|001>")),
+        State(-np.sin(beta) * b_0 + np.cos(beta) * s("|001>")),
+        b_0,
+        b_1,
+        State(np.cos(beta) * b_1 + np.sin(beta) * s("|001>")),
+        State(-np.sin(beta) * b_1 + np.cos(beta) * s("|001>")),
+        s("|111>"),
+    ])
+    print(meas)
+
+
 if __name__ == "__main__":
     test()
-    # test_quantum_processor()
-    # test_light()
-
-    # Partial measurement - I want to measure just the first qubit...
-    # TODO - makes sense???
-    # m0_ = m0 * I
-    # s("|10>").measure_partial(1)
-    # result_0 = post_measurement.measure()
-    # test_measure_partial()
+    test_quantum_processor()
+    test_light()
+    test_krisi_measurement()