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

# In[1]:


# Tutorial 2.8.2. 3D example: zincblende structure
# ================================================
#
# Physical background
# -------------------
#  - 3D Bravais lattices
#
# Kwant features highlighted
# --------------------------
#  - demonstrate different ways of plotting in 3D

from matplotlib import pyplot

import kwant


# In[2]:


import matplotlib
import matplotlib.pyplot
from matplotlib_inline.backend_inline import set_matplotlib_formats

matplotlib.rcParams['figure.figsize'] = matplotlib.pyplot.figaspect(1) * 2
set_matplotlib_formats('svg')


# In[3]:


lat = kwant.lattice.general([(0, 0.5, 0.5), (0.5, 0, 0.5), (0.5, 0.5, 0)],
                            [(0, 0, 0), (0.25, 0.25, 0.25)],
                            norbs=1)
a, b = lat.sublattices


# In[4]:


def make_cuboid(a=15, b=10, c=5):
    def cuboid_shape(pos):
        x, y, z = pos
        return 0 <= x < a and 0 <= y < b and 0 <= z < c

    syst = kwant.Builder()
    syst[lat.shape(cuboid_shape, (0, 0, 0))] = None
    syst[lat.neighbors()] = None

    return syst


# In[5]:


syst = make_cuboid()

kwant.plot(syst);


# In[6]:


kwant.plotter.set_engine('plotly') # Set to plotly engine

kwant.plot(syst)

kwant.plotter.get_engine() # Get the current engine

kwant.plotter.set_engine('matplotlib') # Set to matplotlib engine


# In[7]:


syst = make_cuboid(a=1.5, b=1.5, c=1.5)

def family_colors(site):
    return 'r' if site.family == a else 'g'

kwant.plot(syst, site_size=0.18, site_lw=0.01, hop_lw=0.05,
           site_color=family_colors);