(LearnerND) add advanced usage example

We should add the following "real world usage" code to the tutorial as an "Advanced example" once we merged !127 (merged) and !124 (merged).

This as a downloadable file kwant_functions.py

from functools import lru_cache
import numpy as np
import scipy.linalg
import scipy.spatial
import kwant


@lru_cache()
def create_syst(unit_cell):
    lat = kwant.lattice.Polyatomic(unit_cell, [(0, 0, 0)])
    syst = kwant.Builder(kwant.TranslationalSymmetry(*lat.prim_vecs))
    syst[lat.shape(lambda _: True, (0, 0, 0))] = 6
    syst[lat.neighbors()] = -1
    return kwant.wraparound.wraparound(syst).finalized()


def get_brillouin_zone(unit_cell):
    syst = create_syst(unit_cell)
    A = get_A(syst)
    neighbours = kwant.linalg.lll.voronoi(A)
    lattice_points = np.concatenate(([[0, 0, 0]], neighbours))
    lattice_points = 2 * np.pi * (lattice_points @ A.T)
    vor = scipy.spatial.Voronoi(lattice_points)
    brillouin_zone = vor.vertices[vor.regions[vor.point_region[0]]]
    return scipy.spatial.ConvexHull(brillouin_zone)


def momentum_to_lattice(k, syst):
    A = get_A(syst)
    k, residuals = scipy.linalg.lstsq(A, k)[:2]
    if np.any(abs(residuals) > 1e-7):
        raise RuntimeError("Requested momentum doesn't correspond"
                           " to any lattice momentum.")
    return k


def get_A(syst):
    B = np.asarray(syst._wrapped_symmetry.periods).T
    return np.linalg.pinv(B).T


def energies(k, unit_cell):
    syst = create_syst(unit_cell)
    k_x, k_y, k_z = momentum_to_lattice(k, syst)
    params = {'k_x': k_x, 'k_y': k_y, 'k_z': k_z}
    H = syst.hamiltonian_submatrix(params=params)
    energies = np.linalg.eigvalsh(H)
    return min(energies)

This in the tutorial:

from functools import partial

from ipywidgets import interact_manual
import numpy as np

import adaptive
from kwant_functions import get_brillouin_zone, energies

adaptive.notebook_extension()

# Define the lattice vectors of some common unit cells
lattices = dict(
    hexegonal=(
        (0, 1, 0),
        (np.cos(-np.pi / 6), np.sin(-np.pi / 6), 0),
        (0, 0, 1)
    ),
    simple_cubic=(
        (1, 0, 0),
        (0, 1, 0),
        (0, 0, 1)
    ),
    fcc=(
        (0, .5, .5),
        (.5, .5, 0),
        (.5, 0, .5)
    ),
    bcc=(
        (-.5, .5, .5),
        (.5, -.5, .5),
        (.5, .5, -.5)
    )
)


learners = []
for name, unit_cell in lattices.items():
    hull = get_brillouin_zone(unit_cell)
    learner = adaptive.LearnerND(partial(energies, unit_cell=unit_cell), hull)
    learner.fname = name
    learners.append(learner)

learner = adaptive.BalancingLearner(learners, strategy='npoints')
adaptive.runner.simple(learner, goal=lambda l: l.learners[0].npoints > 20)

# XXX: maybe this could even be a `BalancingLearner` method.
def select(name, learner=learner):
    return next(l for l in learner.learners if l.fname == name)

def iso(unit_cell, level=8.5):
    l = select(unit_cell)
    return l.plot_isosurface(level=level)

def plot_tri(unit_cell):
    l = select(unit_cell)
    return l.plot_3D()

interact_manual(plot_tri, unit_cell=lattices.keys())  # this won't work, but something along these lines
interact_manual(iso, level=(-6, 9, 0.1), unit_cell=lattices.keys())