Add support for calculating the trace distance

trbromley commented 1 year ago

Feature details

The trace distance is one measure of distance between two quantum states rho and sigma. This is a commonly-calculated quantity in quantum information science and will be useful to add to PennyLane. Specifically, we'd like to implement:

Given two QNodes and two sets of parameters to call them with, a user can evaluate and differentiate the trace distance between the two QNodes.
Given two states rho and sigma, represented as a numeric array either as a pure or mixed state, a user can calculate and differentiate the trace distance between the two states.

Implementation

QNode-level interface:

qml.qinfo.trace_distance(qnode_1, qnode_2, wires0=[0], wires1=[2])

Assume each QNode returns a state
Wires arguments can be used to specify a subset of the wires in each qnode
Require number of wires to be the same in each case

qml.math-level interface:

qml.math.trace_distance(rho, sigma, check_state=False, c_dtype='complex128')

rho can be pure or mixed
sigma can be pure or mixed
check_state checks if each state is valid (assume False)
c_dtype is the complex floating point precision

How important would you say this feature is?

1: Not important. Would be nice to have.

Additional information

Check out the relative entropy functionality in PennyLane for an example of how this can be done.

tnemoz commented 1 year ago

Hi! I'd like to work on this as part of the UnitaryHack Hackathon. Should I ask for this issue to be assigned to me, or should I directly fork the repository and issue a PR ASAP? Thanks!

trbromley commented 1 year ago

Hi @tnemoz, welcome! You can go ahead and open a PR and we'll provide feedback once the PR is open. Once a PR is merged, the author is assigned to this issue, which let's the UnitaryHack team know that a bounty has been earned.

You can check out our development guide or just post here if you have any questions. Good luck!

tnemoz commented 1 year ago

Would it make sense to allow the user to indicate that both states are pure, if they're given as density matrices? Something along the lines of:

def trace_distance(rho, sigma, check_state=False, c_dtype='complex128', both_pure=False):
    # Check states
    ...

    if rho.shape == (rho.shape[0], ) and sigma.shape == (sigma.shape[0], ):
        # Compute for pure states as statevectors
        ...
        return res

    elif both_pure:
        return ...

    trace_distance = ...
    return trace_distance

The rationale behind this is that computing the trace distance for Statevectors is way faster, and for pure states given as density matrices, it's still faster than the general case. This could thus save non-negligible computation resources. What do you think about it? WOuld that make sense?

eddddddy commented 1 year ago

Hi @tnemoz, we actually want to move away from having the qml.math functions accept both a state vector and a density matrix. This is due to complications in the dispatch logic when we consider a batch dimension. A rework of this is currently in progress.

For this issue, it suffices to provide an implementation of qml.math.trace_distance that accepts only density matrices, contrary to what @trbromley wrote in the original issue. If the demand comes later (and if it's convenient), it might be possible to implement another more efficient function, e.g. qml.math.trace_distance_pure, that accepts only state vectors. This is not needed here.

tnemoz commented 1 year ago

Hi @tnemoz, we actually want to move away from having the qml.math functions accept both a state vector and a density matrix. This is due to complications in the dispatch logic when we consider a batch dimension. A rework of this is currently in progress.

For this issue, it suffices to provide an implementation of qml.math.trace_distance that accepts only density matrices, contrary to what @trbromley wrote in the original issue. If the demand comes later (and if it's convenient), it might be possible to implement another more efficient function, e.g. qml.math.trace_distance_pure, that accepts only state vectors. This is not needed here.

So if I understand correctly, the function will be given either:

two density matrices (with size (2**N, 2**N) each), in which case it returns the trace distance between these
two batches of the same size batch_dim, in which case both inputs should have shape (batch_dim, 2**N, 2**N), in which case the function will return a (batch_dim,) array containing the trace distances between each pair
one density matrix of size (2**N, 2**N) and one batch of density matrices of the same size, in which case the function returns a (batch_dim,) array containing the tarce distance between each density matrix of the batch and the other one

Is that correct? I mostly infered these from the new _check_density_matrix function, but I could be wrong about the way the function will be used

eddddddy commented 1 year ago

Yep, everything you say is correct.

trbromley commented 1 year ago

Thanks @tnemoz, congrats on winning the bounty!! :100: :tada:

tnemoz commented 1 year ago

Thank you very much!

trbromley commented 1 year ago

Hi @tnemoz!

As part of our 0.31 release, we're planning to promote UnitaryHack contributions to PennyLane on social media. This is completely optional, but would you like to be mentioned in our social media posts? Specifically, we'd need either your Twitter profile, your LinkedIn profile, or your name.

tnemoz commented 1 year ago

Hi @trbromley!

Sure why not! My Twitter is @TNemoz and my LinkedIn profile can be found here. I'm fine with both if that's possible, but would probably prefer Twitter if I had to choose one.

Thanks for the opportunity!

PennyLaneAI / pennylane