quantum/lib_q_computer.py

import numpy as np

# Raw matrixes to be used for initialization of qubits and gates
__0 = [[1],
       [0]]

__1 = [[0],
       [1]]

_I = [
    [1, 0],
    [0, 1],
]

_X = [
    [0, 1],
    [1, 0],
]

_Y = [
    [0, complex(0, -1)],
    [complex(0, 1), 0],
]

_Z = [
    [1, 0],
    [0, -1],
]

_H = [
    [1 / np.sqrt(2), 1 / np.sqrt(2)],
    [1 / np.sqrt(2), -1 / np.sqrt(2)]
]

_CNOT = [
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 0, 1],
    [0, 0, 1, 0],
]


class State(object):
    """Represents a quantum state"""
    def __init__(self, matrix_state):
        """
        Can be initialized with a matrix, e.g. for |0> this is [[0],[1]]
        :param matrix_state: a matrix representing the quantum state
        """
        self.matrix_state = np.array(matrix_state)

    def __getitem__(self, item):
        """
        Kind of hacky way to store a substate of an item for use in Gate operations
        so that one can use state[0] and the encompassing operator can access this.
        :param item:
        :return:
        """
        if item >= len(self):
            raise IndexError
        self.item = item
        return self

    def __len__(self):
        return int(np.log2(self.matrix_state.shape[0]))

    def __repr__(self):
        return str(self.matrix_state)


def kron(_list):
    """Calculates a Kronicker product of a list of matrices
    This is essentially a np.kron(*args) since np.kron takes (a,b)"""
    if type(_list) != list or len(_list) <= 1:
        return np.array([])
    rv = np.array(_list[0])
    for item in _list[1:]:
        rv = np.kron(rv, item)
    return rv


class Gate(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 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_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(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): ...


# List of gates
I = Gate(_I)
X = Gate(_X)
Y = Gate(_Y)
Z = Gate(_Z)
H = Gate(_H)

CNOT = Gate(_CNOT)