Calculate scattered intensity

calc_scatter.py

import numpy as np
from numpy import pi, sin, cos, exp
from scipy.interpolate import interp1d
from multiprocessing import Pool

RAD2ARCSEC = 180. * 3600. / pi  # convert to arcsec for better scale


class CalcScatter(object):

    def __init__(self, alpha, x, y, scale, lam):
        self.alpha = alpha
        self.x = x
        self.y = y
        self.scale = scale
        self.lam = lam

    def __call__(self, theta):
        fac = 2 * pi * 1j / self.lam
        d_sin = fac * (sin(self.alpha) - sin(theta))
        d_cos = fac * (cos(self.alpha) + cos(theta)) * self.scale

        ampl = np.sum(exp(self.x * d_sin - self.y * d_cos),
                      axis=0)
        ampl_sum = np.sum(abs(ampl) ** 2)

        return ampl_sum


def calc_scatter(displ, dx=500, graze_angle=2.0, scale=np.sqrt(2), lam=1.24e-3,
                 thetas=None, n_theta=1000, theta_max=60, n_x=1000, n_proc=4):
    """
    lam = 1.24e-3 um <=> 1 keV
    """
    alpha = (90 - graze_angle) * pi / 180
    theta_max /= RAD2ARCSEC
    n_ax, n_az = displ.shape
    x = np.arange(n_ax) * dx
    x = x - x.mean()
    if thetas is None:
        thetas = np.linspace(alpha - theta_max, alpha + theta_max, n_theta)
    else:
        thetas = thetas + alpha
    I_scatter = []

    interp = interp1d(x, displ, kind='cubic', axis=0)
    x_int = np.linspace(x[0], x[-1], n_x)
    y = interp(x_int)
    x = x_int.reshape(-1, 1)

    calc_func = CalcScatter(alpha, x, y, scale, lam)
    pool = Pool(processes=n_proc)
    I_scatter = pool.map(calc_func, list(thetas))

    return (thetas - thetas.mean()) * RAD2ARCSEC, np.array(I_scatter)

plot_data.py

import time
import numpy as np
import calc_scatter as cs

RAD2ARCSEC = 180 * 3600 / np.pi  # convert to arcsec for better scale
THETA_MAX = 2.55e-4 * RAD2ARCSEC

if 'case' not in globals():
    case = 'uncorr'

if 'displ' not in globals():
    displ = np.loadtxt('{}.dat'.format(case))

if 'scatter_ref' not in globals():
    scatter_ref = np.loadtxt('scatter_{}.dat'.format(case))[:, 1]
    theta_ref = np.linspace(-THETA_MAX, THETA_MAX, 10001)

dx = 200. / 205 * 1000  # spacing in um (200 mm high over 205 points)
t0 = time.time()
theta, scatter = cs.calc_scatter(displ, dx=dx, graze_angle=1.428,
                                 thetas=theta_ref / RAD2ARCSEC,
                                 # n_theta=1000,
                                 # theta_max=40,
                                 n_x=1851, 
                                 n_proc=4,
                                 )
print 'Run time: {:.3f}'.format(time.time() - t0)

scatter = scatter_ref.max() * scatter / scatter.max()

clf()
plot(theta_ref, scatter_ref)
plot(theta, scatter)
grid()
title('Scatter intensity ({}ected)'.format(case))
xlabel('Scatter angle (arcsec)')

i_mid = len(theta) // 2
i1 = 2 * i_mid - 1
angle = theta[i_mid:i1]

sym_scatter = scatter[i_mid:i1] + scatter[i_mid - 1:0:-1]
sym_scatter /= sym_scatter.sum()

ee = np.cumsum(sym_scatter)
i_hpr = np.searchsorted(ee, 0.5)
angle_hpd = angle[i_hpr] * 2
print 'angle_hpd', angle_hpd

i99 = np.searchsorted(ee, 0.99)
angle_rmsd = 2 * np.sqrt(np.sum(angle[:i99] ** 2 * sym_scatter[:i99])
                         / np.sum(sym_scatter[:i99]))
print 'angle_rmsd', angle_rmsd