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

        if not isinstance(dims[0], (list, tuple)):
            # This handles the case where only the hilbert space dimensions are sent in.
            if data.shape[-1] == 1:
                dims = (tuple(dims), tuple([1 for _ in dims]))
            elif data.shape[-2] == 1:
                dims = (tuple([1 for _ in dims]), tuple(dims))
            else:
                dims = (tuple(dims), tuple(dims))
        else:
            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,
    )

QuantumStateTomography

Source code in jaxquantum/core/measurements.py

class QuantumStateTomography:
    def __init__(
        self,
        rho_guess: Qarray,
        measurement_basis: Qarray,
        measurement_results: jnp.ndarray,
        complete_basis: Optional[Qarray] = None,
        true_rho: Optional[Qarray] = None,
    ):
        """
        Reconstruct a quantum state from measurement results using quantum state tomography.
        The tomography can be performed either by direct inversion or by maximum likelihood estimation.

        Args:
            rho_guess (Qarray): The initial guess for the quantum state.
            measurement_basis (Qarray): The basis in which measurements are performed.
            measurement_results (jnp.ndarray): The results of the measurements.
            complete_basis (Optional[Qarray]): The complete basis for state 
            reconstruction used when using direct inversion. 
            Defaults to the measurement basis if not provided.
            true_rho (Optional[Qarray]): The true quantum state, if known.

        """
        self.rho_guess = rho_guess.data
        self.measurement_basis = measurement_basis.data
        self.measurement_results = measurement_results
        self.complete_basis = (
            complete_basis.data
            if (complete_basis is not None)
            else measurement_basis.data
        )
        self.true_rho = true_rho
        self._result = None

    @property
    def result(self) -> Optional[MLETomographyResult]:
        return self._result


    def quantum_state_tomography_mle(
        self, L1_reg_strength: float = 0.0, epochs: int = 10000, lr: float = 5e-3
    ) -> MLETomographyResult:
        """Perform quantum state tomography using maximum likelihood 
        estimation (MLE).

        This method reconstructs the quantum state from measurement results 
        by optimizing
        a likelihood function using gradient descent. The optimization 
        ensures the 
        resulting density matrix is positive semi-definite with trace 1.

        Args:
            L1_reg_strength (float, optional): Strength of L1 
            regularization. Defaults to 0.0.
            epochs (int, optional): Number of optimization iterations. 
            Defaults to 10000.
            lr (float, optional): Learning rate for the Adam optimizer. 
            Defaults to 5e-3.

        Returns:
            MLETomographyResult: Named tuple containing:
                - rho: Reconstructed quantum state as Qarray
                - params_history: List of parameter values during optimization
                - loss_history: List of loss values during optimization
                - grads_history: List of gradient values during optimization
                - infidelity_history: List of infidelities if true_rho was 
                provided, None otherwise
        """

        dim = self.rho_guess.shape[0]
        optimizer = optax.adam(lr)

        # Initialize parameters from the initial guess for the density matrix
        params = _parametrize_density_matrix(self.rho_guess, dim)
        opt_state = optimizer.init(params)

        compute_infidelity_flag = self.true_rho is not None

        # Provide a dummy array if no true_rho is available. It won't be used.
        true_rho_data_or_dummy = (
            self.true_rho.data
            if compute_infidelity_flag
            else jnp.empty((dim, dim), dtype=jnp.complex64)
        )

        final_carry, history = _run_tomography_scan(
            initial_params=params,
            initial_opt_state=opt_state,
            true_rho_data=true_rho_data_or_dummy,
            measurement_basis=self.measurement_basis,
            measurement_results=self.measurement_results,
            dim=dim,
            epochs=epochs,
            optimizer=optimizer,
            compute_infidelity=compute_infidelity_flag,
            L1_reg_strength=L1_reg_strength,
        )

        final_params, _ = final_carry

        rho = Qarray.create(_reconstruct_density_matrix(final_params, dim))

        self._result = MLETomographyResult(
            rho=rho,
            params_history=history["params"],
            loss_history=history["loss"],
            grads_history=history["grads"],
            infidelity_history=history["infidelity"]
            if compute_infidelity_flag
            else None,
        )
        return self._result

    def quantum_state_tomography_direct(
        self,
    ) -> Qarray:

        """Perform quantum state tomography using direct inversion.

        This method reconstructs the quantum state from measurement results by 
        directly solving a system of linear equations. The method assumes that
        the measurement basis is complete and the measurement results are 
        noise-free.

        Returns:
            Qarray: Reconstructed quantum state.
        """

    # Compute overlaps of measurement and complete operator bases
        A = jnp.einsum("ijk,ljk->il", self.complete_basis, self.measurement_basis)
        # Solve the linear system to find the coefficients
        coefficients = jnp.linalg.solve(A, self.measurement_results)
        # Reconstruct the density matrix
        rho = jnp.einsum("i, ijk->jk", coefficients, self.complete_basis)

        return Qarray.create(rho)

    def plot_results(self):
        if self._result is None:
            raise ValueError(
                "No results to plot. Run quantum_state_tomography_mle first."
            )

        fig, ax = plt.subplots(1, figsize=(5, 4))
        if self._result.infidelity_history is not None:
            ax2 = ax.twinx()

        ax.plot(self._result.loss_history, color="C0")
        ax.set_xlabel("Epoch")
        ax.set_ylabel("$\\mathcal{L}$", color="C0")
        ax.set_yscale("log")

        if self._result.infidelity_history is not None:
            ax2.plot(self._result.infidelity_history, color="C1")
            ax2.set_yscale("log")
            ax2.set_ylabel("$1-\\mathcal{F}$", color="C1")
            plt.grid(False)

        plt.show()

__init__(rho_guess, measurement_basis, measurement_results, complete_basis=None, true_rho=None)

Reconstruct a quantum state from measurement results using quantum state tomography. The tomography can be performed either by direct inversion or by maximum likelihood estimation.

Parameters:

Name	Type	Description	Default
rho_guess	Qarray	The initial guess for the quantum state.	required
measurement_basis	Qarray	The basis in which measurements are performed.	required
measurement_results	ndarray	The results of the measurements.	required
complete_basis	Optional[Qarray]	The complete basis for state	None
true_rho	Optional[Qarray]	The true quantum state, if known.	None

