circuits

Quantum Circuits

Gate

Source code in jaxquantum/circuits/gates.py

@struct.dataclass
class Gate:
    dims: List[int] = struct.field(pytree_node=False)
    _U: Optional[Array] # Unitary
    _H: Optional[Array] # Hamiltonian
    _KM: Optional[Qarray] # Kraus map
    _params: Dict[str, Any]
    _ts: Array
    _name: str = struct.field(pytree_node=False)
    num_modes: int = struct.field(pytree_node=False)

    @classmethod
    def create(
        cls,
        dims: Union[int, List[int]],
        name: str = "Gate",
        params: Optional[Dict[str, Any]] = None,
        ts: Optional[Array] = None,
        gen_U: Optional[Callable[[Dict[str, Any]], Qarray]] = None,
        gen_H: Optional[Callable[[Dict[str, Any]], Qarray]] = None,
        gen_KM: Optional[Callable[[Dict[str, Any]], List[Qarray]]] = None,
        num_modes: int = 1,
    ):
        """ Create a gate. 

        Args:
            dims: Dimensions of the gate.
            name: Name of the gate.
            params: Parameters of the gate.
            ts: Times of the gate.
            gen_U: Function to generate the unitary of the gate.
            gen_H: Function to generate the Hamiltonian of the gate.
            gen_KM: Function to generate the Kraus map of the gate.
            num_modes: Number of modes of the gate.
        """

        # TODO: add params to device?

        if isinstance(dims, int):
            dims = [dims]

        assert len(dims) == num_modes, "Number of dimensions must match number of modes."


        # Unitary
        _U = gen_U(params) if gen_U is not None else None 
        _H = gen_H(params) if gen_H is not None else None 

        if gen_KM is not None:
            _KM = gen_KM(params)
        elif _U is not None:
            _KM = Qarray.from_list([_U])

        return Gate(
            dims = dims,
            _U = _U,
            _H = _H,
            _KM = _KM,
            _params = params if params is not None else {},
            _ts=ts if ts is not None else jnp.array([]),
            _name = name,
            num_modes = num_modes
        )

    def __str__(self):
        return self._name

    def __repr__(self):
        return self._name

    @property
    def U(self):
        return self._U

    @property
    def H(self):
        return self._H

    @property
    def KM(self):
        return self._KM

create(dims, name='Gate', params=None, ts=None, gen_U=None, gen_H=None, gen_KM=None, num_modes=1) classmethod

Create a gate.

Parameters:

Name	Type	Description	Default
dims	Union[int, List[int]]	Dimensions of the gate.	required
name	str	Name of the gate.	'Gate'
params	Optional[Dict[str, Any]]	Parameters of the gate.	None
ts	Optional[Array]	Times of the gate.	None
gen_U	Optional[Callable[[Dict[str, Any]], Qarray]]	Function to generate the unitary of the gate.	None
gen_H	Optional[Callable[[Dict[str, Any]], Qarray]]	Function to generate the Hamiltonian of the gate.	None
gen_KM	Optional[Callable[[Dict[str, Any]], List[Qarray]]]	Function to generate the Kraus map of the gate.	None
num_modes	int	Number of modes of the gate.	1

Source code in jaxquantum/circuits/gates.py

@classmethod
def create(
    cls,
    dims: Union[int, List[int]],
    name: str = "Gate",
    params: Optional[Dict[str, Any]] = None,
    ts: Optional[Array] = None,
    gen_U: Optional[Callable[[Dict[str, Any]], Qarray]] = None,
    gen_H: Optional[Callable[[Dict[str, Any]], Qarray]] = None,
    gen_KM: Optional[Callable[[Dict[str, Any]], List[Qarray]]] = None,
    num_modes: int = 1,
):
    """ Create a gate. 

    Args:
        dims: Dimensions of the gate.
        name: Name of the gate.
        params: Parameters of the gate.
        ts: Times of the gate.
        gen_U: Function to generate the unitary of the gate.
        gen_H: Function to generate the Hamiltonian of the gate.
        gen_KM: Function to generate the Kraus map of the gate.
        num_modes: Number of modes of the gate.
    """

    # TODO: add params to device?

    if isinstance(dims, int):
        dims = [dims]

    assert len(dims) == num_modes, "Number of dimensions must match number of modes."


    # Unitary
    _U = gen_U(params) if gen_U is not None else None 
    _H = gen_H(params) if gen_H is not None else None 

    if gen_KM is not None:
        _KM = gen_KM(params)
    elif _U is not None:
        _KM = Qarray.from_list([_U])

    return Gate(
        dims = dims,
        _U = _U,
        _H = _H,
        _KM = _KM,
        _params = params if params is not None else {},
        _ts=ts if ts is not None else jnp.array([]),
        _name = name,
        num_modes = num_modes
    )

