# Prepare lists
thetas = [0,10,20,30,40,50,60,70,80,90,100,110,120,130,140,150,160,170,180]
x = []
y = []
z = []

# Read data
for theta in thetas:
    # Data file
    filename = 'theta-%i.hdf5' % theta

    # Read in the STT analysis object
    stt = nlread(filename, SpinTransferTorque)[0]

    # Get the STT 3D vector grid (units Bohr**-3)
    array = stt.toArray()

    # Get the index of middle position along C
    sh = numpy.shape(array)
    k = sh[2]/2

    # Get the volume element of the STT grid.
    dX, dY, dZ = stt.volumeElement()
    volume = numpy.abs(numpy.dot(dX, numpy.cross(dY, dZ)))

    # Integrate the vector components in the right-hand part of the device
    stt_x = numpy.sum(array[:,:,k:,:,0])*volume*stt.unit()
    stt_y = numpy.sum(array[:,:,k:,:,1])*volume*stt.unit()
    stt_z = numpy.sum(array[:,:,k:,:,2])*volume*stt.unit()

    # append to lists
    x.append(stt_x)
    y.append(stt_y)
    z.append(stt_z)

# Convert lists to arrays
x = numpy.array(x)
y = numpy.array(y)
z = numpy.array(z)

# Save data for future convenience
import pickle
f = open('angles_plot.pckl', 'wb')
pickle.dump((thetas,x,y,z), f)
f.close()

# Plot results
import pylab
pylab.figure(figsize=(10,6))
pylab.plot(thetas, x*1e6, label='x')
pylab.plot(thetas, y*1e6, label='y')
pylab.plot(thetas, z*1e6, label='z')
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('angles_plot.png')
pylab.show()