remove num_orbitals, fix examples/square

examples/square.py

@@ -72,10 +72,6 @@ class System(kwant.system.FiniteSystem):
         self.leads = [Lead(shape[1], hopping, lead_potentials[i])
                       for i in range(2)]
 
-    def num_orbitals(self, site):
-        """Return the number of orbitals of a site."""
-        return 1
-
     def hamiltonian(self, i, j):
         """Return an submatrix of the tight-binding Hamiltonian."""
         if i == j:

examples/tests/test_square.py

-from kwant import square
+import square
 from nose.tools import assert_equal, assert_raises
 from numpy.testing import assert_almost_equal
 
@@ -18,7 +18,7 @@ def test_hamiltonian():
     for i in xrange(sys.graph.num_nodes):
         shape = sys.hamiltonian(i, i).shape
         assert_equal(len(shape), 2)
-        assert_equal(shape[0], sys.num_orbitals(i))
+        assert_equal(shape[0], 1)
         for j in sys.graph.out_neighbors(i):
             m = sys.hamiltonian(i, j)
             shape = m.shape
@@ -29,9 +29,7 @@ def test_hamiltonian():
 def test_selfenergy():
     sys = square.System((2, 4), 1)
     for lead in xrange(len(sys.lead_interfaces)):
-        n_orb = sum(
-            sys.num_orbitals(site) for site in sys.lead_interfaces[lead])
-        se = sys.selfenergy(lead, 0)
+        se = sys.leads[lead].selfenergy(0)
         assert_equal(len(se.shape), 2)
         assert_equal(se.shape[0], se.shape[1])
-        assert_equal(se.shape[0], n_orb)
+        assert_equal(se.shape[0], len(sys.lead_interfaces[lead]))

kwant/system.py

@@ -36,15 +36,6 @@ class System(object):
     """
     __metaclass__ = abc.ABCMeta
 
-    def num_orbitals(self, site, *args):
-        """Return the number of orbitals of a site.
-
-        This is an inefficient general implementation.  It should be
-        overridden, if a more efficient way to calculate is available.
-        """
-        ham = self.hamiltonian(site, site, *args)
-        return 1 if np.isscalar(ham) else ham.shape[0]
-
     @abc.abstractmethod
     def hamiltonian(self, i, j, *args):
         """Return the hamiltonian matrix element for sites `i` and `j`.
@@ -81,7 +72,7 @@ class FiniteSystem(System):
     The length of `leads` must be equal to the length of `lead_interfaces`.
 
     For lead ``n``, the method leads[n].selfenergy must return a square matrix
-    whose size is ``sum(self.num_orbitals(neighbor)`` for neighbor in
+    whose size is ``sum(len(self.hamiltonian(site, site)) for site in
     self.lead_interfaces[n])``. The output format for ``leads[n].modes`` has to
     be as described in `~kwant.physics.ModesTuple`.
 
@@ -218,7 +209,7 @@ class InfiniteSystem(System):
         """Return self-energy of a lead.
 
         The returned matrix has the shape (s, s), where s is
-        ``sum(self.num_orbitals(i)
+        ``sum(len(self.hamiltonian(i, i))
               for i in range(self.graph.num_nodes - self.slice_size))``.
         """
         ham = self.slice_hamiltonian(args=args)