Qarray

Source code in jaxquantum/core/qarray.py

@struct.dataclass # this allows us to send in and return Qarray from jitted functions
class Qarray:
    _data: Array
    _qdims: Qdims = struct.field(pytree_node=False)
    _bdims: Tuple[int] = struct.field(pytree_node=False)


    # Initialization ----
    @classmethod
    def create(cls, data, dims=None, bdims=None):

        # Step 1: Prepare data ----
        data = jnp.asarray(data)

        if len(data.shape) == 1 and data.shape[0] > 0:
            data = data.reshape(data.shape[0], 1)

        if len(data.shape) >= 2:
            if data.shape[-2] != data.shape[-1] and not (data.shape[-2] == 1 or data.shape[-1] == 1):
                data = data.reshape(*data.shape[:-1], data.shape[-1], 1)

        if bdims is not None:
            if len(data.shape) - len(bdims) == 1:
                data = data.reshape(*data.shape[:-1], data.shape[-1], 1)
        # ----

        # Step 2: Prepare dimensions ----
        if bdims is None:
            bdims = tuple(data.shape[:-2])

        if dims is None:
            dims = ((data.shape[-2],), (data.shape[-1],))

        dims = (tuple(dims[0]), tuple(dims[1]))

        check_dims(dims, bdims, data.shape)

        qdims = Qdims(dims)

        # NOTE: Constantly tidying up on Qarray creation might be a bit overkill.
        # It increases the compilation time, but only very slightly 
        # increased the runtime of the jit compiled function.
        # We could instead use this tidy_up where we think we need it.
        data = tidy_up(data, SETTINGS["auto_tidyup_atol"])

        return cls(data, qdims, bdims)

    # ----

    @classmethod
    def from_list(cls, qarr_list: List[Qarray]) -> Qarray:
        """ Create a Qarray from a list of Qarrays. """

        data = jnp.array([qarr.data for qarr in qarr_list])


        if len(qarr_list) == 0:
            dims = ((),())
            bdims = ()
        else:
            dims = qarr_list[0].dims
            bdims = qarr_list[0].bdims

        if not all(qarr.dims == dims and qarr.bdims == bdims for qarr in qarr_list):
            raise ValueError("All Qarrays in the list must have the same dimensions.")

        bdims = (len(qarr_list),) + bdims

        return cls.create(data, dims=dims, bdims=bdims)

    @classmethod
    def from_array(cls, qarr_arr) -> Qarray:
        """ Create a Qarray from a nested list of Qarrays. 

        Args:
            qarr_arr (list): nested list of Qarrays

        Returns:
            Qarray: Qarray object
        """
        if isinstance(qarr_arr, Qarray):
            return qarr_arr

        bdims = ()
        lvl = qarr_arr
        while not isinstance(lvl, Qarray):
            bdims = bdims + (len(lvl),)
            if len(lvl) > 0:
                lvl = lvl[0]
            else:
                break

        depth = len(bdims)

        def flat(lis):
            flatList = []
            for element in lis:
                if type(element) is list:
                    flatList += flat(element)
                else:
                    flatList.append(element)
            return flatList

        qarr_list = flat(qarr_arr)
        qarr = cls.from_list(qarr_list)
        qarr = qarr.reshape_bdims(*bdims)
        return qarr


    # Properties ----
    @property
    def qtype(self):
        return self._qdims.qtype

    @property
    def dtype(self):
        return self._data.dtype

    @property
    def dims(self):
        return self._qdims.dims

    @property
    def bdims(self):
        return self._bdims

    @property
    def qdims(self):
        return self._qdims

    @property
    def space_dims(self):
        if self.qtype in [Qtypes.oper, Qtypes.ket]:
            return self.dims[0]
        elif self.qtype == Qtypes.bra:
            return self.dims[1]
        else:
            # TODO: not reached for some reason
            raise ValueError("Unsupported qtype.")

    @property
    def data(self):
        return self._data

    @property
    def shaped_data(self):
        return self._data.reshape(self.bdims + self.dims[0] + self.dims[1])

    @property 
    def shape(self):
        return self.data.shape

    @property
    def is_batched(self):
        return len(self.bdims) > 0

    def __getitem__(self, index):
        if len(self.bdims) > 0:
            return Qarray.create(
                self.data[index],
                dims=self.dims,
            )
        else:
            raise ValueError("Cannot index a non-batched Qarray.")

    def reshape_bdims(self, *args):
        """ Reshape the batch dimensions of the Qarray. """
        new_bdims = tuple(args)

        if prod(new_bdims) == 0:
            new_shape = new_bdims 
        else:
            new_shape = new_bdims + self.dims[0] + self.dims[1]
        return Qarray.create(
            self.data.reshape(new_shape),
            dims=self.dims,
            bdims=new_bdims,
        )

    def space_to_qdims(self, space_dims: List[int]):

        if isinstance(space_dims[0], (list, tuple)):
            return space_dims

        if self.qtype in [Qtypes.oper, Qtypes.ket]:
            return (tuple(space_dims), tuple([1 for _ in range(len(space_dims))]))
        elif self.qtype == Qtypes.bra:
            return (tuple([1 for _ in range(len(space_dims))]), tuple(space_dims))
        else:
            raise ValueError("Unsupported qtype for space_to_qdims conversion.")

    def reshape_qdims(self, *args):
        """ Reshape the quantum dimensions of the Qarray. 

        Note that this does not take in qdims but rather the new Hilbert space dimensions.

        Args:
            *args: new Hilbert dimensions for the Qarray.

        Returns:
            Qarray: reshaped Qarray.
        """

        new_space_dims = tuple(args)
        current_space_dims = self.space_dims
        assert prod(new_space_dims) == prod(current_space_dims)


        new_qdims = self.space_to_qdims(new_space_dims)
        new_bdims = self.bdims

        return Qarray.create(
            self.data,
            dims=new_qdims,
            bdims=new_bdims
        )

    def resize(self, new_shape):
        """ Resize the Qarray to a new shape. 

        TODO: review and maybe deprecate this method.
        """
        dims = self.dims
        data = jnp.resize(self.data, new_shape)
        return Qarray.create(
            data,
            dims=dims,
        )

    def __len__(self):
        """ Length of the Qarray. """
        if len(self.bdims) > 0:
            return self.data.shape[0]
        else:
            raise ValueError("Cannot get length of a non-batched Qarray.")


    def __eq__(self, other):
        if not isinstance(other, Qarray):
            raise ValueError("Cannot calculate equality of a Qarray with a non-Qarray.")

        if self.dims != other.dims:
            return False

        if self.bdims != other.bdims:
            return False

        return jnp.all(self.data == other.data)

    def __ne__(self, other):
        return not self.__eq__(other)
    # ----


    # Elementary Math ----
    def __matmul__(self, other):
        if not isinstance(other, Qarray):
            return NotImplemented
        _qdims_new = self._qdims @ other._qdims
        return Qarray.create(
            self.data @ other.data,
            dims=_qdims_new.dims,
        )

    # NOTE: not possible to reach this.
    # def __rmatmul__(self, other):        
    #     if not isinstance(other, Qarray):
    #         return NotImplemented

    #     _qdims_new = other._qdims @ self._qdims
    #     return Qarray.create(
    #         other.data @ self.data,
    #         dims=_qdims_new.dims,
    #     )


    def __mul__(self, other):
        if isinstance(other, Qarray):
            return self.__matmul__(other)


        other = other + 0.0j
        if not robust_isscalar(other) and len(other.shape) > 0: # not a scalar
            other = other.reshape(other.shape + (1,1))

        return Qarray.create(
            other * self.data,
            dims=self._qdims.dims,
        )

    def __rmul__(self, other):

        # NOTE: not possible to reach this.
        # if isinstance(other, Qarray):
        #     return self.__rmatmul__(other)

        return self.__mul__(other)

    def __neg__(self):
        return self.__mul__(-1)

    def __truediv__(self, other):
        """ For Qarray's, this only really makes sense in the context of division by a scalar. """

        if isinstance(other, Qarray):
            raise ValueError("Cannot divide a Qarray by another Qarray.")

        return self.__mul__(1/other)

    def __add__(self, other):
        if isinstance(other, Qarray):
            if self.dims != other.dims:
                msg = (
                    "Dimensions are incompatible: "
                    + repr(self.dims) + " and " + repr(other.dims)
                )
                raise ValueError(msg)
            return Qarray.create(self.data + other.data, dims=self.dims)

        if robust_isscalar(other) and other == 0:
            return self.copy()

        if self.data.shape[-2] == self.data.shape[-1]:
            other = other + 0.0j
            if not robust_isscalar(other) and len(other.shape) > 0: # not a scalar
                other = other.reshape(other.shape + (1,1))
            other = Qarray.create(other * jnp.eye(self.data.shape[-2], dtype=self.data.dtype), dims=self.dims)
            return self.__add__(other)

        return NotImplemented

    def __radd__(self, other):
        return self.__add__(other)

    def __sub__(self, other):    
        if isinstance(other, Qarray):
            if self.dims != other.dims:
                msg = (
                    "Dimensions are incompatible: "
                    + repr(self.dims) + " and " + repr(other.dims)
                )
                raise ValueError(msg)
            return Qarray.create(self.data - other.data, dims=self.dims)

        if robust_isscalar(other) and other == 0:
            return self.copy()

        if self.data.shape[-2] == self.data.shape[-1]:
            other = other + 0.0j
            if not robust_isscalar(other) and len(other.shape) > 0: # not a scalar
                other = other.reshape(other.shape + (1,1))
            other = Qarray.create(other * jnp.eye(self.data.shape[-2], dtype=self.data.dtype), dims=self.dims)
            return self.__sub__(other)

        return NotImplemented

    def __rsub__(self, other):
        return self.__neg__().__add__(other)

    def __xor__(self, other):
        if not isinstance(other, Qarray):
            return NotImplemented
        return tensor(self, other)

    def __rxor__(self, other):
        if not isinstance(other, Qarray):
            return NotImplemented
        return tensor(other, self)

    def __pow__(self, other):
        if not isinstance(other, int):
            return NotImplemented

        return powm(self, other)

    # ----

    # String Representation ----
    def _str_header(self):
        out = ", ".join([
            "Quantum array: dims = " + str(self.dims),
            "bdims = " + str(self.bdims),
            "shape = " + str(self._data.shape),
            "type = " + str(self.qtype),
        ])
        return out

    def __str__(self):
        return self._str_header() + "\nQarray data =\n" + str(self.data)

    @property
    def header(self):
        """ Print the header of the Qarray. """
        return self._str_header()

    def __repr__(self):
        return self.__str__()

    # ----

    # Utilities ----
    def copy(self, memo=None):
        # return Qarray.create(deepcopy(self.data), dims=self.dims)
        return self.__deepcopy__(memo)

    def __deepcopy__(self, memo):
        """ Need to override this when defininig __getattr__. """

        return Qarray(
            _data = deepcopy(self._data, memo=memo),
            _qdims = deepcopy(self._qdims, memo=memo),
            _bdims = deepcopy(self._bdims, memo=memo)
        )

    def __getattr__(self, method_name):

        if "__" == method_name[:2]:
            # NOTE: we return NotImplemented for binary special methods logic in python, plus things like __jax_array__
            return lambda *args, **kwargs: NotImplemented

        modules = [jnp, jnp.linalg, jsp, jsp.linalg]

        method_f = None 
        for mod in modules:
            method_f = getattr(mod, method_name, None)
            if method_f is not None:
                break

        if method_f is None:
            raise NotImplementedError(f"Method {method_name} does not exist. No backup method found in {modules}.")

        def func(*args, **kwargs): 
            res = method_f(self.data, *args, **kwargs)

            if getattr(res, "shape", None) is None or res.shape != self.data.shape:
                return res
            else:
                return Qarray.create(
                    res,
                    dims=self._qdims.dims
                )
        return func

    # ----

    # Conversions / Reshaping ----
    def dag(self):
        return dag(self)

    def to_dm(self):
        return ket2dm(self)

    def is_dm(self):
        return self.qtype == Qtypes.oper

    def to_ket(self):
        return to_ket(self)

    def transpose(self, *args):
        return transpose(self, *args)

    def keep_only_diag_elements(self):
        return keep_only_diag_elements(self)

    # ----

    # Math Functions ----
    def unit(self):
        return unit(self)

    def norm(self):
        return norm(self)

    def expm(self):
        return expm(self)

    def powm(self, n):
        return powm(self, n)

    def cosm(self):
        return cosm(self)

    def sinm(self):
        return sinm(self)

    def tr(self, **kwargs):
        return tr(self, **kwargs)

    def trace(self, **kwargs):
        return tr(self, **kwargs)

    def ptrace(self, indx):
        return ptrace(self, indx)

    def eigenstates(self):
        return eigenstates(self)

    def eigenenergies(self):
        return eigenenergies(self)

    def collapse(self, mode="sum"):
        return collapse(self, mode=mode)    

