Allow operators to act on sequences of wavefunctions (matrices)
Currently operators can only act on a single vector at a time, but the following is very common (e.g. if you want to get the current that would flow when applying an infinitesimal voltage to a lead):
wf = kwant.wave_function(syst, ...)
J = kwant.operator.Current(syst)
J_0 = sum(J(psi) for psi in wf(0))
This is sub-optimal because of the explicit loop:
- Users don't expect that they have to do this (at least 2 messages on the mailing list about this)
- We pay the Python overhead for not vectorizing the loop
Proposal
Allow operators to allow input of arbitrary rank.
Challenges
We have to define what the output should be for all possible inputs. In the simplest case this is relatively straightforward:
wfs = kwant.wave_function(syst, ...)
rho = kwant.operator.Density(syst, ...)
wfs_0 = wfs(0) # shape: (nmodes, norbs)
rho(wf(0)) # shape: (nmodes, nsites)
all_wfs = np.array([wfs(i) for i in range(num_leads)]) # (nleads, nmodes, norbs)
rho(all_wfs) # (nleads, nmodes, nsites)
but what if we provide bra
as well as ket
?
# (nmodes, nmodes, norbs) ?
# (nmodes, norbs) ?
rho(wfs(0), wfs(0))
Implementation details
Even though the input may be of arbitrary rank, the lowest layers of kwant.operator
only need to operate on rank-2 quantities, which means that we can still use a fixed number of C loops.
The python layer can simply reshape the input and output to be in the correct format.
In the above example:
all_wfs = np.array([wfs(i) for i in range(num_leads)]) # (nleads, nmodes, norbs)
rho(all_wfs) # (nleads, nmodes, nsites)
We would internally reshape all_wfs
to shape (nleads * nmodes, norbs)
, pass it down to the lower layers, and then reshape the output of the lower layers to (nleads, nmodes, nsites)
before returning it to the user.