Source code in jaxquantum/core/measurements.py

def __init__(
    self,
    rho_guess: Qarray,
    measurement_basis: Qarray,
    measurement_results: jnp.ndarray,
    complete_basis: Optional[Qarray] = None,
    true_rho: Optional[Qarray] = None,
):
    """
    Reconstruct a quantum state from measurement results using quantum state tomography.
    The tomography can be performed either by direct inversion or by maximum likelihood estimation.

    Args:
        rho_guess (Qarray): The initial guess for the quantum state.
        measurement_basis (Qarray): The basis in which measurements are performed.
        measurement_results (jnp.ndarray): The results of the measurements.
        complete_basis (Optional[Qarray]): The complete basis for state 
        reconstruction used when using direct inversion. 
        Defaults to the measurement basis if not provided.
        true_rho (Optional[Qarray]): The true quantum state, if known.

    """
    self.rho_guess = rho_guess.data
    self.measurement_basis = measurement_basis.data
    self.measurement_results = measurement_results
    self.complete_basis = (
        complete_basis.data
        if (complete_basis is not None)
        else measurement_basis.data
    )
    self.true_rho = true_rho
    self._result = None

quantum_state_tomography_direct()

Perform quantum state tomography using direct inversion.

This method reconstructs the quantum state from measurement results by directly solving a system of linear equations. The method assumes that the measurement basis is complete and the measurement results are noise-free.

Returns:

Name	Type	Description
Qarray	Qarray	Reconstructed quantum state.

Source code in jaxquantum/core/measurements.py

def quantum_state_tomography_direct(
    self,
) -> Qarray:

    """Perform quantum state tomography using direct inversion.

    This method reconstructs the quantum state from measurement results by 
    directly solving a system of linear equations. The method assumes that
    the measurement basis is complete and the measurement results are 
    noise-free.

    Returns:
        Qarray: Reconstructed quantum state.
    """

# Compute overlaps of measurement and complete operator bases
    A = jnp.einsum("ijk,ljk->il", self.complete_basis, self.measurement_basis)
    # Solve the linear system to find the coefficients
    coefficients = jnp.linalg.solve(A, self.measurement_results)
    # Reconstruct the density matrix
    rho = jnp.einsum("i, ijk->jk", coefficients, self.complete_basis)

    return Qarray.create(rho)

quantum_state_tomography_mle(L1_reg_strength=0.0, epochs=10000, lr=0.005)

Perform quantum state tomography using maximum likelihood estimation (MLE).

This method reconstructs the quantum state from measurement results by optimizing a likelihood function using gradient descent. The optimization ensures the resulting density matrix is positive semi-definite with trace 1.

Parameters:

Name	Type	Description	Default
L1_reg_strength	float	Strength of L1	0.0
epochs	int	Number of optimization iterations.	10000
lr	float	Learning rate for the Adam optimizer.	0.005

Returns:

Name	Type	Description
MLETomographyResult	MLETomographyResult	Named tuple containing: - rho: Reconstructed quantum state as Qarray - params_history: List of parameter values during optimization - loss_history: List of loss values during optimization - grads_history: List of gradient values during optimization - infidelity_history: List of infidelities if true_rho was provided, None otherwise

Source code in jaxquantum/core/measurements.py

def quantum_state_tomography_mle(
    self, L1_reg_strength: float = 0.0, epochs: int = 10000, lr: float = 5e-3
) -> MLETomographyResult:
    """Perform quantum state tomography using maximum likelihood 
    estimation (MLE).

    This method reconstructs the quantum state from measurement results 
    by optimizing
    a likelihood function using gradient descent. The optimization 
    ensures the 
    resulting density matrix is positive semi-definite with trace 1.

    Args:
        L1_reg_strength (float, optional): Strength of L1 
        regularization. Defaults to 0.0.
        epochs (int, optional): Number of optimization iterations. 
        Defaults to 10000.
        lr (float, optional): Learning rate for the Adam optimizer. 
        Defaults to 5e-3.

    Returns:
        MLETomographyResult: Named tuple containing:
            - rho: Reconstructed quantum state as Qarray
            - params_history: List of parameter values during optimization
            - loss_history: List of loss values during optimization
            - grads_history: List of gradient values during optimization
            - infidelity_history: List of infidelities if true_rho was 
            provided, None otherwise
    """

    dim = self.rho_guess.shape[0]
    optimizer = optax.adam(lr)

    # Initialize parameters from the initial guess for the density matrix
    params = _parametrize_density_matrix(self.rho_guess, dim)
    opt_state = optimizer.init(params)

    compute_infidelity_flag = self.true_rho is not None

    # Provide a dummy array if no true_rho is available. It won't be used.
    true_rho_data_or_dummy = (
        self.true_rho.data
        if compute_infidelity_flag
        else jnp.empty((dim, dim), dtype=jnp.complex64)
    )

    final_carry, history = _run_tomography_scan(
        initial_params=params,
        initial_opt_state=opt_state,
        true_rho_data=true_rho_data_or_dummy,
        measurement_basis=self.measurement_basis,
        measurement_results=self.measurement_results,
        dim=dim,
        epochs=epochs,
        optimizer=optimizer,
        compute_infidelity=compute_infidelity_flag,
        L1_reg_strength=L1_reg_strength,
    )

    final_params, _ = final_carry

    rho = Qarray.create(_reconstruct_density_matrix(final_params, dim))

    self._result = MLETomographyResult(
        rho=rho,
        params_history=history["params"],
        loss_history=history["loss"],
        grads_history=history["grads"],
        infidelity_history=history["infidelity"]
        if compute_infidelity_flag
        else None,
    )
    return self._result

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)

cf_wigner(psi, xvec, yvec)

Wigner function for a state vector or density matrix at points xvec + i * yvec.

Parameters

Qarray

A state vector or density matrix.

array_like

x-coordinates at which to calculate the Wigner function.

array_like

y-coordinates at which to calculate the Wigner function.

