partial application of multiqubit gates

lib.py

@@ -13,6 +13,7 @@ ListOrNdarray = Union[list, np.ndarray]
 
 REPR_EMPTY_SET = "Ø"
 REPR_TENSOR_OP = "⊗"
+REPR_TARGET = "⊕"
 REPR_GREEK_BETA = "β"
 REPR_GREEK_PSI = "ψ"
 REPR_GREEK_PHI = "φ"
@@ -26,7 +27,8 @@ UNIVERSE_STATES = []
 
 
 class Matrix(object):
-    """Wraps a Matrix... it's for my understanding, this could easily probably be np.array"""
+    """Wraps a Matrix... it's for my understanding, this could easily
+    probably be np.array"""
 
     def __init__(self, m: ListOrNdarray = None, *args, **kwargs):
         """
@@ -52,8 +54,6 @@ class Matrix(object):
     def __eq__(self, other):
         if isinstance(other, Complex):
             return bool(np.allclose(self.m, other))
-        elif isinstance(other, PartialQubit):
-            return False
         return np.allclose(self.m, other.m)
 
     def __mul_or_kron__(self, other):
@@ -182,7 +182,9 @@ class State(Vector):
     def to_bloch_angles(self):
         """Returns the angles of this state on the Bloch sphere"""
         if not self.m.shape == (2, 1):
-            raise Exception("State needs to describe only 1 qubit on the bloch sphere (2x1 matrix)")
+            raise Exception(
+                "State needs to describe only 1 qubit on the bloch sphere ("
+                "2x1 matrix)")
         m0, m1 = self.m[0][0], self.m[1][0]
 
         # theta is between 0 and pi
@@ -193,7 +195,8 @@ class State(Vector):
 
         div = np.sin(theta / 2)
         if div == 0:
-            # here is doesn't matter what phi is as phase at the poles is arbitrary
+            # here is doesn't matter what phi is as phase at the poles is
+            # arbitrary
             phi = 0
         else:
             exp = m1 / div
@@ -234,10 +237,13 @@ class State(Vector):
                             return repr(used_state)
                     except:
                         ...
-                next_state_sub = ''.join([REPR_MATH_SUBSCRIPT_NUMBERS[int(d)] for d in str(len(UNIVERSE_STATES))])
+                next_state_sub = ''.join(
+                    [REPR_MATH_SUBSCRIPT_NUMBERS[int(d)] for d in
+                     str(len(UNIVERSE_STATES))])
                 self.name = '{}{}'.format(REPR_GREEK_PSI, next_state_sub)
         UNIVERSE_STATES.append(self)
-        # matrix_rep = "{}".format(self.m).replace('[', '').replace(']', '').replace('\n', '|').strip()
+        # matrix_rep = "{}".format(self.m).replace('[', '').replace(']',
+        # '').replace('\n', '|').strip()
         # state_name = '|{}> = {}'.format(self.name, matrix_rep)
         state_name = '|{}>'.format(self.name)
         return state_name
@@ -251,7 +257,8 @@ class State(Vector):
         return np.sqrt((self.inner(self)).m).item(0)
 
     def is_orthogonal(self, other):
-        """If the inner (dot) product is zero, this vector is orthogonal to other"""
+        """If the inner (dot) product is zero, this vector is orthogonal to
+        other"""
         return self.inner(other) == 0
 
     def get_prob_from_measurement_op(self, m_op):
@@ -284,8 +291,10 @@ class State(Vector):
               What about if measured twice but in different bases?
               E.g. measure1 -> computation -> A
                    measure2 -> +/- basis -> B
-                   measure3 -> computation -> should it return A again or random weighted?
+                   measure3 -> computation -> should it return A again or
-                   measure4 -> +/- basis -> should it return B again or random weighted?
+                   random weighted?
+                   measure4 -> +/- basis -> should it return B again or
+                   random weighted?
         :return: binary representation of the measured qubit (e.g. "011")
         """
         # if self.measurement_result:
@@ -320,7 +329,8 @@ class State(Vector):
             weights = list(np.array(weights) / sum(weights))
 
         format_str = self.get_fmt_of_element()
-        choices = empty_choices + [format_str.format(i) for i in range(len(weights))]
+        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]
 
@@ -329,7 +339,8 @@ class State(Vector):
     def measure_with_op(self, mo):
         """
         Measures with a measurement operator mo
-        TODO: Can't define `mo: MeasurementOperator` because in python you can't declare classes before defining them
+        TODO: Can't define `mo: MeasurementOperator` because in python you
+        can't declare classes before defining them
         """
         m = mo.on(self).m / np.sqrt(mo.get_prob(self))
         return State(m)
