# Load STT data
import pickle
f = open('angles_plot.pckl', 'rb')
(thetas,x,y,z) = pickle.load(f)
f.close()

# Process STT data
inplane = (x**2 + z**2)**0.5
outplane = y

# Analytical in-plane STT
transmissions = numpy.zeros(2)
for j,theta in enumerate([0,180]):
    # Data file
    filename = 'theta-%i.hdf5' % theta
    # Read the transmission spectrum
    transmission = nlread(filename, TransmissionSpectrum)[0]
    # Get the energies
    energies = transmission.energies().inUnitsOf(eV)
    # Find the index of the Fermi level
    i_Ef = numpy.argmin(abs(energies))
    # Calculate the spin-transmission (Up - Down)
    T = transmission.evaluate(spin=Spin.Up)[i_Ef]-transmission.evaluate(spin=Spin.Down)[i_Ef]
    # Append to list
    transmissions[j] = T
analytical = abs(transmissions[0]-transmissions[1])/2*numpy.sin(numpy.array(thetas)*numpy.pi/180)/(2*numpy.pi)

# Plot results
import pylab
pylab.figure(figsize=(10,6))
pylab.plot(thetas, inplane*1e6,         label=r'$\tau_\parallel$')
pylab.plot(thetas, analytical*1e6, 'o', label=r'$\tau_\parallel$'+'(analytical)')
pylab.plot(thetas, outplane*1e6,        label=r'$\tau_\perp$')
pylab.axhline(0, color='k', linestyle=':')
pylab.legend()
pylab.xlabel(r'$\theta$ (degrees)')
pylab.ylabel(r'$\tau_\mathrm{right} \, \, \, (\mu eV/V)$')
ax = pylab.gca()
ax.set_xlim((-5,185))
ax.set_xticks(thetas)
pylab.savefig('analytical.png')
pylab.show()