Returns

array

Values representing the Wigner function calculated over the specified range [xvec,yvec].

Source code in jaxquantum/core/cfunctions.py

def cf_wigner(psi, xvec, yvec):
    """Wigner function for a state vector or density matrix at points
    `xvec + i * yvec`.

    Parameters
    ----------

    state : Qarray
        A state vector or density matrix.

    xvec : array_like
        x-coordinates at which to calculate the Wigner function.

    yvec : array_like
        y-coordinates at which to calculate the Wigner function.


    Returns
    -------

    W : array
        Values representing the Wigner function calculated over the specified
        range [xvec,yvec].


    """
    N = psi.dims[0][0]
    x, y = jnp.meshgrid(xvec, yvec)
    alpha = x + 1.0j * y
    displacement = jqt.displace(N, alpha)

    vmapped_overlap = [vmap(vmap(jqt.overlap, in_axes=(None, 0)), in_axes=(
        None, 0))]
    for _ in psi.bdims:
        vmapped_overlap.append(vmap(vmapped_overlap[-1], in_axes=(0, None)))

    cf = vmapped_overlap[-1](psi, displacement)
    return cf

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

concatenate(qarr_list, axis=0)

Concatenate a list of Qarrays along a specified axis.

Parameters:

Name	Type	Description	Default
qarr_list	List[Qarray]	List of Qarrays to concatenate.	required
axis	int	Axis along which to concatenate. Default is 0.	0

Returns:

Name	Type	Description
Qarray	Qarray	Concatenated Qarray.

Source code in jaxquantum/core/qarray.py

def concatenate(qarr_list: List[Qarray], axis: int = 0) -> Qarray:
    """Concatenate a list of Qarrays along a specified axis.

    Args:
        qarr_list (List[Qarray]): List of Qarrays to concatenate.
        axis (int): Axis along which to concatenate. Default is 0.

    Returns:
        Qarray: Concatenated Qarray.
    """

    non_empty_qarr_list = [qarr for qarr in qarr_list if len(qarr.data) != 0]

    if len(non_empty_qarr_list) == 0:
        return Qarray.from_list([])

    concatenated_data = jnp.concatenate(
        [qarr.data for qarr in non_empty_qarr_list], axis=axis
    )

    dims = non_empty_qarr_list[0].dims
    return Qarray.create(concatenated_data, dims=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

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, force_positivity=False)

Fidelity between two states.

Parameters:

Name	Type	Description	Default
rho	Qarray	state.	required
sigma	Qarray	state.	required
force_positivity	bool	force the states to be positive semidefinite	False

Returns:

Type	Description
ndarray	Fidelity between rho and sigma.

Source code in jaxquantum/core/measurements.py

def fidelity(rho: Qarray, sigma: Qarray, force_positivity: bool=False) -> (
        jnp.ndarray):
    """Fidelity between two states.

    Args:
        rho: state.
        sigma: state.
        force_positivity: force the states to be positive semidefinite

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

    sqrt_rho = powm(rho, 0.5, clip_eigvals=force_positivity)

    return jnp.real(((powm(sqrt_rho @ sigma @ sqrt_rho, 0.5,
                           clip_eigvals=force_positivity)).tr())
                    ** 2)

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, saveat_tlist=None, 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
saveat_tlist	Optional[Array]	list of times at which to save the state. If None, save at all times in tlist. Default: None.	None
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,
    saveat_tlist: Optional[Array] = None,
    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
        saveat_tlist: list of times at which to save the state. If None, save at all times in tlist. Default: None.
        c_ops: qarray list of collapse operators
        solver_options: SolverOptions with solver options

    Returns:
        list of states
    """

    saveat_tlist = saveat_tlist if saveat_tlist is not None else tlist

    saveat_tlist = jnp.atleast_1d(saveat_tlist)

    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, saveat_tlist, c_ops,
                       solver_options=solver_options)

    return jnp2jqt(ys, dims=dims)

multi_mode_basis_set(Ns)

Creates a multi-mode basis set.

Parameters:

Name	Type	Description	Default
Ns	List[int]	List of Hilbert space dimensions for each mode.	required

Returns:

Type	Description
Qarray	Multi-mode basis set.

Source code in jaxquantum/core/operators.py

def multi_mode_basis_set(Ns: List[int]) -> Qarray:
    """Creates a multi-mode basis set.

    Args:
        Ns: List of Hilbert space dimensions for each mode.

    Returns:
        Multi-mode basis set.
    """
    data = jnp.eye(prod(Ns))
    dims = (tuple(Ns), tuple([1 for _ in Ns]))
    return Qarray.create(data, dims=dims, bdims=(prod(Ns),))

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_cf(state, pts_x, pts_y=None, axs=None, contour=True, qp_type=WIGNER, cbar_label='', axis_scale_factor=1, plot_cbar=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None)

Plot characteristic function.

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions, result will	required
pts_x		x points to evaluate quasi-probability distribution at	required
pts_y		y points to evaluate quasi-probability distribution at	None
axs		matplotlib axes to plot on	None
contour		make the plot use contouring	True
qp_type		type of quasi probability distribution ("wigner")	WIGNER
cbar_label		labels for the real and imaginary cbar	''
axis_scale_factor		scale of the axes labels relative	1
plot_cbar		whether to plot cbar	True
x_ticks		tick position for the x-axis	None
y_ticks		tick position for the y-axis	None
z_ticks		tick position for the z-axis	None
subtitles		subtitles for the subplots	None
figtitle		figure title	None

Returns:

Type	Description
	axis on which the plot was plotted.

Source code in jaxquantum/core/visualization.py