header property

Print the header of the Qarray.

__deepcopy__(memo)

Need to override this when defininig getattr.

Source code in jaxquantum/core/qarray.py

def __deepcopy__(self, memo):
    """ Need to override this when defininig __getattr__. """

    return Qarray(
        _data = deepcopy(self._data, memo=memo),
        _qdims = deepcopy(self._qdims, memo=memo),
        _bdims = deepcopy(self._bdims, memo=memo)
    )

__len__()

Length of the Qarray.

Source code in jaxquantum/core/qarray.py

def __len__(self):
    """ Length of the Qarray. """
    if len(self.bdims) > 0:
        return self.data.shape[0]
    else:
        raise ValueError("Cannot get length of a non-batched Qarray.")

__truediv__(other)

For Qarray's, this only really makes sense in the context of division by a scalar.

Source code in jaxquantum/core/qarray.py

def __truediv__(self, other):
    """ For Qarray's, this only really makes sense in the context of division by a scalar. """

    if isinstance(other, Qarray):
        raise ValueError("Cannot divide a Qarray by another Qarray.")

    return self.__mul__(1/other)

from_array(qarr_arr) classmethod

Create a Qarray from a nested list of Qarrays.

Parameters:

Name	Type	Description	Default
qarr_arr	list	nested list of Qarrays	required

