test_plotter.py 2.19 KB
 Kwant authors committed Jan 10, 2011 1 ``````import tempfile, os `````` Christoph Groth committed May 03, 2012 2 ``````from nose.tools import assert_raises `````` Kwant authors committed Jan 10, 2011 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 ``````import numpy as np import kwant from kwant import plotter lat = kwant.lattice.Square() def make_ribbon(width, dir, E, t): b = kwant.Builder(kwant.TranslationalSymmetry([(dir, 0)])) # Add sites to the builder. for y in xrange(width): b[lat(0, y)] = E # Add hoppings to the builder. for y in xrange(width): b[lat(0, y), lat(1, y)] = t if y+1 < width: b[lat(0, y), lat(0, y+1)] = t return b def make_rectangle(length, width, E, t): b = kwant.Builder() # Add sites to the builder. for x in xrange(length): for y in xrange(width): b[lat(x, y)] = E # Add hoppings to the builder. for x in xrange(length): for y in xrange(width): if x+1 < length: b[lat(x, y), lat(x+1, y)] = t if y+1 < width: b[lat(x, y), lat(x, y+1)] = t return b def test_plot(): E = 4.0 t = -1.0 length = 5 width = 5 b = make_rectangle(length, width, E, t) b.attach_lead(make_ribbon(width, -1, E, t)) b.attach_lead(make_ribbon(width, 1, E, t)) directory = tempfile.mkdtemp() filename = os.path.join(directory, "test.pdf") kwant.plot(b.finalized(), filename=filename, symbols=plotter.Circle(r=0.25, fcol=plotter.red), lines=plotter.Line(lw=0.1, lcol=plotter.red), lead_symbols=plotter.Circle(r=0.25, fcol=plotter.black), lead_lines=plotter.Line(lw=0.1, lcol=plotter.black), lead_fading=[0, 0.2, 0.4, 0.6, 0.8]) os.unlink(filename) os.rmdir(directory) `````` Christoph Groth committed May 03, 2012 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 `````` def test_non_2d_fails(): directory = tempfile.mkdtemp() filename = os.path.join(directory, "test.pdf") for d in [1, 2, 3, 15]: b = kwant.Builder() lat = kwant.make_lattice(np.identity(d)) site = kwant.builder.Site(lat, (0,) * d) b[site] = 0 if d == 2: kwant.plot(b, filename=filename) plotter.interpolate(b, b) else: assert_raises(ValueError, kwant.plot, b) assert_raises(ValueError, plotter.interpolate, b, b) os.unlink(filename) os.rmdir(directory)``````