jaxquantum

jaxquantum

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

        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 + (prod(self.dims[0]),) + (-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 is_vec(self):
        return self.qtype == Qtypes.ket or self.qtype == Qtypes.bra

    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

    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 + (prod(self.dims[0]),) + (-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))

basis_like(A, ks)

Creates a |k> (i.e. fock state) ket with the same space dims as A.

Parameters:

Name	Type	Description	Default
A	Qarray	state or operator.	required
k		fock number.	required

Returns:

Type	Description
Qarray	Fock State |k> with the same space dims as A.

Source code in jaxquantum/core/operators.py

def basis_like(A: Qarray, ks: List[int]) -> Qarray:
    """Creates a |k> (i.e. fock state) ket with the same space dims as A.

    Args:
        A: state or operator.
        k: fock number.

    Returns:
        Fock State |k> with the same space dims as A.
    """
    space_dims = A.space_dims
    assert len(space_dims) == len(ks), "len(ks) must be equal to len(space_dims)"

    kets = []
    for j, k in enumerate(ks):
        kets.append(basis(space_dims[j], k))
    return tensor(*kets)

coherent(N, α)

Coherent state.

Parameters:

Name	Type	Description	Default
N	int	Hilbert Space Size.	required
α	complex	coherent state amplitude.	required

Return

Coherent state |α⟩.

Source code in jaxquantum/core/operators.py

def coherent(N: int, α: complex) -> Qarray:
    """Coherent state.

    Args:
        N: Hilbert Space Size.
        α: coherent state amplitude.

    Return:
        Coherent state |α⟩.
    """
    return displace(N, α) @ basis(N, 0)

collapse(qarr, mode='sum')

Collapse the Qarray.

Parameters:

Name	Type	Description	Default
qarr	Qarray	quantum array array	required

Returns:

Type	Description
Qarray	Collapsed quantum array

Source code in jaxquantum/core/qarray.py

def collapse(qarr: Qarray, mode="sum") -> Qarray:
    """Collapse the Qarray.

    Args:
        qarr (Qarray): quantum array array

    Returns:
        Collapsed quantum array
    """
    if mode == "sum":
        if len(qarr.bdims) == 0:
            return qarr

        batch_axes = list(range(len(qarr.bdims)))
        return Qarray.create(jnp.sum(qarr.data, axis=batch_axes), dims=qarr.dims)

comb(N, k)

NCk

TODO: replace with jsp.special.comb once issue is closed:

https://github.com/google/jax/issues/9709

Parameters:

Name	Type	Description	Default
N		total items	required
k		of items to choose	required

Returns:

Name	Type	Description
NCk		N choose k

Source code in jaxquantum/utils/utils.py

def comb(N, k):
    """
    NCk

    #TODO: replace with jsp.special.comb once issue is closed:
    https://github.com/google/jax/issues/9709

    Args:
        N: total items
        k: # of items to choose

    Returns:
        NCk: N choose k
    """
    one = 1
    N_plus_1 = lax.add(N, one)
    k_plus_1 = lax.add(k, one)
    return lax.exp(
        lax.sub(
            gammaln(N_plus_1), lax.add(gammaln(k_plus_1), gammaln(lax.sub(N_plus_1, k)))
        )
    )

cosm(qarr)

Matrix cosine wrapper.

Parameters:

Name	Type	Description	Default
qarr	Qarray	quantum array	required

Returns:

Type	Description
Qarray	matrix cosine

Source code in jaxquantum/core/qarray.py

def cosm(qarr: Qarray) -> Qarray:
    """Matrix cosine wrapper.

    Args:
        qarr (Qarray): quantum array

    Returns:
        matrix cosine
    """
    dims = qarr.dims
    data = cosm_data(qarr.data)
    return Qarray.create(data, dims=dims)

cosm_data(data, **kwargs)

Matrix cosine wrapper.

Returns:

Type	Description
Array	matrix cosine

Source code in jaxquantum/core/qarray.py

def cosm_data(data: Array, **kwargs) -> Array:
    """Matrix cosine wrapper.

    Returns:
        matrix cosine
    """
    return (expm_data(1j * data) + expm_data(-1j * data)) / 2

create(N)

creation operator

Parameters:

Name	Type	Description	Default
N		Hilbert space size	required

Returns:

Type	Description
Qarray	creation operator in Hilber Space of size N

Source code in jaxquantum/core/operators.py