Returns:

Name	Type	Description
Qarray	Qarray	Qarray object

Source code in jaxquantum/core/qarray.py

@classmethod
def from_array(cls, qarr_arr) -> Qarray:
    """ Create a Qarray from a nested list of Qarrays. 

    Args:
        qarr_arr (list): nested list of Qarrays

    Returns:
        Qarray: Qarray object
    """
    if isinstance(qarr_arr, Qarray):
        return qarr_arr

    bdims = ()
    lvl = qarr_arr
    while not isinstance(lvl, Qarray):
        bdims = bdims + (len(lvl),)
        if len(lvl) > 0:
            lvl = lvl[0]
        else:
            break

    depth = len(bdims)

    def flat(lis):
        flatList = []
        for element in lis:
            if type(element) is list:
                flatList += flat(element)
            else:
                flatList.append(element)
        return flatList

    qarr_list = flat(qarr_arr)
    qarr = cls.from_list(qarr_list)
    qarr = qarr.reshape_bdims(*bdims)
    return qarr

from_list(qarr_list) classmethod

Create a Qarray from a list of Qarrays.

Source code in jaxquantum/core/qarray.py

@classmethod
def from_list(cls, qarr_list: List[Qarray]) -> Qarray:
    """ Create a Qarray from a list of Qarrays. """

    data = jnp.array([qarr.data for qarr in qarr_list])


    if len(qarr_list) == 0:
        dims = ((),())
        bdims = ()
    else:
        dims = qarr_list[0].dims
        bdims = qarr_list[0].bdims

    if not all(qarr.dims == dims and qarr.bdims == bdims for qarr in qarr_list):
        raise ValueError("All Qarrays in the list must have the same dimensions.")

    bdims = (len(qarr_list),) + bdims

    return cls.create(data, dims=dims, bdims=bdims)

