Feature Entanglement Fidelity Speedup and Bugfix #263
Conversation
Improves the performance of the entanglement fidelity calculation for the case where the process matrices are TP and the target it unitary by using the magic of einsum to only calculate the diagonal elements of a matrix we want to calculate the trace of.
The same trick we used to speed up the entanglement fidelity calculation can be used to speed up the unitarity calculation.
pygsti/tools/optools.py
Outdated
| @@ -427,7 +427,7 @@ def jtracedist(a, b, mx_basis='pp'): # Jamiolkowski trace distance: Tr(|J(a)-J | |||
| return tracedist(JA, JB) | |||
|
|
|||
|
|
|||
| def entanglement_fidelity(a, b, mx_basis='pp'): | |||
| def entanglement_fidelity(a, b, mx_basis='pp', is_tp_flag=None, is_unitary_flag=None): | |||
There was a problem hiding this comment.
I'd rather have these kwargs be is_tp and is_unitary and we rename the helper functions below with a _fn suffix.
There was a problem hiding this comment.
Gotcha, I renamed the tp check function as suggested, but opted to inline the unitarity check (it was only called once in any case). Kwargs have also been updated as per your suggestion.
pygsti/tools/optools.py
Outdated
| @@ -450,20 +450,51 @@ def entanglement_fidelity(a, b, mx_basis='pp'): | |||
| The basis of the matrices. Allowed values are Matrix-unit (std), | |||
| Gell-Mann (gm), Pauli-product (pp), and Qutrit (qt) | |||
| (or a custom basis object). | |||
|
|
|||
| is_tp_flag : bool, optional (default None) | |||
There was a problem hiding this comment.
I originally didn't like the dual purpose of one variable to denote (i) whether to do the check and (ii) the value in lieu of the check, but I understand why you've done it this way. I think I'll be OK with it, so long as this docstring is a little more clear. Maybe something like:
Flag indicating both matrices are TP. If None (the default), an explicit check is performed. If True/False, the check is skipped and provided value is used (faster, but should only be used when the user is certain this is true apriori).
IDK exactly, but I want it to be clear that None is not just a dummy default - it actually changes the behavior of the function.
There was a problem hiding this comment.
Great suggestion, wording was definitely unclear. I more or less verbatim copied your language directly since I thought it was very clear and concise.
pygsti/tools/optools.py
Outdated
| #Use einsum black magic to only calculate the diagonal elements | ||
| #if the basis is either pp or qm we know the elements are real-valued, so we | ||
| #don't need to take the conjugate | ||
| if mx_basis=='pp' or mx_basis=='qm': |
There was a problem hiding this comment.
Is this line correct? From the comment, which jives with my understanding, .conjugate() can be dropped when the basis is 'pp' or 'qm', not added. I think this may be a "sign" error ;)
There was a problem hiding this comment.
Also, while we're using einsum magic, we could swap ji to ij and drop the .T in the three einsum lines, e.g. _np.einsum('ij,ij->',a, b). This may be even slightly faster?
There was a problem hiding this comment.
Thanks for catching that, they were indeed switched. Incidentally, I also had a typo in the basis label for the gell-mann matrices I hadn't noticed.
As for dropping the transpose, I just gave that a shot and found that it gives a 7% speedup for the 2-qubit case, but is within error bars of the version with the transpose for the 3-qubit case. Since numpy implements transposes using views, I'm guessing that the overhead of constructing that view is still measurable for 2-qubits, but is negligible for 3. Given that I'm going to opt to keep the transpose in for readability.
Fixes a typo where I had the inclusion of the conjugates switched and another typo where I mislabeled the basis key for the gell-man matrices.
Renames the kwargs for manually specifying if the input matrices are unitary and TP. Reworks the docstrings to make clear that the default value of None for the flags means that a built-in check is performed. Also in-lines the unitarity check (only used once and doesn't actually enhance readability).
|
RE: Unit tests. I'm not sure if there are currently, if not I'll add some. |
|
There is one unit test that covers the general use case. It would be nice to add one or two more that cover the other code paths. They could be pretty simple: generate a test case that passes the TP/unitary condition, and compare the results of that against the function called with explicit is_tp/is_unitary=False to force it through the (known good) general use case. Similar thing could be done with the mx_basis by doing a change_basis to non-pp/gm basis (maybe std?) and comparing to the default pp version? Just some thoughts. |
Adds a check for the special case when the inputs are TP and the target is unitary.
|
Got it, I've added a unit test for the TP/unitary case and can add some additional ones to cover changing the mx_basis and comparing the results for the special case to the jamiolkowski isomorphism. Thoughts on adding the new kwargs to functions which call entanglement_fidelity as a sub-routine, to pass through? Skimming through that includes: average_gate_fidelity, average_gate_infidelity, entanglement_infidelity and gateset_infidelity (POVM fidelity does too, but the POVM "gates" are most certainly not TP or unitary). |
Adds additional entanglement fidelity unit tests.
|
Added additional unit tests for all of the test cases I could think of. |
|
This looks good to me. As for plumbing through the options, I think that sounds like a good idea but not essential. I'll leave the decision up to you - now that I've approved, I think you would be able to merge now if you didn't feel like putting in that effort, or you can add the extra plumbing and merge afterwards. |
Plumbs the new entanglement fidelity kwargs to functions which call it as a subroutine.
Some requests were made to look into making the calculation of the entanglement fidelity faster. This version speeds up the calculation when using TP process matrices, and when the target is unitary. Same formula as before for this case, but this uses the magic that is Numpy's einsum to calculate the trace by only constructing the diagonal of the requisite matrix product. In practice I found this to be 130-280X faster than the previous implementation, depending on the number of qubits.
I have also added optional flags that allow a user to manually specify whether the inputs are TP and unitary to short circuit the built in tests, which add time since the unitarity check at least involves doing a full matrix multiplication. Code still falls back to the built in tests when these flags aren't specified.
Also adds a minor bugfix to the TP check code to fix one of the ranges. This typo was causing the TP check to always fail which means that the fast implementation for the special case described above was never invoked and we were always falling back to the use of the Jamiolkowski isomorphism (which is orders of magnitude slower).
As a BOGO special I also used the same einsum trick to speedup the unitarity calculation. Not requested, but I noticed it'd apply there too.