def create(N) -> Qarray:
    """creation operator

    Args:
        N: Hilbert space size

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

dag(qarr)

Conjugate transpose.

Parameters:

Name	Type	Description	Default
qarr	Qarray	quantum array	required

Returns:

Type	Description
Qarray	conjugate transpose of qarr

Source code in jaxquantum/core/qarray.py

def dag(qarr: Qarray) -> Qarray:
    """Conjugate transpose.

    Args:
        qarr (Qarray): quantum array

    Returns:
        conjugate transpose of qarr
    """
    dims = qarr.dims[::-1]

    data = dag_data(qarr.data)

    return Qarray.create(data, dims=dims)

dag_data(arr)

Conjugate transpose.

Parameters:

Name	Type	Description	Default
arr	Array	operator	required

Returns:

Type	Description
Array	conjugate of op, and transposes last two axes

Source code in jaxquantum/core/qarray.py

def dag_data(arr: Array) -> Array:
    """Conjugate transpose.

    Args:
        arr: operator

    Returns:
        conjugate of op, and transposes last two axes
    """
    # TODO: revisit this case...
    if len(arr.shape) == 1:
        return jnp.conj(arr)

    return jnp.moveaxis(
        jnp.conj(arr), -1, -2
    )  # transposes last two axes, good for batching

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()

eigenenergies(qarr)

Eigenvalues of a quantum array.

Parameters:

Name	Type	Description	Default
qarr	Qarray	quantum array	required

Returns:

Type	Description
Array	eigenvalues

Source code in jaxquantum/core/qarray.py

def eigenenergies(qarr: Qarray) -> Array:
    """Eigenvalues of a quantum array.

    Args:
        qarr (Qarray): quantum array

    Returns:
        eigenvalues
    """

    evals = jnp.linalg.eigvalsh(qarr.data)
    return evals

eigenstates(qarr)

Eigenstates of a quantum array.

Parameters:

Name	Type	Description	Default
qarr	Qarray	quantum array	required

Returns:

Type	Description
Qarray	eigenvalues and eigenstates

Source code in jaxquantum/core/qarray.py

def eigenstates(qarr: Qarray) -> Qarray:
    """Eigenstates of a quantum array.

    Args:
        qarr (Qarray): quantum array

    Returns:
        eigenvalues and eigenstates
    """

    evals, evecs = jnp.linalg.eigh(qarr.data)
    idxs_sorted = jnp.argsort(evals, axis=-1)

    dims = ket_from_op_dims(qarr.dims)

    evals = jnp.take_along_axis(evals, idxs_sorted, axis=-1)
    evecs = jnp.take_along_axis(evecs, idxs_sorted[..., None, :], axis=-1)

    evecs = Qarray.create(
        evecs,
        dims=dims,
        bdims=evecs.shape[:-1],
    )

    return evals, evecs

expm(qarr, **kwargs)

Matrix exponential wrapper.

Returns:

Type	Description
Qarray	matrix exponential

Source code in jaxquantum/core/qarray.py

def expm(qarr: Qarray, **kwargs) -> Qarray:
    """Matrix exponential wrapper.

    Returns:
        matrix exponential
    """
    dims = qarr.dims
    data = expm_data(qarr.data, **kwargs)
    return Qarray.create(data, dims=dims)

expm_data(data, **kwargs)

Matrix exponential wrapper.

Returns:

Type	Description
Array	matrix exponential

Source code in jaxquantum/core/qarray.py

def expm_data(data: Array, **kwargs) -> Array:
    """Matrix exponential wrapper.

    Returns:
        matrix exponential
    """
    return jsp.linalg.expm(data, **kwargs)

extract_dims(arr, dims=None)

Extract dims from a JAX array or Qarray.

Parameters:

Name	Type	Description	Default
arr	Array	JAX array or Qarray.	required
dims	Optional[Union[DIMS_TYPE, List[int]]]	Qarray dims.	None

Returns:

Type	Description
	Qarray dims.

Source code in jaxquantum/core/conversions.py

def extract_dims(arr: Array, dims: Optional[Union[DIMS_TYPE, List[int]]] = None):
    """Extract dims from a JAX array or Qarray.

    Args:
        arr: JAX array or Qarray.
        dims: Qarray dims.

    Returns:
        Qarray dims.
    """
    if isinstance(dims[0], Number):
        is_op = arr.shape[-2] == arr.shape[-1]
        if is_op:
            dims = [dims, dims]
        else:
            dims = [dims, [1] * len(dims)]  # defaults to ket
    return dims

fidelity(rho, sigma)

Fidelity between two states.

Parameters:

Name	Type	Description	Default
rho	Qarray	state.	required
sigma	Qarray	state.	required

Returns:

Type	Description
float	Fidelity between rho and sigma.

Source code in jaxquantum/core/measurements.py

def fidelity(rho: Qarray, sigma: Qarray) -> float:
    """Fidelity between two states.

    Args:
        rho: state.
        sigma: state.

    Returns:
        Fidelity between rho and sigma.
    """
    rho = rho.to_dm()
    sigma = sigma.to_dm()

    sqrt_rho = powm(rho, 0.5)

    return ((powm(sqrt_rho @ sigma @ sqrt_rho, 0.5)).tr()) ** 2

get_logical_basis(n_qubits)

Compute a complete operator basis of a system composed of logical qubits.

        Args:
            n_qubits: number of qubits

        Returns:
            List containing the complete operator basis.

Source code in jaxquantum/core/measurements.py

def get_logical_basis(n_qubits: int) -> Qarray:
    """Compute a complete operator basis of a system composed of logical
    qubits.

                Args:
                    n_qubits: number of qubits

                Returns:
                    List containing the complete operator basis.
    """
    if n_qubits < 1:
        raise ValueError("n_qubits must be at least 1.")

    n_qubits -= 1

    operators = Qarray.from_list(
        [identity(2) / 2, sigmax() / 2, sigmay() / 2, sigmaz() / 2]
    )

    if n_qubits == 0:
        return operators

    sub_basis = get_logical_basis(n_qubits)
    basis = []

    sub_basis_size = sub_basis.bdims[-1]

    for i in range(4):
        for j in range(sub_basis_size):
            basis.append(operators[i] ^ sub_basis[j])

    return Qarray.from_list(basis)

get_physical_basis(qubits)

Compute a complete operator basis of a QEC code on a physical system specified by a number of qubits.

    Args:
        qubits: list of qubit codes, must have
        common_gates and params attributes.

    Returns:
        List containing the complete operator basis.

Source code in jaxquantum/core/measurements.py

def get_physical_basis(qubits: List) -> Qarray:
    """Compute a complete operator basis of a QEC code on a
    physical system specified by a number of qubits.

            Args:
                qubits: list of qubit codes, must have
                common_gates and params attributes.

            Returns:
                List containing the complete operator basis.
    """

    qubit = qubits[0]
    qubits = qubits[1:]
    try:
        operators = Qarray.from_list(
            [
                identity(qubit.params["N"]),
                qubit.common_gates["X"],
                qubit.common_gates["Y"],
                qubit.common_gates["Z"],
            ]
        )
    except KeyError:
        print("QEC code must have common_gates for all three axes.")
    except AttributeError:
        print("QEC code must have common_gates and params attribute.")

    if len(qubits) == 0:
        return operators

    sub_basis = get_physical_basis(qubits)
    basis = []

    sub_basis_size = sub_basis.bdims[-1]

    for i in range(4):
        for j in range(sub_basis_size):
            basis.append(operators[i] ^ sub_basis[j])

    return Qarray.from_list(basis)

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))

identity_like(A)

Identity matrix with the same shape as A.

Parameters:

Name	Type	Description	Default
A		Matrix.	required

Returns:

Type	Description
Qarray	Identity matrix with the same shape as A.

Source code in jaxquantum/core/operators.py

def identity_like(A) -> Qarray:
    """Identity matrix with the same shape as A.

    Args:
        A: Matrix.

    Returns:
        Identity matrix with the same shape as A.
    """
    space_dims = A.space_dims
    total_dim = prod(space_dims)
    return Qarray.create(jnp.eye(total_dim, total_dim), dims=[space_dims, space_dims])

is_dm_data(data)

Check if data is a density matrix.

Parameters:

Name	Type	Description	Default
data	Array	matrix	required

Returns: True if data is a density matrix

Source code in jaxquantum/core/qarray.py

def is_dm_data(data: Array) -> bool:
    """Check if data is a density matrix.

    Args:
        data: matrix
    Returns:
        True if data is a density matrix
    """
    return data.shape[-2] == data.shape[-1]

jnp2jqt(arr, dims=None)

JAX array -> QuTiP state.

Parameters:

Name	Type	Description	Default
jnp_obj		JAX array.	required
dims	Optional[Union[DIMS_TYPE, List[int]]]	Qarray dims.	None

Returns:

Type	Description
	QuTiP state.

Source code in jaxquantum/core/conversions.py

def jnp2jqt(arr: Array, dims: Optional[Union[DIMS_TYPE, List[int]]] = None):
    """JAX array -> QuTiP state.

    Args:
        jnp_obj: JAX array.
        dims: Qarray dims.

    Returns:
        QuTiP state.
    """
    dims = extract_dims(arr, dims) if dims is not None else None
    return Qarray.create(arr, dims=dims)

jqt2qt(jqt_obj)

Qarray -> QuTiP state.

Parameters:

Name	Type	Description	Default
jqt_obj		Qarray.	required
dims		QuTiP dims.	required

Returns:

Type	Description
	QuTiP state.

Source code in jaxquantum/core/conversions.py

def jqt2qt(jqt_obj):
    """Qarray -> QuTiP state.

    Args:
        jqt_obj: Qarray.
        dims: QuTiP dims.

    Returns:
        QuTiP state.
    """
    if isinstance(jqt_obj, Qobj) or jqt_obj is None:
        return jqt_obj

    if jqt_obj.is_batched:
        res = []
        for i in range(len(jqt_obj)):
            res.append(jqt2qt(jqt_obj[i]))
        return res

    dims = [list(jqt_obj.dims[0]), list(jqt_obj.dims[1])]
    return Qobj(np.array(jqt_obj.data), dims=dims)

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()

mesolve(H, rho0, tlist, c_ops=None, solver_options=None)

Quantum Master Equation solver.

Parameters:

Name	Type	Description	Default
H	Union[Qarray, Callable[[float], Qarray]]	time dependent Hamiltonian function or time-independent Qarray.	required
rho0	Qarray	initial state, must be a density matrix. For statevector evolution, please use sesolve.	required
tlist	Array	time list	required
c_ops	Optional[Qarray]	qarray list of collapse operators	None
solver_options	Optional[SolverOptions]	SolverOptions with solver options	None

Returns:

Type	Description
Qarray	list of states

Source code in jaxquantum/core/solvers.py

def mesolve(
    H: Union[Qarray, Callable[[float], Qarray]],
    rho0: Qarray,
    tlist: Array,
    c_ops: Optional[Qarray] = None,
    solver_options: Optional[SolverOptions] = None,
) -> Qarray:
    """Quantum Master Equation solver.

    Args:
        H: time dependent Hamiltonian function or time-independent Qarray.
        rho0: initial state, must be a density matrix. For statevector evolution, please use sesolve.
        tlist: time list
        c_ops: qarray list of collapse operators
        solver_options: SolverOptions with solver options

    Returns:
        list of states
    """

    c_ops = c_ops if c_ops is not None else Qarray.from_list([])

    # if isinstance(H, Qarray):

    if len(c_ops) == 0 and rho0.qtype != Qtypes.oper:
        logging.warning(
            "Consider using `jqt.sesolve()` instead, as `c_ops` is an empty list and the initial state is not a density matrix."
        )

    ρ0 = rho0.to_dm()
    dims = ρ0.dims
    ρ0 = ρ0.data

    c_ops = c_ops.data

    if isinstance(H, Qarray):
        Ht_data = lambda t: H.data
    else:
        Ht_data = lambda t: H(t).data if H is not None else None

    ys = _mesolve_data(Ht_data, ρ0, tlist, c_ops, solver_options=solver_options)

    return jnp2jqt(ys, dims=dims)

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)))

overlap(rho, sigma)

Overlap between two states or operators.

Parameters:

Name	Type	Description	Default
rho	Qarray	state/operator.	required
sigma	Qarray	state/operator.	required

Returns:

Type	Description
Array	Overlap between rho and sigma.

Source code in jaxquantum/core/measurements.py

def overlap(rho: Qarray, sigma: Qarray) -> Array:
    """Overlap between two states or operators.

    Args:
        rho: state/operator.
        sigma: state/operator.

    Returns:
        Overlap between rho and sigma.
    """

    if rho.is_vec() and sigma.is_vec():
        return jnp.abs(((rho.to_ket().dag() @ sigma.to_ket()).trace())) ** 2
    elif rho.is_vec():
        rho = rho.to_ket()
        res = (rho.dag() @ sigma @ rho).data
        return res.squeeze(-1).squeeze(-1)
    elif sigma.is_vec():
        sigma = sigma.to_ket()
        res = (sigma.dag() @ rho @ sigma).data
        return res.squeeze(-1).squeeze(-1)
    else:
        return (rho.dag() @ sigma).trace()

plot_qp(state, pts, ax=None, contour=True, qp_type=WIGNER, cbar_label='', axis_scale_factor=1, plot_cbar=True)

Plot quasi-probability distribution.

TODO: decouple this from qutip.

Parameters:

Name	Type	Description	Default
state		statevector	required
pts		points to evaluate quasi-probability distribution on	required
dim		dimensions of state	required
ax		matplotlib axis to plot on	None
contour		make the plot use contouring	True
qp_type		type of quasi probability distribution ("wigner", "qfunc")	WIGNER

Returns:

Type	Description
	axis on which the plot was plotted.

Source code in jaxquantum/core/visualization.py

def plot_qp(
    state,
    pts,
    ax=None,
    contour=True,
    qp_type=WIGNER,
    cbar_label="",
    axis_scale_factor=1,
    plot_cbar=True,
):
    """Plot quasi-probability distribution.

    TODO: decouple this from qutip.

    Args:
        state: statevector
        pts: points to evaluate quasi-probability distribution on
        dim: dimensions of state
        ax: matplotlib axis to plot on
        contour: make the plot use contouring
        qp_type: type of quasi probability distribution ("wigner", "qfunc")

    Returns:
        axis on which the plot was plotted.
    """
    pts = np.array(pts)
    state = jqt2qt(state)
    if ax is None:
        _, ax = plt.subplots(1, figsize=(4, 3), dpi=200)
    # fig = ax.get_figure()

    if qp_type == WIGNER:
        vmin = -1
        vmax = 1
        scale = np.pi / 2
        cmap = "seismic"
    elif qp_type == QFUNC:
        vmin = 0
        vmax = 1
        scale = np.pi
        cmap = "jet"

    QP = scale * getattr(qt, qp_type)(state, pts, pts, g=2)

    pts = pts * axis_scale_factor

    if contour:
        im = ax.contourf(
            pts,
            pts,
            QP,
            cmap=cmap,
            vmin=vmin,
            vmax=vmax,
            levels=np.linspace(vmin, vmax, 101),
        )
    else:
        im = ax.pcolormesh(
            pts,
            pts,
            QP,
            cmap=cmap,
            vmin=vmin,
            vmax=vmax,
        )
    ax.axhline(0, linestyle="-", color="black", alpha=0.7)
    ax.axvline(0, linestyle="-", color="black", alpha=0.7)
    ax.grid()
    ax.set_aspect("equal", adjustable="box")

    if plot_cbar:
        cbar = plt.colorbar(
            im, ax=ax, orientation="vertical", ticks=np.linspace(-1, 1, 11)
        )
        cbar.ax.set_title(cbar_label)

    ax.set_xlabel(r"Re[$\alpha$]")
    ax.set_ylabel(r"Im[$\alpha$]")

    fig = ax.get_figure()
    fig.tight_layout()
    return ax, im

powm(qarr, n)

Matrix power.

Parameters:

Name	Type	Description	Default
qarr	Qarray	quantum array	required
n	int	power	required

Returns:

Type	Description
Qarray	matrix power

Source code in jaxquantum/core/qarray.py

def powm(qarr: Qarray, n: Union[int, float]) -> Qarray:
    """Matrix power.

    Args:
        qarr (Qarray): quantum array
        n (int): power

    Returns:
        matrix power
    """
    if isinstance(n, int):
        data_res = jnp.linalg.matrix_power(qarr.data, n)
    else:
        evalues, evectors = jnp.linalg.eig(qarr.data)
        if not (evalues >= 0).all():
            raise ValueError(
                "Non-integer power of a matrix can only be "
                "computed if the matrix is positive semi-definite."
                "Got a matrix with a negative eigenvalue."
            )
        data_res = evectors * jnp.pow(evalues, n) @ jnp.linalg.inv(evectors)
    return Qarray.create(data_res, dims=qarr.dims)

powm_data(data, n)

Matrix power.

Parameters:

Name	Type	Description	Default
data	Array	matrix	required
n	int	power	required

Returns:

Type	Description
Array	matrix power

Source code in jaxquantum/core/qarray.py

def powm_data(data: Array, n: int) -> Array:
    """Matrix power.

    Args:
        data: matrix
        n: power

    Returns:
        matrix power
    """
    return jnp.linalg.matrix_power(data, n)

ptrace(qarr, indx)

Partial Trace.

Parameters:

Name	Type	Description	Default
rho		density matrix	required
indx		index of quantum object to keep, rest will be partial traced out	required

Returns:

Type	Description
Qarray	partial traced out density matrix

TODO: Fix weird tracing errors that arise with reshape

Source code in jaxquantum/core/qarray.py

def ptrace(qarr: Qarray, indx) -> Qarray:
    """Partial Trace.

    Args:
        rho: density matrix
        indx: index of quantum object to keep, rest will be partial traced out

    Returns:
        partial traced out density matrix

    TODO: Fix weird tracing errors that arise with reshape
    """

    qarr = ket2dm(qarr)
    rho = qarr.shaped_data
    dims = qarr.dims

    Nq = len(dims[0])

    indxs = [indx, indx + Nq]
    for j in range(Nq):
        if j == indx:
            continue
        indxs.append(j)
        indxs.append(j + Nq)

    bdims = qarr.bdims
    len_bdims = len(bdims)
    bdims_indxs = list(range(len_bdims))
    indxs = bdims_indxs + [j + len_bdims for j in indxs]
    rho = rho.transpose(indxs)

    for j in range(Nq - 1):
        rho = jnp.trace(rho, axis1=2 + len_bdims, axis2=3 + len_bdims)

    return Qarray.create(rho)

qt2jqt(qt_obj, dtype=jnp.complex128)

QuTiP state -> Qarray.

Parameters:

Name	Type	Description	Default
qt_obj		QuTiP state.	required
dtype		JAX dtype.	complex128

Returns:

Type	Description
	Qarray.

Source code in jaxquantum/core/conversions.py

def qt2jqt(qt_obj, dtype=jnp.complex128):
    """QuTiP state -> Qarray.

    Args:
        qt_obj: QuTiP state.
        dtype: JAX dtype.

    Returns:
        Qarray.
    """
    if isinstance(qt_obj, Qarray) or qt_obj is None:
        return qt_obj
    return Qarray.create(jnp.array(qt_obj.full(), dtype=dtype), dims=qt_obj.dims)

quantum_state_tomography(rho, physical_basis, logical_basis)

Perform quantum state tomography to retrieve the density matrix in the logical basis.

Args:
    rho: state expressed in the physical Hilbert space basis.
    physical_basis: list of logical operators expressed in the physical
    Hilbert space basis forming a complete logical operator basis.
    logical_basis: list of logical operators expressed in the
    logical Hilbert space basis forming a complete operator basis.


Returns:
    Density matrix of state rho expressed in the logical basis.

Source code in jaxquantum/core/measurements.py

def quantum_state_tomography(
    rho: Qarray, physical_basis: Qarray, logical_basis: Qarray
) -> Qarray:
    """Perform quantum state tomography to retrieve the density matrix in
    the logical basis.

        Args:
            rho: state expressed in the physical Hilbert space basis.
            physical_basis: list of logical operators expressed in the physical
            Hilbert space basis forming a complete logical operator basis.
            logical_basis: list of logical operators expressed in the
            logical Hilbert space basis forming a complete operator basis.


        Returns:
            Density matrix of state rho expressed in the logical basis.
    """
    dm = jnp.zeros_like(logical_basis[0].data)
    rho = rho.to_dm()

    if physical_basis.bdims[-1] != logical_basis.bdims[-1]:
        raise ValueError(
            f"The two bases should have the same size for the "
            f"last batch dimension. Received "
            f"{physical_basis.bdims} and {logical_basis.bdims} "
            f"instead."
        )

    space_size = physical_basis.bdims[-1]

    for i in tqdm(range(space_size), total=space_size):
        p_i = (rho @ physical_basis[i]).trace()
        dm += p_i * logical_basis[i].data

    return Qarray.create(dm, dims=logical_basis.dims, bdims=physical_basis[0].bdims)

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()
    )

sesolve(H, rho0, tlist, solver_options=None)

Schrödinger Equation solver.

Parameters:

Name	Type	Description	Default
H	Union[Qarray, Callable[[float], Qarray]]	time dependent Hamiltonian function or time-independent Qarray.	required
rho0	Qarray	initial state, must be a density matrix. For statevector evolution, please use sesolve.	required
tlist	Array	time list	required
solver_options	Optional[SolverOptions]	SolverOptions with solver options	None

Returns:

Type	Description
Qarray	list of states

Source code in jaxquantum/core/solvers.py

def sesolve(
    H: Union[Qarray, Callable[[float], Qarray]],
    rho0: Qarray,
    tlist: Array,
    solver_options: Optional[SolverOptions] = None,
) -> Qarray:
    """Schrödinger Equation solver.

    Args:
        H: time dependent Hamiltonian function or time-independent Qarray.
        rho0: initial state, must be a density matrix. For statevector evolution, please use sesolve.
        tlist: time list
        solver_options: SolverOptions with solver options

    Returns:
        list of states
    """

    ψ = rho0

    if ψ.qtype == Qtypes.oper:
        raise ValueError(
            "Please use `jqt.mesolve` for initial state inputs in density matrix form."
        )

    ψ = ψ.to_ket()
    dims = ψ.dims
    ψ = ψ.data

    if isinstance(H, Qarray):
        Ht_data = lambda t: H.data
    else:
        Ht_data = lambda t: H(t).data if H is not None else None

    ys = _sesolve_data(Ht_data, ψ, tlist, solver_options=solver_options)

    return jnp2jqt(ys, dims=dims)

sigmam()

σ-

Returns:

Type	Description
Qarray	σ- Pauli Operator

Source code in jaxquantum/core/operators.py

def sigmam() -> Qarray:
    """σ-

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