def plot_cf(
        state,
        pts_x,
        pts_y=None,
        axs=None,
        contour=True,
        qp_type=WIGNER,
        cbar_label="",
        axis_scale_factor=1,
        plot_cbar=True,
        x_ticks=None,
        y_ticks=None,
        z_ticks=None,
        subtitles=None,
        figtitle=None,
):
    """Plot characteristic function.


    Args:
        state: state with arbitrary number of batch dimensions, result will
        be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate quasi-probability distribution at
        pts_y: y points to evaluate quasi-probability distribution at
        axs: matplotlib axes to plot on
        contour: make the plot use contouring
        qp_type: type of quasi probability distribution ("wigner")
        cbar_label: labels for the real and imaginary cbar
        axis_scale_factor: scale of the axes labels relative
        plot_cbar: whether to plot cbar
        x_ticks: tick position for the x-axis
        y_ticks: tick position for the y-axis
        z_ticks: tick position for the z-axis
        subtitles: subtitles for the subplots
        figtitle: figure title

    Returns:
        axis on which the plot was plotted.
    """
    if pts_y is None:
        pts_y = pts_x
    pts_x = jnp.array(pts_x)
    pts_y = jnp.array(pts_y)

    bdims = state.bdims
    added_baxes = 0

    if subtitles is not None:
        if subtitles.shape != bdims:
            raise ValueError(
                f"labels must have same shape as bdims, "
                f"got shapes {subtitles.shape} and {bdims}"
            )

    if len(bdims) == 0:
        bdims = (1,)
        added_baxes += 1
    if len(bdims) == 1:
        bdims = (1, bdims[0])
        added_baxes += 1

    extra_dims = bdims[2:]
    if extra_dims != ():
        state = state.reshape_bdims(
            bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1]
        )
        if subtitles is not None:
            subtitles = subtitles.reshape(
                bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1]
            )
        bdims = state.bdims

    if axs is None:
        _, axs = plt.subplots(
            bdims[0],
            bdims[1]*2,
            figsize=(4 * bdims[1]*2, 3 * bdims[0]),
            dpi=200,
        )


    if qp_type == WIGNER:
        vmin = -1
        vmax = 1
        scale = 1
        cmap = "seismic"
        cbar_label = [r"$\mathcal{Re}(\chi_W(\alpha))$", r"$\mathcal{"
                                                         r"Im}(\chi_W("
                                                         r"\alpha))$"]
        QP = scale * cf_wigner(state, pts_x, pts_y)

    for _ in range(added_baxes):
        QP = jnp.array([QP])
        axs = np.array([axs])
        if subtitles is not None:
            subtitles = np.array([subtitles])

    if added_baxes==2:
        axs = axs[0] # When the input state is zero-dimensional, remove an
                     # axis that is automatically added due to the subcolumns


    pts_x = pts_x * axis_scale_factor
    pts_y = pts_y * axis_scale_factor

    x_ticks = (
        jnp.linspace(jnp.min(pts_x), jnp.max(pts_x),
                     5) if x_ticks is None else x_ticks
    )
    y_ticks = (
        jnp.linspace(jnp.min(pts_y), jnp.max(pts_y),
                     5) if y_ticks is None else y_ticks
    )
    z_ticks = jnp.linspace(vmin, vmax, 11) if z_ticks is None else z_ticks
    print(axs.shape)
    for row in range(bdims[0]):
        for col in range(bdims[1]):
            for subcol in range(2):
                ax = axs[row, 2 * col + subcol]
                if contour:
                    im = ax.contourf(
                        pts_x,
                        pts_y,
                        jnp.real(QP[row, col]) if subcol==0 else jnp.imag(QP[
                                                                           row, col]),
                        cmap=cmap,
                        vmin=vmin,
                        vmax=vmax,
                        levels=np.linspace(vmin, vmax, 101),
                    )
                else:
                    im = ax.pcolormesh(
                        pts_x,
                        pts_y,
                        jnp.real(QP[row, col]) if subcol == 0 else jnp.imag(QP[
                                                                                row, col]),
                        cmap=cmap,
                        vmin=vmin,
                        vmax=vmax,
                    )
                ax.set_xticks(x_ticks)
                ax.set_yticks(y_ticks)
                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[subcol])
                    cbar.set_ticks(z_ticks)

                ax.set_xlabel(r"Re[$\alpha$]")
                ax.set_ylabel(r"Im[$\alpha$]")
                if subtitles is not None:
                    ax.set_title(subtitles[row, col])

    fig = ax.get_figure()
    fig.tight_layout()
    if figtitle is not None:
        fig.suptitle(figtitle, y=1.04)
    return axs, im

plot_cf_wigner(state, pts_x, pts_y=None, axs=None, contour=True, cbar_label='', axis_scale_factor=1, plot_cbar=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None)

Plot the Wigner characteristic function of the state.

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions, result will	required
pts_x		x points to evaluate quasi-probability distribution at	required
pts_y		y points to evaluate quasi-probability distribution at	None
axs		matplotlib axes to plot on	None
contour		make the plot use contouring	True
cbar_label		label for the cbar	''
axis_scale_factor		scale of the axes labels relative	1
plot_cbar		whether to plot cbar	True
x_ticks		tick position for the x-axis	None
y_ticks		tick position for the y-axis	None
z_ticks		tick position for the z-axis	None
subtitles		subtitles for the subplots	None
figtitle		figure title	None

Returns:

Type	Description
	axis on which the plot was plotted.

Source code in jaxquantum/core/visualization.py

def plot_cf_wigner(
    state,
    pts_x,
    pts_y=None,
    axs=None,
    contour=True,
    cbar_label="",
    axis_scale_factor=1,
    plot_cbar=True,
    x_ticks=None,
    y_ticks=None,
    z_ticks=None,
    subtitles=None,
    figtitle=None,
):
    """Plot the Wigner characteristic function of the state.


    Args:
        state: state with arbitrary number of batch dimensions, result will
        be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate quasi-probability distribution at
        pts_y: y points to evaluate quasi-probability distribution at
        axs: matplotlib axes to plot on
        contour: make the plot use contouring
        cbar_label: label for the cbar
        axis_scale_factor: scale of the axes labels relative
        plot_cbar: whether to plot cbar
        x_ticks: tick position for the x-axis
        y_ticks: tick position for the y-axis
        z_ticks: tick position for the z-axis
        subtitles: subtitles for the subplots
        figtitle: figure title

    Returns:
        axis on which the plot was plotted.
    """
    return plot_cf(
        state=state,
        pts_x=pts_x,
        pts_y=pts_y,
        axs=axs,
        contour=contour,
        qp_type=WIGNER,
        cbar_label=cbar_label,
        axis_scale_factor=axis_scale_factor,
        plot_cbar=plot_cbar,
        x_ticks=x_ticks,
        y_ticks=y_ticks,
        z_ticks=z_ticks,
        subtitles=subtitles,
        figtitle=figtitle,
    )