@@ -341,7 +352,10 @@ class State(Vector):
         """
         max_qubits = int(np.log2(len(self)))
         if not (0 < qubit_n <= max_qubits):
-            raise Exception("Partial measurement of qubit_n must be between 1 and {}".format(max_qubits))
+            raise Exception(
+                "Partial measurement of qubit_n must be between 1 and {"
+                "}".format(
+                    max_qubits))
         format_str = self.get_fmt_of_element()
         # e.g. for state |000>:
         # ['000', '001', '010', '011', '100', '101', '110', '111']
@@ -352,12 +366,15 @@ class State(Vector):
         # [0, 1, 4, 5]
         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]
+            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]]))
+        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)))
 
@@ -375,7 +392,9 @@ class State(Vector):
             bin_length = 2 ** len(bin_repr)
             arr = [[1] if i == dec_bin else [0] for i in range(bin_length)]
         except:
-            raise Exception("State from string should be of the form |00..01> with all numbers either 0 or 1")
+            raise Exception(
+                "State from string should be of the form |00..01> with all "
+                "numbers either 0 or 1")
         return cls(arr)
 
     @staticmethod
@@ -405,7 +424,8 @@ def test_measure_partial():
 
 
 def normalize_state(vector: Vector):
-    """Normalize a state by dividing by the square root of sum of the squares of states"""
+    """Normalize a state by dividing by the square root of sum of the squares
+    of states"""
     norm_coef = np.sqrt(np.sum(np.array(vector.m) ** 2))
     if norm_coef == 0:
         raise TypeError("zero-sum vector")
@@ -464,11 +484,11 @@ class Operator(object):
         """An Operator turns one function into another"""
         self.func = func
 
-    def on(self, *args):
+    def on(self, *args, **kwargs):
-        return self(*args)
+        return self(*args, **kwargs)
 
-    def __call__(self, *args):
+    def __call__(self, *args, **kwargs):
-        return self.func(*args)
+        return self.func(*args, **kwargs)
 
 
 class LinearOperator(Operator):
@@ -480,8 +500,10 @@ class LinearOperator(Operator):
 
     def _is_linear(self):
         # A function is (jointly) linear in a given set of variables
-        # if all second-order derivatives are identically zero (including mixed ones).
+        # if all second-order derivatives are identically zero (including
-        # https://stackoverflow.com/questions/36283548/check-if-an-equation-is-linear-for-a-specific-set-of-variables
+        # mixed ones).
+        # https://stackoverflow.com/questions/36283548/check-if-an-equation
+        # -is-linear-for-a-specific-set-of-variables
         expr, vars_ = self.func.expr, self.func.variables
         for x in vars_:
             for y in vars_:
@@ -501,7 +523,8 @@ def test_linear_operator():
     assert _times_2.on(5) == 10
     assert _times_2(5) == 10
 
-    assert_raises(TypeError, "Not a linear operator", LinearOperator, sp.Lambda(_x, _x ** 2))
+    assert_raises(TypeError, "Not a linear operator", LinearOperator,
+                  sp.Lambda(_x, _x ** 2))
 
 
 class SquareMatrix(Matrix):
@@ -515,9 +538,12 @@ class SquareMatrix(Matrix):
 
 
 class LinearTransformation(LinearOperator, Matrix):
-    def __init__(self, m: ListOrNdarray, func=None, name: str = '', *args, **kwargs):
+    def __init__(self, m: ListOrNdarray, func=None, name: str = '', *args,
-        """LinearTransformation (or linear map) inherits from both LinearOperator and a Matrix
+                 **kwargs):
-        It is used to act on a State vector by defining the operator to be the dot product"""
+        """LinearTransformation (or linear map) inherits from both
+        LinearOperator and a Matrix
+        It is used to act on a State vector by defining the operator to be
+        the dot product"""
         self.name = name
         Matrix.__init__(self, m=m, *args, **kwargs)
         self.func = func or self.operator_func
@@ -525,7 +551,8 @@ class LinearTransformation(LinearOperator, Matrix):
 
     def _is_linear(self):
         # Every matrix transformation is a linear transformation
-        # https://www.mathbootcamps.com/proof-every-matrix-transformation-is-a-linear-transformation/
+        # https://www.mathbootcamps.com/proof-every-matrix-transformation-is
+        # -a-linear-transformation/
         return True
 
     def operator_func(self, other):
@@ -555,7 +582,8 @@ class UnitaryMatrix(SquareMatrix):
             raise TypeError("Not a Unitary matrix")
 
     def _is_unitary(self):
-        """Checks if the matrix product of itself with conjugate transpose is Identity UU+ = I"""
+        """Checks if the matrix product of itself with conjugate transpose is
+        Identity UU+ = I"""
         UU_ = np.dot(self._conjugate_transpose(), self.m)
         I = np.eye(self.m.shape[0])
         return np.isclose(UU_, I).all()
@@ -574,25 +602,67 @@ class HermitianMatrix(SquareMatrix):
 
 
 class UnitaryOperator(LinearTransformation, UnitaryMatrix):
-    def __init__(self, m: ListOrNdarray, name: str = '', *args, **kwargs):
+    def __init__(self, m: ListOrNdarray, name: str = '', partials=None, *args,
-        """UnitaryOperator inherits from both LinearTransformation and a Unitary matrix
+                 **kwargs):
-        It is used to act on a State vector by defining the operator to be the dot product"""
+        """UnitaryOperator inherits from both LinearTransformation and a
+        Unitary matrix
+        It is used to act on a State vector by defining the operator to be
+        the dot product
+
+        Partials - a list of partial operators that are used in a multi-qubit
+        UnitaryOperator to define the operator on each qubit. Each element of
+        the list is the Nth partial that is used, i.e. for the first - |0><0|,
+        for the second - |1><1|
+        """
+        if np.shape(m) != (2, 2) and partials is None:
+            raise Exception("Please define partials in a non-single operator")
         UnitaryMatrix.__init__(self, m=m, *args, **kwargs)
-        LinearTransformation.__init__(self, m=m, func=self.operator_func, *args, **kwargs)
+        LinearTransformation.__init__(self, m=m, func=self.operator_func, *args,
+                                      **kwargs)
         self.name = name
+        self.partials = partials
 
     def operator_func(self, other, which_qbit=None):
         this_cols, other_rows = np.shape(self.m)[1], np.shape(other.m)[0]
-        if this_cols != other_rows:
+        if this_cols == other_rows:
-            if which_qbit is None:
+            return State(np.dot(self.m, other.m))
-                raise Exception("Operating dim-{} operator on a dim-{} state. "
+        if which_qbit is None:
-                      "Please specify which_qubit to operate on".format(
+            raise Exception("Operating dim-{} operator on a dim-{} state. "
-                    this_cols, other_rows))
+                            "Please specify which_qubit to operate on".format(
+                this_cols, other_rows))
+        total_qbits = int(np.log2(other_rows))
+        if type(which_qbit) is int:
+            # single qubit-gate
+            assert this_cols == 2
+            assert 0 <= which_qbit < total_qbits
+            extended_m = [self if i == which_qbit else I for i in
+                          range(total_qbits)]
+            new_m = np.prod(extended_m).m
+        elif type(which_qbit) is list:
+            # single or multiple qubit-gate
+            assert 1 <= len(which_qbit) < total_qbits
+            assert len(which_qbit) == len(self.partials)
+            assert all([q < total_qbits for q in which_qbit])
+            extended_m, next_partial = [[], []], 0
+            for qbit in range(total_qbits):
+                if qbit not in which_qbit:
+                    extended_m[0].append(I)
+                    extended_m[1].append(I)
+                else:
+                    this_partial = self.partials[next_partial]
+                    if this_partial == C_partial:
+                        extended_m[0].append(s("|0><0|"))
+                        extended_m[1].append(s("|1><1|"))
+                    else:
+                        extended_m[0].append(I)
+                        extended_m[1].append(this_partial)
+                    next_partial += 1
+            new_m = sum([np.prod(e).m for e in extended_m])
+        else:
+            raise Exception("which_qubit needs to be either an int of N-th qubit or list")
 
-            extended_m = [self if i == which_qbit else I for i in range(int(np.log2(other_rows)))]
+        extended_op = UnitaryOperator(new_m, name=self.name, partials=self.partials)
-            extended_op = UnitaryOperator(np.prod(extended_m).m)
+        return State(np.dot(extended_op.m, other.m))
-            return State(np.dot(extended_op.m, other.m))
-        return State(np.dot(self.m, other.m))
 
     def __repr__(self):
         if self.name:
@@ -608,14 +678,17 @@ class Gate(UnitaryOperator):
 
 class HermitianOperator(LinearTransformation, HermitianMatrix):
     def __init__(self, m: ListOrNdarray, name: str = '', *args, **kwargs):
-        """HermitianOperator inherits from both LinearTransformation and a Hermitian matrix
+        """HermitianOperator inherits from both LinearTransformation and a
-        It is used to act on a State vector by defining the operator to be the dot product"""
+        Hermitian matrix
+        It is used to act on a State vector by defining the operator to be
+        the dot product"""
         self.name = name
         HermitianMatrix.__init__(self, m=m, *args, **kwargs)
         LinearOperator.__init__(self, func=self.operator_func, *args, **kwargs)
 
     def operator_func(self, other):
-        """This might not return a normalized state vector so don't wrap it in State"""
+        """This might not return a normalized state vector so don't wrap it
+        in State"""
         return Vector(np.dot(self.m, other.m))
 
     def __repr__(self):