reshape_bdims(*args)

Reshape the batch dimensions of the Qarray.

Source code in jaxquantum/core/qarray.py

def reshape_bdims(self, *args):
    """ Reshape the batch dimensions of the Qarray. """
    new_bdims = tuple(args)

    if prod(new_bdims) == 0:
        new_shape = new_bdims 
    else:
        new_shape = new_bdims + self.dims[0] + self.dims[1]
    return Qarray.create(
        self.data.reshape(new_shape),
        dims=self.dims,
        bdims=new_bdims,
    )

reshape_qdims(*args)

Reshape the quantum dimensions of the Qarray.

Note that this does not take in qdims but rather the new Hilbert space dimensions.

Parameters:

Name	Type	Description	Default
*args		new Hilbert dimensions for the Qarray.	()

Returns:

Name	Type	Description
Qarray		reshaped Qarray.

Source code in jaxquantum/core/qarray.py

def reshape_qdims(self, *args):
    """ Reshape the quantum dimensions of the Qarray. 

    Note that this does not take in qdims but rather the new Hilbert space dimensions.

    Args:
        *args: new Hilbert dimensions for the Qarray.

    Returns:
        Qarray: reshaped Qarray.
    """

    new_space_dims = tuple(args)
    current_space_dims = self.space_dims
    assert prod(new_space_dims) == prod(current_space_dims)


    new_qdims = self.space_to_qdims(new_space_dims)
    new_bdims = self.bdims

    return Qarray.create(
        self.data,
        dims=new_qdims,
        bdims=new_bdims
    )

