How to plot cost function of maximization of purity by using qml.SpecialUnitary

I would like to plot the optimization of a cost of maximum purity of density matrix. Here are my codes

dev = qml.device("default.mixed", wires = n)
def circuit(rho,thetas):
  qml.QubitDensityMatrix(rho, wires = wireList)
  qml.adjoint(qml.SpecialUnitary(thetas, wires=wireList))
  qml.SpecialUnitary(thetas, wires = wireList)
  return qml.density_matrix([0])
def purity(thetas):
    Initializes a density matrix and takes its purity.
    qml.QubitDensityMatrix(circuit(rho,thetas), wires=0)
    return qml.purity(0)

# def fidi(thetas,pure_density_matrix):
#   '''
#   Initializes a the reference state and evaluates its purity of the desired state
#   '''
#   return,qml.QubitDensityMatrix(circuit(rho,thetas), wires=0)) 

def cost(thetas):
    Computes linear entropy from purity.
    return 1 - purity(thetas)

I would like to maximize the purity of denstiy matrix so minimizing the cost as follows:

thetas = jnp.array(thetas)
minimize(purity, thetas, method='BFGS', jac=jax.grad(cost, argnums=0))

How to show this in the plot that cost function minimizes.


Just want to make sure I understand your question; do you want to plot the value of the cost function as a function of the iteration / step number? If so, you can do so with matplotlib :slight_smile: : Getting started — Matplotlib 3.8.3 documentation

1 Like