@@ -643,71 +716,6 @@ class DensityMatrix(HermitianOperator):
         return np.isclose(np.trace(self.m), 1.0)
 
 
-# How to add a CNOT gate to the Quantum Processor?
-# Imagine if I have to act on a 3-qubit computer and CNOT(q1, q3)
-#
-# 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
-#
-# CNOT(q1, I, q2):
-# |0><0| x I_2 x I_2 + |1><1| x I_2 x X
-# np.kron(np.kron(np.outer(_0.m, _0.m), np.eye(2)), np.eye(2)) + np.kron(np.kron(np.outer(_1.m, _1.m), np.eye(2)), X.m)
-#
-#
-# CNOT(q1, q2):
-# |0><0| x I + |1><1| x X
-# np.kron(np.outer(_0.m, _0.m), np.eye(2)) + np.kron(np.outer(_1.m, _1.m), X.m)
-# _0.x(_0) * Matrix(I.m) + _1.x(_1) * Matrix(X.m)
-
-class PartialQubit(object):
-    def __init__(self, rpr):
-        self.rpr = rpr
-        self.operator = None
-
-    def __repr__(self):
-        return str("-{}-".format(self.rpr))
-
-
-C_partial = PartialQubit("C")
-x_partial = PartialQubit("x")
-
-
-class TwoQubitOperator(UnitaryOperator):
-    def __init__(self, m: ListOrNdarray, A: PartialQubit, B: PartialQubit,
-                 A_p: UnitaryOperator, B_p: UnitaryOperator, *args, **kwargs):
-        super().__init__(m, *args, **kwargs)
-        A.operator, B.operator = self, self
-        self.A = A
-        self.B = B
-        self.A_p = A_p
-        self.B_p = B_p
-
-        A.operator = self
-        B.operator = self
-
-    def verify_step(self, step):
-        if not (step.count(self.A) == 1 and step.count(self.B) == 1):
-            raise RuntimeError("Both {} and {} need to be defined in the same step exactly once".format(
-                self.A, self.B
-            ))
-
-    def compose(self, step, state):
-        # _0.x(_0) * Matrix(I.m) + _1.x(_1) * Matrix(X.m)
-        outer_0, outer_1 = [], []
-        for s in step:
-            if s == self.A:
-                outer_0.append(_0.x(_0))
-                outer_1.append(_1.x(_1))
-            elif s == self.B:
-                outer_0.append(Matrix(self.A_p.m))
-                outer_1.append(Matrix(self.B_p.m))
-            else:
-                outer_0.append(Matrix(s.m))
-                outer_1.append(Matrix(s.m))
-        return reduce((lambda x, y: x * y), outer_0) + reduce((lambda x, y: x * y), outer_1)
-
-
 """
 Define States and Operators 
 """
