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
Allow operators to allow input of arbitrary rank.
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
# (nmodes, nmodes, norbs) ? # (nmodes, norbs) ? rho(wfs(0), wfs(0))
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.