1. States¶

class QuantumState:
    """
    The QuantumState class is the unified API for accessing state representation backends.
    It contains one state representation.
    """

    def __init__(self, data, rep_type=None, mixed=False):
        """
        Creates the QuantumState class with one initial representation

        :param data: valid data input for "rep_type".
            Density matrices representations support np.ndarray or int inputs
            Stabilizer representations take int or StabilizerTableau
            Graph representations take networkx.Graph
        :type data: list OR numpy.ndarray OR Graph OR nx.Graph or CliffordTableau
        :param rep_type: selected representation to initialize;
            if not specified, the default choice is the density matrix if the number of qubits is
            less than the threshold value or stabilizer otherwise.
        :type rep_type: str
        :param mixed: boolean flag to initialize as a mixed state or not (mainly used for Stabilizer rep_type)
        :type mixed: boolean
        :return: nothing
        :rtype: None
        """

        self.mixed = mixed
        self._rep_type = self._get_rep_type_name(rep_type)
        valid, self.n_qubits = self.validate_data(data)

        if not valid:
            raise TypeError(
                f"Data's type is incorrect. Input data is {data}, which is of type {type(data)}"
            )
        if rep_type is None:
            if (
                self.n_qubits < DENSITY_MATRIX_QUBIT_THRESH
                and DensityMatrix.valid_datatype(data)
            ):
                self._rep_data = self._initialize_dm(data)
            elif Stabilizer.valid_datatype(data):
                self._rep_data = self._initialize_stabilizer(data)
            elif Graph.valid_datatype(data):
                self._rep_data = self._initialize_graph(data)
            elif DensityMatrix.valid_datatype(data):
                raise ValueError(
                    "Data's type is correct for density matrix representation, but state size exceeds the "
                    "recommended size for density matrix representation"
                )
            else:
                raise ValueError(
                    "Data's type is invalid for initializing a QuantumState"
                )

        elif isinstance(rep_type, str):
            self._rep_data = self._initialize_representation(rep_type, data)

        else:
            raise ValueError("passed rep_type argument must be a String")

    @property
    def rep_data(self):
        """
        Representation data

        :return: representation data
        :rtype: DensityMatrix or Stabilizer or Graph
        """
        return self._rep_data

    @rep_data.setter
    def rep_data(self, data):
        self._rep_data = self._initialize_representation(self._rep_type, data)

    @property
    def rep_type(self):
        """
        Representation type

        :return: representation type
        :rtype: str
        """
        return self._rep_type

    @rep_type.setter
    def rep_type(self, new_rep_type):
        self._rep_type = self._get_rep_type_name(new_rep_type)

    def partial_trace(self, keep, dims):
        r"""
        Calculates the partial trace on all state representations which are currently defined

        :param keep:  An array of indices of the spaces to keep. For instance, if the space is
                    $A \times B \times C \times D$, and we want to trace out B and D, keep = [0,2]
        :type keep: list OR numpy.ndarray
        :param dims: An array of the dimensions of each space. For instance,
                    if the space is $A \times B \times C \times D$,
                    dims = [$dim_A$, $dim_B$, $dim_C$, $dim_D$]
        :type dims: list OR numpy.ndarray
        :return: nothing
        :rtype: None
        """

        if self._rep_type == "g":
            raise NotImplementedError(
                "Partial trace not yet implemented on graph state"
            )
        else:
            self._rep_data.partial_trace(keep, dims)

    def show(self, show=True, ax=None):
        """
        Plots the state representation using matplotlib formatting

        :param show: if True, the state representation is plotted. Otherwise, it is drawn but not plotted
        :type show: bool
        :param ax: axis/axes on which to plot the state representation
        :type ax: matplotlib.Axis
        :return: fig, ax (the figure and axes on which data was plotted)
        :rtype: matplotlib.Figure, matplotlib.Axis
        """
        # Use visualizers module
        if self._rep_type == "dm":
            fig, ax = self._rep_data.draw(show=show, style="heat")
        elif self._rep_type == "g":
            fig, ax = self._rep_data.draw(show=show, ax=ax, with_labels=True)
        elif self._rep_type == "s":
            raise NotImplementedError(
                "No visualization tool is implemented for stabilizer backend."
            )
        else:
            raise ValueError("Invalid representation type.")

        return fig, ax

    def _initialize_dm(self, data):
        """
        Initializes a density matrix based on the data

        :param data: either a graph or ndarray matrix
        :type data: Graph OR nx.Graph OR numpy.ndarray
        :raises AssertionError: if the density matrix being initialized does not have self.n_qubits
        """
        if isinstance(data, Graph) or isinstance(data, nx.Graph):
            dm = DensityMatrix.from_graph(data)
        else:
            dm = DensityMatrix(data)

        assert dm.data.shape[0] == dm.data.shape[1] == 2**self.n_qubits
        return dm

    def _initialize_graph(self, data):
        """
        Initializes a graph state based on the data

        :param data: data to construct the Graph representation
        :type data: networkX.Graph
        :raises AssertionError: if the graph being initialized does not have self.n_qubits
        """
        graph = Graph(data)

        assert graph.n_qubits == self.n_qubits, (
            f"Expected {self.n_qubits} qubits, " f"graph rep_type has {graph.n_qubits}"
        )
        return graph

    def _initialize_stabilizer(self, data):
        """
        Initializes a stabilizer state based on the data

        :param data: data to construct the stabilizer state representation
        :type data: Stabilizer or int or CliffordTableau or MixedStabilizer or list[(float, CliffordTableau)]
        """
        if not self.mixed:
            if isinstance(data, Stabilizer):
                return data
            else:
                stabilizer = Stabilizer(data)
        else:
            if isinstance(data, MixedStabilizer):
                return data
            else:
                stabilizer = MixedStabilizer(data)

        assert stabilizer.n_qubits == self.n_qubits, (
            f"Expected {self.n_qubits} qubits, "
            f"Stabilizer representation has {stabilizer.n_qubits}"
        )
        return stabilizer

    def _initialize_representation(self, rep_type, data):
        """
        Helper function to initialize any given representation

        :param rep_type: representation type to initialize
        :type rep_type: str
        :param data: data with which the rep_type should be initialized
        :type data: int OR Graph OR nx.Graph OR numpy.ndarray
        :raises ValueError: if rep_type is invalid
        """
        if rep_type in ("dm", "density matrix"):
            if self.n_qubits > DENSITY_MATRIX_QUBIT_THRESH:
                warnings.warn(
                    UserWarning(
                        "Density matrix is not recommended for a state of this size"
                    )
                )
            return self._initialize_dm(data)
        elif rep_type in ("g", "graph"):
            return self._initialize_graph(data)
        elif rep_type in ("s", "stab", "stabilizer"):
            return self._initialize_stabilizer(data)
        else:
            raise ValueError("Passed rep_type is invalid")

    def _density_to_graph(self, rep):
        """
        Helper function. Convert a density matrix representation to a graph representation.

        :param rep: density matrix representation
        :type rep: DensityMatrix
        :return: graph representation
        :rtype: Graph or MixedGraph
        """
        new_data = rc.density_to_graph(rep.data)
        if isinstance(new_data, list):
            new_rep = MixedGraph(new_data)
        else:
            new_rep = Graph(new_data)
        return new_rep

    def _density_to_stabilizer(self, rep):
        """
        Helper function. Convert a density matrix representation to a stabilizer representation

        :param rep: density matrix representation
        :type rep: DensityMatrix
        :return: stabilizer representation
        :rtype: Stabilizer or MixedStabilizer
        """
        new_data = rc.density_to_stabilizer(rep.data)
        # new_data is of type list[(float, StabilizerTableau)]
        if self.mixed:
            new_tab_list = []
            for p_i, s_i in new_data:
                new_tab_list.append((p_i, CliffordTableau(s_i)))

            new_rep = MixedStabilizer(new_tab_list)
        else:
            new_tableau = CliffordTableau(new_data[0][1])
            new_rep = Stabilizer(new_tableau)
        return new_rep

    def _stabilizer_to_density(self, rep):
        """
        Helper function. Convert a stabilizer representation to density matrix

        :param rep: initial representation
        :type rep: Stabilizer or MixedStabilizer
        :return: a density matrix
        :rtype: DensityMatrix
        """
        if self.mixed:
            data_list = []
            assert isinstance(rep, MixedStabilizer)
            for p_i, t_i in rep.mixture:
                data_list.append((p_i, t_i.to_stabilizer()))
            rho = rc.stabilizer_to_density(data_list)
        else:
            rho = rc.stabilizer_to_density(rep.data.to_stabilizer())
        return DensityMatrix(rho)

    def _stabilizer_to_graph(self, rep):
        """
        Convert a stabilizer representation to a graph representation

        :param rep: stabilizer representation
        :type rep: Stabilizer or MixedStabilizer
        :return: graph representation
        :rtype: Graph or MixedGraph
        """
        if self.mixed:
            assert isinstance(rep, MixedStabilizer)
            data_list = []
            for p_i, t_i in rep.mixture:
                data_list.append((p_i, t_i.to_stabilizer()))
            graph_list = rc.stabilizer_to_graph(data_list)
            return MixedGraph(graph_list)
        else:
            graph_list = rc.stabilizer_to_graph(rep.data)
            return Graph(graph_list[0][1])

    def _graph_to_density(self, rep):
        """
        Helper function. Convert a graph representation to a density matrix representation

        :param rep: graph representation
        :type rep: Graph or MixedGraph
        :return: density matrix representation
        :rtype: DensityMatrix
        """
        if self.mixed:
            assert isinstance(rep, MixedGraph)
            rho = 0
            for p_i, t_i in rep.mixture:
                rho += p_i * rc.graph_to_density(t_i.data)
        else:
            rho = rc.graph_to_density(rep.data)
        return DensityMatrix(rho)

    def _graph_to_stabilizer(self, rep):
        """
        Helper function. Convert a graph representation to a stabilizer representation

        :param rep: graph representation
        :type rep: Graph or MixedGraph
        :return: stabilizer representation
        :rtype: Stabilizer or MixedStabilizer
        """
        new_data = rc.graph_to_stabilizer(rep.data)
        # new_data is of type list[(float, StabilizerTableau)]
        if self.mixed:
            new_tab_list = []
            for p_i, s_i in new_data:
                new_tab_list.append((p_i, CliffordTableau(s_i)))

            new_rep = MixedStabilizer(new_tab_list)
        else:
            new_tableau = CliffordTableau(new_data[0][1])
            new_rep = Stabilizer(new_tableau)
        return new_rep

    def _identity_fun(self, rep):
        return rep

    @staticmethod
    def _get_rep_type_name(rep_type):
        if rep_type in ("s", "stab", "stabilizer"):
            return "s"
        elif rep_type in ("dm", "density matrix"):
            return "dm"
        elif rep_type in ("g", "graph"):
            return "g"
        elif rep_type is None:
            return None
        else:
            raise ValueError(
                f"QuantumState does not support the representation of type {rep_type}"
            )

    def convert_representation(self, new_rep_type):
        """
        Convert to a representation specified by new_rep_type

        :param new_rep_type: new representation type
        :type new_rep_type: str
        :return: nothing
        :rtype: None
        """
        rep_type = self._get_rep_type_name(new_rep_type)
        if rep_type is None:
            raise ValueError("Cannot convert representation to None type")

        conversion_dict = {
            ("dm", "g"): self._density_to_graph,
            ("dm", "s"): self._density_to_stabilizer,
            ("s", "dm"): self._stabilizer_to_density,
            ("s", "g"): self._stabilizer_to_graph,
            ("g", "dm"): self._graph_to_density,
            ("g", "s"): self._graph_to_stabilizer,
            ("dm", "dm"): self._identity_fun,
            ("s", "s"): self._identity_fun,
            ("g", "g"): self._identity_fun,
        }
        if self._rep_type != rep_type:
            if rep_type == "dm" and self.n_qubits > DENSITY_MATRIX_QUBIT_THRESH:
                warnings.warn(
                    UserWarning(
                        "Density matrix is not recommended for a state of this size"
                    )
                )

            tmp_data = self._rep_data
            conversion_func = conversion_dict[(self._rep_type, rep_type)]

            self._rep_data = conversion_func(tmp_data)
            self._rep_type = rep_type

    @classmethod
    def validate_data(cls, data):
        """
        Validate data type for input data

        :param data: input data
        :type data: int or np.ndarray or CliffordTableau or nx.Graph
        :return: True and the number of qubits if the data type is valid
        :rtype: bool, int
        """
        valid = True
        if isinstance(data, int):
            n_qubits = data
        elif isinstance(data, np.ndarray):
            assert (
                data.shape[0] == data.shape[1]
            ), "Input data is a matrix but it is not a square matrix."
            n_qubits = int(np.log2(data.shape[0]))
        elif isinstance(data, CliffordTableau):
            n_qubits = data.n_qubits
        elif isinstance(data, nx.Graph):
            n_qubits = data.number_of_nodes()
        elif isinstance(data, list):
            if isinstance(data[0][1], CliffordTableau) or isinstance(
                data[0][1], nx.Graph
            ):
                n_qubits = data[0][1].n_qubits
            else:
                return False, None
        else:
            valid = False
            n_qubits = None
        return valid, n_qubits

    def copy(self):
        """
        Make a copy of this QuantumState object

        :return: a copy of the current QuantumState object
        :rtype: QuantumState
        """
        return copy.deepcopy(self)