@@ -767,73 +775,92 @@ bell_basis = [b_phi_p, b_psi_p, b_phi_m, b_psi_m]
 well_known_states = [_p, _m, b_phi_p, b_psi_p, b_phi_m, b_psi_m]
 
 _ = I = Gate([[1, 0],
-             [0, 1]],
+              [0, 1]],
-            name="-")
+             name="-")
 
 X = Gate([[0, 1],
-         [1, 0]],
+          [1, 0]],
-        name="X")
+         name="X")
 
 Y = Gate([[0, -1j],
-         [1j, 0]],
+          [1j, 0]],
-        name="Y")
+         name="Y")
 
 Z = Gate([[1, 0],
-         [0, -1]],
+          [0, -1]],
-        name="Z")
+         name="Z")
-
 
 # These are rotations that are specified commonly e.g. in
 #       https://www.quantum-inspire.com/kbase/rotation-operators/
-#       http://www.vcpc.univie.ac.at/~ian/hotlist/qc/talks/bloch-sphere-rotations.pdf
+#       http://www.vcpc.univie.ac.at/~ian/hotlist/qc/talks/bloch-sphere
+#       -rotations.pdf
 #       and elsewhere DO NOT equate X, Y and Z for theta=np.pi
-# However - they are correct up to a global phase which is all that matters for measurement purposes
+# However - they are correct up to a global phase which is all that matters
+# for measurement purposes
 #
 
 H = Gate([[1 / np.sqrt(2), 1 / np.sqrt(2)],
-         [1 / np.sqrt(2), -1 / np.sqrt(2)], ],
+          [1 / np.sqrt(2), -1 / np.sqrt(2)], ],
          name="H")
 
 # Rotation operators
 # https://www.quantum-inspire.com/kbase/rx-gate/
 Rx = lambda theta: Gate([[np.cos(theta / 2), -1j * np.sin(theta / 2)],
-                [-1j * np.sin(theta / 2), np.cos(theta / 2)]],
+                         [-1j * np.sin(theta / 2), np.cos(theta / 2)]],
-               name="Rx")
+                        name="Rx")
-
 
 Ry = lambda theta: Gate([[np.cos(theta / 2), -np.sin(theta / 2)],
-                [np.sin(theta / 2), np.cos(theta / 2)]],
+                         [np.sin(theta / 2), np.cos(theta / 2)]],
-               name="Ry")
+                        name="Ry")
-
 
 Rz = lambda theta: Gate([[np.power(np.e, -1j * theta / 2), 0],
-                [0, np.power(np.e, 1j * theta / 2)]],
+                         [0, np.power(np.e, 1j * theta / 2)]],
-               name="Rz")
+                        name="Rz")
-
 
 # See [T-Gate](https://www.quantum-inspire.com/kbase/t-gate/)
 T = lambda phi: Gate([[1, 0],
-                     [0, np.power(np.e, 1j * phi / 2)]],
+                      [0, np.power(np.e, 1j * phi / 2)]],
                      name="T")
 
