diff --git a/doc/source/images/3-closed_system.py b/doc/source/images/3-closed_system.py
index c4dd0dc33b5b55e48591a2a6c362a62f3c873336..33a9929793390fad084b422c54b789530916b43b 100644
--- a/doc/source/images/3-closed_system.py
+++ b/doc/source/images/3-closed_system.py
@@ -64,7 +64,7 @@ def plot_spectrum(sys, Bfields):
B = Bfield
# Obtain the Hamiltonian as a dense matrix
- ham_mat = sys.hamiltonian_submatrix()[0]
+ ham_mat = sys.hamiltonian_submatrix()
ev = la.eigh(ham_mat, eigvals_only=True)
diff --git a/doc/source/whatsnew/0.2.rst b/doc/source/whatsnew/0.2.rst
index 9ae614d6636b2b3ef563fcff849b91f7d7e85d29..4682b061e2bbd46435dee7c5208d54334e0b680f 100644
--- a/doc/source/whatsnew/0.2.rst
+++ b/doc/source/whatsnew/0.2.rst
@@ -40,11 +40,6 @@ Return value of sparse solver
`~kwant.solvers.sparse.BlockResult`. The latter has been generalized to
include more information for leads defined as infinite systems.
-Return value of `~kwant.system.System.hamiltonian_submatrix`
-------------------------------------------------------------
-`~kwant.system.System.hamiltonian_submatrix` now always returns a 3-tuple
-``hamiltonian_part, to_norb, from_norb``.
-
Return value of `~kwant.solvers.sparse.make_linear_sys` has changed
-------------------------------------------------------------------
A namedtuple is used for more clarity.
diff --git a/kwant/_system.pyx b/kwant/_system.pyx
index 2cb98eb47394075a1dc8870e26359a55f9cc4cf2..7523864c4df833c5b6b20bfe4ccaff4e73205423 100644
--- a/kwant/_system.pyx
+++ b/kwant/_system.pyx
@@ -204,7 +204,7 @@ def make_dense_full(ham, CGraph gr, diag,
def hamiltonian_submatrix(self, to_sites=None, from_sites=None,
- sparse=False):
+ sparse=False, return_norb=False):
"""Return a submatrix of the system Hamiltonian.
Parameters
@@ -213,15 +213,19 @@ def hamiltonian_submatrix(self, to_sites=None, from_sites=None,
from_sites : sequence of sites or None (default)
sparse : bool
Whether to return a sparse or a dense matrix. Defaults to `False`.
+ return_norb : bool
+ Whether to return arrays of numbers of orbitals. Defaults to `False`.
Returns
-------
hamiltonian_part : numpy.ndarray or scipy.sparse.coo_matrix
Submatrix of Hamiltonian of the system.
to_norb : array of integers
- Numbers of orbitals on each site in to_sites.
+ Numbers of orbitals on each site in to_sites. Only returned when
+ `return_norb` is true.
from_norb : array of integers
- Numbers of orbitals on each site in from_sites.
+ Numbers of orbitals on each site in from_sites. Only returned when
+ `return_norb` is true.
Notes
-----
@@ -295,4 +299,4 @@ def hamiltonian_submatrix(self, to_sites=None, from_sites=None,
func = make_sparse if sparse else make_dense
mat = func(ham, self.graph, diag, from_sites, n_by_to_site,
to_norb, to_off, from_norb, from_off)
- return mat, to_norb, from_norb
+ return (mat, to_norb, from_norb) if return_norb else mat
diff --git a/kwant/linalg/tests/test_mumps.py b/kwant/linalg/tests/test_mumps.py
index db9842f19c0ff3e0fa7ae25e0f3ca783c5d6f9a9..e0e062a107531c015aa9181a262545d44bae84b0 100644
--- a/kwant/linalg/tests/test_mumps.py
+++ b/kwant/linalg/tests/test_mumps.py
@@ -70,7 +70,7 @@ def test_error_minus_9(r=10):
for hopping in hoppings:
sys[sys.possible_hoppings(*hopping)] = - 1
- ham = sys.finalized().hamiltonian_submatrix(sparse=True)[0]
+ ham = sys.finalized().hamiltonian_submatrix(sparse=True)
# No need to check result, it's enough if no exception is raised
MUMPSContext().factor(ham)
diff --git a/kwant/solvers/common.py b/kwant/solvers/common.py
index efcbf164d10ebc9796e993e9c8400dfe1506af52..44dc9548d8cb5da514216bd61eba9870edabedb0 100644
--- a/kwant/solvers/common.py
+++ b/kwant/solvers/common.py
@@ -137,7 +137,8 @@ class SparseSolver(object):
if not sys.lead_neighbor_seqs:
raise ValueError('System contains no leads.')
- lhs, norb = sys.hamiltonian_submatrix(sparse=True)[:2]
+ lhs, norb = sys.hamiltonian_submatrix(
+ sparse=True, return_norb=True)[:2]
lhs = getattr(lhs, 'to' + self.lhsformat)()
lhs = lhs - energy * sp.identity(lhs.shape[0], format=self.lhsformat)
diff --git a/kwant/system.py b/kwant/system.py
index 23fc3053ab1397dcfa1ed592463352483783e845..f2d130ab5901d17b156e52a08cdb6cb9f124c53d 100644
--- a/kwant/system.py
+++ b/kwant/system.py
@@ -127,14 +127,14 @@ class InfiniteSystem(System):
"""Hamiltonian of a single slice of the infinite system."""
slice_sites = xrange(self.slice_size)
return self.hamiltonian_submatrix(slice_sites, slice_sites,
- sparse=sparse)[0]
+ sparse=sparse)
def inter_slice_hopping(self, sparse=False):
"""Hopping Hamiltonian between two slices of the infinite system."""
slice_sites = xrange(self.slice_size)
neighbor_sites = xrange(self.slice_size, self.graph.num_nodes)
return self.hamiltonian_submatrix(slice_sites, neighbor_sites,
- sparse=sparse)[0]
+ sparse=sparse)
def self_energy(self, energy):
"""Return self-energy of a lead.
diff --git a/kwant/tests/test_system.py b/kwant/tests/test_system.py
index b2bb7fc4f8fd4b29854b163ffeb3567af279ad2e..eea7869aab543905c14a24ebc7c147c48e77f3ef 100644
--- a/kwant/tests/test_system.py
+++ b/kwant/tests/test_system.py
@@ -12,7 +12,7 @@ def test_hamiltonian_submatrix():
sys[(i,), (i + 1,)] = 1j * (i + 1)
sys2 = sys.finalized()
- mat = sys2.hamiltonian_submatrix()[0]
+ mat = sys2.hamiltonian_submatrix()
assert mat.shape == (3, 3)
# Sorting is required due to unknown compression order of builder.
perm = np.argsort(sys2.onsite_hamiltonians)
@@ -22,17 +22,17 @@ def test_hamiltonian_submatrix():
mat = mat[:, perm]
np.testing.assert_array_equal(mat, mat_should_be)
- mat = sys2.hamiltonian_submatrix(sparse=True)[0]
+ mat = sys2.hamiltonian_submatrix(sparse=True)
assert sparse.isspmatrix_coo(mat)
mat = mat.todense()
mat = mat[perm, :]
mat = mat[:, perm]
np.testing.assert_array_equal(mat, mat_should_be)
- mat = sys2.hamiltonian_submatrix(perm[[0, 1]], perm[[2]])[0]
+ mat = sys2.hamiltonian_submatrix(perm[[0, 1]], perm[[2]])
np.testing.assert_array_equal(mat, mat_should_be[:2, 2:3])
- mat = sys2.hamiltonian_submatrix(perm[[0, 1]], perm[[2]], sparse=True)[0]
+ mat = sys2.hamiltonian_submatrix(perm[[0, 1]], perm[[2]], sparse=True)
mat = mat.todense()
np.testing.assert_array_equal(mat, mat_should_be[:2, 2:3])
@@ -45,8 +45,8 @@ def test_hamiltonian_submatrix():
sys[(1,), (0,)] = np.array([[1, 2j]])
sys[(2,), (1,)] = np.array([[3j]])
sys2 = sys.finalized()
- mat_dense = sys2.hamiltonian_submatrix()[0]
- mat_sp = sys2.hamiltonian_submatrix(sparse=True)[0].todense()
+ mat_dense = sys2.hamiltonian_submatrix()
+ mat_sp = sys2.hamiltonian_submatrix(sparse=True).todense()
np.testing.assert_array_equal(mat_sp, mat_dense)
# Test for shape errors.
diff --git a/tutorial/3-closed_system.py b/tutorial/3-closed_system.py
index f09603abff97a85a50c96c883b69effbee9678e4..2cb44f58bf57fbd7026f9246766e6d954fbf552d 100644
--- a/tutorial/3-closed_system.py
+++ b/tutorial/3-closed_system.py
@@ -62,7 +62,7 @@ def plot_spectrum(sys, Bfields):
B = Bfield
# Obtain the Hamiltonian as a dense matrix
- ham_mat = sys.hamiltonian_submatrix()[0]
+ ham_mat = sys.hamiltonian_submatrix()
ev = la.eigh(ham_mat, eigvals_only=True)