A suggestion for a new implementation of kwant.operator.current

current-implementation.pyx

cdef class Current(_LocalOperator):
    # ...
    def _operate(self, complex[:] out_data, complex[:] bra, complex[:] ket,
                 args, operation op, *, params=None):
        # prepare onsite matrices and hamiltonians
        cdef int unique_onsite = not callable(self.onsite)
        cdef complex[:, :] _tmp_mat
        cdef complex *M_a = NULL
        cdef complex *H_ab = NULL
        cdef BlockSparseMatrix M_a_blocks, H_ab_blocks

        if unique_onsite:
            _tmp_mat = self.onsite
            M_a = <complex*> &_tmp_mat[0, 0]
        elif self._bound_onsite:
            M_a_blocks = self._bound_onsite
        else:
            M_a_blocks = self._eval_onsites(args, params)

        if self._bound_hamiltonian:
            H_ab_blocks = self._bound_hamiltonian
        else:
            H_ab_blocks = self._eval_hamiltonian(args, params)

        # main loop
        cdef gint a, a_s, a_norbs, b, b_s, b_norbs
        cdef gint i, j, k, w
        cdef complex tmp
        for w in range(self.where.shape[0]):
            ### get the next hopping's start orbitals and numbers of orbitals
            a_s = H_ab_blocks.block_offsets[w, 0]
            b_s = H_ab_blocks.block_offsets[w, 1]
            a_norbs = H_ab_blocks.block_shapes[w, 0]
            b_norbs = H_ab_blocks.block_shapes[w, 1]
            ### get the next onsite and Hamiltonian matrices
            H_ab = H_ab_blocks.get(w)
            if not unique_onsite:
                M_a = M_a_blocks.get(w)
            ### do the actual calculation
            if op == MAT_ELS:
                tmp = 0
                for i in range(b_norbs):
                    for j in range(a_norbs):
                        for k in range(a_norbs):
                            tmp += (bra[b_s + i].conjugate() *
                                    H_ab[j * b_norbs + i].conjugate() *
                                    M_a[j * a_norbs + k] * ket[a_s + k]
                                  - bra[a_s + j].conjugate() *
                                    M_a[j * a_norbs + k] *
                                    H_ab[k * b_norbs + i] * ket[b_s + i])
                out_data[w] = 1j * tmp
            elif op == ACT:

suggestion.pyx

cdef class Current(_LocalOperator):
    # ...
    def _operate(self, complex[:] out_data, complex[:] bra, complex[:] ket,
                 args, operation op, *, params=None):
        # prepare onsite matrices and hamiltonians
        cdef int unique_onsite = not callable(self.onsite)      

        cdef gint[:, :] offsets, norbs = _get_all_orbs(self.where, self._site_ranges)
        
        # MathVector is an "accesser" class
        cdef MathVector wf = MathVector(ket, offsets, norbs)

        # main loop
        cdef gint a b
        for w in range(self.where.shape[0]):           
            if op == MAT_ELS:
                a = self.where[w, 0]
                b = self.where[w, 1]
                
                cdef MathMatrix H_ab = MathMatrix(self.syst.hamiltonian(a, b, *args, params=params))              

                if not unique_onsite:
                    cdef MathMatrix M_aa = MathMatrix(self.onsite(a, *args, params=params))                   
                    out_data[w] = (wf[a].dagger() * M_aa * H_ab * wf[b]).imag()  
                else:
                    out_data[w] = (wf[a].dagger() * H_ab * wf[b]).imag()