-CNOT = TwoQubitOperator([
+# How to add a CNOT gate to the Quantum Processor?
+# Imagine if I have to act on a 3-qubit computer and CNOT(q1, q3)
+#
+# 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
+#
+# CNOT(q1, I, q2):
+# |0><0| x I_2 x I_2 + |1><1| x I_2 x X
+# np.kron(np.kron(np.outer(_0.m, _0.m), np.eye(2)), np.eye(2)) + np.kron(
+# np.kron(np.outer(_1.m, _1.m), np.eye(2)), X.m)
+#
+#
+# CNOT(q1, q2):
+# |0><0| x I + |1><1| x X
+# np.kron(np.outer(_0.m, _0.m), np.eye(2)) + np.kron(np.outer(_1.m, _1.m), X.m)
+# _0.x(_0) * Matrix(I.m) + _1.x(_1) * Matrix(X.m)
+
+C_partial = Gate(I.m, name="C")
+x_partial = Gate(X.m, name=REPR_TARGET)
+
+CNOT = Gate([
     [1, 0, 0, 0],
     [0, 1, 0, 0],
     [0, 0, 0, 1],
     [0, 0, 1, 0],
-], C_partial, x_partial, I, X)
+], name="CNOT", partials=[C_partial, x_partial])
 
-
+TOFF = Gate([
-# TOFFOLLI_GATE = ThreeQubitOperator([
+    [1, 0, 0, 0, 0, 0, 0, 0],
-#     [1, 0, 0, 0, 0, 0, 0, 0],
+    [0, 1, 0, 0, 0, 0, 0, 0],
-#     [0, 1, 0, 0, 0, 0, 0, 0],
+    [0, 0, 1, 0, 0, 0, 0, 0],
-#     [0, 0, 1, 0, 0, 0, 0, 0],
+    [0, 0, 0, 1, 0, 0, 0, 0],
-#     [0, 0, 0, 1, 0, 0, 0, 0],
+    [0, 0, 0, 0, 1, 0, 0, 0],
-#     [0, 0, 0, 0, 1, 0, 0, 0],
+    [0, 0, 0, 0, 0, 1, 0, 0],
-#     [0, 0, 0, 0, 0, 1, 0, 0],
+    [0, 0, 0, 0, 0, 0, 0, 1],
-#     [0, 0, 0, 0, 0, 0, 0, 1],
+    [0, 0, 0, 0, 0, 0, 1, 0],
-#     [0, 0, 0, 0, 0, 0, 1, 0],
+], partials=[C_partial, C_partial, x_partial])
-# ], C_partial, C_partial, x_partial, I, I, X)
 
 
 def assert_raises(exception, msg, callable, *args, **kwargs):
@@ -857,12 +884,15 @@ def test_unitary_hermitian():
     # Unitary is UU+ = I; Hermitian is U = U+
     # Matrixes could be either, neither or both
     # Quantum operators (gates) are described *only* by unitary transformations
-    # Hermitian operators are used for measurement operators - https://towardsdatascience.com/understanding-basics-of-measurements-in-quantum-computation-4c885879eba0
+    # Hermitian operators are used for measurement operators -
+    # https://towardsdatascience.com/understanding-basics-of-measurements-in
+    # -quantum-computation-4c885879eba0
     h_not_u = [
         [1, 0],
         [0, 2],
     ]
-    assert_not_raises(TypeError, "Not a Hermitian matrix", HermitianMatrix, h_not_u)
+    assert_not_raises(TypeError, "Not a Hermitian matrix", HermitianMatrix,
+                      h_not_u)
     assert_raises(TypeError, "Not a Unitary matrix", UnitaryMatrix, h_not_u)
 
     u_not_h = [
@@ -876,14 +906,16 @@ def test_unitary_hermitian():
         [0, 1],
         [1, 0],
     ]