resize(new_shape)

Resize the Qarray to a new shape.

TODO: review and maybe deprecate this method.

Source code in jaxquantum/core/qarray.py

def resize(self, new_shape):
    """ Resize the Qarray to a new shape. 

    TODO: review and maybe deprecate this method.
    """
    dims = self.dims
    data = jnp.resize(self.data, new_shape)
    return Qarray.create(
        data,
        dims=dims,
    )

basis(N, k)

Creates a |k> (i.e. fock state) ket in a specified Hilbert Space.

Parameters:

Name	Type	Description	Default
N	int	Hilbert space dimension	required
k	int	fock number	required

Returns:

Type	Description
	Fock State |k>

Source code in jaxquantum/core/operators.py

def basis(N: int, k: int):
    """Creates a |k> (i.e. fock state) ket in a specified Hilbert Space.

    Args:
        N: Hilbert space dimension
        k: fock number

    Returns:
        Fock State |k>
    """
    return Qarray.create(one_hot(k, N).reshape(N, 1))

destroy(N)

annihilation operator

Parameters:

Name	Type	Description	Default
N		Hilbert space size	required

Returns:

Type	Description
Qarray	annilation operator in Hilber Space of size N

Source code in jaxquantum/core/operators.py

def destroy(N) -> Qarray:
    """annihilation operator

    Args:
        N: Hilbert space size

    Returns:
        annilation operator in Hilber Space of size N
    """
    return Qarray.create(jnp.diag(jnp.sqrt(jnp.arange(1, N)), k=1))

displace(N, α)

Displacement operator

Parameters:

Name	Type	Description	Default
N		Hilbert Space Size	required
α		Phase space displacement	required

Returns:

Type	Description
Qarray	Displace operator D(α)

Source code in jaxquantum/core/operators.py

