#!/usr/bin/env python
# coding: utf-8

# In[1]:


from __future__ import print_function, division

import kwant

import numpy as np
from numpy import sqrt
get_ipython().run_line_magic('matplotlib', 'inline')


# # Define the system

# In[2]:


HZ = sqrt(6) / 6.0
h = 1 / (sqrt(3) * 4)
lat = kwant.lattice.general([(1, 0, 0), (0.5, -sqrt(3) * 0.5, 0), (0, 1 / sqrt(3), 2 * HZ)],
                            [(0, 0, 0), (0.5, 0, 0), (0.25, -0.25 * sqrt(3), 0), (0.25, -h, HZ)])

def make_latticeperiodic_pyrochlore_slab(t_1, t_2, t_perp, lambda_1, lambda_2, shape=(5, 4, 2)):
    """
    create a pyrochlore slab in a parallelepiped shape,a s given by the lattice vectors
    with periodic boundary conditions along the second lattice vector
    and leads attached at tle left and right (x direction)
    """
    a, b, c, d = lat.sublattices

    def initialize_hopping(sys):
        # construct all hoppings with SO explicit
        sys[kwant.HoppingKind((0, 0, 0), b, a)] = t_1 + 1j * lambda_1
        sys[kwant.HoppingKind((-1, 0, 0), b, a)] = t_1 - 1j * lambda_1

        sys[kwant.HoppingKind((0, 0, 0), c, a)] = t_1 - 1j * lambda_1
        sys[kwant.HoppingKind((0, -1, 0), c, a)] = t_1 - 1j * lambda_1

        sys[kwant.HoppingKind((0, 0, 0), c, b)] = t_1 + 1j * lambda_1
        sys[kwant.HoppingKind((1, -1, 0), c, b)] = t_1 + 1j * lambda_1

        # construct all NNN hoppings with SO explicit
        sys[kwant.HoppingKind((0, -1, 0), b, a)] = t_2 - 1j * lambda_2
        sys[kwant.HoppingKind((-1, 1, 0), b, a)] = t_2 - 1j * lambda_2

        sys[kwant.HoppingKind((-1, 0, 0), c, a)] = t_2 + 1j * lambda_2
        sys[kwant.HoppingKind((1, -1, 0), c, a)] = t_2 + 1j * lambda_2

        sys[kwant.HoppingKind((0, -1, 0), c, b)] = t_2 - 1j * lambda_2
        sys[kwant.HoppingKind((1, 0, 0), c, b)] = t_2 - 1j * lambda_2

        nn_hoppings = lat.neighbors(n=1)
        for hop in nn_hoppings:
            # hoppings in z-direction
            if hop.family_a == d or hop.family_b == d:
                sys[hop] = t_perp

    # Build ancillary system with y-periodic BCs.  This system cannot be
    # finalized in Kwant 1.0.
    sym = kwant.TranslationalSymmetry(lat.vec((0, shape[1], 0)))

    anc = kwant.Builder(sym)
    anc[lat.shape(lambda p: 0 <= p[0] < shape[0] and 0 <= p[2] < shape[2], (0, 0, 0))] = None
    initialize_hopping(anc)

    sys = kwant.Builder()

    sys[anc.sites()] = 0
    for s1, s2 in anc.hoppings():
        sys[sym.to_fd(s1), sym.to_fd(s2)] = anc[s1, s2]

    # leads
    # TRICK ancestor only periodic in Y-direction to catch also the x-direction hoppings later on
    anc_sym = kwant.TranslationalSymmetry(lat.vec((0, shape[1], 0)))

    anc_lead = kwant.Builder(anc_sym)
    anc_lead[lat.shape(lambda p: -2 <= p[0] <= 2 and 0 <= p[2] < shape[2], (0, 0, 0))] = 0
    initialize_hopping(anc_lead)

    # at the left
    lead_symmetry = kwant.TranslationalSymmetry(lat.vec([-1, 0, 0]))
    lead = kwant.Builder(lead_symmetry)

    lead[anc_lead.sites()] = 0
    for s1 in anc_lead.sites():
        lead[s1] = 1

    for s1, s2 in anc_lead.hoppings():
        lead[s1, anc_sym.to_fd(s2)] = anc_lead[s1, s2]

    sys.attach_lead(lead)
    # and right
    sys.attach_lead(lead.reversed())

    sys.parameter = dict(t_1=t_1,
                         t_2=t_2,
                         t_perp=t_perp,
                         lambda_1=lambda_1,
                         lambda_2=lambda_2,
                         shape=shape,
                         peridoic=True)
    return sys, lat