-    assert_not_raises(TypeError, "Not a Hermitian matrix", HermitianMatrix, u_and_h)
+    assert_not_raises(TypeError, "Not a Hermitian matrix", HermitianMatrix,
+                      u_and_h)
     assert_not_raises(TypeError, "Not a Unitary matrix", UnitaryMatrix, u_and_h)
 
     not_u_not_h = [
         [1, 2],
         [0, 1],
     ]
-    assert_raises(TypeError, "Not a Hermitian matrix", HermitianMatrix, not_u_not_h)
+    assert_raises(TypeError, "Not a Hermitian matrix", HermitianMatrix,
+                  not_u_not_h)
     assert_raises(TypeError, "Not a Unitary matrix", UnitaryMatrix, not_u_not_h)
 
 
@@ -947,32 +979,69 @@ def test_eigenstuff():
            [(1.0, HorizontalVector([1., 0.])), (0., HorizontalVector([0., 1.]))]
 
 
+def test_partials():
+    # normal 2 qbit state
+    assert CNOT.on(s("|00>")) == s("|00>")
+    assert CNOT.on(s("|10>")) == s("|11>")
+    assert_raises(Exception,
+                  "Operating dim-4 operator on a dim-8 state. Please specify "
+                  "which_qubit to operate on",
+                  CNOT.on, s("|100>"))
+    # apply on 0, 1 of 3qbit state
+    assert CNOT.on(s("|000>"), which_qbit=[0, 1]) == s("|000>")
+    assert CNOT.on(s("|100>"), which_qbit=[0, 1]) == s("|110>")
+    # apply on 1, 2 of 3qbit state
+    assert CNOT.on(s("|000>"), which_qbit=[1, 2]) == s("|000>")
+    assert CNOT.on(s("|010>"), which_qbit=[1, 2]) == s("|011>")
+    # apply on 0, 2 of 3qbit state
+    assert CNOT.on(s("|000>"), which_qbit=[0, 2]) == s("|000>")
+    assert CNOT.on(s("|100>"), which_qbit=[0, 2]) == s("|101>")
+    # apply on 0, 2 of 4qbit state
+    assert CNOT.on(s("|1000>"), which_qbit=[0, 2]) == s("|1010>")
+    # apply on 0, 3 of 4qbit state
+    assert CNOT.on(s("|1000>"), which_qbit=[0, 3]) == s("|1001>")
+
+    # test Toffoli gate
+    # assert TOFF.on(s("|000>")) == s("|000>")
+    # assert TOFF.on(s("|100>")) == s("|100>")
+    # assert TOFF.on(s("|110>")) == s("|111>")
+    # assert TOFF.on(s("|1100>"), which_qbit=[0, 1, 3]) == s("|1101>")
+
+
 def test():
     # Test properties of Hilbert vector space
     # The four postulates of Quantum Mechanics
     # I: States | Associated to any physical system is a complex vector space
-    #             known as the state space of the system. If the system is closed
+    #             known as the state space of the system. If the system is
+    #             closed
     #             then the system is described completely by its state vector
     #             which is a unit vector in the space.
     # Mathematically, this vector space is also a function space
     assert _0 + _1 == _1 + _0  # commutativity of vector addition
     assert _0 + (_1 + _p) == (_0 + _1) + _p  # associativity of vector addition
-    assert 8 * (_0 + _1) == 8 * _0 + 8 * _1  # Linear when multiplying by constants
+    assert 8 * (
+            _0 + _1) == 8 * _0 + 8 * _1  # Linear when multiplying by
+    # constants
     assert _0.inner(_0) == 1  # parallel have 1 product
     assert _0.inner(_1) == 0  # orthogonal have 0 product
     assert _0.is_orthogonal(_1)
-    assert _1.inner(8 * _0) == 8 * _1.inner(_0)  # Inner product is linear multiplied by constants
+    assert _1.inner(8 * _0) == 8 * _1.inner(
-    assert _p.inner(_1 + _0) == _p.inner(_1) + _p.inner(_0)  # Inner product is linear in superpos of vectors
+        _0)  # Inner product is linear multiplied by constants
+    assert _p.inner(_1 + _0) == _p.inner(_1) + _p.inner(
+        _0)  # Inner product is linear in superpos of vectors
 
-    assert np.isclose(_1.length(), 1.0)  # all of the vector lengths are normalized
+    assert np.isclose(_1.length(),
+                      1.0)  # all of the vector lengths are normalized
     assert np.isclose(_0.length(), 1.0)
     assert np.isclose(_p.length(), 1.0)
 
