restructure BlockResult

kwant/solvers/common.py

@@ -342,12 +342,9 @@ class SparseSolver(object):
 
         flhs = self._factorized(linsys.lhs)
         data = self._solve_linear_sys(flhs, linsys.rhs, linsys.kept_vars)
-        result = BlockResult(data, lead_info)
 
-        result.in_leads = in_leads
-        result.out_leads = out_leads
+        return BlockResult(data, lead_info, out_leads, in_leads)
 
-        return result
 
     def ldos(self, fsys, energy=0):
         """
@@ -444,12 +441,11 @@ class WaveFunc(object):
         return result.transpose()
 
 
-class BlockResult(namedtuple('BlockResultTuple', ['data', 'lead_info'])):
+class BlockResult(object):
     """
     Solution of a transport problem, subblock of retarded Green's function.
 
-    This class is derived from ``namedtuple('BlockResultTuple', ['data',
-    'lead_info'])``. In addition to direct access to `data` and `lead_info`,
+    In addition to direct access to `data` and `lead_info`,
     this class also supports a higher level interface via its methods.
 
     Instance Variables
@@ -460,28 +456,53 @@ class BlockResult(namedtuple('BlockResultTuple', ['data', 'lead_info'])):
     lead_info : list of data
         a list with output of `kwant.physics.modes` for each lead defined as a
         builder, and self-energy for each lead defined as self-energy term.
