import pylab

# Filenames
filename_out_of_plane = 'W8-Fe2-theta-0-phi-0.hdf5'
filename_in_plane = 'W8-Fe2-theta-90-phi-150.hdf5'

# Read the calculated LDOS
ldos_out_of_plane = nlread(filename_out_of_plane, LocalDensityOfStates)[0]
ldos_in_plane = nlread(filename_in_plane, LocalDensityOfStates)[0]

# Get the configuration.
conf = ldos_out_of_plane._configuration()

# Get the coordinate of the surface Fe atom
R0 = conf.cartesianCoordinates()[-1, :]

# Set the shift in the z-direction, i.e. tip distance above Fe atom
z_shift = 2*Ang

R = R0.copy()
R[2] += z_shift

nlprint('Tip position')
nlprint(R)

# Calculate the LDOS at the tip position for both in- and out of plane
n_out = ldos_out_of_plane.evaluate(x=R[0], y=R[1], z=R[2], spin=Spin.Sum)
n_in = ldos_in_plane.evaluate(x=R[0], y=R[1], z=R[2], spin=Spin.Sum)
energies = ldos_out_of_plane.energies().inUnitsOf(eV)

# Calculate the TAMR
tamr_ldos = (n_out - n_in) / n_out * 100

# Plot the results
pylab.figure()
pylab.plot(energies, tamr_ldos, 'k-')
pylab.plot([-0.2, 0.5], [0, 0], 'k--')
pylab.plot([0, 0], [-30, 20], 'k--')
pylab.xlabel('Energy  (eV)')
pylab.ylabel('TAMR (%)')
pylab.xlim([-0.2, 0.5])
pylab.ylim([-30, 20])
pylab.show()