Improve tokwant() export for multi-orbital models

And improve test coverage.

changelog.md

@@ -2,7 +2,7 @@
 
 ### v0.9.3 | In development
 
-* Now compatible with Kwant v1.3.x.
+* Added support for Kwant v1.3.x and improved `Model.tokwant()` exporting of multi-orbital models.
 
 * Fixed errors when compiling with GCC 6.
 

pybinding/support/kwant.py

@@ -13,8 +13,9 @@
 def _warn_if_not_empty(args, params):
     if args or params:
         warnings.warn(
-            "Additional `args/params` are ignored because pybinding's Hamiltonian is immutable.\n"
-            "Complete the model with all parameters before calling tokwant() conversion."
+            "Additional `args/params` are ignored because pybinding's Hamiltonian is immutable. "
+            "Complete the model with all parameters before calling the `tokwant()` conversion.",
+            stacklevel=3
         )
 
 
@@ -23,8 +24,8 @@ class Graph:
 
     Only the `num_nodes` attribute seems to be required, at least for `smatrix`.
     """
-    def __init__(self, hamiltonian):
-        self.num_nodes = hamiltonian.shape[0]
+    def __init__(self, num_nodes):
+        self.num_nodes = num_nodes
 
 
 class KwantFiniteSystem(FiniteSystem):
@@ -34,8 +35,8 @@ class KwantFiniteSystem(FiniteSystem):
     however it seems to work well for `smatrix`.
     """
     def __init__(self, pb_model):
-        self.hamiltonian = pb_model.hamiltonian
-        self.graph = Graph(self.hamiltonian)
+        self.pb_model = pb_model
+        self.graph = Graph(pb_model.system.num_sites)
         self._pos = np.array(pb_model.system.positions[:pb_model.lattice.ndim]).T
         self.leads = [KwantInfiniteSystem(l) for l in pb_model.leads]
         self.lead_interfaces = [l.indices for l in pb_model.leads]
@@ -45,21 +46,23 @@ def pos(self, index):
 
     def hamiltonian(self, i, j, *args, params=None):
         _warn_if_not_empty(args, params)
-        return self.hamiltonian[i, j]
+        return self.pb_model.hamiltonian[i, j]
 
     def hamiltonian_submatrix(self, args=(), to_sites=None, from_sites=None,
                               sparse=False, return_norb=False, *, params=None):
         if to_sites is not None or from_sites is not None:
-            raise RuntimeError("Not supported")
+            raise RuntimeError("The `to_sites` and `from_sites` arguments are not supported")
         _warn_if_not_empty(args, params)
 
-        matrix = self.hamiltonian.tocoo() if sparse else self.hamiltonian.todense()
+        ham = self.pb_model.hamiltonian
+        matrix = ham.tocoo() if sparse else ham.todense()
         if not return_norb:
             return matrix
         else:
-            size = self.hamiltonian.shape[0]
-            to_norb = np.ones(size)
-            from_norb = np.ones(size)
+            subs = self.pb_model.system.sublattices
+            norb = self.pb_model.system.impl.compressed_sublattices.orbital_counts
+            to_norb = norb[subs]
+            from_norb = to_norb
             return matrix, to_norb, from_norb
 
 

tests/test_kwant.py

@@ -4,7 +4,7 @@
 
 import pybinding as pb
 from pybinding.support.kwant import kwant_installed
-from pybinding.repository import graphene
+from pybinding.repository import graphene, group6_tmd
 
 
 if not kwant_installed:
@@ -72,3 +72,38 @@ def test_kwant():
         pb_result = np.array([calc_transmission(pb_model(v).tokwant(), energy) for v in vs])
 
         assert pytest.fuzzy_equal(pb_result, kwant_result)
+
+
+    @pytest.mark.parametrize("lattice, norb", [
+        (graphene.monolayer(), 1),
+        (group6_tmd.monolayer_3band("MoS2"), 3),
+    ])
+    def test_hamiltonian_submatrix_orbitals(lattice, norb):
+        """Return the number of orbitals at each site in addition to the Hamiltonian"""
+        model = pb.Model(lattice, pb.rectangle(1, 1))
+        kwant_sys = model.tokwant()
+
+        matrix, to_norb, from_norb = kwant_sys.hamiltonian_submatrix(sparse=True, return_norb=True)
+        assert matrix.shape == model.hamiltonian.shape
+        assert to_norb.size == model.system.num_sites
+        assert from_norb.size == model.system.num_sites
+        assert np.all(to_norb == norb)
+        assert np.all(from_norb == norb)
+
+
+    def test_hamiltonian_submatrix_sites():
+        """The `to_sites` and `from_sites` arguments are not supported"""
+        kwant_sys = pb.Model(graphene.monolayer(), pb.rectangle(1, 1)).tokwant()
+
+        with pytest.raises(RuntimeError) as excinfo:
+            kwant_sys.hamiltonian_submatrix(to_sites=1, from_sites=1)
+        assert "not supported" in str(excinfo.value)
+
+
+    def test_warnings():
+        """Extra arguments and ignored by pybinding -- warn users"""
+        kwant_sys = pb.Model(graphene.monolayer(), pb.rectangle(1, 1)).tokwant()
+        with pytest.warns(UserWarning):
+            kwant_sys.hamiltonian_submatrix(sparse=True, args=(1,))
+        with pytest.warns(UserWarning):
+            kwant_sys.hamiltonian_submatrix(sparse=True, params=dict(v=1))