Supporting counts from raw samples

[Jaybsoni]

Feature details

Given a quantum circuit we can sample the final state it prepares in the computation basis. For example:

dev = qml.device('default.qubit', wires=3, shots=10)

@qml.qnode(dev)
def my_circ():
    qml.Hadamard(wires=0)
    qml.CNOT(wires=[0,1])
    qml.PauliX(wires=2) 

    return qml.sample()

>>> my_samples = my_circ()
>>> print(my_samples)
[[1 1 1]
 [0 0 1]
 [1 1 1]
 [0 0 1]
 [1 1 1]
 [1 1 1]
 [0 0 1]
 [1 1 1]
 [1 1 1]
 [0 0 1]]
>>>

Samples produces an array containing the results of sampling the circuit. Its shape is (num_shots, num_qubits) in this case it would be (10, 3). In general it could be the case that someone asks for a very larger number of shots on a system with many qubits. In this case, it would be difficult to work with this full array of samples.

To address this we would like to support the ability to pass in a counts=True option to qml.sample such that the QNode upon evaluation returns a dictionary containing each unique state and the number of times it was sampled.

dev = qml.device('default.qubit', wires=3, shots=10)

@qml.qnode(dev)
def my_circ():
    qml.Hadamard(wires=0)
    qml.CNOT(wires=[0,1])
    qml.PauliX(wires=2) 

    return qml.sample(counts=True)

>>> my_counts = my_circ()
>>> print(my_counts)
{'111':6, '001':4}

Implementation

How important would you say this feature is?

1: Not important. Would be nice to have.

Additional information

Once the counts are available, we can also return a vector representation of the counts:

my_counts = my_circ()
wires = 3
count_vector = np.zeros(2 ** wires)
for k, v in my_counts.items():
    idx = int(k, 2)
    count_vector[idx] = v

>>> print(count_vector)
array([0., 4., 0., 0., 0., 0., 0., 6.])

In this case the list would have 2^n elements where n is the number of qubits. The element at index i would correspond to the number of samples for the ith basis state.

[antalszava]

This change will likely require engineering not just qml.sample, but also QubitDevice and the scope of the change may become non-trivial.

[antalszava]

This change could include creating a new Counts observable return type (in the measurements.py file). The logic for computing statistics in the QubitDevice base class could then be extended such that the Counts observable return type is handled too. The existing machinery for the Sample return type could be utilized.

[theodotk]

Hello PennyLane team,

I submitted PR 2686 that adds this feature. I based my implementation on editing the sample function in _qubit_device.py, though it was possible to base it on estimate_probability function (or to create something new that seems superfluous to me). Choice of editing the sample function seemed more natural for this task, but please tell me if that's not the best approach.

As for the count_vector idea, it seems to be already realized as qml.probs() (up to the number of shots factor), so I didn't change anything in this regard.

@antalszava , @Jaybsoni , could you please review the PR?

Thanks!

[antalszava]

Closing with #2686 merged.