from matplotlib import pyplot as plt

# Unit conversion from eV to uJ/m^2
surface = 2.866 * 2.866 * 1e-20 # in m * m
conversion = (1 * eV).inUnitsOf(Joule) * 1e6 / surface

filename = 'stt-90.hdf5'
initial_free_layer_site = 22

# Retrieve the saved quantities
torkance_linear_response = nlread(filename, object_id='torkance')[0]
torque = nlread(filename, object_id='torque')[0]
biases = nlread(filename, object_id='biases')[0]

# Plot the torque obtained from the linear repsonse torkance.
total_torkance = numpy.sum(torkance_linear_response[:, initial_free_layer_site:], axis=1)
torque_from_torkance_x = [- total_torkance[0] * bias * conversion for bias in biases.inUnitsOf(Volt)]
plt.plot(biases, torque_from_torkance_x, '--', label='In-plane torque (linear response)')

total_torque_x = [numpy.sum(x[0, initial_free_layer_site:]) * conversion for x in torque.inUnitsOf(eV)]
plt.plot(biases, total_torque_x, 'o-', linewidth=2, label='In-plane torque (finite bias)')

total_torque_y = [numpy.sum(x[1, initial_free_layer_site:]) * conversion for x in torque.inUnitsOf(eV)]
plt.plot(biases, total_torque_y, 'o-', linewidth=2, label='Out-of-plane torque (finite bias)')

plt.xlabel('Voltage [V]')
plt.ylabel('Torque [uJ/m^2]')
plt.title(r'$\theta = \pi/2$')
plt.legend()
plt.grid()
plt.show()

# Plot the torque per atomic site.
single_bias_point = - 0.6
idx = 4

plt.plot(- single_bias_point * torkance_linear_response[0, :] * conversion, 'o--', label='In-plane torque (linear response)')
plt.plot(torque.inUnitsOf(eV)[idx][0, :] * conversion, '*-', label='In-plane torque (finite bias)')

plt.title('Torque at {} V'.format(single_bias_point))
plt.ylabel('Torque [uJ/m^2]')
plt.xlabel('Atom no.')
plt.legend()
plt.grid()
plt.show()