diff --git a/lib_q_computer.py b/lib_q_computer.py index 0085040..fc6c2a6 100644 --- a/lib_q_computer.py +++ b/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): ...