hamiltonian_submatrix: return only the matrix by default

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)
 

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.

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

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)

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)
 

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.

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.

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)