-    assert _0.inner(_1) == _1.inner(_0).complex_conjugate()  # non-commutative inner product
+    assert _0.inner(_1) == _1.inner(
+        _0).complex_conjugate()  # non-commutative inner product
 
     test_to_from_angles()
 
-    # II: Dynamics | The evolution of a closed system is described by a unitary transformation
+    # II: Dynamics | The evolution of a closed system is described by a
+    # unitary transformation
     #
     test_linear_operator()
 
@@ -994,14 +1063,23 @@ def test():
     # Test Pauli rotation gates
     testRotPauli()
 
-    # III: Measurement | A quantum measurement is described by an orthonormal basis |e_j>
+    # Test 2+ qbit gates and partials
-    #                    for state space. If the initial state of the system is |ψ>
+    test_partials()
-    #                    then we get outcome j with probability pr(j) = |<e_j|ψ>|^2
+
+    # III: Measurement | A quantum measurement is described by an orthonormal
+    # basis |e_j>
+    #                    for state space. If the initial state of the system
+    #                    is |ψ>
+    #                    then we get outcome j with probability pr(j) =
+    #                    |<e_j|ψ>|^2
     # Note: The postulates are applicable on closed, isolated systems.
-    #       Systems that are closed and are described by unitary time evolution by a Hamiltonian
+    #       Systems that are closed and are described by unitary time
-    #       can be measured by projective measurements. Systems are not closed in reality and hence
+    #       evolution by a Hamiltonian
+    #       can be measured by projective measurements. Systems are not
+    #       closed in reality and hence
     #       are immeasurable using projective measurements.
-    #       POVM (Positive Operator-Valued Measure) is a restriction on the projective measurements,
+    #       POVM (Positive Operator-Valued Measure) is a restriction on the
+    #       projective measurements,
     #       such that it encompasses everything except the environment.
 
     assert _0.get_prob(0) == 1  # Probability for |0> in 0 is 1
@@ -1052,7 +1130,8 @@ def test():
     assert CNOT.on(_1 * _0) == _1 * _1
     assert CNOT.on(_1 * _1) == _1 * _0
 
-    # ALL FOUR NOW - Create a Bell state (I) that has H|0> (II), measures (III) and composition in CNOT (IV)
+    # ALL FOUR NOW - Create a Bell state (I) that has H|0> (II), measures (
+    # III) and composition in CNOT (IV)
     # Bell state
     # First - create a superposition
     superpos = H.on(_0)
@@ -1065,10 +1144,14 @@ def test():
                           [0.],
                           [1 / np.sqrt(2)], ])
 