class DensityMatrix(StateRepresentationBase):
    """
    Density matrix of a state
    """

    def __init__(self, data, normalized=True, *args, **kwargs):
        """
        Construct a DensityMatrix object from a numpy.ndarray or from the number of qubits. If an integer is specified,
        then the state is initialized as a product state of :math:`|0\\rangle` with the given number of qubits.

        :param data: density matrix or the number of qubits
        :type data: numpy.ndarray or int
        :param normalized: whether the state is normalized
        :type normalized: bool
        :return: nothing
        :rtype: None
        """

        if isinstance(data, np.ndarray):
            if not dmf.is_psd(data):
                # check if state_data is positive semi-definite
                raise ValueError("The input matrix is not a valid density matrix")

            if normalized and not np.equal(np.trace(data), 1):
                data = data / np.trace(data)
        elif isinstance(data, int):
            # initialize as a tensor product of |0> state
            data = dmf.create_n_product_state(data, dmf.state_ketz0())
        else:
            raise TypeError("Input must be a numpy.ndarray or an integer")

        super().__init__(data, *args, **kwargs)

    @classmethod
    def from_graph(cls, graph):
        """
        Builds a density matrix representation from a graph (either nx.Graph or a Graph representation)

        :param graph: the graph from which we will build a density matrix
        :type graph: networkx.Graph OR Graph
        :raises TypeError: if the input graph is neither nx.Graph or Graph
        :return: a DensityMatrix representation with the data contained by graph
        :rtype: DensityMatrix
        """
        if isinstance(graph, Graph):
            return cls(graph_to_density(graph.data))
        else:
            return cls(graph_to_density(graph))

    @classmethod
    def valid_datatype(cls, data):
        return isinstance(data, (int, np.ndarray))

    @property
    def trace(self):
        """
        Return the trace of the state

        :return: the trace of the state
        :rtype: float
        """
        return np.trace(self.data)

    @property
    def normalized(self):
        """
        Return whether the state is normalized, that is, trace is 1

        :return: whether the state is normalized
        :rtype: bool
        """
        return np.allclose(self.trace, 1.0)

    def apply_unitary(self, unitary):
        """
        Apply a unitary to the state.
        Assumes the dimensions match; Otherwise, raise ValueError

        :param unitary: unitary matrix to apply
        :type unitary: numpy.ndarray
        :raises ValueError: if the density matrix of the state has a different size from the unitary gate to be applied
        :return: nothing
        :rtype: None
        """
        if self._data.shape == unitary.shape:
            self._data = unitary @ self._data @ np.transpose(np.conjugate(unitary))
            # to avoid small numerical error that causes the state non-Hermitian
            self._data = dmf.hermitianize(self._data)
        else:
            raise ValueError(
                "The density matrix of the state has a different size from the unitary gate to be applied."
            )

    def apply_channel(self, kraus_ops):
        """
        Apply a quantum channel on the state where the quantum channel is described by Kraus representation.
        Assumes the dimensions match; Otherwise, raise ValueError

        :param kraus_ops: a list of Kraus operators of the channel
        :type kraus_ops: list[numpy.ndarray]
        :raises ValueError: if Kraus operators have wrong dimensions.
        :return: nothing
        :rtype: None
        """
        tmp_state = 0
        if len(kraus_ops) == 0:
            return
        if self._data.shape[0] == kraus_ops[0].shape[1]:
            for i in range(len(kraus_ops)):
                tmp_state = tmp_state + kraus_ops[i] @ self._data @ np.conjugate(
                    kraus_ops[i].T
                )
            # to avoid small numerical error that causes the state non-Hermitian
            self._data = dmf.hermitianize(tmp_state)
        else:
            raise ValueError("Kraus operators have wrong dimensions.")

    def apply_measurement(self, projectors, measurement_determinism="probabilistic"):
        """
        Apply a measurement, either deterministically (with a certain outcome) or probabilistically

        :param projectors: a list of projective measurements in the computational basis
        :type projectors: list[numpy.ndarray]
        :param measurement_determinism: if "probabilistic", measurement results are probabilistically selected
                                    if 1, measurement results default to 1 unless the probability of measuring p(1) = 0
                                    if 0, measurement results default to 0 unless the probability of measuring p(0) = 0
        :type measurement_determinism: str/int
        :return: the measurement outcome
        :rtype: int
        """

        if self._data.shape == projectors[0].shape:
            probs = []
            for m in projectors:
                prob = np.real(np.trace(self._data @ m))
                if prob < 0:
                    prob = 0
                probs.append(prob)
            probs = np.array(probs)
            if measurement_determinism == "probabilistic":
                outcome = numpy.random.choice([0, 1], p=probs / np.sum(probs))
            elif measurement_determinism == 1:
                if probs[1] > 0:
                    outcome = 1
                else:
                    outcome = 0

            elif measurement_determinism == 0:
                if probs[1] < 1:
                    outcome = 0
                else:
                    outcome = 1
            else:
                raise ValueError(
                    f'measurement_determinism parameter must be "probabilistic", 0, or 1'
                )

            m, norm = projectors[outcome], probs[outcome]

            # this assumes that the projector, m, has the properties: m = sqrt(m) and m = m.dag()
            self._data = (m @ self._data @ np.transpose(np.conjugate(m))) / norm

        else:
            raise ValueError(
                "The density matrix of the state has a different size from the POVM elements."
            )
        return outcome

    def apply_measurement_controlled_gate(
        self, projectors, target_gate, measurement_determinism=1
    ):
        """
        Apply a measurement, either deterministically (with a certain outcome) or probabilistically
        and conditioned on the measurement outcome, apply the target_gate

        :param projectors: a list of projective measurements in the computational basis
        :type projectors: list[numpy.ndarray]
        :param target_gate: the gate to be applied if the measurement outcome is 1
        :type target_gate: numpy.ndarray
        :param measurement_determinism: if "probabilistic", measurement results are probabilistically selected
                                    if 1, measurement results default to 1 unless the probability of measuring p(1) = 0
                                    if 0, measurement results default to 0 unless the probability of measuring p(0) = 0
        :type measurement_determinism: str/int
        :raises AssertionError: if target_gate has different dimensions from the density matrix of the state
        :return: the measurement outcome
        :rtype: int
        """
        assert self._data.shape == target_gate.shape
        outcome = self.apply_measurement(projectors, measurement_determinism)
        if outcome == 1:
            self.apply_unitary(target_gate)
        return outcome

    def partial_trace(self, keep, dims):
        """
        Take the partial trace of the state

        :param keep:  An array of indices of the spaces to keep. For instance, if the space is
                    :math:`A \\times B \\times C \\times D` and we want to trace out B and D, keep = [0,2]
        :type keep: list OR numpy.ndarray
        :param dims: An array of the dimensions of each space. For instance,
                    if the space is :math:`A \\times B \\times C \\times D`,
                    dims = [dim_A, dim_B, dim_C, dim_D]
        :type dims: list OR numpy.ndarray
        :return:
        :rtype:
        """
        self.data = dmf.partial_trace(self.data, keep, dims)

    def draw(self, style="bar", show=True):
        """
        Draw a bar graph or heatmap of the DensityMatrix representation data

        :param style: 'bar' for bar plot, 'heat' for heatmap
        :type style: str
        :param show: if True, show the density matrix plot. Otherwise, draw the density matrix plot but do not show
        :type show: bool
        :return: fig, axes on which the state is drawn
        :rtype: matplotlib.Figure, matplotlib.Axes

        """
        # TODO: add a "ax" parameter to match the other viewing utils
        if style == "bar":
            fig, axs = density_matrix_bars(self.data)
        else:
            fig, axs = density_matrix_heatmap(self.data)

        if show:
            plt.show()

        return fig, axs

    def __eq__(self, other):
        """
        Compare two DensityMatrix objects and return True if the underlying density matrices are equal
        (up to precision)

        :param other: another DensityMatrix object
        :type other: DensityMatrix
        :return: True if they are equal; False otherwise
        :rtype: bool
        """
        return np.allclose(self._data, other.data)