from __future__ import print_function, division
import time
import sys
from multiprocessing import Pool
from haslib.science_basics import A, mel
from haslib.icc_2d_2 import iCC_2D_2 as icc
from numpy import *
from numpy import pi
import numpy as np
[docs]def worker(p_arr):
"""
Simulate a specific point using the given arguments
Args:
p_arr: tuple containing the following parameters(in order)
* energy: energy of the beam
* phi: azimuth angle
* g_x: reciprical vector 1
* g_y: recpiprical vector 2
* pot: potential class (Morse)
* theta: incident angle
* kwargs: extra arguments for iCC_2D_2
Returns:
array: array with one row containing energy, azimuthal angle and
scattered intensity
"""
energy, phi, g_x, g_y, pot, theta, kwargs = p_arr
# Calculate the specular intenstiy of the reflection via eCC
res = icc(pot.h,0,pot.a,pot.D,pot.xi,energy,theta,g_x,g_y,**kwargs)[:,[0,1,6,5]]
index = np.where(res[:, 1].astype(int) != 0)[0]
res = delete(res, index, 0)
index2 = np.where(res[:, 0].astype(int) == 0)[0]
i = res[index2, 2][0]
# Extract corresponding Intensity from the output (for the chosen G)
return [energy / mel, phi * 180. / pi, i]
[docs]class AzimuthalScan:
# Define Vector rotation function
[docs] def rotate(self, ang, vec):
"""
rotate the 2D vector vec by angle ang
Args:
ang: angle in radians
vec: 2D vector
Returns:
array: rotated vector
"""
r = array([[cos(ang), -sin(ang)], [sin(ang), cos(ang)]])
erg = dot(r, vec)
return erg
[docs] def run(self, e, phi, theta, pot, cpus=1, **ccargs):
"""
calculate the spec intensity for a grid of e und phi
Args:
e: energy
phi: phi in degrees
theta: theta in degrees
pot: Potential (class:Morse)
cpus: cpus to use for Pool call
**ccargs: optional iCC parameters(see icc_2d_2)
Returns:
array: array with energy, azimuthal angle and intensity.
::
[[energy_1,angle_1,intensity_1],
[energy_2,angle_2,intensity_2],
...
[energy_n,angle_n,intensity_n]]
every row ist calculated by the worker method
"""
phi_rad = phi * (pi / 180)
# Read G vectors in zero position
g_x_zero = pot.crystal.b1[0:2]
g_y_zero = pot.crystal.b2[0:2]
# Loop though possibilities
Point_array = []
for e_now in e:
for phi_now in phi_rad:
# Write current values in container
# Calculate local azimuthal angle
gx = self.rotate(phi_now, g_x_zero)
gy = self.rotate(phi_now, g_y_zero)
Point_array.append([e_now * mel, phi_now, gx, gy, pot, theta, ccargs])
#from ipyparallel import Client
#c = Client(profile="ssh")
#v = c.load_balanced_view()
#res = v.map(worker, Point_array)
if cpus >= 1:
pool = Pool(cpus)
res = pool.imap(worker, Point_array)
pool.close()
old_complete = -1
while True:
#completed = res.progress
completed = res._index
if completed%1 == 0 and completed != old_complete:
print("\r{}/{}".format(completed, len(Point_array)), end="")
sys.stdout.flush()
if completed == len(Point_array):
print("done")
break
old_complete = completed
time.sleep(5)
else:
res = map(worker, Point_array)
dummy = array([0, 0, 0])
for r in res:
dummy = vstack([dummy, r])
container = delete(dummy, 0, 0)
return container