diff --git a/kwant/kpm.py b/kwant/kpm.py index 3c335bdcca2cf3194f7e2ec4eb9ac2bb52805015..e30401ee1958d123cda43ce7ecc2a05618658595 100644 --- a/kwant/kpm.py +++ b/kwant/kpm.py @@ -1051,10 +1051,11 @@ def _velocity(hamiltonian, params, op_type, positions): elif isinstance(op_type, str): direction = directions[op_type] if isinstance(hamiltonian, system.System): - operator = hamiltonian.hamiltonian_submatrix(params=params, - sparse=True) - positions = np.array([site.pos for site in hamiltonian.sites - for iorb in range(site.family.norbs)]) + operator, norbs, norbs = hamiltonian.hamiltonian_submatrix( + params=params, sparse=True, return_norb=True + ) + positions = np.vstack([[hamiltonian.pos(i)] * norb + for i, norb in enumerate(norbs)]) elif positions is not None: operator = coo_matrix(hamiltonian, copy=True) displacements = positions[operator.col] - positions[operator.row]