Skip to content
Snippets Groups Projects
Commit 32c0972b authored by Christoph Groth's avatar Christoph Groth
Browse files

replace verbose by a printable builder

parent 184c7f3f
Branches
No related tags found
No related merge requests found
--- original
+++ modified
@@ -9,6 +9,8 @@
@@ -9,6 +9,7 @@
# --------------------------
# - discretizer module from kwant.continuum
+import _defs
+from contextlib import redirect_stdout
import kwant
import scipy.sparse.linalg
@@ -20,9 +22,19 @@
@@ -20,10 +21,19 @@
from matplotlib import pyplot as plt
......@@ -23,14 +22,14 @@
+
def stadium_system(L=20, W=20):
hamiltonian = "k_x**2 + k_y**2 + V(x, y)"
- template = kwant.continuum.discretize(hamiltonian, verbose=True)
template = kwant.continuum.discretize(hamiltonian)
- print(template)
+ with open('discretizer_verbose.txt', 'w') as f:
+ with redirect_stdout(f):
+ template = kwant.continuum.discretize(hamiltonian, verbose=True)
+ print(template, file=f)
def stadium(site):
(x, y) = site.pos
@@ -43,7 +55,7 @@
@@ -44,7 +54,7 @@
ham = syst.hamiltonian_submatrix(params=dict(V=potential), sparse=True)
evecs = scipy.sparse.linalg.eigsh(ham, k=10, which='SM')[1]
kwant.plotter.map(syst, abs(evecs[:, n])**2, show=False)
......@@ -39,7 +38,7 @@
def qsh_system(a=20, L=2000, W=1000):
@@ -90,7 +102,8 @@
@@ -91,7 +101,8 @@
plt.ylim(-0.05, 0.05)
plt.xlabel('momentum [1/A]')
plt.ylabel('energy [eV]')
......@@ -49,7 +48,7 @@
# get scattering wave functions at E=0
wf = kwant.wave_function(syst, energy=0, params=params)
@@ -118,7 +131,7 @@
@@ -119,7 +130,7 @@
ax1.set_title('Probability density')
ax2.set_title('Spin density')
......@@ -58,7 +57,7 @@
def lattice_spacing():
@@ -159,7 +172,7 @@
@@ -160,7 +171,7 @@
plot(ax1, a=1)
plot(ax2, a=.25)
......@@ -67,27 +66,25 @@
def substitutions():
@@ -172,14 +185,20 @@
@@ -173,15 +184,18 @@
sympify('k_x**2 * sz + alpha * k_x * sx', subs=subs),
)
- print(e[0] == e[1] == e[2])
+ with open('discretizer_subs_1.txt', 'w') as f:
+ with redirect_stdout(f):
+ print(e[0] == e[1] == e[2])
+ print(e[0] == e[1] == e[2], file=f)
subs = {'A': 'A(x) + B', 'V': 'V(x) + V_0', 'C': 5}
- print(sympify('k_x * A * k_x + V + C', subs=subs))
+ with open('discretizer_subs_2.txt', 'w') as f:
+ with redirect_stdout(f):
+ print(sympify('k_x * A * k_x + V + C', subs=subs))
+ print(sympify('k_x * A * k_x + V + C', subs=subs), file=f)
def main():
- template = kwant.continuum.discretize('k_x * A(x) * k_x', verbose=True)
template = kwant.continuum.discretize('k_x * A(x) * k_x')
- print(template)
+ with open('discretizer_intro_verbose.txt', 'w') as f:
+ with redirect_stdout(f):
+ kwant.continuum.discretize('k_x * A(x) * k_x', verbose=True)
+ print(template, file=f)
syst = stadium_system()
plot_eigenstate(syst)
......@@ -23,7 +23,8 @@ from matplotlib import pyplot as plt
def stadium_system(L=20, W=20):
#HIDDEN_BEGIN_template
hamiltonian = "k_x**2 + k_y**2 + V(x, y)"
template = kwant.continuum.discretize(hamiltonian, verbose=True)
template = kwant.continuum.discretize(hamiltonian)
print(template)
#HIDDEN_END_template
#HIDDEN_BEGIN_fill
......@@ -204,7 +205,8 @@ def substitutions():
def main():
#HIDDEN_BEGIN_symbolic_discretization
template = kwant.continuum.discretize('k_x * A(x) * k_x', verbose=True)
template = kwant.continuum.discretize('k_x * A(x) * k_x')
print(template)
#HIDDEN_END_symbolic_discretization
syst = stadium_system()
......
......@@ -68,9 +68,7 @@ It is worth noting that ``discretize`` treats ``k_x`` and ``x`` as
non-commuting operators, and so their order is preserved during the
discretization process.
Setting the ``verbose`` parameter to ``True`` prints extra information about the
onsite and hopping functions assigned to the ``Builder`` produced
by ``discretize``:
The builder produced by ``discretize`` may be printed to show the source code of its onsite and hopping functions (this is a special feature of builders returned by ``discretize``):
.. literalinclude:: ../images/discretizer_intro_verbose.txt
......
......@@ -7,6 +7,7 @@
# http://kwant-project.org/authors.
from collections import defaultdict
import itertools
import numpy as np
import tinyarray as ta
......@@ -25,18 +26,62 @@ from ._common import (sympify, gcd, position_operators, momentum_operators,
__all__ = ['discretize']
################ Globals variables and definitions
_wf = sympy.Function('_internal_unique_name', commutative=False)
_momentum_operators = {s.name: s for s in momentum_operators}
_position_operators = {s.name: s for s in position_operators}
_displacements = {s: sympy.Symbol('_internal_a_{}'.format(s)) for s in 'xyz'}
class _DiscretizedBuilder(builder.Builder):
"""A builder that is made from a discretized model and knows how to
pretty-print itself."""
def __init__(self, symmetry=None, discrete_coords=[], **kwargs):
super().__init__(symmetry, **kwargs)
self._discrete_coords = discrete_coords
def __str__(self):
result = []
sv = list(s for s in self.site_value_pairs())
if len(sv) != 1:
raise ValueError("Cannot pretty-print _DiscretizedBuilder: "
"must contain a single site.")
site, site_value = sv[0]
if any(e != 0 for e in site.tag):
raise ValueError("Cannot pretty-print _DiscretizedBuilder: "
"site must be located at origin.")
result.extend(["# Discrete coordinates: ",
" ".join(self._discrete_coords),
"\n\n"])
for key, val in itertools.chain(self.site_value_pairs(),
self.hopping_value_pairs()):
if isinstance(key, builder.Site):
result.append("# Onsite element:\n")
else:
a, b = key
assert a is site
result.extend(["# Hopping in direction ",
str(tuple(b.tag)),
":\n"])
result.append(val._source if callable(val) else repr(val))
result.append('\n\n')
result.pop()
return "".join(result)
# For the Jupyter notebook:
def __repr_html__(self):
return self.__str__()
################ Interface functions
def discretize(hamiltonian, discrete_coords=None, *, grid_spacing=1,
subs=None, verbose=False):
subs=None):
"""Construct a tight-binding model from a continuum Hamiltonian.
This is a convenience function that is equivalent to first calling
......@@ -66,8 +111,6 @@ def discretize(hamiltonian, discrete_coords=None, *, grid_spacing=1,
proceeding further. For example:
``subs={'k': 'k_x + I * k_y'}`` or
``subs={'s_z': [[1, 0], [0, -1]]}``.
verbose : bool, default: False
If ``True`` additional information will be printed.
Returns
-------
......@@ -75,14 +118,12 @@ def discretize(hamiltonian, discrete_coords=None, *, grid_spacing=1,
with translational symmetry, which can be used as a template.
"""
tb, coords = discretize_symbolic(hamiltonian, discrete_coords,
subs=subs, verbose=verbose)
subs=subs)
return build_discretized(tb, coords, grid_spacing=grid_spacing,
verbose=verbose)
return build_discretized(tb, coords, grid_spacing=grid_spacing)
def discretize_symbolic(hamiltonian, discrete_coords=None, *,
subs=None, verbose=False):
def discretize_symbolic(hamiltonian, discrete_coords=None, *, subs=None):
"""Discretize a continuous Hamiltonian into a tight-binding representation.
The two objects returned by this function may be used directly as the first
......@@ -109,8 +150,6 @@ def discretize_symbolic(hamiltonian, discrete_coords=None, *,
proceeding further. For example:
``subs={'k': 'k_x + I * k_y'}`` or
``subs={'s_z': [[1, 0], [0, -1]]}``.
verbose : bool, default: False
If ``True`` additional information will be printed.
Returns
-------
......@@ -146,10 +185,6 @@ def discretize_symbolic(hamiltonian, discrete_coords=None, *,
"your input. You can use the 'discrete_coords'"
"parameter to provide them.")
if verbose:
print('Discrete coordinates set to: ',
discrete_coords, end='\n\n')
onsite_zeros = (0,) * len(discrete_coords)
if not isinstance(hamiltonian, sympy.matrices.MatrixBase):
......@@ -179,7 +214,7 @@ def discretize_symbolic(hamiltonian, discrete_coords=None, *,
def build_discretized(tb_hamiltonian, discrete_coords, *,
grid_spacing=1, subs=None, verbose=False):
grid_spacing=1, subs=None):
"""Create a template Builder from a symbolic tight-binding Hamiltonian.
This return values of `~kwant.continuum.discretize_symbolic` may be used
......@@ -204,8 +239,6 @@ def build_discretized(tb_hamiltonian, discrete_coords, *,
proceeding further. For example:
``subs={'k': 'k_x + I * k_y'}`` or
``subs={'s_z': [[1, 0], [0, -1]]}``.
verbose : bool, default: False
If ``True`` additional information will be printed.
Returns
-------
......@@ -222,15 +255,7 @@ def build_discretized(tb_hamiltonian, discrete_coords, *,
discrete_coords = sorted(discrete_coords)
tb = {}
first = True
for n, (offset, hopping) in enumerate(tb_hamiltonian.items()):
if verbose:
if first:
first = False
else:
print('\n')
print("Function generated for {}:".format(offset))
onsite = all(i == 0 for i in offset)
if onsite:
......@@ -238,9 +263,8 @@ def build_discretized(tb_hamiltonian, discrete_coords, *,
else:
name = 'hopping_{}'.format(n)
tb[offset] = _value_function(hopping, discrete_coords,
grid_spacing, onsite, name,
verbose=verbose)
tb[offset] = _builder_value(hopping, discrete_coords,
grid_spacing, onsite, name)
dim = len(discrete_coords)
onsite_zeros = (0,) * dim
......@@ -256,7 +280,8 @@ def build_discretized(tb_hamiltonian, discrete_coords, *,
hoppings = {builder.HoppingKind(tuple(-i for i in d), lat): val
for d, val in tb.items()}
syst = builder.Builder(lattice.TranslationalSymmetry(*prim_vecs))
syst = _DiscretizedBuilder(lattice.TranslationalSymmetry(*prim_vecs),
discrete_coords)
syst[lat(*onsite_zeros)] = onsite
for hop, val in hoppings.items():
syst[hop] = val
......@@ -507,9 +532,9 @@ def _assign_symbols(map_func_calls, grid_spacing,
return lines
def _value_function(expr, discrete_coords, grid_spacing, onsite,
name='_anonymous_func', verbose=False):
"""Generate a numeric function from a sympy expression.
def _builder_value(expr, discrete_coords, grid_spacing, onsite,
name='_anonymous_func'):
"""Generate a builder value from a sympy expression.
Parameters
----------
......@@ -519,14 +544,15 @@ def _value_function(expr, discrete_coords, grid_spacing, onsite,
List of coodinates present in the system.
grid_spacing : int or float
Lattice spacing of the system
verbose : bool, default: False
If True, the function body is printed.
Returns
-------
numerical function that can be used with Kwant.
`expr` transformed into an object that can be used as a
`kwant.builder.Builder` value. Either a numerical value
(``tinyarray.array`` instance or complex number) or a value function. In
the case of a function, the source code is available in its `_source`
attribute.
"""
expr = expr.subs({sympy.Symbol('a'): grid_spacing})
return_string, map_func_calls, const_symbols, _cache = _return_string(
expr, discrete_coords=discrete_coords)
......@@ -545,14 +571,9 @@ def _value_function(expr, discrete_coords, grid_spacing, onsite,
if (not required_kwargs) and (discrete_coords is None):
# we can just use a constant value instead of a value function
if isinstance(expr, sympy.MatrixBase):
output = ta.array(expr.tolist(), complex)
return ta.array(expr.tolist(), complex)
else:
output = complex(expr)
if verbose:
print("\n{}".format(output))
return output
return complex(expr)
lines = _assign_symbols(map_func_calls, onsite=onsite,
grid_spacing=grid_spacing,
......@@ -573,12 +594,13 @@ def _value_function(expr, discrete_coords, grid_spacing, onsite,
namespace = {'pi': np.pi}
namespace.update(_cache)
if verbose:
for k, v in _cache.items():
print("\n{} = (\n{})".format(k, repr(np.array(v))))
print('\n' + func_code)
source = []
for k, v in _cache.items():
source.append("{} = (\n{})\n".format(k, repr(np.array(v))))
source.append(func_code)
exec(func_code, namespace)
f = namespace[name]
f._source = "".join(source)
return f
......@@ -335,7 +335,7 @@ def test_numeric_functions_basic_symbolic():
@pytest.mark.parametrize('commutative', [ True, False])
def test_numeric_function_coords_from_site(commutative):
tb = {(0,): sympy.symbols('x', commutative=commutative)}
builder = build_discretized(tb, 'x', verbose=True)
builder = build_discretized(tb, 'x')
lat = next(iter(builder.sites()))[0]
onsite = builder[lat(0)]
......@@ -515,30 +515,3 @@ def test_numeric_functions_with_parameter():
rhs = f_num
assert np.allclose(lhs, rhs)
def test_basic_verbose(capsys): # or use "capfd" for fd-level
discretize('k_x * A(x) * k_x', verbose=True)
out, err = capsys.readouterr()
assert "Discrete coordinates set to" in out
assert "Function generated for (0,)" in out
def test_that_verbose_covers_all_hoppings(capsys):
discretize('k_x**2 + k_y**2 + k_x*k_y', verbose=True)
out, err = capsys.readouterr()
for tag in [(0, 1), (0, 0), (1, -1), (1, 1)]:
assert "Function generated for {}".format(tag) in out
def test_verbose_cache(capsys):
discretize('[[k_x * A(x) * k_x]]', verbose=True)
out, err = capsys.readouterr()
assert '_cache_0' in out
def test_no_output_when_verbose_false(capsys):
discretize('[[k_x * A(x) * k_x]]', verbose=False)
out, err = capsys.readouterr()
assert out == ''
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment