Get number of calls to np.dot down to at most 1000 (replace with ``@``)

[rileyjmurray]

Meta-comment

It turns out that doing this automatically isn't practical. We need to do the replacements manually. I think the best way to handle this is to keep the title as a moving target. Right now there between 1100 and 1200 calls to np.dot, and we can say we've made material progress if we can get that down to 1000 calls.

Original comment

Numpy's dot function executes matrix-matrix multiplication. There are many places where we use this function in pyGSTi. Python 3.5 introduced the @ operator for matrix multiplication. Since then the preferred way to express C = np.dot(A, B) is C = A @ B. Using @ has a big advantage of working with matrix/array datatypes from libraries other than numpy (like PyTorch Tensors or Dask Arrays). In view of this, I suggest we replace calls to np.dot with an appropriate use of @.

How do @sandialabs/pygsti-maintainers feel about this?

[sserita]

I'm generally for replacing np.dot with @ in places where we mean matrix multiplication. That sort of modernization has been a low priority but would certainly be worthwhile IMO.

[eendebakpt]

There is a slight performance penalty for @:

In [167]: %timeit _np.dot(L, R)
1.49 µs ± 65 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

In [168]: %timeit L @ R
1.84 µs ± 137 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

(L and R dimension (16,16) from one of the inner loops in 2-qubit GST)

(maybe the overhead is due to @ being able to dispatch to libraries like dask and pytorch, I did not check)

[coreyostrove]

I suspect that any performance differences are likely going to be system configuration dependent and hard to predict. I did a quick check on my machine and found that the two were within error bars of each other.

import numpy as np
from numpy.random import default_rng
rng = default_rng()
L = rng.random((16,16))
R = rng.random((16,16))

%%timeit
L@R

1.67 µs ± 81.7 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

%%timeit
np.dot(L,R)

1.61 µs ± 16.1 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

That said, even if it had been the case that I'd found a performance hit on my machine, I think it would still be generally worth it for the sake of readability (if we spot any particularly hot bits of the codebase where this difference could add up we can always selectively keep dot there).

[rileyjmurray]

Hi all! After taking a stab at making this replacement in an automated way, I've come to the (tentative) conclusion that this just isn't practical. Calls to np.dot are too nested and/or heterogeneous to write regular expressions that are "sane" and also comprehensive across pyGSTi's codebase. See below for an overview of my findings and a proposed course of action.

CC: @coreyostrove, @sserita, @enielse.

A complicated attempt that didn't get me far

PR #433 shows a branch where I tried to automate things. The branch was made off my branch from #432 instead of from develop, so most of the changes GitHub mentions aren't worth looking at.

Instead, just look at this script. It contains code to generate the following regex pattern

'_np\\.dot\\((?P<left>[ ]*[A-Za-z0-9]+(\\.T)?(\\.conj\\(\\))?[ ]*),[ ]*(?P<right>[A-Za-z0-9]+(\\.T)?(\\.conj\\(\\))?[ ]*)\\)'

along with the substitution function

    def simple_replacer(match):
        return f"{match.group('left')} @ {match.group('right')}"

I can use this with Python's regex.sub to correctly handle the following test cases

    print('\nInstances where we want to replace ... ')
    out = npdotdot.sub(simple_replacer, '_np.dot(aAa, B11), nice to see you. _np.dot(aAa, B12)'); print(out)
    out = npdotdot.sub(simple_replacer, '_np.dot(aAa, B11), nice to see you. _np.dot( aAa, B12 )'); print(out)
    out = npdotdot.sub(simple_replacer, '_np.dot(aAa, B11), nice to see you. _np.dot(aAa, B12.T)'); print(out)
    out = npdotdot.sub(simple_replacer, '_np.dot(aAa, B11), nice to see you. _np.dot(aAa.T, B12)'); print(out)

    out = npdotdot.sub(simple_replacer, '_np.dot(aAa.T, B11), nice to see you. _np.dot(aAa, B12)'); print(out)
    out = npdotdot.sub(simple_replacer, '_np.dot(aAa.T, B11.T), nice to see you. _np.dot(aAa, B12.T)'); print(out)
    out = npdotdot.sub(simple_replacer, '_np.dot(aAa.conj(), B11), nice to see you. _np.dot(aAa.T, B12)'); print(out)

    print('\nInstances where we do NOT expect to replace ...')
    out = npdotdot.sub(simple_replacer, '_np.dot(aAa.I, B11), nice to see you. _np.dot(aAa, B12.J)'); print(out)
    out = npdotdot.sub(simple_replacer, '_np.dot(aAa.H, B11.K), nice to see you. _np.dot(aAa.T, B12.A)'); print(out)

which prints out

Instances where we want to replace ... 
aAa @ B11, nice to see you. aAa @ B12
aAa @ B11, nice to see you.  aAa @ B12 
aAa @ B11, nice to see you. aAa @ B12.T
aAa @ B11, nice to see you. aAa.T @ B12
aAa.T @ B11, nice to see you. aAa @ B12
aAa.T @ B11.T, nice to see you. aAa @ B12.T
aAa.conj() @ B11, nice to see you. aAa.T @ B12

Instances where we do NOT expect to replace ...
_np.dot(aAa.I, B11), nice to see you. _np.dot(aAa, B12.J)
_np.dot(aAa.H, B11.K), nice to see you. _np.dot(aAa.T, B12.A)

But that only satisfied ~150 of the 725 matches in the codebase! You can see the changes for yourself by looking at this commit on the relevant branch.

Other considerations

On top of the difficulty in writing regex that's comprehensive, it's impossible to fully automate this in safe way, since there are a few instances where np.dot is used on arrays with n.dim > 2 (where np.matmul and np.dot can have different behavior).

Proposal

I think the best way forward here is to split up this work across the various pyGSTi developers. Maybe each week we can have one of us replace 50 calls to _np.dot(... with equivalent uses of @. We'd only do this in the pygsti/ package folder; no need to look at notebooks or tests. If we take this approach then we'd be done in a few months.

If we want, we can start by making the ~150 or so changes in my automated approach. We'd just need to audit them to make sure they didn't touch cases with arrays with more than two dimensions.

Alternatively, we could just punt on this for a little bit, so I can ruminate on other possibilities for regex'ing.

What do folks think?

[coreyostrove]

Thanks for the excellent work on this, @rileyjmurray!

I'm on-board with dividing and conquering on this. Sometimes it is nice have semi-mindless work to fall back on as a palettepilates cleanser (thanks for catching the typo, @robinbk!) between other tasks, so I am sure we can knock this out without too much difficulty. As for starting with the 150 that your regex caught, I'm slightly in favor of not doing so. If we're going to need to carefully audit these changes anyway, then I'd rather bite the bullet and do the auditing in situ.