diff --git a/kwant/continuum/discretizer.py b/kwant/continuum/discretizer.py index bd05342b615f85fbeb7dcf3e8e927f5fd2bbbd5d..baef963c4c3ed4647852053d92a61f560f566424 100644 --- a/kwant/continuum/discretizer.py +++ b/kwant/continuum/discretizer.py @@ -582,16 +582,13 @@ def _return_string(expr, coords): return 'return {}'.format(output), map_func_calls, const_symbols, _cache -def _assign_symbols(map_func_calls, grid_spacing, - coords, onsite): +def _assign_symbols(map_func_calls, coords, onsite): """Generate a series of assignments. Parameters ---------- map_func_calls : dict mapping of function calls to assigned constants. - grid_spacing : int or float - Used to get site.pos from site.tag coords : sequence of strings If left as None coordinates will not be read from a site. onsite : bool @@ -606,8 +603,8 @@ def _assign_symbols(map_func_calls, grid_spacing, if coords: site = 'site' if onsite else 'site1' - args = ', '.join(coords), str(grid_spacing), site - lines.append('({}, ) = {} * {}.tag'.format(*args)) + args = ', '.join(coords), site + lines.append('({}, ) = {}.pos'.format(*args)) for k, v in map_func_calls.items(): lines.append("{} = {}".format(v, _print_sympy(k))) @@ -667,10 +664,7 @@ def _builder_value(expr, coords, grid_spacing, onsite, else: return complex(expr) - lines = _assign_symbols(map_func_calls, onsite=onsite, - grid_spacing=grid_spacing, - coords=coords) - + lines = _assign_symbols(map_func_calls, onsite=onsite, coords=coords) lines.append(return_string) separator = '\n ' diff --git a/kwant/continuum/tests/test_discretizer.py b/kwant/continuum/tests/test_discretizer.py index 19ec7454bce8309b67bbc923fb60fcd76a9a7cfe..54c8da657e6ad30ad7a9c74d9d61e516b23e3618 100644 --- a/kwant/continuum/tests/test_discretizer.py +++ b/kwant/continuum/tests/test_discretizer.py @@ -563,6 +563,16 @@ def test_grid_input(ham, grid_offset, offset, norbs): assert tmp.lattice.norbs == norbs +def test_grid_offset_passed_to_functions(): + V = lambda x: x + grid = Monatomic([[1, ]], offset=[0.5, ]) + tb = discretize('V(x)', 'x', grid=grid) + onsite = tb[tb.lattice(0)] + bools = [np.allclose(onsite(tb.lattice(i), V), V(tb.lattice(i).pos)) + for i in [0, 1, 5]] + assert all(bools) + + @pytest.mark.parametrize("ham, coords, grid", [ ("k_x", None, Monatomic([[1, 0]])), ("k_x", 'xy', Monatomic([[1, 0]])),