-    assert np.isclose(bell.get_prob(0b00), 0.5)  # Probability for bell in 00 is 0.5
+    assert np.isclose(bell.get_prob(0b00),
-    assert np.isclose(bell.get_prob(0b01), 0)  # Probability for bell in 01 is 0.
+                      0.5)  # Probability for bell in 00 is 0.5
-    assert np.isclose(bell.get_prob(0b10), 0)  # Probability for bell in 10 is 0.
+    assert np.isclose(bell.get_prob(0b01),
-    assert np.isclose(bell.get_prob(0b11), 0.5)  # Probability for bell in 11 is 0.5
+                      0)  # Probability for bell in 01 is 0.
+    assert np.isclose(bell.get_prob(0b10),
+                      0)  # Probability for bell in 10 is 0.
+    assert np.isclose(bell.get_prob(0b11),
+                      0.5)  # Probability for bell in 11 is 0.5
 
     ################################
     assert _0.x(_0) == Matrix([[1, 0],
@@ -1108,16 +1191,17 @@ def naive_load_test(N):
     from time import time
     from sys import getsizeof
 
-    print("{:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10}".format(
+    print(
-        "qbits",
+        "{:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10}".format(
-        "kron_len",
+            "qbits",
-        "mem_used",
+            "kron_len",
-        "mem_per_q",
+            "mem_used",
-        "getsizeof",
+            "mem_per_q",
-        "getsiz/len",
+            "getsizeof",
-        "nbytes",
+            "getsiz/len",
-        "nbytes/len",
+            "nbytes",
-        "time"))
+            "nbytes/len",
+            "time"))
 
     _0 = State([[1], [0]], name='0')
 
@@ -1131,16 +1215,18 @@ def naive_load_test(N):
         len_m = len(m)
         elapsed = time() - start
         mem_b = process.memory_info().rss - mem_init
-        print("{:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10}".format(
+        print(
-            i,
+            "{:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {:>10} {"
-            len(m),
+            ":>10}".format(
-            sizeof_fmt(mem_b),
+                i,
-            sizeof_fmt(mem_b / len_m),
+                len(m),
-            sizeof_fmt(getsizeof(m)),
+                sizeof_fmt(mem_b),
-            sizeof_fmt(getsizeof(m) / len_m),
+                sizeof_fmt(mem_b / len_m),
-            sizeof_fmt(m.m.nbytes),
+                sizeof_fmt(getsizeof(m)),
-            sizeof_fmt(m.m.nbytes / len_m),
+                sizeof_fmt(getsizeof(m) / len_m),
-            np.round(elapsed, 2)))
+                sizeof_fmt(m.m.nbytes),
+                sizeof_fmt(m.m.nbytes / len_m),
+                np.round(elapsed, 2)))
 
         gc.collect()
 
@@ -1160,12 +1246,15 @@ class QuantumCircuit(object):
 
     def add_row_step(self, row: int, step: int, qbit_state):
         if len(self.steps) <= step:
-            self.steps += [[I for _ in range(self.n_qubits)] for _ in range(len(self.steps) - step + 1)]
+            self.steps += [[I for _ in range(self.n_qubits)] for _ in
+                           range(len(self.steps) - step + 1)]
         self.steps[step][row] = qbit_state
 
     def add_step(self, step_data: list):
         if len(step_data) != self.n_qubits:
-            raise RuntimeError("Length of step is: {}, should be: {}".format(len(step_data), self.n_qubits))
+            raise RuntimeError(
+                "Length of step is: {}, should be: {}".format(len(step_data),
+                                                              self.n_qubits))
         step_i = len(step_data)
         for row, qubit_state in enumerate(step_data):
             self.add_row_step(row, step_i, qubit_state)
@@ -1187,7 +1276,8 @@ class QuantumCircuit(object):
 
     def compose_quantum_state(self, step):
         partials = [op for op in step if isinstance(op, PartialQubit)]
-        # TODO: No more than 1 TwoQubitGate **OR** UnitaryOperator can be used in a step
+        # TODO: No more than 1 TwoQubitGate **OR** UnitaryOperator can be
+        #  used in a step
         for partial in partials:
             two_qubit_op = partial.operator
             two_qubit_op.verify_step()
@@ -1205,7 +1295,8 @@ class QuantumCircuit(object):
         if not self.current_quantum_state:
             self.current_quantum_state = step_quantum_state
         else:
-            self.current_quantum_state = step_quantum_state.inner(self.current_quantum_state)
+            self.current_quantum_state = step_quantum_state.inner(
+                self.current_quantum_state)
         self.current_step += 1
 
     def run(self):
@@ -1235,9 +1326,11 @@ class QuantumProcessor(object):
 
     def __init__(self, circuit: QuantumCircuit):
         if circuit.rows != circuit.n_qubits:
-            raise Exception("Declared circuit with n_qubits: {} but called add_row: {}".format(
+            raise Exception(
-                circuit.n_qubits, circuit.rows
+                "Declared circuit with n_qubits: {} but called add_row: {"
-            ))
+                "}".format(
+                    circuit.n_qubits, circuit.rows
+                ))
         self.circuit = circuit
         self.c_step = 0
         self.c_q_state = None
@@ -1245,13 +1338,6 @@ class QuantumProcessor(object):
         self.reset()
 
     def compose_quantum_state(self, step):
-        if any([type(s) is PartialQubit for s in step]):
-            two_qubit_gates = filter(lambda s: type(s) is PartialQubit, step)
-            state = []
-            for two_qubit_gate in two_qubit_gates:
-                two_qubit_gate.operator.verify_step(step)
-                state = two_qubit_gate.operator.compose(step, state)
-            return state
         return reduce((lambda x, y: x * y), step)
 
     def step(self):
@@ -1327,7 +1413,8 @@ def test_light():
     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)]])
+        random_pol = Vector(
+            [[np.random.uniform(0, 1)], [np.random.uniform(0, 1)]])
         return normalize_state(random_pol)
 
     def experiment(id, random_ls, filters):
@@ -1354,12 +1441,22 @@ def test_light():
     experiment(2, random_lights, [ver_filter, hor_filter])
 
     # TODO: Something is wrong here...
-    # Vertical after diagonal after horizontal - should result in ~ 50% compared to only Horizontal
+    # Vertical after diagonal after horizontal - should result in ~ 50%
+    # compared to only Horizontal
     experiment(3, random_lights, [ver_filter, diag_filter, hor_filter])
 
 
+def get_bin_fmt(count):
+    return "{:0" + str(int(np.ceil(np.log2(count)))) + "b}"
+
+
+def generate_bin(i, count):
+    format_str = get_bin_fmt(count)
+    return format_str.format(i)
+
+
 def generate_bins(count):
-    format_str = "{:0" + str(int(np.ceil(np.log2(count)))) + "b}"
+    format_str = get_bin_fmt(count)
     return [format_str.format(i) for i in range(count)]