def make_periodic_pyrochlore_slab(t_1, t_2, t_perp, lambda_1, lambda_2, shape=(5, 4, 2)):
    """
    create a pyrochlore slab in a rectangular shape
    with periodic boundary conditions in y-direction
    and leads attached at tle left and right (x direction)
    """
    a, b, c, d = lat.sublattices

    def initialize_hopping(sys):
        # construct all hoppings with SO explicit
        sys[kwant.HoppingKind((0, 0, 0), b, a)] = t_1 + 1j * lambda_1
        sys[kwant.HoppingKind((-1, 0, 0), b, a)] = t_1 - 1j * lambda_1

        sys[kwant.HoppingKind((0, 0, 0), c, a)] = t_1 - 1j * lambda_1
        sys[kwant.HoppingKind((0, -1, 0), c, a)] = t_1 - 1j * lambda_1

        sys[kwant.HoppingKind((0, 0, 0), c, b)] = t_1 + 1j * lambda_1
        sys[kwant.HoppingKind((1, -1, 0), c, b)] = t_1 + 1j * lambda_1

        # construct all NNN hoppings with SO explicit
        sys[kwant.HoppingKind((0, -1, 0), b, a)] = t_2 - 1j * lambda_2
        sys[kwant.HoppingKind((-1, 1, 0), b, a)] = t_2 - 1j * lambda_2

        sys[kwant.HoppingKind((-1, 0, 0), c, a)] = t_2 + 1j * lambda_2
        sys[kwant.HoppingKind((1, -1, 0), c, a)] = t_2 + 1j * lambda_2

        sys[kwant.HoppingKind((0, -1, 0), c, b)] = t_2 - 1j * lambda_2
        sys[kwant.HoppingKind((1, 0, 0), c, b)] = t_2 - 1j * lambda_2

        nn_hoppings = lat.neighbors(n=1)
        for hop in nn_hoppings:
            # hoppings in z-direction
            if hop.family_a == d or hop.family_b == d:
                sys[hop] = t_perp

    # Build ancillary system with y-periodic BCs.  This system cannot be
    # finalized in Kwant 1.0.
    sym = kwant.TranslationalSymmetry(lat.vec((-shape[1] / 2, shape[1], 0)))

    # Define unit cell ('fundamental domain') to be along y-axis and not along lattice vector
    for fam in (a, b, c, d):
        sym.add_site_family(fam, other_vectors=[[-1, 0, 0], [-1, 2, 3]])

    anc = kwant.Builder(sym)
    anc[lat.shape(lambda p: 0 <= p[0] < shape[0] and 0 <= p[2] < shape[2], (0, 0, 0))] = None
    initialize_hopping(anc)

    sys = kwant.Builder()

    sys[anc.sites()] = 0
    for s1, s2 in anc.hoppings():
        sys[sym.to_fd(s1), sym.to_fd(s2)] = anc[s1, s2]

    # leads
    # TRICK ancestor only periodic in Y-direction to catch also the x-direction hoppings later on
    anc_sym = kwant.TranslationalSymmetry(lat.vec((-shape[1] / 2, shape[1], 0)))

    # Define unit cell ('fundamental domain') to be along y-axis and not along lattice vector
    for fam in (a, b, c, d):
        anc_sym.add_site_family(fam, other_vectors=[[-1, 0, 0], [-1, 2, 3]])

    anc_lead = kwant.Builder(anc_sym)
    anc_lead[lat.shape(lambda p: -2 <= p[0] <= 2 and 0 <= p[2] < shape[2], (0, 0, 0))] = 0
    initialize_hopping(anc_lead)

    # at the left
    lead_symmetry = kwant.TranslationalSymmetry(lat.vec([-1, 0, 0]))
    lead = kwant.Builder(lead_symmetry)

    # Define unit cell ('fundamental domain') to be along y-axis and not along lattice vector
    for fam in (a, b, c, d):
        lead_symmetry.add_site_family(fam, other_vectors=[[-1, 2, 0], [-1, 2, 3]])

    lead[anc_lead.sites()] = 0
    for s1 in anc_lead.sites():
        lead[s1] = 1

    for s1, s2 in anc_lead.hoppings():
        lead[s1, anc_sym.to_fd(s2)] = anc_lead[s1, s2]

    sys.attach_lead(lead)
    # and right
    sys.attach_lead(lead.reversed())

    sys.parameter = dict(t_1=t_1,
                         t_2=t_2,
                         t_perp=t_perp,
                         lambda_1=lambda_1,
                         lambda_2=lambda_2,
                         shape=shape,
                         peridoic=True)
    return sys, lat


# In[3]:


#lattice
sys, lat = make_latticeperiodic_pyrochlore_slab(
    t_1=-1, t_2=0.3, t_perp=1, lambda_1=0.3, lambda_2=0.2, shape=(5,8,(4*2-1)*HZ))

sys2, lat2 = make_periodic_pyrochlore_slab(
    t_1=-1, t_2=0.3, t_perp=1, lambda_1=0.3, lambda_2=0.2, shape=(5,8,(4*2-1)*HZ))


# In[4]:


lattice_sys = sys.finalized()


# In[5]:


rectangle_sys = sys2.finalized()


# # QZ iteration fails

# In[6]:


kwant.smatrix(rectangle_sys, 1.5).transmission(0, 1)


# But it works on the lattice-periodic system

# In[7]:


kwant.smatrix(lattice_sys, 1.5).transmission(0, 1)


# # Persists over large range:

# List all energies for which we get the QZ Error

# In[8]:


energies = np.linspace(-5, 5, 100)
conductances = []
failures = []
for en in energies:
    try:
        conductances.append((en, kwant.smatrix(rectangle_sys, en).transmission(0, 1)))
    except kwant.linalg.lapack.LinAlgError:
        failures.append(en)
    


# In[9]:


failures


# In[10]:


conductances


# In[ ]: