In [1]:
import sys

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

init_notebook()
import itertools
import warnings

warnings.simplefilter("ignore", UserWarning)

# Onsite and hoppings terms for BHZ and QAH model
def onsite_BHZ(site, p):
    return (
        (p.M - 4 * p.B) * pauli.s0sz
        - 4 * p.D * pauli.s0s0
        + p.field * site.pos[1] * pauli.s0s0
    )


def hopx_BHZ(site1, site2, p):
    return p.B * pauli.s0sz + p.D * pauli.s0s0 + 0.5j * p.A * pauli.szsx


def hopy_BHZ(site1, site2, p):
    return p.B * pauli.s0sz + p.D * pauli.s0s0 - 0.5j * p.A * pauli.s0sy


def weak_hopping_BHZ(site1, site2, p):
    return p.t_inter * np.eye(4)


def onsite_QAH(site, p):
    return (p.mu - 4 * p.B) * pauli.sz + p.field * site.pos[1] * pauli.s0


def hopx_QAH(site1, site2, p):
    return p.B * pauli.sz + 0.5j * p.A * pauli.sx


def hopy_QAH(site1, site2, p):
    return p.B * pauli.sz + 0.5j * p.A * pauli.sy


def weak_hopping_QAH(site1, site2, p):
    return p.t_inter * np.eye(2)


# Systems
def create_screw_system(L, W, H, xs=None, ys=None, ye=None, pbc=True, model="BHZ"):
    """ Create system with screw dislocation.

    Function creates system with a screw dislocation. 
    L, W, H are dimension of scattering region.
    L, W are dimension of cross section. 
    Leads are attached from top and bottom (0,0,1) direction.

    xs, ys, ye describes where disloaction is placed.
    """
    if model == "BHZ":
        onsite, hopx, hopy, weak_hopping = (
            onsite_BHZ,
            hopx_BHZ,
            hopy_BHZ,
            weak_hopping_BHZ,
        )
    elif model == "QAH":
        onsite, hopx, hopy, weak_hopping = (
            onsite_QAH,
            hopx_QAH,
            hopy_QAH,
            weak_hopping_QAH,
        )

    def shape(pos):
        (x, y, z) = pos
        return (0 <= x < L) and (0 <= y < W) and (0 <= z < H)

    def lead_shape(pos):
        (x, y, z) = pos
        return (0 <= x < L) and (0 <= y < W)

    # calling kwant
    lat = kwant.lattice.general(np.identity(3))
    syst = kwant.Builder()

    sym = kwant.TranslationalSymmetry((0, 0, 1))
    lead = kwant.Builder(sym)

    # scattering system
    syst[lat.shape(shape, (0, 0, 0))] = onsite
    syst[kwant.HoppingKind((1, 0, 0), lat)] = hopx
    syst[kwant.HoppingKind((0, 1, 0), lat)] = hopy
    syst[kwant.HoppingKind((0, 0, 1), lat)] = weak_hopping

    # lead system
    lead[lat.shape(lead_shape, (0, 0, 0))] = onsite
    lead[kwant.HoppingKind((1, 0, 0), lat)] = hopx
    lead[kwant.HoppingKind((0, 1, 0), lat)] = hopy
    lead[kwant.HoppingKind((0, 0, 1), lat)] = weak_hopping

    # defining dislocation
    if xs is not None:
        for y in range(ys, ye):
            for z in range(H):
                del syst[lat(xs + 1, y, z), lat(xs, y, z)]

            del lead[lat(xs + 1, y, 0), lat(xs, y, 0)]
            lead[lat(xs + 1, y, z + 1), lat(xs, y, z)] = hopx

        for y, z in itertools.product(range(ys, ye), range(H - 1)):
            syst[lat(xs + 1, y, z + 1), lat(xs, y, z)] = hopx

    # adding periodic boundary conditions
    if pbc:
        for x in range(L):
            lead[lat(x, 0, 0), lat(x, W - 1, 0)] = hopy
            for z in range(H):
                syst[lat(x, 0, z), lat(x, W - 1, z)] = hopy

        for y in range(W):
            lead[lat(0, y, 0), lat(L - 1, y, 0)] = hopx
            for z in range(H):
                syst[lat(0, y, z), lat(L - 1, y, z)] = hopx

    syst.attach_lead(lead)
    syst.attach_lead(lead.reversed())

    return syst.finalized(), lead


