diff --git a/kwant/kpm.py b/kwant/kpm.py index db19f4ea6504f3170cbe438be4904864dd405121..e30401ee1958d123cda43ce7ecc2a05618658595 100644 --- a/kwant/kpm.py +++ b/kwant/kpm.py @@ -778,6 +778,8 @@ def conductivity(hamiltonian, alpha='x', beta='x', positions=None, **kwargs): Parameters ---------- + hamiltonian : `~kwant.system.FiniteSystem` or matrix Hamiltonian + If a system is passed, it should contain no leads. alpha, beta : str, or operators If ``hamiltonian`` is a kwant system, or if the ``positions`` are provided, ``alpha`` and ``beta`` can be the directions of the @@ -1049,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] diff --git a/kwant/tests/test_kpm.py b/kwant/tests/test_kpm.py index 7ae187339e51163549637cf44306643b7f816bf4..524cc57d78a97434ed6605f36b72d5833c5fd410 100644 --- a/kwant/tests/test_kpm.py +++ b/kwant/tests/test_kpm.py @@ -80,6 +80,13 @@ def test_conductivity(): cond = kwant.kpm.conductivity(syst, alpha=alpha, beta=beta, positions=None, **kpm_params) + # test when 'norbs' not provided + n = lat.norbs + lat.norbs = None + cond = kwant.kpm.conductivity(syst, alpha=alpha, beta=beta, + positions=None, **kpm_params) + lat.norbs = n + # test system or hamiltonian, no positions, but velocity operators cond_xx = kwant.kpm.conductivity(syst, alpha='x', beta='x', positions=None, **kpm_params)