Postselection statistics in pennylane

Hi everyone, I’m new to pennylane and I’m trying to apply post-selection to my VQE and get the success rate of the post selection. I’m doing it as the following with qml.Tracker but I’m not sure whether it’s correct. Can anybody help me to verify this is the correct way to collect the ratio of postselected/all runs?

# Define the VQE routine with ancilla qubit control
def perform_vqe_with_ancilla_control(num_system_qubits, depth):
    # Initialize device
    dev = qml.device('default.qubit', wires=num_qubits)
    
    # Initialize the Hamiltonian
    H = get_hamiltonian(num_rows, num_cols)
    
    # Define the cost function
 
    @qml.qnode(dev)
    def cost_fn(params):
        ansatz_with_ancilla(params, num_system_qubits, depth)
        qml.measure(num_system_qubits, postselect=0, reset=True) 
        return qml.expval(H)

    # Initialize the parameters
    num_params = depth * (2 * (num_system_qubits - 1) + num_system_qubits * 3)  # Including parameters for the ancilla rotation
    params = np.random.random(num_params, requires_grad=True)
    
    # Initialize the optimizer
    optimizer = NesterovMomentumOptimizer(stepsize=0.005)
    # Initialize counters for success rate
  
    # Perform optimization
    total_run=0
    max_iter = 200
    print(cost_fn(params,shots=max_iter))
    for n in range(max_iter):
        with qml.Tracker(dev) as tracker:  # track the number of executions
            params, energy = optimizer.step_and_cost(cost_fn, params)
            if n%40==0:
                print(f"Iteration = {n}, Energy = {energy}")
        total_run+= tracker.totals['executions']
    print(f"Success Rate:{max_iter/total_run}")
    return energy, params

Hi @Haxxx_Akxxx ,

Thank you for your question!

It looks like you want to do something like this:

import pennylane as qml

dev = qml.device("default.qubit")

@qml.qnode(dev)
def f():
    qml.Hadamard(0)
    m = qml.measure(0, postselect=0)
    ...
    return qml.probs(op=m)

Right now, PennyLane would give you:

>>> f()
tensor([1., 0.], requires_grad=True)

This is expected, since you are doing the postselection so the probability should be one.

For running your algorithm however, it looks like you’ll need a slightly different functionality that PennyLane currently doesn’t have – but it’s on our radar! If we had a separate qml.postselect() operation then it’d look like:

import pennylane as qml

dev = qml.device("default.qubit")

@qml.qnode(dev)
def f():
    qml.Hadamard(0)
    m = qml.measure(0)
    qml.postselect(m, 0)
    ...
    return qml.probs(op=m)

>>> f()
tensor([0.5, 0.5], requires_grad=True)

Is this what you were looking for?
I hope this isn’t too confusing.
Let me know if you have any additional questions!