# Tutorial 2.3.1. Matrix structure of on-site and hopping elements
# ================================================================
#
# Physics background
# ------------------
#  Gaps in quantum wires with spin-orbit coupling and Zeeman splititng,
#  as theoretically predicted in
#   http://prl.aps.org/abstract/PRL/v90/i25/e256601
#  and (supposedly) experimentally oberved in
#   http://www.nature.com/nphys/journal/v6/n5/abs/nphys1626.html
#
# Kwant features highlighted
# --------------------------
#  - Numpy matrices as values in Builder

import kwant

# For plotting
from matplotlib import pyplot

# For matrix support
import tinyarray
from numpy import *

# define Pauli-matrices for convenience
sigma_0 = tinyarray.array([[1, 0], [0, 1]])
sigma_x = tinyarray.array([[0, 1], [1, 0]])
sigma_y = tinyarray.array([[0, -1j], [1j, 0]])
sigma_z = tinyarray.array([[1, 0], [0, -1]])


def make_system(a=1, t=1.0, alpha=0.5, e_z=0.0, W=10, L=40):
    # Start with an empty tight-binding system and a single square lattice.
    # `a` is the lattice constant (by default set to 1 for simplicity).
    lat = kwant.lattice.square(a)

    sys = kwant.Builder()

    #### Define the scattering region. ####
    sys[(lat(x, y) for x in range(L) for y in range(W))] = \
        4 * t * sigma_0 + e_z * sigma_z
    # hoppings in x-direction
    sys[kwant.builder.HoppingKind((1, 0), lat, lat)] = \
        -t * sigma_0 - 1j * alpha * sigma_y
    # hoppings in y-directions
    sys[kwant.builder.HoppingKind((0, 1), lat, lat)] = \
        -t * sigma_0 + 1j * alpha * sigma_x

    #### Define the left lead. ####
    lead = kwant.Builder(kwant.TranslationalSymmetry((-a, 0)))

    lead[(lat(0, j) for j in xrange(W))] = 4 * t * sigma_0 + e_z * sigma_z
    # hoppings in x-direction
    lead[kwant.builder.HoppingKind((1, 0), lat, lat)] = \
        -t * sigma_0 - 1j * alpha * sigma_y
    # hoppings in y-directions
    lead[kwant.builder.HoppingKind((0, 1), lat, lat)] = \
        -t * sigma_0 + 1j * alpha * sigma_x

    #### Attach the leads and return the finalized system. ####
    sys.attach_lead(lead)
    sys.attach_lead(lead.reversed())

    return sys


def plot_conductance(sys, energies):
    # Compute conductance
    data = []
    for energy in energies:
        smatrix = kwant.smatrix(sys, energy)
        data.append(smatrix.transmission(1, 0))

    pyplot.figure()
    pyplot.plot(energies, data)
    pyplot.xlabel("energy [t]")
    pyplot.ylabel("conductance [e^2/h]")
    pyplot.show()
	
def plot_ldos(sys, energy):
    # Compute local dos
    local_dos = kwant.ldos(sys, energy)
    # Calculate ldos per site, by summing spin up and spin down components
    local_dos = sum(local_dos.reshape(-1, 2), axis=1)
    kwant.plotter.map(sys, local_dos, num_lead_cells=5)
	
def plot_Si(sys, energy, n_lead):
    # Compute wave function
	psi = kwant.wave_function(sys, energy)
	psi_n = psi(n_lead).sum(axis=0)
	# Calculate <Si> per site
	Nsites = psi_n.shape[0] / 2
	sx = zeros((Nsites,))
	sy = zeros((Nsites,))
	sz = zeros((Nsites,))
	psi_i = array([[0.0], [0.0]])
	
	for i in xrange(Nsites):
		psi_i = array([[psi_n[i * 2]], [psi_n[i * 2 + 1]]])
		sx[i] = abs(dot(psi_i.conj().T, dot(sigma_x, psi_i)))[0, 0]
		sy[i] = abs(dot(psi_i.conj().T, dot(sigma_y, psi_i)))[0, 0]
		sz[i] = abs(dot(psi_i.conj().T, dot(sigma_z, psi_i)))[0, 0]
	# Plot
	kwant.plotter.map(sys, sx, num_lead_cells=5)
	kwant.plotter.map(sys, sy, num_lead_cells=5)
	kwant.plotter.map(sys, sz, num_lead_cells=5)
	kwant.plotter.map(sys, sx+sy+sz, num_lead_cells=5)

def main():
    sys = make_system()

    # Check that the system looks as intended.
    kwant.plot(sys)

    # Finalize the system.
    sys = sys.finalized()

    # We should see non-monotonic conductance steps.
    plot_conductance(sys, energies=[0.01 * i - 0.3 for i in xrange(100)])
    plot_ldos(sys, energy=0.0)
    plot_Si(sys, energy=0.0, n_lead=0)


# Call the main function if the script gets executed (as opposed to imported).
# See <http://docs.python.org/library/__main__.html>.
if __name__ == '__main__':
    main()