def create_edge_dislocation_system(
    L, W, H, xs=None, ys=None, ye=None, pbc=True, model="BHZ"
):
    """ Create system with edge dislocation.

    Function creates system with an edge dislocation. 
    L, W, H are dimension of scattering region.
    L, W are dimension of cross section. 
    Leads are attached from top and bottom (0,0,1) direction.

    xs, ys, ye describes where disloaction is placed.
    """
    if model == "BHZ":
        onsite, hopx, hopy, weak_hopping = (
            onsite_BHZ,
            hopx_BHZ,
            hopy_BHZ,
            weak_hopping_BHZ,
        )
    elif model == "QAH":
        onsite, hopx, hopy, weak_hopping = (
            onsite_QAH,
            hopx_QAH,
            hopy_QAH,
            weak_hopping_QAH,
        )

    def shape(pos):
        (x, y, z) = pos
        return (0 <= x < L) and (0 <= y < W) and (0 <= z < H)

    def lead_shape(pos):
        (x, y, z) = pos
        return (0 <= x < L) and (0 <= y < W)

    # Calling kwant
    lat = kwant.lattice.general(np.identity(3))
    syst = kwant.Builder()

    sym = kwant.TranslationalSymmetry((0, 0, 1))
    lead = kwant.Builder(sym)

    # scattering system
    syst[lat.shape(shape, (0, 0, 0))] = onsite
    syst[kwant.HoppingKind((1, 0, 0), lat)] = weak_hopping
    syst[kwant.HoppingKind((0, 1, 0), lat)] = hopy
    syst[kwant.HoppingKind((0, 0, 1), lat)] = hopx

    # lead system
    lead[lat.shape(lead_shape, (0, 0, 0))] = onsite
    lead[kwant.HoppingKind((1, 0, 0), lat)] = weak_hopping
    lead[kwant.HoppingKind((0, 1, 0), lat)] = hopy
    lead[kwant.HoppingKind((0, 0, 1), lat)] = hopx

    # defining disclocation
    if xs != None:
        for y in range(ys, ye):
            del lead[lat(xs, y, 0)]
            lead[lat(xs + 1, y, 0), lat(xs - 1, y, 0)] = weak_hopping

            for z in range(H):
                del syst[lat(xs, y, z)]
                syst[lat(xs + 1, y, z), lat(xs - 1, y, z)] = weak_hopping

    # periodic boundary conditions
    if pbc:
        for x in range(L):
            lead[lat(x, 0, 0), lat(x, W - 1, 0)] = hopy
            for z in range(H):
                syst[lat(x, 0, z), lat(x, W - 1, z)] = hopy

        for y in range(W):
            lead[lat(0, y, 0), lat(L - 1, y, 0)] = weak_hopping
            for z in range(H):
                syst[lat(0, y, z), lat(L - 1, y, z)] = weak_hopping

    # attaching leads
    syst.attach_lead(lead)
    syst.attach_lead(lead.reversed())

    return syst.finalized(), lead


def get_densities(lead, momentum, p, model, sorting_mid=0.0):
    """ Calculate density of states in the lead at a given momentum. """
    coord = np.array([i.pos for i in lead.sites])
    xy = coord[coord[:, 2] == 0][:, :-1]
    indxs_xy = np.lexsort(xy.T)
    xy = xy[indxs_xy, :]

    h, t = lead.cell_hamiltonian(params=dict(p=p)), lead.inter_cell_hopping(params=dict(p=p))
    h_k = lambda k: h + t * np.exp(-1j * k) + t.T.conj() * np.exp(1j * k)

    vals, vecs = np.linalg.eigh(h_k(momentum))

    if model == "BHZ":
        num_orbital = 4
    if model == "QAH":
        num_orbital = 2

    densities = np.linalg.norm(vecs.reshape(-1, num_orbital, len(vecs)), axis=1) ** 2

    indxs = np.argsort(abs(vals - sorting_mid))
    vals = vals[indxs]
    densities = densities[:, indxs]
    densities = densities[indxs_xy, :]

    L, W = int(np.max(xy[:, 0]) + 1), int(np.max(xy[:, 1]) + 1)
    twod_densities = np.zeros((W, L, densities.shape[1]))

    for coord, val in zip(xy, densities):
        i, j = np.array(coord, dtype=int)
        twod_densities[j, i, :] = val

    return twod_densities, vals


def get_spectrum_and_densities(sys_func, p, model, momentum, kwargs):
    syst, lead = sys_func(L=11, W=21, H=5, xs=4, ys=5, ye=16, model=model)
    spect = spectrum(lead, p, **kwargs)
    densities = get_densities(syst.leads[0], momentum, p, model)
    return spect, densities


def create_hm(sys_func, momentum, kwargs):
    parameters = {
        "BHZ": {"A": 1.0, "B": 1.0, "D": 0.0, "M": 0.8},
        "QAH": {"A": 1.0, "B": 1.0, "D": 0.0, "mu": 0.8},
    }

    p_BHZ = SimpleNamespace(field=0.005, t_inter=-0.1, **parameters["BHZ"])
    p_QAH = SimpleNamespace(field=0.01, t_inter=-0.1, **parameters["QAH"])

    spectrum_QAH, densities_QAH = get_spectrum_and_densities(
        sys_func, p_QAH, "QAH", momentum, kwargs
    )
    spectrum_BHZ, densities_BHZ = get_spectrum_and_densities(
        sys_func, p_BHZ, "BHZ", momentum, kwargs
    )

    hm_dict = {}
    for n in range(7):
        for model, density, spect in zip(
            ["BHZ", "QAH"], [densities_BHZ, densities_QAH], [spectrum_BHZ, spectrum_QAH]
        ):
            hm_dict[n, model] = (
                spect
                * holoviews.Points((momentum, density[1][n]))
                * holoviews.VLine(momentum)
            ).relabel("Bandstructure") + holoviews.Raster(
                density[0][:, :, n], label=r"$\left|\psi\right|^2$"
            )

    return holoviews.HoloMap(hm_dict, kdims=["n", "model"])
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:41.154670.