plot_qfunc(state, pts_x, pts_y=None, g=2, axs=None, contour=True, cbar_label='', axis_scale_factor=1, plot_cbar=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None)

Plot the husimi function of the state.

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions, result will	required
pts_x		x points to evaluate quasi-probability distribution at	required
pts_y		y points to evaluate quasi-probability distribution at	None
g		float, default: 2	2
related	to the value of	math:\hbar in the commutation relation	required
		math:[x,\,y] = i\hbar via :math:\hbar=2/g^2.	required
axs		matplotlib axes to plot on	None
contour		make the plot use contouring	True
cbar_label		label for the cbar	''
axis_scale_factor		scale of the axes labels relative	1
plot_cbar		whether to plot cbar	True
x_ticks		tick position for the x-axis	None
y_ticks		tick position for the y-axis	None
z_ticks		tick position for the z-axis	None
subtitles		subtitles for the subplots	None
figtitle		figure title	None

Returns:

Type	Description
	axis on which the plot was plotted.

Source code in jaxquantum/core/visualization.py

def plot_qfunc(
    state,
    pts_x,
    pts_y=None,
    g=2,
    axs=None,
    contour=True,
    cbar_label="",
    axis_scale_factor=1,
    plot_cbar=True,
    x_ticks=None,
    y_ticks=None,
    z_ticks=None,
    subtitles=None,
    figtitle=None,
):
    """Plot the husimi function of the state.


    Args:
        state: state with arbitrary number of batch dimensions, result will
        be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate quasi-probability distribution at
        pts_y: y points to evaluate quasi-probability distribution at
        g : float, default: 2
        Scaling factor for ``a = 0.5 * g * (x + iy)``.  The value of `g` is
        related to the value of :math:`\\hbar` in the commutation relation
        :math:`[x,\,y] = i\\hbar` via :math:`\\hbar=2/g^2`.
        axs: matplotlib axes to plot on
        contour: make the plot use contouring
        cbar_label: label for the cbar
        axis_scale_factor: scale of the axes labels relative
        plot_cbar: whether to plot cbar
        x_ticks: tick position for the x-axis
        y_ticks: tick position for the y-axis
        z_ticks: tick position for the z-axis
        subtitles: subtitles for the subplots
        figtitle: figure title

    Returns:
        axis on which the plot was plotted.
    """
    return plot_qp(
        state=state,
        pts_x=pts_x,
        pts_y=pts_y,
        g=g,
        axs=axs,
        contour=contour,
        qp_type=HUSIMI,
        cbar_label=cbar_label,
        axis_scale_factor=axis_scale_factor,
        plot_cbar=plot_cbar,
        x_ticks=x_ticks,
        y_ticks=y_ticks,
        z_ticks=z_ticks,
        subtitles=subtitles,
        figtitle=figtitle,
    )

plot_qp(state, pts_x, pts_y=None, g=2, axs=None, contour=True, qp_type=WIGNER, cbar_label='', axis_scale_factor=1, plot_cbar=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None)

Plot quasi-probability distribution.

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions, result will	required
pts_x		x points to evaluate quasi-probability distribution at	required
pts_y		y points to evaluate quasi-probability distribution at	None
g		float, default: 2	2
related	to the value of	math:\hbar in the commutation relation	required
		math:[x,\,y] = i\hbar via :math:\hbar=2/g^2.	required
axs		matplotlib axes to plot on	None
contour		make the plot use contouring	True
qp_type		type of quasi probability distribution ("wigner", "qfunc")	WIGNER
cbar_label		label for the cbar	''
axis_scale_factor		scale of the axes labels relative	1
plot_cbar		whether to plot cbar	True
x_ticks		tick position for the x-axis	None
y_ticks		tick position for the y-axis	None
z_ticks		tick position for the z-axis	None
subtitles		subtitles for the subplots	None
figtitle		figure title	None

Returns:

Type	Description
	axis on which the plot was plotted.

Source code in jaxquantum/core/visualization.py