sigmap()

σ+

Returns:

Type	Description
Qarray	σ+ Pauli Operator

Source code in jaxquantum/core/operators.py

def sigmap() -> Qarray:
    """σ+

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

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]]))

sinm_data(data, **kwargs)

Matrix sine wrapper.

Parameters:

Name	Type	Description	Default
data	Array	matrix	required

Returns:

Type	Description
Array	matrix sine

Source code in jaxquantum/core/qarray.py

def sinm_data(data: Array, **kwargs) -> Array:
    """Matrix sine wrapper.

    Args:
        data: matrix

    Returns:
        matrix sine
    """
    return (expm_data(1j * data) - expm_data(-1j * data)) / (2j)

solve(f, ρ0, tlist, args, solver_options=None)

Gets teh desired solver from diffrax.

Parameters:

Name	Type	Description	Default
solver_options	Optional[SolverOptions]	dictionary with solver options	None

Returns:

Type	Description
	solution

Source code in jaxquantum/core/solvers.py

def solve(f, ρ0, tlist, args, solver_options: Optional[SolverOptions] = None):
    """Gets teh desired solver from diffrax.

    Args:
        solver_options: dictionary with solver options

    Returns:
        solution
    """

    # f and ts
    term = ODETerm(f)
    saveat = SaveAt(ts=tlist)

    # solver
    solver_options = solver_options or SolverOptions.create()

    solver_name = solver_options.solver
    solver = getattr(diffrax, solver_name)()
    stepsize_controller = PIDController(rtol=1e-6, atol=1e-6)

    # solve!
    with warnings.catch_warnings():
        warnings.simplefilter(
            "ignore", UserWarning
        )  # NOTE: suppresses complex dtype warning in diffrax
        sol = diffeqsolve(
            term,
            solver,
            t0=tlist[0],
            t1=tlist[-1],
            dt0=tlist[1] - tlist[0],
            y0=ρ0,
            saveat=saveat,
            stepsize_controller=stepsize_controller,
            args=args,
            max_steps=solver_options.max_steps,
            progress_meter=CustomProgressMeter()
            if solver_options.progress_meter
            else NoProgressMeter(),
        )

    return sol

tensor(*args, **kwargs)

Tensor product.

Parameters:

Name	Type	Description	Default
*args	Qarray	tensors to take the product of	()
parallel	bool	if True, use parallel einsum for tensor product true: [A,B] ^ [C,D] = [A^C, B^D] false: [A,B] ^ [C,D] = [A^C, A^D, B^C, B^D]	required

Returns:

Type	Description
Qarray	Tensor product of given tensors

Source code in jaxquantum/core/qarray.py

def tensor(*args, **kwargs) -> Qarray:
    """Tensor product.

    Args:
        *args (Qarray): tensors to take the product of
        parallel (bool): if True, use parallel einsum for tensor product
            true: [A,B] ^ [C,D] = [A^C, B^D]
            false: [A,B] ^ [C,D] = [A^C, A^D, B^C, B^D]

    Returns:
        Tensor product of given tensors

    """

    parallel = kwargs.pop("parallel", False)

    data = args[0].data
    dims = deepcopy(args[0].dims)
    dims_0 = dims[0]
    dims_1 = dims[1]
    for arg in args[1:]:
        if parallel:
            a = data
            b = arg.data

            if len(a.shape) > len(b.shape):
                batch_dim = a.shape[:-2]
            elif len(a.shape) == len(b.shape):
                if prod(a.shape[:-2]) > prod(b.shape[:-2]):
                    batch_dim = a.shape[:-2]
                else:
                    batch_dim = b.shape[:-2]
            else:
                batch_dim = b.shape[:-2]

            data = jnp.einsum("...ij,...kl->...ikjl", a, b).reshape(
                *batch_dim, a.shape[-2] * b.shape[-2], -1
            )
        else:
            data = jnp.kron(data, arg.data, **kwargs)

        dims_0 = dims_0 + arg.dims[0]
        dims_1 = dims_1 + arg.dims[1]

    return Qarray.create(data, dims=(dims_0, dims_1))

thermal(N, beta)

Thermal state.

Parameters:

Name	Type	Description	Default
N	int	Hilbert Space Size.	required
beta	float	thermal state inverse temperature.	required

Return

Thermal state.

Source code in jaxquantum/core/operators.py

def thermal(N: int, beta: float) -> Qarray:
    """Thermal state.

    Args:
        N: Hilbert Space Size.
        beta: thermal state inverse temperature.

    Return:
        Thermal state.
    """

    return Qarray.create(
        jnp.where(
            jnp.isposinf(beta),
            basis(N, 0).to_dm().data,
            jnp.diag(jnp.exp(-beta * jnp.linspace(0, N - 1, N))),
        )
    ).unit()

tr(qarr, **kwargs)

Full trace.

Parameters:

Name	Type	Description	Default
qarr	Qarray	quantum array	required

Returns:

Type	Description
Array	Full trace.

Source code in jaxquantum/core/qarray.py

def tr(qarr: Qarray, **kwargs) -> Array:
    """Full trace.

    Args:
        qarr (Qarray): quantum array

    Returns:
        Full trace.
    """
    axis1 = kwargs.get("axis1", -2)
    axis2 = kwargs.get("axis2", -1)
    return jnp.trace(qarr.data, axis1=axis1, axis2=axis2, **kwargs)

trace(qarr, **kwargs)

Full trace.

Parameters:

Name	Type	Description	Default
qarr	Qarray	quantum array	required

Returns:

Type	Description
Array	Full trace.

Source code in jaxquantum/core/qarray.py

def trace(qarr: Qarray, **kwargs) -> Array:
    """Full trace.

    Args:
        qarr (Qarray): quantum array

    Returns:
        Full trace.
    """
    return tr(qarr, **kwargs)

transpose(qarr, indices)

Transpose the quantum array.

Parameters:

Name	Type	Description	Default
qarr	Qarray	quantum array	required
*args		axes to transpose	required

Returns:

Type	Description
Qarray	tranposed Qarray

Source code in jaxquantum/core/qarray.py

def transpose(qarr: Qarray, indices: List[int]) -> Qarray:
    """Transpose the quantum array.

    Args:
        qarr (Qarray): quantum array
        *args: axes to transpose

    Returns:
        tranposed Qarray
    """

    indices = list(indices)

    shaped_data = qarr.shaped_data
    dims = qarr.dims
    bdims_indxs = list(range(len(qarr.bdims)))

    reshape_indices = indices + [j + len(dims[0]) for j in indices]

    reshape_indices = bdims_indxs + [j + len(bdims_indxs) for j in reshape_indices]

    shaped_data = shaped_data.transpose(reshape_indices)
    new_dims = (
        tuple([dims[0][j] for j in indices]),
        tuple([dims[1][j] for j in indices]),
    )

    full_dims = prod(dims[0])
    full_data = shaped_data.reshape(*qarr.bdims, full_dims, -1)
    return Qarray.create(full_data, dims=new_dims)

unit(qarr)

Normalize the quantum array.

Parameters:

Name	Type	Description	Default
qarr	Qarray	quantum array	required

Returns:

Type	Description
Qarray	Normalized quantum array

Source code in jaxquantum/core/qarray.py

def unit(qarr: Qarray) -> Qarray:
    """Normalize the quantum array.

    Args:
        qarr (Qarray): quantum array

    Returns:
        Normalized quantum array
    """
    data = qarr.data
    data = data / qarr.norm()
    return Qarray.create(data, dims=qarr.dims)