In [1]:
import sys

sys.path.append("../../code")
from init_mooc_nb import *

init_notebook()
import scipy

%output size=150 fig='png'
pi_ticks = [(-np.pi, r"$-\pi$"), (0, "0"), (np.pi, r"$\pi$")]


def haldane(w=20, boundary="zigzag"):
    def ribbon_shape_zigzag(pos):
        return -0.5 / np.sqrt(3) - 0.1 <= pos[1] < np.sqrt(3) * w / 2 + 0.01

    def ribbon_shape_armchair(pos):
        return -1 <= pos[0] < w

    def onsite(site, p):
        if site.family == a:
            return p.m
        else:
            return -p.m

    def nn_hopping(site1, site2, p):
        return p.t

    def nnn_hopping(site1, site2, p):
        return 1j * p.t_2

    lat = kwant.lattice.honeycomb()
    a, b = lat.sublattices
    nnn_hoppings_a = (((-1, 0), a, a), ((0, 1), a, a), ((1, -1), a, a))
    nnn_hoppings_b = (((1, 0), b, b), ((0, -1), b, b), ((-1, 1), b, b))

    if boundary == "zigzag":
        syst = kwant.Builder(kwant.TranslationalSymmetry((1, 0)))
        syst[lat.shape(ribbon_shape_zigzag, (0, 0))] = onsite
    elif boundary == "armchair":
        syst = kwant.Builder(kwant.TranslationalSymmetry((0, np.sqrt(3))))
        syst[lat.shape(ribbon_shape_armchair, (0, 0))] = onsite
    else:
        syst = kwant.Builder(kwant.TranslationalSymmetry(*lat.prim_vecs))
        syst[lat.shape(lambda pos: True, (0, 0))] = onsite

    syst[lat.neighbors()] = nn_hopping
    syst[
        [kwant.builder.HoppingKind(*hopping) for hopping in nnn_hoppings_a]
    ] = nnn_hopping
    syst[
        [kwant.builder.HoppingKind(*hopping) for hopping in nnn_hoppings_b]
    ] = nnn_hopping

    return syst


def Qi_Wu_Zhang():
    def onsite(site, p):
        return -p.mu * pauli.sz

    def hopx(site1, site2, p):
        return -0.5j * p.delta * pauli.sy - p.t * pauli.sz

    def hopy(site1, site2, p):
        return -1j * p.gamma * pauli.sx - p.gamma * pauli.sz

    lat = kwant.lattice.square()

    syst = kwant.Builder(kwant.TranslationalSymmetry(*lat.prim_vecs))
    syst[lat.shape(lambda pos: True, (0, 0))] = onsite
    syst[kwant.HoppingKind((1, 0), lat)] = hopx
    syst[kwant.HoppingKind((0, 1), lat)] = hopy
    return syst


def berry_curvature(syst, p, ks, num_filled_bands=1):
    """Berry curvature of a system.
    
    Parameters:
    -----------
    sys : kwant.Builder
        A 2D infinite system.
    p : SimpleNamespace
        The arguments expected by the system.
    ks : 1D array-like
        Values of momentum grid to be used for Berry curvature calculation.
    num_filled_bands : int
        The number of filled bands.

    Returns:
    --------
    bc : 2D array
        Berry curvature on each square in a `ks x ks` grid.
    """
    # Calculate an array of eigenvectors.
    B = np.array(syst.symmetry.periods).T
    A = B @ np.linalg.inv(B.T @ B)

    syst = kwant.wraparound.wraparound(syst).finalized()

    def energy(kx, ky):
        k = np.array([kx, ky])
        kx, ky = np.linalg.solve(A, k)
        H = syst.hamiltonian_submatrix(params=dict(p=p, k_x=kx, k_y=ky), sparse=False)
        return scipy.linalg.eigh(H)[1]

    vectors = np.array(
        [[energy(kx, ky)[:, :num_filled_bands] for kx in ks] for ky in ks]
    )

    # The actual Berry curvature calculation
    vectors_x = np.roll(vectors, 1, 0)
    vectors_xy = np.roll(vectors_x, 1, 1)
    vectors_y = np.roll(vectors, 1, 1)

    shifted_vecs = [vectors, vectors_x, vectors_xy, vectors_y]

    v_shape = vectors.shape

    shifted_vecs = [i.reshape(-1, v_shape[-2], v_shape[-1]) for i in shifted_vecs]

    dets = np.ones(len(shifted_vecs[0]), dtype=complex)
    for vec, shifted in zip(shifted_vecs, np.roll(shifted_vecs, 1, 0)):
        dets *= [np.linalg.det(a.T.conj() @ b) for a, b in zip(vec, shifted)]
    bc = np.angle(dets).reshape(int(np.sqrt(len(dets))), -1)

    bc = (bc + np.pi / 2) % (np.pi) - np.pi / 2

    return bc


def plot_berry_curvature(syst, p, ks=None, title=None):
    if ks is None:
        ks = np.linspace(-np.pi, np.pi, 150, endpoint=False)
    bc = berry_curvature(syst, p, ks)[1:-1, 1:-1]
    vmax = max(np.abs(bc).min(), np.abs(bc).max())
    kwargs = {
        "bounds": (ks.min(), ks.min(), ks.max(), ks.max()),
        "kdims": [r"$k_x$", r"$k_y$"],
    }

    if callable(title):
        kwargs["label"] = title(p)

    plot = {"xticks": pi_ticks, "yticks": pi_ticks}
    style = {"clims": [-vmax, vmax]}
    return holoviews.Image(bc, **kwargs).opts(plot=plot, style=style)


def title(p):
    title = r"$t={:.2}$, $t_2={:.2}$, $M={:.2}$"
    return title.format(p.t, p.t_2, p.m)
Populated the namespace with:
np, matplotlib, kwant, holoviews, init_notebook, SimpleNamespace, pprint_matrix, scientific_number, pretty_fmt_complex, plt, pf, display_html
from code/edx_components:
MoocVideo, MoocDiscussion, MoocCheckboxesAssessment, MoocMultipleChoiceAssessment, MoocSelfAssessment
from code/functions:
spectrum, hamiltonian_array, h_k, pauli
Using kwant 1.4.2 and holoviews 1.13.2
Executed on 2020-06-07 at 08:23:10.983708.