Commit 5c0fc3a9 authored by Christoph Groth's avatar Christoph Groth
Browse files

add 'lattice' attribute to builders returned by discretize

parent d2c1401c
...@@ -36,8 +36,9 @@ class _DiscretizedBuilder(builder.Builder): ...@@ -36,8 +36,9 @@ class _DiscretizedBuilder(builder.Builder):
"""A builder that is made from a discretized model and knows how to """A builder that is made from a discretized model and knows how to
pretty-print itself.""" pretty-print itself."""
def __init__(self, symmetry=None, coords=[], **kwargs): def __init__(self, coords, lattice, symmetry=None, **kwargs):
super().__init__(symmetry, **kwargs) super().__init__(symmetry, **kwargs)
self.lattice = lattice
self._coords = coords self._coords = coords
def __str__(self): def __str__(self):
...@@ -120,8 +121,11 @@ def discretize(hamiltonian, coords=None, *, grid_spacing=1, ...@@ -120,8 +121,11 @@ def discretize(hamiltonian, coords=None, *, grid_spacing=1,
------- -------
model : `~kwant.builder.Builder` model : `~kwant.builder.Builder`
The translationally symmetric builder that corresponds to the provided The translationally symmetric builder that corresponds to the provided
Hamiltonian. Hamiltonian. This builder instance belongs to a subclass of the
standard builder that may be printed to obtain the source code of the
value functions. It also holds the discretization lattice (a
`~kwant.lattice.Monatomic` instance with lattice constant
`grid_spacing`) in the ``lattice`` attribute.
""" """
tb, coords = discretize_symbolic(hamiltonian, coords, locals=locals) tb, coords = discretize_symbolic(hamiltonian, coords, locals=locals)
...@@ -262,7 +266,11 @@ def build_discretized(tb_hamiltonian, coords, *, grid_spacing=1, locals=None): ...@@ -262,7 +266,11 @@ def build_discretized(tb_hamiltonian, coords, *, grid_spacing=1, locals=None):
------- -------
model : `~kwant.builder.Builder` model : `~kwant.builder.Builder`
The translationally symmetric builder that corresponds to the provided The translationally symmetric builder that corresponds to the provided
Hamiltonian. Hamiltonian. This builder instance belongs to a subclass of the
standard builder that may be printed to obtain the source code of the
value functions. It also holds the discretization lattice (a
`~kwant.lattice.Monatomic` instance with lattice constant
`grid_spacing`) in the ``lattice`` attribute.
""" """
if len(coords) == 0: if len(coords) == 0:
...@@ -301,8 +309,8 @@ def build_discretized(tb_hamiltonian, coords, *, grid_spacing=1, locals=None): ...@@ -301,8 +309,8 @@ def build_discretized(tb_hamiltonian, coords, *, grid_spacing=1, locals=None):
hoppings = {builder.HoppingKind(tuple(-i for i in d), lat): val hoppings = {builder.HoppingKind(tuple(-i for i in d), lat): val
for d, val in tb.items()} for d, val in tb.items()}
syst = _DiscretizedBuilder(lattice.TranslationalSymmetry(*prim_vecs), syst = _DiscretizedBuilder(coords, lat,
coords) lattice.TranslationalSymmetry(*prim_vecs))
syst[lat(*onsite_zeros)] = onsite syst[lat(*onsite_zeros)] = onsite
for hop, val in hoppings.items(): for hop, val in hoppings.items():
syst[hop] = val syst[hop] = val
......
...@@ -308,6 +308,7 @@ def test_numeric_functions_basic_symbolic(): ...@@ -308,6 +308,7 @@ def test_numeric_functions_basic_symbolic():
for i in [0, 1, 3, 5]: for i in [0, 1, 3, 5]:
builder = discretize(i, 'x') builder = discretize(i, 'x')
lat = next(iter(builder.sites()))[0] lat = next(iter(builder.sites()))[0]
assert builder.lattice is lat
assert builder[lat(0)] == i assert builder[lat(0)] == i
p = dict(t=i) p = dict(t=i)
......
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