def plot_qp(
    state,
    pts_x,
    pts_y=None,
    g=2,
    axs=None,
    contour=True,
    qp_type=WIGNER,
    cbar_label="",
    axis_scale_factor=1,
    plot_cbar=True,
    x_ticks=None,
    y_ticks=None,
    z_ticks=None,
    subtitles=None,
    figtitle=None,
):
    """Plot quasi-probability distribution.


    Args:
        state: state with arbitrary number of batch dimensions, result will
        be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate quasi-probability distribution at
        pts_y: y points to evaluate quasi-probability distribution at
        g : float, default: 2
        Scaling factor for ``a = 0.5 * g * (x + iy)``.  The value of `g` is
        related to the value of :math:`\\hbar` in the commutation relation
        :math:`[x,\,y] = i\\hbar` via :math:`\\hbar=2/g^2`.
        axs: matplotlib axes to plot on
        contour: make the plot use contouring
        qp_type: type of quasi probability distribution ("wigner", "qfunc")
        cbar_label: label for the cbar
        axis_scale_factor: scale of the axes labels relative
        plot_cbar: whether to plot cbar
        x_ticks: tick position for the x-axis
        y_ticks: tick position for the y-axis
        z_ticks: tick position for the z-axis
        subtitles: subtitles for the subplots
        figtitle: figure title

    Returns:
        axis on which the plot was plotted.
    """
    if pts_y is None:
        pts_y = pts_x
    pts_x = jnp.array(pts_x)
    pts_y = jnp.array(pts_y)

    if len(state.bdims)==1 and state.bdims[0]==1:
        state = state[0]


    bdims = state.bdims
    added_baxes = 0

    if subtitles is not None:
        if subtitles.shape != bdims:
            raise ValueError(
                f"labels must have same shape as bdims, "
                f"got shapes {subtitles.shape} and {bdims}"
            )

    if len(bdims) == 0:
        bdims = (1,)
        added_baxes += 1
    if len(bdims) == 1:
        bdims = (1, bdims[0])
        added_baxes += 1

    extra_dims = bdims[2:]
    if extra_dims != ():
        state = state.reshape_bdims(
            bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1]
        )
        if subtitles is not None:
            subtitles = subtitles.reshape(
                bdims[0] * int(jnp.prod(jnp.array(extra_dims))), bdims[1]
            )
        bdims = state.bdims

    if axs is None:
        _, axs = plt.subplots(
            bdims[0],
            bdims[1],
            figsize=(4 * bdims[1], 3 * bdims[0]),
            dpi=200,
        )

    if qp_type == WIGNER:
        vmin = -1
        vmax = 1
        scale = np.pi / 2
        cmap = "seismic"
        cbar_label = r"$\mathcal{W}(\alpha)$"
        QP = scale * wigner(state, pts_x, pts_y, g=g)

    elif qp_type == HUSIMI:
        vmin = 0
        vmax = 1
        scale = np.pi
        cmap = "jet"
        cbar_label = r"$\mathcal{Q}(\alpha)$"
        QP = scale * qfunc(state, pts_x, pts_y, g=g)



    for _ in range(added_baxes):
        QP = jnp.array([QP])
        axs = np.array([axs])
        if subtitles is not None:
            subtitles = np.array([subtitles])




    pts_x = pts_x * axis_scale_factor
    pts_y = pts_y * axis_scale_factor

    x_ticks = (
        jnp.linspace(jnp.min(pts_x), jnp.max(pts_x), 5) if x_ticks is None else x_ticks
    )
    y_ticks = (
        jnp.linspace(jnp.min(pts_y), jnp.max(pts_y), 5) if y_ticks is None else y_ticks
    )
    z_ticks = jnp.linspace(vmin, vmax, 11) if z_ticks is None else z_ticks

    for row in range(bdims[0]):
        for col in range(bdims[1]):
            ax = axs[row, col]
            if contour:
                im = ax.contourf(
                    pts_x,
                    pts_y,
                    QP[row, col],
                    cmap=cmap,
                    vmin=vmin,
                    vmax=vmax,
                    levels=np.linspace(vmin, vmax, 101),
                )
            else:
                im = ax.pcolormesh(
                    pts_x,
                    pts_y,
                    QP[row, col],
                    cmap=cmap,
                    vmin=vmin,
                    vmax=vmax,
                )
            ax.set_xticks(x_ticks)
            ax.set_yticks(y_ticks)
            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)
                cbar.set_ticks(z_ticks)

            ax.set_xlabel(r"Re[$\alpha$]")
            ax.set_ylabel(r"Im[$\alpha$]")
            if subtitles is not None:
                ax.set_title(subtitles[row, col])

    fig = ax.get_figure()
    fig.tight_layout()
    if figtitle is not None:
        fig.suptitle(figtitle, y=1.04)
    return axs, im

plot_wigner(state, pts_x, pts_y=None, g=2, axs=None, contour=True, cbar_label='', axis_scale_factor=1, plot_cbar=True, x_ticks=None, y_ticks=None, z_ticks=None, subtitles=None, figtitle=None)

Plot the wigner function of the state.

Parameters:

Name	Type	Description	Default
state		state with arbitrary number of batch dimensions, result will	required
pts_x		x points to evaluate quasi-probability distribution at	required
pts_y		y points to evaluate quasi-probability distribution at	None
g		float, default: 2	2
related	to the value of	math:\hbar in the commutation relation	required
		math:[x,\,y] = i\hbar via :math:\hbar=2/g^2.	required
axs		matplotlib axes to plot on	None
contour		make the plot use contouring	True
cbar_label		label for the cbar	''
axis_scale_factor		scale of the axes labels relative	1
plot_cbar		whether to plot cbar	True
x_ticks		tick position for the x-axis	None
y_ticks		tick position for the y-axis	None
z_ticks		tick position for the z-axis	None
subtitles		subtitles for the subplots	None
figtitle		figure title	None

Returns:

Type	Description
	axis on which the plot was plotted.

Source code in jaxquantum/core/visualization.py

def plot_wigner(
    state,
    pts_x,
    pts_y=None,
    g=2,
    axs=None,
    contour=True,
    cbar_label="",
    axis_scale_factor=1,
    plot_cbar=True,
    x_ticks=None,
    y_ticks=None,
    z_ticks=None,
    subtitles=None,
    figtitle=None,
):
    """Plot the wigner function of the state.


    Args:
        state: state with arbitrary number of batch dimensions, result will
        be flattened to a 2d grid to allow for plotting
        pts_x: x points to evaluate quasi-probability distribution at
        pts_y: y points to evaluate quasi-probability distribution at
        g : float, default: 2
        Scaling factor for ``a = 0.5 * g * (x + iy)``.  The value of `g` is
        related to the value of :math:`\\hbar` in the commutation relation
        :math:`[x,\,y] = i\\hbar` via :math:`\\hbar=2/g^2`.
        axs: matplotlib axes to plot on
        contour: make the plot use contouring
        cbar_label: label for the cbar
        axis_scale_factor: scale of the axes labels relative
        plot_cbar: whether to plot cbar
        x_ticks: tick position for the x-axis
        y_ticks: tick position for the y-axis
        z_ticks: tick position for the z-axis
        subtitles: subtitles for the subplots
        figtitle: figure title

    Returns:
        axis on which the plot was plotted.
    """
    return plot_qp(
        state=state,
        pts_x=pts_x,
        pts_y=pts_y,
        g=g,
        axs=axs,
        contour=contour,
        qp_type=WIGNER,
        cbar_label=cbar_label,
        axis_scale_factor=axis_scale_factor,
        plot_cbar=plot_cbar,
        x_ticks=x_ticks,
        y_ticks=y_ticks,
        z_ticks=z_ticks,
        subtitles=subtitles,
        figtitle=figtitle,
    )

powm(qarr, n, clip_eigvals=False)

Matrix power.

Parameters:

