Commit eae1b77b authored by Joseph Weston's avatar Joseph Weston

Merge branch 'feature/nd-systems' into 'master'

First steps in ND vectorized systems

Closes #326

See merge request !341
parents e9ea4e8c ac264fb5
Pipeline #26956 passed with stages
in 36 minutes and 53 seconds
......@@ -23,7 +23,6 @@ Abstract base classes
.. autosummary::
:toctree: generated/
Symmetry
Lead
Functions
......
......@@ -28,6 +28,14 @@ Sites
SiteArray
SiteFamily
Symmetries
----------
.. autosummary::
:toctree: generated/
Symmetry
NoSymmetry
Systems
-------
.. autosummary::
......
......@@ -548,13 +548,12 @@ def vectorized_cell_hamiltonian(self, args=(), sparse=False, *, params=None):
# Site array where next cell starts
next_cell = bisect.bisect(site_offsets, self.cell_size) - 1
def inside_cell(term):
(to_which, from_which), _= self.subgraphs[term.subgraph]
return to_which < next_cell and from_which < next_cell
def inside_fd(term):
return all(s == 0 for s in term.symmetry_element)
cell_terms = [
n for n, t in enumerate(self.terms)
if inside_cell(t)
if inside_fd(t)
]
subgraphs = [
......@@ -593,34 +592,38 @@ def vectorized_cell_hamiltonian(self, args=(), sparse=False, *, params=None):
def vectorized_inter_cell_hopping(self, args=(), sparse=False, *, params=None):
"""Hopping Hamiltonian between two cells of the infinite system.
This method returns a complex matrix that represents the hopping from
the *interface sites* of unit cell ``n - 1`` to *all* the sites of
unit cell ``n``. It is therefore generally a *rectangular* matrix of
shape ``(N_uc, N_iface)`` where ``N_uc`` is the number of orbitals
in the unit cell, and ``N_iface`` is the number of orbitals on the
*interface* sites (i.e. the sites with hoppings *to* the next unit cell).
Providing positional arguments via 'args' is deprecated,
instead, provide named parameters as a dictionary via 'params'.
"""
# Take advantage of the fact that there are distinct entries in
# onsite_terms for sites inside the unit cell, and distinct entries
# in onsite_terms for hoppings between the unit cell sites and
# interface sites. This way we only need to check the first entry
# in onsite/hopping_terms
site_offsets = np.cumsum([0] + [len(arr) for arr in self.site_arrays])
# Site array where next cell starts
next_cell = bisect.bisect(site_offsets, self.cell_size) - 1
# This method is only meaningful for systems with a 1D translational
# symmetry, and we use this fact in several places
assert all(len(t.symmetry_element) == 1 for t in self.terms)
# Symmetry element -1 means hoppings *from* the *previous*
# unit cell. These are directly the hoppings we wish to return.
inter_cell_hopping_terms = [
n for n, t in enumerate(self.terms)
# *from* site is in interface
if (self.subgraphs[t.subgraph][0][1] >= next_cell
and self.subgraphs[t.subgraph][0][0] < next_cell)
if t.symmetry_element[0] == -1
]
# Symmetry element +1 means hoppings *from* the *next* unit cell.
# These are related by translational symmetry to hoppings *to*
# the *previous* unit cell. We therefore need the *reverse*
# (and conjugate) of these hoppings.
reversed_inter_cell_hopping_terms = [
n for n, t in enumerate(self.terms)
# *to* site is in interface
if (self.subgraphs[t.subgraph][0][0] >= next_cell
and self.subgraphs[t.subgraph][0][1] < next_cell)
if t.symmetry_element[0] == +1
]
# Evaluate inter-cell hoppings only
inter_cell_hams = [
self.hamiltonian_term(n, args=args, params=params)
for n in inter_cell_hopping_terms
......@@ -642,18 +645,21 @@ def vectorized_inter_cell_hopping(self, args=(), sparse=False, *, params=None):
for n in reversed_inter_cell_hopping_terms
]
# All the 'from' sites are in the previous domain, but to build a
# matrix we need to get the equivalent sites in the fundamental domain.
# We set the 'from' site array to the one from the fundamental domain.
subgraphs = [
((to_sa, from_sa - next_cell), (to_off, from_off))
for (to_sa, from_sa), (to_off, from_off) in subgraphs
]
_, norbs, orb_offsets = self.site_ranges.transpose()
shape = (orb_offsets[next_cell],
orb_offsets[len(self.site_arrays) - next_cell])
# SiteArrays containing interface sites appear before SiteArrays
# containing non-interface sites, so the max of the site array
# indices that appear in the 'from' site arrays of the inter-cell
# hoppings allows us to get the number of interface orbitals.
last_iface_site_array = max(
from_site_array for (_, from_site_array), _ in subgraphs
)
iface_norbs = orb_offsets[last_iface_site_array + 1]
fd_norbs = orb_offsets[-1]
# TODO: return a square matrix when we no longer need to maintain
# backwards compatibility with unvectorized systems.
shape = (fd_norbs, iface_norbs)
func = _vectorized_make_sparse if sparse else _vectorized_make_dense
mat = func(subgraphs, hams, norbs, orb_offsets, site_offsets, shape=shape)
return mat
This diff is collapsed.
......@@ -18,6 +18,7 @@ import sympy
import kwant.lattice
import kwant.builder
import kwant.system
import kwant.continuum
import kwant.continuum._common
......@@ -144,7 +145,7 @@ def discretize_landau(hamiltonian, N, momenta=None, grid_spacing=1):
if _has_coordinate(normal_coordinate, hamiltonian):
sym = kwant.lattice.TranslationalSymmetry([grid_spacing, 0])
else:
sym = kwant.builder.NoSymmetry()
sym = kwant.system.NoSymmetry()
lat = LandauLattice(grid_spacing, norbs=norbs)
syst = kwant.Builder(sym)
......
......@@ -13,8 +13,8 @@ import sympy
import pytest
import itertools
import kwant.builder
import kwant.lattice
import kwant.system
from .._common import position_operators, momentum_operators, sympify
from ..landau_levels import (
......@@ -118,7 +118,7 @@ def test_discretize_landau():
# test a basic Hamiltonian with no normal coordinate dependence
syst = discretize_landau("k_x**2 + k_y**2", N=n_levels)
lat = LandauLattice(1, norbs=1)
assert isinstance(syst.symmetry, kwant.builder.NoSymmetry)
assert isinstance(syst.symmetry, kwant.system.NoSymmetry)
syst = syst.finalized()
assert set(syst.sites) == {lat(0, j) for j in range(n_levels)}
assert np.allclose(
......
......@@ -153,7 +153,7 @@ class Polyatomic:
algorithm finds and yields all the lattice sites inside the specified
shape starting from the specified position.
A `~kwant.builder.Symmetry` or `~kwant.builder.Builder` may be passed as
A `~kwant.system.Symmetry` or `~kwant.builder.Builder` may be passed as
sole argument when calling the function returned by this method. This
will restrict the flood-fill to the fundamental domain of the symmetry
(or the builder's symmetry). Note that unless the shape function has
......@@ -174,8 +174,8 @@ class Polyatomic:
Site = system.Site
if symmetry is None:
symmetry = builder.NoSymmetry()
elif not isinstance(symmetry, builder.Symmetry):
symmetry = system.NoSymmetry()
elif not isinstance(symmetry, system.Symmetry):
symmetry = symmetry.symmetry
def fd_site(lat, tag):
......@@ -259,7 +259,7 @@ class Polyatomic:
center = ta.array(center, float)
def wire_sites(sym):
if not isinstance(sym, builder.Symmetry):
if not isinstance(sym, system.Symmetry):
sym = sym.symmetry
if not isinstance(sym, TranslationalSymmetry):
raise ValueError('wire shape only works with '
......@@ -514,7 +514,7 @@ class Monatomic(system.SiteFamily, Polyatomic):
# The following class is designed such that it should avoid floating
# point precision issues.
class TranslationalSymmetry(builder.Symmetry):
class TranslationalSymmetry(system.Symmetry):
"""A translational symmetry defined in real space.
An alias exists for this common name: ``kwant.TranslationalSymmetry``.
......@@ -568,7 +568,7 @@ class TranslationalSymmetry(builder.Symmetry):
return TranslationalSymmetry(*ta.dot(generators, self.periods))
def has_subgroup(self, other):
if isinstance(other, builder.NoSymmetry):
if isinstance(other, system.NoSymmetry):
return True
elif not isinstance(other, TranslationalSymmetry):
raise ValueError("Unknown symmetry type.")
......@@ -718,8 +718,9 @@ class TranslationalSymmetry(builder.Symmetry):
raise ValueError("must provide a single group element when "
"acting on single sites.")
if (len(element.shape) == 1 and not is_site):
raise ValueError("must provide a sequence of group elements "
"when acting on site arrays.")
# We can act on a whole SiteArray with a single group element
# using numpy broadcasting.
element = element.reshape(1, -1)
m_part = self._get_site_family_data(a.family)[0]
try:
delta = array_mod.dot(m_part, element)
......
......@@ -6,7 +6,8 @@ import pytest
import kwant
from ... import lattice
from ...builder import HoppingKind, Builder, NoSymmetry, Site
from ...builder import HoppingKind, Builder, Site
from ...system import NoSymmetry
from .. import gauge
......
This diff is collapsed.
......@@ -153,7 +153,7 @@ def test_site_families():
assert fam1 < fam2 # string '1' is lexicographically less than '2'
class VerySimpleSymmetry(builder.Symmetry):
class VerySimpleSymmetry(system.Symmetry):
def __init__(self, period):
self.period = period
......@@ -162,7 +162,7 @@ class VerySimpleSymmetry(builder.Symmetry):
return 1
def has_subgroup(self, other):
if isinstance(other, builder.NoSymmetry):
if isinstance(other, system.NoSymmetry):
return True
elif isinstance(other, VerySimpleSymmetry):
return not other.period % self.period
......@@ -321,10 +321,10 @@ def random_hopping_integral(rng):
def check_onsite(fsyst, sites, subset=False, check_values=True):
vectorized = isinstance(fsyst, (system.FiniteVectorizedSystem, system.InfiniteVectorizedSystem))
vectorized = system.is_vectorized(fsyst)
if vectorized:
site_offsets = np.cumsum([0] + [len(s) for s in fsyst.site_arrays])
site_offsets, _, _ = fsyst.site_ranges.transpose()
freq = {}
for node in range(fsyst.graph.num_nodes):
......@@ -353,10 +353,10 @@ def check_onsite(fsyst, sites, subset=False, check_values=True):
def check_hoppings(fsyst, hops):
vectorized = isinstance(fsyst, (system.FiniteVectorizedSystem, system.InfiniteVectorizedSystem))
vectorized = system.is_vectorized(fsyst)
if vectorized:
site_offsets = np.cumsum([0] + [len(s) for s in fsyst.site_arrays])
site_offsets, _, _ = fsyst.site_ranges.transpose()
assert fsyst.graph.num_edges == 2 * len(hops)
for edge_id, edge in enumerate(fsyst.graph):
......@@ -366,6 +366,12 @@ def check_hoppings(fsyst, hops):
if vectorized:
term = fsyst._hopping_term_by_edge_id[edge_id]
# Map sites in previous cell to fundamental domain; vectorized
# infinite systems only store sites in the FD. We already know
# that this hopping is between unit cells because this is encoded
# in the edge_id and term id.
if system.is_infinite(fsyst):
i, j = i % fsyst.cell_size, j % fsyst.cell_size
if term < 0: # Hermitian conjugate
assert (head, tail) in hops
else:
......@@ -464,7 +470,8 @@ def test_finalization(vectorize):
assert tuple(fsyst.sites) == tuple(sorted(fam(*site) for site in sr_sites))
# Build lead from blueprint and test it.
lead = builder.Builder(kwant.TranslationalSymmetry((size, 0)))
lead = builder.Builder(kwant.TranslationalSymmetry((size, 0)),
vectorize=vectorize)
for site, value in lead_sites.items():
shift = rng.randrange(-5, 6) * size
site = site[0] + shift, site[1]
......@@ -757,10 +764,6 @@ def test_vectorized_hamiltonian_evaluation():
syst_simple[(lat(2, 1), lat(1, 1))] = hopping
fsyst_simple = syst_simple.finalized()
assert np.allclose(
fsyst_vectorized.hamiltonian_submatrix(),
fsyst_simple.hamiltonian_submatrix(),
)
assert np.allclose(
fsyst_vectorized.cell_hamiltonian(),
fsyst_simple.cell_hamiltonian(),
......@@ -812,7 +815,7 @@ def test_vectorized_value_normalization():
@pytest.mark.parametrize("sym", [
builder.NoSymmetry(),
system.NoSymmetry(),
kwant.TranslationalSymmetry([-1]),
])
def test_vectorized_requires_norbs(sym):
......@@ -921,7 +924,7 @@ def test_fill():
## Test that copying a builder by "fill" preserves everything.
for sym, func in [(kwant.TranslationalSymmetry(*np.diag([3, 4, 5])),
lambda pos: True),
(builder.NoSymmetry(),
(system.NoSymmetry(),
lambda pos: ta.dot(pos, pos) < 17)]:
cubic = kwant.lattice.general(ta.identity(3), norbs=1)
......@@ -1642,7 +1645,7 @@ def test_subs():
lat = kwant.lattice.chain(norbs=1)
def make_system(sym=kwant.builder.NoSymmetry(), n=3):
def make_system(sym=system.NoSymmetry(), n=3):
syst = kwant.Builder(sym)
syst[(lat(i) for i in range(n))] = onsite
syst[lat.neighbors()] = hopping
......
......@@ -11,7 +11,7 @@ from math import sqrt
import numpy as np
import tinyarray as ta
from pytest import raises
from kwant import lattice, builder
from kwant import lattice, builder, system
from kwant._common import ensure_rng
import pytest
......@@ -285,7 +285,7 @@ def test_symmetry_has_subgroup():
## test whether actual subgroups are detected as such
vecs = rng.randn(3, 3)
sym1 = lattice.TranslationalSymmetry(*vecs)
ns = builder.NoSymmetry()
ns = system.NoSymmetry()
assert ns.has_subgroup(ns)
assert sym1.has_subgroup(sym1)
assert sym1.has_subgroup(ns)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment