unnecessary high multiplication of kron space

lib_q_computer.py

@@ -80,27 +80,26 @@ def kron(_list):
 
 
 class Gate(State):
-    def on(self, state):
+    def on(self, other_state):
         """
         Applies a gate operation on a state.
         If applying to a substate, use the index of the substate, e.g. `H.on(bell[0])` will apply the Hadamard gate
         on the 0th qubit of the `bell` state.
-        :param state: another state (e.g. `H.on(q1)` or a list of states (e.g. for `CNOT.on([q1, q2])`)
+        :param other_state: another state (e.g. `H.on(q1)` or a list of states (e.g. for `CNOT.on([q1, q2])`)
         :return: the state after the application of the Gate
         """
-        this_state = self.matrix_state
-        if type(state) == list:
-            other_state = kron([e.matrix_state for e in state])
-        else:
-            other_state = state.matrix_state
-        if this_state.shape[1] != other_state.shape[0]:
+        this_state_m = self.matrix_state
+        if type(other_state) == list:
+            other_state = State(kron([e.matrix_state for e in other_state]))
+        other_state_m = other_state.matrix_state
+        if this_state_m.shape[1] != other_state_m.shape[0]:
             # The two arrays are different sizes
             # Use the Kronicker product of Identity with the state.item where
             # state.item is the substate
-            larger_side = max(this_state.shape[1], this_state.shape[0])
-            _list = [this_state if i == state.item else _I for i in range(larger_side)]
-            this_state = kron(_list)
-        return State(this_state.dot(other_state))
+            larger_side = max(len(self), len(other_state))
+            _list = [this_state_m if i == other_state.item else _I for i in range(larger_side)]
+            this_state_m = kron(_list)
+        return State(this_state_m.dot(other_state_m))
 
 
 class Qubit(State): ...