qarray

QArray.

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

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)

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

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

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)

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]

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

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)

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)

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

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)