Name	Type	Description	Default
qarr	Qarray	quantum array	required
n	int	power	required
clip_eigvals	bool	clip eigenvalues to always be able to compute	False

Returns:

Type	Description
Qarray	matrix power

Source code in jaxquantum/core/qarray.py

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

    Args:
        qarr (Qarray): quantum array
        n (int): power
        clip_eigvals (bool): clip eigenvalues to always be able to compute
        non-integer powers

    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 clip_eigvals:
            evalues = jnp.maximum(evalues, 0)
        else:
            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)

propagator(H, ts, solver_options=None)

Generate the propagator for a time dependent Hamiltonian.

Parameters:

Name	Type	Description	Default
H	Qarray or callable	A Qarray static Hamiltonian OR a function that takes a time argument and returns a Hamiltonian.	required
ts	float or Array	A single time point or an Array of time points.	required

Returns:

Type	Description
	Qarray or List[Qarray]: The propagator for the Hamiltonian at time t. OR a list of propagators for the Hamiltonian at each time in t.

Source code in jaxquantum/core/solvers.py

def propagator(
    H: Union[Qarray, Callable[[float], Qarray]],
    ts: Union[float, Array],
    solver_options=None
):
    """ Generate the propagator for a time dependent Hamiltonian.

    Args:
        H (Qarray or callable):
            A Qarray static Hamiltonian OR
            a function that takes a time argument and returns a Hamiltonian.
        ts (float or Array):
            A single time point or
            an Array of time points.

    Returns:
        Qarray or List[Qarray]:
            The propagator for the Hamiltonian at time t.
            OR a list of propagators for the Hamiltonian at each time in t.

    """


    ts_is_scalar = robust_isscalar(ts)
    H_is_qarray = isinstance(H, Qarray)

    if H_is_qarray:
        return (-1j * H * ts).expm()
    else:

        if ts_is_scalar:
            H_first = H(0.0)
            if ts == 0:
                return identity_like(H_first)
            ts = jnp.array([0.0, ts])
        else:
            H_first = H(ts[0])

        basis_states = multi_mode_basis_set(H_first.space_dims)
        results = sesolve(H, basis_states, ts)
        propagators_data = results.data.squeeze(-1)
        propagators = Qarray.create(propagators_data, dims=H_first.space_dims)

        return propagators

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)

qfunc(psi, xvec, yvec, g=2)

Husimi-Q function of a given state vector or density matrix at phase-space points 0.5 * g * (xvec + i*yvec).

Parameters

state : Qarray A state vector or density matrix. This cannot have tensor-product structure.

xvec, yvec : array_like x- and y-coordinates at which to calculate the Husimi-Q function.

float, default: 2

Scaling factor for a = 0.5 * g * (x + iy). The value of g is related to the value of :math:\hbar in the commutation relation :math:[x,\,y] = i\hbar via :math:\hbar=2/g^2.

Returns

jnp.ndarray Values representing the Husimi-Q function calculated over the specified range [xvec, yvec].

Source code in jaxquantum/core/qp_distributions.py

def qfunc(psi, xvec, yvec, g=2):
    r"""
    Husimi-Q function of a given state vector or density matrix at phase-space
    points ``0.5 * g * (xvec + i*yvec)``.

    Parameters
    ----------
    state : Qarray
        A state vector or density matrix. This cannot have tensor-product
        structure.

    xvec, yvec : array_like
        x- and y-coordinates at which to calculate the Husimi-Q function.

    g : float, default: 2
        Scaling factor for ``a = 0.5 * g * (x + iy)``.  The value of `g` is
        related to the value of :math:`\hbar` in the commutation relation
        :math:`[x,\,y] = i\hbar` via :math:`\hbar=2/g^2`.

    Returns
    -------
    jnp.ndarray
        Values representing the Husimi-Q function calculated over the specified
        range ``[xvec, yvec]``.

    """

    alpha_grid, prefactor = _qfunc_coherent_grid(xvec, yvec, g)

    if psi.is_vec():
        psi = psi.to_ket()

        def _compute_qfunc(psi, alpha_grid, prefactor, g):
            out = _qfunc_iterative_single(psi, alpha_grid, prefactor, g)
            out /= jnp.pi
            return out
    else:

        def _compute_qfunc(psi, alpha_grid, prefactor, g):
            values, vectors = jnp.linalg.eigh(psi)
            vectors = vectors.T
            out = values[0] * _qfunc_iterative_single(
                vectors[0], alpha_grid, prefactor, g
            )
            for value, vector in zip(values[1:], vectors[1:]):
                out += value * _qfunc_iterative_single(vector, alpha_grid, prefactor, g)
            out /= jnp.pi

            return out

    psi = psi.data

    vmapped_compute_qfunc = [_compute_qfunc]

    for _ in psi.shape[:-2]:
        vmapped_compute_qfunc.append(
            vmap(
                vmapped_compute_qfunc[-1],
                in_axes=(0, None, None, None),
                out_axes=0,
            )
        )
    return vmapped_compute_qfunc[-1](psi, alpha_grid, prefactor, g)

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)

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, saveat_tlist=None, 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
saveat_tlist	Optional[Array]	list of times at which to save the state. If None, save at all times in tlist. Default: None.	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 sesolve(
    H: Union[Qarray, Callable[[float], Qarray]],
    rho0: Qarray,
    tlist: Array,
    saveat_tlist: Optional[Array] = None,
    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
        saveat_tlist: list of times at which to save the state. If None, save at all times in tlist. Default: None.
        solver_options: SolverOptions with solver options

    Returns:
        list of states
    """

    saveat_tlist = saveat_tlist if saveat_tlist is not None else tlist

    saveat_tlist = jnp.atleast_1d(saveat_tlist)

    ψ = 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, saveat_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, saveat_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, saveat_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)
    if saveat_tlist.shape[0] == 1 and saveat_tlist == tlist[-1]:
        saveat = SaveAt(t1=True)
    else:
        saveat = SaveAt(ts=saveat_tlist)

    # solver
    solver_options = solver_options or SolverOptions.create()

    solver_name = solver_options.solver
    solver = getattr(diffrax, solver_name)()
    stepsize_controller = PIDController(rtol=solver_options.rtol, atol=solver_options.atol)

    # solve!
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore",
                                message="Complex dtype support in Diffrax",
                                category=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 (default): [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 (default): [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))

tensor_basis(single_basis, n)

Construct n-fold tensor product basis from a single-system basis.

Parameters:

Name	Type	Description	Default
single_basis	Qarray	The single-system operator basis as a Qarray.	required
n	int	Number of tensor copies to construct.	required

Returns:

Type	Description
Qarray	Qarray containing the n-fold tensor product basis operators.
Qarray	The resulting basis has b^n elements where b is the number
Qarray	of operators in the single-system basis.

Source code in jaxquantum/core/measurements.py

def tensor_basis(single_basis: Qarray, n: int) -> Qarray:
    """Construct n-fold tensor product basis from a single-system basis.

    Args:
        single_basis: The single-system operator basis as a Qarray.
        n: Number of tensor copies to construct.

    Returns:
        Qarray containing the n-fold tensor product basis operators.
        The resulting basis has b^n elements where b is the number
        of operators in the single-system basis.
    """

    dims = single_basis.dims

    single_basis = single_basis.data
    b, d, _ = single_basis.shape
    indices = jnp.stack(jnp.meshgrid(*[jnp.arange(b)] * n, indexing="ij"),
                        axis=-1).reshape(-1, n)  # shape (b^n, n)

    # Select the operators based on indices: shape (b^n, n, d, d)
    selected = single_basis[indices]  # shape: (b^n, n, d, d)

    # Vectorized Kronecker products
    full_basis = vmap(lambda ops: reduce(jnp.kron, ops))(selected)

    new_dims = tuple(tuple(x**n for x in row) for row in dims)

    return Qarray.create(full_basis, dims=new_dims, bdims=(b**n,))

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)

wigner(psi, xvec, yvec, method='clenshaw', g=2)

Wigner function for a state vector or density matrix at points xvec + i * yvec.

Parameters

Qarray

A state vector or density matrix.

array_like

x-coordinates at which to calculate the Wigner function.

array_like

y-coordinates at which to calculate the Wigner function.

float, default: 2

Scaling factor for a = 0.5 * g * (x + iy), default g = 2. The value of g is related to the value of hbar in the commutation relation [x, y] = i * hbar via hbar=2/g^2.

string {'clenshaw', 'iterative', 'laguerre', 'fft'}, default: 'clenshaw'

Only 'clenshaw' is currently supported. Select method 'clenshaw' 'iterative', 'laguerre', or 'fft', where 'clenshaw' and 'iterative' use an iterative method to evaluate the Wigner functions for density matrices :math:|m><n|, while 'laguerre' uses the Laguerre polynomials in scipy for the same task. The 'fft' method evaluates the Fourier transform of the density matrix. The 'iterative' method is default, and in general recommended, but the 'laguerre' method is more efficient for very sparse density matrices (e.g., superpositions of Fock states in a large Hilbert space). The 'clenshaw' method is the preferred method for dealing with density matrices that have a large number of excitations (>~50). 'clenshaw' is a fast and numerically stable method.

Returns

array

Values representing the Wigner function calculated over the specified range [xvec,yvec].

References

Ulf Leonhardt, Measuring the Quantum State of Light, (Cambridge University Press, 1997)

Source code in jaxquantum/core/qp_distributions.py

def wigner(psi, xvec, yvec, method="clenshaw", g=2):
    """Wigner function for a state vector or density matrix at points
    `xvec + i * yvec`.

    Parameters
    ----------

    state : Qarray
        A state vector or density matrix.

    xvec : array_like
        x-coordinates at which to calculate the Wigner function.

    yvec : array_like
        y-coordinates at which to calculate the Wigner function.

    g : float, default: 2
        Scaling factor for `a = 0.5 * g * (x + iy)`, default `g = 2`.
        The value of `g` is related to the value of `hbar` in the commutation
        relation `[x, y] = i * hbar` via `hbar=2/g^2`.

    method : string {'clenshaw', 'iterative', 'laguerre', 'fft'}, default: 'clenshaw'
        Only 'clenshaw' is currently supported.
        Select method 'clenshaw' 'iterative', 'laguerre', or 'fft', where 'clenshaw'
        and 'iterative' use an iterative method to evaluate the Wigner functions for density
        matrices :math:`|m><n|`, while 'laguerre' uses the Laguerre polynomials
        in scipy for the same task. The 'fft' method evaluates the Fourier
        transform of the density matrix. The 'iterative' method is default, and
        in general recommended, but the 'laguerre' method is more efficient for
        very sparse density matrices (e.g., superpositions of Fock states in a
        large Hilbert space). The 'clenshaw' method is the preferred method for
        dealing with density matrices that have a large number of excitations
        (>~50). 'clenshaw' is a fast and numerically stable method.

    Returns
    -------

    W : array
        Values representing the Wigner function calculated over the specified
        range [xvec,yvec].


    References
    ----------

    Ulf Leonhardt,
    Measuring the Quantum State of Light, (Cambridge University Press, 1997)

    """

    if not (psi.is_vec() or psi.is_dm()):
        raise TypeError("Input state is not a valid operator.")

    if method == "fft":
        raise NotImplementedError("Only the 'clenshaw' method is implemented.")

    if method == "iterative":
        raise NotImplementedError("Only the 'clenshaw' method is implemented.")

    elif method == "laguerre":
        raise NotImplementedError("Only the 'clenshaw' method is implemented.")

    elif method == "clenshaw":
        rho = psi.to_dm()
        rho = rho.data

        vmapped_wigner_clenshaw = [_wigner_clenshaw]

        for _ in rho.shape[:-2]:
            vmapped_wigner_clenshaw.append(
                vmap(
                    vmapped_wigner_clenshaw[-1],
                    in_axes=(0, None, None, None),
                    out_axes=0,
                )
            )
        return vmapped_wigner_clenshaw[-1](rho, xvec, yvec, g)

    else:
        raise TypeError("method must be either 'iterative', 'laguerre', or 'fft'.")