Commit f10f3de2 authored by Joseph Weston's avatar Joseph Weston

use an infinite cycle of colors when plotting site families

Previously if there were more site families than colours in the
matplotlib cycle, then plotting would raise an exception. Now
we cycle through the colours indefinitely.

Closes #257
parent 7234365b
Pipeline #13604 passed with stages
in 20 minutes and 2 seconds
......@@ -48,6 +48,12 @@ def _sample_array(array, n_samples, rng=None):
# matplotlib helper functions.
def _color_cycle():
"""Infinitely cycle through colors from the matplotlib color cycle."""
props = _p.matplotlib.rcParams['axes.prop_cycle']
return itertools.cycle(x['color'] for x in props)
def _make_figure(dpi, fig_size, use_pyplot=False):
if 'matplotlib.backends' not in sys.modules:
warnings.warn(
......@@ -944,7 +950,7 @@ def plot(sys, num_lead_cells=2, unit='nn',
fancy_indexing = False
if site_color is None:
cycle = (x['color'] for x in _p.matplotlib.rcParams['axes.prop_cycle'])
cycle = _color_cycle()
if isinstance(syst, (builder.FiniteSystem, builder.InfiniteSystem)):
# Skipping the leads for brevity.
families = sorted({site.family for site in syst.sites})
......
......@@ -155,6 +155,20 @@ def test_plot():
# test 2D projections of 3D systems
plot(syst3d, file=out, pos_transform=lambda pos: pos[:2])
@pytest.mark.skipif(not _plotter.mpl_available, reason="Matplotlib unavailable.")
def test_plot_more_site_families_than_colors():
# 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) for i in range(ncolors + 1)]
for i, lat in enumerate(lattices):
syst[lat(i, 0)] = None
with tempfile.TemporaryFile('w+b') as out:
plotter.plot(syst, file=out)
def good_transform(pos):
x, y = pos
return y, x
......
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