def displace(N, α) -> Qarray:
    """Displacement operator

    Args:
        N: Hilbert Space Size
        α: Phase space displacement

    Returns:
        Displace operator D(α)
    """
    a = destroy(N)
    return (α * a.dag() - jnp.conj(α) * a).expm()

hadamard()

H

Returns:

Name	Type	Description
H	Qarray	Hadamard gate

Source code in jaxquantum/core/operators.py

def hadamard() -> Qarray:
    """H

    Returns:
        H: Hadamard gate
    """
    return Qarray.create(jnp.array([[1, 1], [1, -1]]) / jnp.sqrt(2))

identity(*args, **kwargs)

Identity matrix.

Returns:

Type	Description
Qarray	Identity matrix.

Source code in jaxquantum/core/operators.py

def identity(*args, **kwargs) -> Qarray:
    """Identity matrix.

    Returns:
        Identity matrix.
    """
    return Qarray.create(jnp.eye(*args, **kwargs))

ket2dm(qarr)

Turns ket into density matrix. Does nothing if already operator.

Parameters:

Name	Type	Description	Default
qarr	Qarray	qarr	required

Returns:

Type	Description
Qarray	Density matrix

Source code in jaxquantum/core/qarray.py

def ket2dm(qarr: Qarray) -> Qarray:
    """Turns ket into density matrix.
    Does nothing if already operator.

    Args:
        qarr (Qarray): qarr

    Returns:
        Density matrix
    """

    if qarr.qtype == Qtypes.oper:
        return qarr

    if qarr.qtype == Qtypes.bra:
        qarr = qarr.dag()

    return qarr @ qarr.dag()

num(N)

Number operator

Parameters:

Name	Type	Description	Default
N		Hilbert Space size	required

Returns:

Type	Description
Qarray	number operator in Hilber Space of size N

Source code in jaxquantum/core/operators.py

def num(N) -> Qarray:
    """Number operator

    Args:
        N: Hilbert Space size

    Returns:
        number operator in Hilber Space of size N
    """
    return Qarray.create(jnp.diag(jnp.arange(N)))

qubit_rotation(theta, nx, ny, nz)

Single qubit rotation.

Parameters:

Name	Type	Description	Default
theta	float	rotation angle.	required
nx		rotation axis x component.	required
ny		rotation axis y component.	required
nz		rotation axis z component.	required

Returns:

Type	Description
Qarray	Single qubit rotation operator.

Source code in jaxquantum/core/operators.py

def qubit_rotation(theta: float, nx, ny, nz) -> Qarray:
    """Single qubit rotation.

    Args:
        theta: rotation angle.
        nx: rotation axis x component.
        ny: rotation axis y component.
        nz: rotation axis z component.

    Returns:
        Single qubit rotation operator.
    """
    return jnp.cos(theta / 2) * identity(2) - 1j * jnp.sin(theta / 2) * (nx * sigmax() + ny * sigmay() + nz * sigmaz())

sigmax()

σx

Returns:

Type	Description
Qarray	σx Pauli Operator

Source code in jaxquantum/core/operators.py

def sigmax() -> Qarray:
    """σx

    Returns:
        σx Pauli Operator
    """
    return Qarray.create(jnp.array([[0.0, 1.0], [1.0, 0.0]]))

sigmay()

σy

Returns:

Type	Description
Qarray	σy Pauli Operator

Source code in jaxquantum/core/operators.py

def sigmay() -> Qarray:
    """σy

    Returns:
        σy Pauli Operator
    """
    return Qarray.create(jnp.array([[0.0, -1.0j], [1.0j, 0.0]]))

sigmaz()

σz

Returns:

Type	Description
Qarray	σz Pauli Operator

Source code in jaxquantum/core/operators.py

def sigmaz() -> Qarray:
    """σz

    Returns:
        σz Pauli Operator
    """
    return Qarray.create(jnp.array([[1.0, 0.0], [0.0, -1.0]]))

simulate(circuit, initial_state, mode=SimulateMode.UNITARY)

Simulates the evolution of a quantum state through a given quantum circuit.

Parameters:

Name	Type	Description	Default
circuit	Circuit	The quantum circuit to simulate. The circuit is composed of layers, each of which can generate unitary or Kraus operators.	required
initial_state	Qarray	The initial quantum state to be evolved. This can be a state vector or a density matrix.	required
mode	SimulateMode	The mode of simulation. It can be either SimulateMode.UNITARY for unitary evolution or SimulateMode.KRAUS for Kraus operator evolution. Defaults to SimulateMode.UNITARY.	UNITARY

Returns:

Name	Type	Description
Results	Results	An object containing the results of the simulation, which includes the quantum states at each step of the circuit.

Source code in jaxquantum/circuits/simulate.py

def simulate(
    circuit: Circuit,
    initial_state: Qarray,
    mode: SimulateMode = SimulateMode.UNITARY
) -> Results:
    """
    Simulates the evolution of a quantum state through a given quantum circuit.

    Args:
        circuit (Circuit): The quantum circuit to simulate. The circuit is composed of layers, 
                           each of which can generate unitary or Kraus operators.
        initial_state (Qarray): The initial quantum state to be evolved. This can be a state vector 
                                or a density matrix.
        mode (SimulateMode, optional): The mode of simulation. It can be either SimulateMode.UNITARY 
                                       for unitary evolution or SimulateMode.KRAUS for Kraus operator 
                                       evolution. Defaults to SimulateMode.UNITARY.

    Returns:
        Results: An object containing the results of the simulation, which includes the quantum states 
                 at each step of the circuit.
    """

    results = Results.create([])
    state = initial_state
    results.append(Qarray.from_list([state]))

    for layer in circuit.layers:
        result = simulate_layer(layer, state, mode=mode)
        results.append(result)
        state = result[-1]

    return results

simulate_layer(layer, initial_state, mode=SimulateMode.UNITARY)

Simulates the evolution of a quantum state through a given layer.

Parameters:

Name	Type	Description	Default
layer	Layer	The layer through which the quantum state evolves. This layer should have methods to generate unitary (gen_U) and Kraus (gen_KM) operators.	required
initial_state	Qarray	The initial quantum state to be evolved. This can be a state vector or a density matrix.	required
mode	SimulateMode	The mode of simulation. It can be either SimulateMode.UNITARY for unitary evolution or SimulateMode.KRAUS for Kraus operator evolution or SimulateMode.DEFAULT to use the default simulate mode in the layer. Defaults to SimulateMode.UNITARY.	UNITARY

Returns: Qarray: The result of the simulation containing the evolved quantum state.

Source code in jaxquantum/circuits/simulate.py

def simulate_layer(layer: Layer, initial_state: Qarray, mode: SimulateMode = SimulateMode.UNITARY) -> Qarray:
    """
    Simulates the evolution of a quantum state through a given layer.

    Args:
        layer (Layer): The layer through which the quantum state evolves. 
                       This layer should have methods to generate unitary (gen_U) 
                       and Kraus (gen_KM) operators.
        initial_state (Qarray): The initial quantum state to be evolved. 
                                This can be a state vector or a density matrix.
        mode (SimulateMode, optional): The mode of simulation. It can be either 
                                       SimulateMode.UNITARY for unitary evolution 
                                       or SimulateMode.KRAUS for Kraus operator evolution
                                       or SimulateMode.DEFAULT to use the default simulate mode in the layer.
                                       Defaults to SimulateMode.UNITARY.
    Returns:
        Qarray: The result of the simulation containing the evolved quantum state.
    """

    state = initial_state 

    if mode == SimulateMode.DEFAULT:
        mode = layer._default_simulate_mode

    if mode == SimulateMode.UNITARY:
        U = layer.gen_U()
        if state.is_dm():
            state = U @ state @ U.dag()
        else:
            state = U @ state 

        result = Qarray.from_list([state])

    elif mode == SimulateMode.KRAUS:
        KM = layer.gen_KM()

        state = ket2dm(state)
        state = (KM @ state @ KM.dag()).collapse()

        # new_state = 0
        # for op_j in range(len(KM)):
        #     op = KM[op_j]
        #     new_state += op @ state @ op.dag()
        # state = new_state

        result = Qarray.from_list([state])

    return result