+    out_leads : list of integers
+    in_leads : list of integers
+        indices of the leads where current is extracted (out) or injected (in).
+        Only those are listed for which BlockResult contains the calculated
+        result.
     """
+
+    def __init__(self, data, lead_info, out_leads, in_leads):
+        self.data = data
+        self.lead_info = lead_info
+        self.out_leads = out_leads
+        self.in_leads = in_leads
+
+        sizes = []
+        for i in self.lead_info:
+            if isinstance(i, tuple):
+                sizes.append(i[2])
+            else:
+                sizes.append(i.shape[0])
+        self._sizes = np.array(sizes)
+
+        self._in_offsets = np.zeros(len(self.in_leads) + 1, int)
+        self._in_offsets[1 :] = np.cumsum(self._sizes[self.in_leads])
+        self._out_offsets = np.zeros(len(self.out_leads) + 1, int)
+        self._out_offsets[1 :] = np.cumsum(self._sizes[self.out_leads])
+
     def block_coords(self, lead_out, lead_in):
         """
         Return slices corresponding to the block from lead_in to lead_out.
         """
+        return self.out_block_coords(lead_out), self.in_block_coords(lead_in)
+
+    def out_block_coords(self, lead_out):
+        """Return a slice corresponding to the rows in the block corresponding
+        to lead_out
+        """
         lead_out = self.out_leads.index(lead_out)
-        lead_in = self.in_leads.index(lead_in)
-        if not hasattr(self, '_sizes'):
-            sizes = []
-            for i in self.lead_info:
-                if isinstance(i, tuple):
-                    sizes.append(i[2])
-                else:
-                    sizes.append(i.shape[0])
-            self._sizes = np.array(sizes)
-            self._in_offsets = np.zeros(len(self.in_leads) + 1, int)
-            self._in_offsets[1 :] = np.cumsum(self._sizes[self.in_leads])
-            self._out_offsets = np.zeros(len(self.out_leads) + 1, int)
-            self._out_offsets[1 :] = np.cumsum(self._sizes[self.out_leads])
         return slice(self._out_offsets[lead_out],
-                     self._out_offsets[lead_out + 1]), \
-               slice(self._in_offsets[lead_in], self._in_offsets[lead_in + 1])
+                     self._out_offsets[lead_out + 1])
+
+    def in_block_coords(self, lead_in):
+        """Return a slice corresponding to the columns in the block
+        corresponding to lead_in
+        """
+        lead_in = self.in_leads.index(lead_in)
+        return slice(self._in_offsets[lead_in],
+                     self._in_offsets[lead_in + 1])
 
     def submatrix(self, lead_out, lead_in):
         """Return the matrix elements from lead_in to lead_out."""
@@ -523,3 +544,8 @@ class BlockResult(namedtuple('BlockResultTuple', ['data', 'lead_info'])):
                 result += 2 * np.trace(np.dot(gamma, gf)).imag + N
 
             return result
+
+    def __repr__(self):
+        return "BlockResult(data=%r, lead_info=%r, " \
+            "out_leads=%r, in_leads=%r)" % (self.data, self.lead_info,
+                                            self.out_leads, self.in_leads)

kwant/solvers/tests/_test_sparse.py

@@ -31,16 +31,17 @@ def test_output(solve):
     fsys = system.finalized()
 
     result1 = solve(fsys)
-    s, modes1 = result1
+    s, modes1 = result1.data, result1.lead_info
     assert s.shape == 2 * (sum(i[2] for i in modes1),)
     s1 = result1.submatrix(1, 0)
-    s2, modes2 = solve(fsys, 0, [1], [0])
+    result2 = solve(fsys, 0, [1], [0])
+    s2, modes2 = result2.data, result2.lead_info
     assert s2.shape == (modes2[1][2], modes2[0][2])
     assert_almost_equal(s1, s2)
     assert_almost_equal(np.dot(s.conjugate().transpose(), s),
                         np.identity(s.shape[0]))
     assert_raises(ValueError, solve, fsys, 0, [])
-    modes = solve(fsys)[1]
+    modes = solve(fsys).lead_info
     h = fsys.leads[0].slice_hamiltonian()
     t = fsys.leads[0].inter_slice_hopping()
     modes1 = kwant.physics.modes(h, t)
@@ -68,7 +69,7 @@ def test_one_lead(solve):
     system.attach_lead(lead)
     fsys = system.finalized()
 
-    s = solve(fsys)[0]
+    s = solve(fsys).data
     assert_almost_equal(np.dot(s.conjugate().transpose(), s),
                         np.identity(s.shape[0]))
 
@@ -96,22 +97,22 @@ def test_smatrix_shape(solve):
 
     lead0_val = 4
     lead1_val = 4
-    s = solve(fsys, energy=1.0, out_leads=[1], in_leads=[0])[0]
+    s = solve(fsys, energy=1.0, out_leads=[1], in_leads=[0]).data
     assert s.shape == (0, 0)
 
     lead0_val = 2
     lead1_val = 2
-    s = solve(fsys, energy=1.0, out_leads=[1], in_leads=[0])[0]
+    s = solve(fsys, energy=1.0, out_leads=[1], in_leads=[0]).data
     assert s.shape == (1, 1)
 
     lead0_val = 4
     lead1_val = 2
-    s = solve(fsys, energy=1.0, out_leads=[1], in_leads=[0])[0]
+    s = solve(fsys, energy=1.0, out_leads=[1], in_leads=[0]).data
     assert s.shape == (1, 0)
 
     lead0_val = 2
     lead1_val = 4
-    s = solve(fsys, energy=1.0, out_leads=[1], in_leads=[0])[0]
+    s = solve(fsys, energy=1.0, out_leads=[1], in_leads=[0]).data
     assert s.shape == (0, 1)
 
 
@@ -120,7 +121,7 @@ def test_smatrix_shape(solve):
 def test_two_equal_leads(solve):
     def check_fsys():
         sol = solve(fsys)
-        s, leads = sol[:2]
+        s, leads = sol.data, sol.lead_info
         assert_almost_equal(np.dot(s.conjugate().transpose(), s),
                             np.identity(s.shape[0]))
         n_modes = leads[0][2]
@@ -177,7 +178,8 @@ def test_graph_system(solve):
     system.attach_lead(lead.reversed())
     fsys = system.finalized()
 
-    s, leads = solve(fsys)[: 2]
+    result = solve(fsys)
+    s, leads = result.data, result.lead_info
     assert_almost_equal(np.dot(s.conjugate().transpose(), s),
                         np.identity(s.shape[0]))
     n_modes = leads[0][2]
@@ -210,7 +212,8 @@ def test_singular_graph_system(solve):
     system.attach_lead(lead.reversed())
     fsys = system.finalized()
 
-    s, leads = solve(fsys)[: 2]
+    result = solve(fsys)
+    s, leads = result.data, result.lead_info
     assert_almost_equal(np.dot(s.conjugate().transpose(), s),
                         np.identity(s.shape[0]))
     n_modes = leads[0][2]
@@ -223,7 +226,7 @@ def test_singular_graph_system(solve):
 
 
 # This test features inside the onslice Hamiltonian a hopping matrix with more
-# zero eigenvalues than the lead hopping matrix. The previous version of the
+# zero eigenvalues than the lead hopping matrix. Older version of the
 # sparse solver failed here.
 def test_tricky_singular_hopping(solve):
     system = kwant.Builder()
@@ -249,7 +252,7 @@ def test_tricky_singular_hopping(solve):
     system.leads.append(kwant.builder.BuilderLead(lead, interface))
     fsys = system.finalized()
 
-    s = solve(fsys, -1.3)[0]
+    s = solve(fsys, -1.3).data
     assert_almost_equal(np.dot(s.conjugate().transpose(), s),
                         np.identity(s.shape[0]))
 
@@ -349,8 +352,8 @@ def test_very_singular_leads(solve):
     sys.attach_lead(left_lead)
     sys.attach_lead(right_lead)
     fsys = sys.finalized()
-    result = solve(fsys)
-    assert [i[2] for i in result[1]] == [0, 2]
+    leads = solve(fsys).lead_info
+    assert [i[2] for i in leads] == [0, 2]
 
 
 def test_ldos(ldos):