diff --git a/changelog.md b/changelog.md index e79b13ea..5992ce3e 100644 --- a/changelog.md +++ b/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. diff --git a/pybinding/support/kwant.py b/pybinding/support/kwant.py index e2c968ce..99e635f9 100644 --- a/pybinding/support/kwant.py +++ b/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 diff --git a/tests/test_kwant.py b/tests/test_kwant.py index bba982b8..9bf93881 100644 --- a/tests/test_kwant.py +++ b/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))