# Source code for mitiq.pec.representations.optimal

# Copyright (C) Unitary Fund
#
# LICENSE file in the root directory of this source tree.

"""Functions for finding optimal representations given a noisy basis."""

import functools
from typing import List, Optional, cast

import numpy as np
import numpy.typing as npt
from cirq import kraus
from scipy.optimize import LinearConstraint, minimize

from mitiq import QPROGRAM
from mitiq.interface import convert_to_mitiq
from mitiq.pec.channels import kraus_to_super
from mitiq.pec.types import NoisyOperation, OperationRepresentation
from mitiq.utils import matrix_to_vector

[docs]
def minimize_one_norm(
ideal_matrix: npt.NDArray[np.complex64],
basis_matrices: List[npt.NDArray[np.complex64]],
tol: float = 1.0e-8,
initial_guess: Optional[npt.NDArray[np.float64]] = None,
) -> npt.NDArray[np.float64]:
r"""
Returns the list of real coefficients :math:[x_0, x_1, \dots],
which minimizes :math:\sum_j |x_j| with the contraint that
the following representation of the input ideal_matrix holds:

.. math::
\text{ideal_matrix} = x_0 A_0 + x_1 A_1 + \cdots

where :math:\{A_j\} are the basis matrices, i.e., the elements of
the input basis_matrices.

This function can be used to compute the optimal representation
of an ideal superoperator (or Choi state) as a linear
combination of real noisy superoperators (or Choi states).

Args:
ideal_matrix: The ideal matrix to represent.
basis_matrices: The list of basis matrices.
tol: The error tolerance for each matrix element
of the represented matrix.
initial_guess: Optional initial guess for the coefficients
:math:[x_0, x_1, \dots].

Returns:
The list of optimal coefficients :math:[x_0, x_1, \dots].
"""

# Map complex matrices to extended real matrices
ideal_matrix_real = np.hstack(
(np.real(ideal_matrix), np.imag(ideal_matrix))
)
basis_matrices_real = [
np.hstack((np.real(mat), np.imag(mat))) for mat in basis_matrices
]

# Express the representation constraint written in the docstring in the
# form of a matrix multiplication applied to the x vector: A @ x == b.
matrix_a = np.array(
[
matrix_to_vector(mat)  # type: ignore[arg-type]
for mat in basis_matrices_real
]
).T
array_b = matrix_to_vector(ideal_matrix_real)  # type: ignore[arg-type]

constraint = LinearConstraint(matrix_a, lb=array_b - tol, ub=array_b + tol)

def one_norm(x: npt.NDArray[np.complex64]) -> float:
return cast(float, np.linalg.norm(x, 1))

if initial_guess is None:
initial_guess = np.zeros(len(basis_matrices))

result = minimize(one_norm, x0=initial_guess, constraints=constraint)

if not result.success:
raise RuntimeError("The search for an optimal representation failed.")

return result.x

[docs]
def find_optimal_representation(
ideal_operation: QPROGRAM,
noisy_operations: List[NoisyOperation],
tol: float = 1.0e-8,
initial_guess: Optional[npt.NDArray[np.float64]] = None,
is_qubit_dependent: bool = True,
) -> OperationRepresentation:
r"""Returns the OperationRepresentation of the input ideal operation
which minimizes the one-norm of the associated quasi-probability
distribution.

More precisely, it solve the following optimization problem:

.. math::
\min_{\eta_\alpha} \left\{\sum_\alpha |\eta_\alpha| \, : \,
\mathcal G = \sum_\alpha \eta_\alpha \mathcal O_\alpha \right\}

where :math:\{\mathcal O_j\} is the input basis of noisy operations,
and :math:\mathcal{G} is the ideal operation to be represented.

Args:
ideal_operation: The ideal operation to represent.
noisy_operations: The basis in which the ideal_operation
should be represented. Must be a list of NoisyOperation objects
which are initialized with a numerical superoperator matrix.
tol: The error tolerance for each matrix element
of the represented operation.
initial_guess: Optional initial guess for the coefficients
:math:\{ \eta_\alpha \}.
is_qubit_dependent: If True, the representation corresponds to the
operation on the specific qubits defined in ideal_operation.
If False, the representation is valid for the same gate even if
acting on different qubits from those specified in
ideal_operation.

Returns: The optimal OperationRepresentation.
"""
ideal_cirq_circuit, _ = convert_to_mitiq(ideal_operation)

ops = list(ideal_cirq_circuit.all_operations())
super_ops = [
kraus_to_super(cast(List[npt.NDArray[np.complex64]], kraus(op)))
for op in ops
]

# Compute super operator of the full ideal_circuit
ideal_matrix = functools.reduce(lambda a, b: a @ b, super_ops)

try:
basis_matrices = [
noisy_op.channel_matrix for noisy_op in noisy_operations
]
except ValueError as err:
if str(err) == "The channel matrix is unknown.":
raise ValueError(
"The input noisy_basis should contain NoisyOperation objects"
" which are initialized with a numerical superoperator matrix."
)
else:
raise err  # pragma no cover

# Run numerical optimization problem
quasi_prob_dist = minimize_one_norm(
ideal_matrix,
basis_matrices,
tol=tol,
initial_guess=initial_guess,
)

return OperationRepresentation(
ideal_operation,
noisy_operations,
quasi_prob_dist.tolist(),
is_qubit_dependent,
)