def test_map(engine):
plotter.set_engine(engine)
syst = syst_2d()
with tempfile.NamedTemporaryFile(
'w+b', suffix=plotter_file_suffix(engine)) as out:
out_filename = out.name
plotter.map(syst,
lambda site: site.tag[0],
pos_transform=good_transform,
file=out_filename,
method='linear',
a=4,
oversampling=4,
cmap='flag',
show=False)
pytest.raises(ValueError,
plotter.map,
syst,
lambda site: site.tag[0],
pos_transform=bad_transform,
file=out_filename)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
plotter.map(syst.finalized(),
range(len(syst.sites())),
file=out_filename,
show=False)
pytest.raises(ValueError,
plotter.map,
syst,
range(len(syst.sites())),
file=out_filename)
def test_plotly_plot():
plotter.set_engine('plotly')
plot = plotter.plot
syst2d = syst_2d()
syst3d = syst_3d()
color_opts = ['black', (lambda site: site.tag[0]),
lambda site: (abs(site.tag[0] / 100),
abs(site.tag[1] / 100), 0)]
engine = plotter.get_engine()
with tempfile.NamedTemporaryFile('w+b', suffix=plotter_file_suffix(engine)) as out:
out_filename = out.name
for color in color_opts:
for syst in (syst2d, syst3d):
plot(syst, site_color=color, cmap='binary', file=out_filename, show=False)
color_opts = ['black', (lambda site, site2: site.tag[0]),
lambda site, site2: (abs(site.tag[0] / 100),
abs(site.tag[1] / 100), 0)]
syst2d.leads = []
plot(syst2d, file=out_filename, show=False)
del syst2d[list(syst2d.hoppings())]
plot(syst2d, file=out_filename, show=False)
plot(syst3d, file=out_filename, show=False)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
plot(syst2d.finalized(), file=out_filename, show=False)
# test 2D projections of 3D systems
plot(syst3d, file=out_filename, pos_transform=lambda pos: pos[:2], show=False)
def test_matplotlib_plot():
plotter.set_engine('matplotlib')
plot = plotter.plot
syst2d = syst_2d()
syst3d = syst_3d()
color_opts = [
'k', (lambda site: site.tag[0]), lambda site:
(abs(site.tag[0] / 100), abs(site.tag[1] / 100), 0)
]
engine = plotter.get_engine()
with tempfile.NamedTemporaryFile(
'w+b', suffix=plotter_file_suffix(engine)) as out:
out_filename = out.name
for color in color_opts:
for syst in (syst2d, syst3d):
fig = plot(syst,
site_color=color,
cmap='binary',
file=out_filename)
if (color != 'k' and isinstance(
color(next(iter(syst2d.sites()))), float)):
assert fig.axes[0].collections[0].get_array() is not None
assert len(fig.axes[0].collections) == 6
color_opts = [
'k', (lambda site, site2: site.tag[0]), lambda site, site2:
(abs(site.tag[0] / 100), abs(site.tag[1] / 100), 0)
]
for color in color_opts:
for syst in (syst2d, syst3d):
fig = plot(syst2d,
hop_color=color,
cmap='binary',
file=out_filename,
fig_size=(2, 10),
dpi=30)
if color != 'k' and isinstance(
color(next(iter(syst2d.sites())), None), float):
assert fig.axes[0].collections[1].get_array() is not None
assert isinstance(
plot(syst3d, file=out_filename).axes[0], mplot3d.axes3d.Axes3D)
syst2d.leads = []
plot(syst2d, file=out_filename)
del syst2d[list(syst2d.hoppings())]
plot(syst2d, file=out_filename)
plot(syst3d, file=out_filename)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
plot(syst2d.finalized(), file=out_filename)
# test 2D projections of 3D systems
plot(syst3d, file=out_filename, pos_transform=lambda pos: pos[:2])
def test_spectrum(engine):
plotter.set_engine(engine)
def ham_1d(a, b, c):
return a**2 + b**2 + c**2
def ham_2d(a, b, c):
return np.eye(2) * (a**2 + b**2 + c**2)
lat = kwant.lattice.chain(norbs=1)
syst = kwant.Builder()
syst[(lat(i) for i in range(3))] = lambda site, a, b: a + b
syst[lat.neighbors()] = lambda site1, site2, c: c
fsyst = syst.finalized()
vals = np.linspace(0, 1, 3)
with tempfile.NamedTemporaryFile('w+b', suffix=plotter_file_suffix(engine)) as out:
out_filename = out.name
for ham in (ham_1d, ham_2d, fsyst):
plotter.spectrum(ham, ('a', vals), params=dict(b=1, c=1), file=out_filename, show=False)
if engine == 'matplotlib':
# test with explicit figsize
plotter.spectrum(ham, ('a', vals), params=dict(b=1, c=1),
fig_size=(10, 10), file=out_filename, show=False)
for ham in (ham_1d, ham_2d, fsyst):
plotter.spectrum(ham, ('a', vals), ('b', 2 * vals),
params=dict(c=1), file=out_filename, show=False)
if engine == 'matplotlib':
# test with explicit figsize
plotter.spectrum(ham, ('a', vals), ('b', 2 * vals),
params=dict(c=1), fig_size=(10, 10), file=out_filename, show=False)
if engine == 'matplotlib':
# test 2D plot and explicitly passing axis
fig = pyplot.figure()
ax = fig.add_subplot(1, 1, 1, projection='3d')
plotter.spectrum(ham_1d, ('a', vals), ('b', 2 * vals),
params=dict(c=1), ax=ax, file=out_filename, show=False)
# explicitly pass axis without 3D support
ax = fig.add_subplot(1, 1, 1)
with pytest.raises(TypeError):
plotter.spectrum(ham_1d, ('a', vals), ('b', 2 * vals),
params=dict(c=1), ax=ax, file=out_filename, show=False)
def mask(a, b):
return a > 0.5
with tempfile.NamedTemporaryFile('w+b', suffix=plotter_file_suffix(engine)) as out:
out_filename = out.name
plotter.spectrum(ham, ('a', vals), ('b', 2 * vals), params=dict(c=1),
mask=mask, file=out_filename, show=False)
def test_plot_raises_on_bad_site_spec(engine):
syst = kwant.Builder()
lat = kwant.lattice.square(norbs=1)
syst[(lat(i, j) for i in range(5) for j in range(5))] = None
# Cannot provide site_size as an array when syst is a Builder
plotter.set_engine(engine)
with pytest.raises(TypeError):
plotter.plot(syst, site_size=[1] * 25)
# Cannot provide site_size as an array when syst is a Builder
with pytest.raises(TypeError):
plotter.plot(syst, site_symbol=['o'] * 25)
def test_current():
plotter.set_engine('matplotlib')
syst = syst_2d().finalized()
J = kwant.operator.Current(syst)
current = J(kwant.wave_function(syst, energy=1)(1)[0])
# Test good codepath
with tempfile.NamedTemporaryFile('w+b') as out:
plotter.current(syst, current, file=out)
fig = pyplot.Figure()
ax = fig.add_subplot(1, 1, 1)
plotter.current(syst, current, ax=ax, file=out)
def test_plot_more_site_families_than_colors(engine):
# test against regression reported in
# https://gitlab.kwant-project.org/kwant/kwant/issues/257
ncolors = len(pyplot.rcParams['axes.prop_cycle'])
syst = kwant.Builder()
lattices = [kwant.lattice.square(name=i, norbs=1)
for i in range(ncolors + 1)]
for i, lat in enumerate(lattices):
syst[lat(i, 0)] = None
plotter.set_engine(engine)
with tempfile.NamedTemporaryFile('w+b', suffix=plotter_file_suffix(engine)) as out:
out_filename = out.name
plotter.plot(syst, file=out_filename, show=False)
def test_bands(engine):
plotter.set_engine(engine)
syst = syst_2d().finalized().leads[0]
with tempfile.NamedTemporaryFile('w+b', suffix=plotter_file_suffix(engine)) as out:
out_filename = out.name
plotter.bands(syst, show=False, file=out_filename)
plotter.bands(syst, show=False, momenta=np.linspace(0, 2 * np.pi), file=out_filename)
if engine == 'matplotlib':
plotter.bands(syst, show=False, fig_size=(10, 10), file=out_filename)
fig = pyplot.Figure()
ax = fig.add_subplot(1, 1, 1)
plotter.bands(syst, show=False, ax=ax, file=out_filename)
