如何使用 numpy 快速找到布拉格反射？

Question

对于 X 射线衍射，需要找到所谓的劳厄方程的解

G_hkl - k_in + |k_in|*(sin(theta) cos(phi) , sin(theta) sin(phi) , cos(theta))= 0

其中 G_hkl 是给定的 3 维向量，k_in 可以选择为 (0,0,1) 并且 theta 和 phi 是满足方程式的自由参数。在一个典型的实验中，G_hkl 然后绕 x 轴旋转，对于旋转的每一步，都需要找到方程的解。方程在给定的旋转下不能有多个解。

我已经编写了这个 python 脚本来找到这些解决方案，但对于我的应用程序来说它不够快。

import numpy as np
import time
# Just initialization of variables. This is fast enough
res_list = []
N_rot = 100
Ghkl = np.array([0,0.7,0.7])
Ghkl_norm = np.linalg.norm(Ghkl)
kin = np.array([0,0,1])
kin_norm = np.linalg.norm(kin)
alpha_step = 2*np.pi/(N_rot-1)
rot_mat = np.array([[1,0,0],[0,np.cos(alpha_step),-np.sin(alpha_step)],[0,np.sin(alpha_step),np.cos(alpha_step)]])

# You can deduce theta from the norm of the vector equation
theta = np.arccos(1-(Ghkl_norm**2/(2*kin_norm**2))) 
sint = np.sin(theta)
cost = np.cos(theta)

# This leaves only phi as paramter to find
# I find phi by introducing a finite test vector
# and testing whether the norm of the vector equation is close
# to zero for any of those test phis
phi_test = np.linspace(0,2*np.pi,200)

kout = kin_norm * np.array([sint * np.cos(phi_test), sint * np.sin(phi_test), cost + 0*phi_test]).T
##############
start_time = time.time()

for j in range(100): # just to make it longer to measure the time
    res_list = []
    for i in range(N_rot): # This loop is too slow
        # Here the norm of the vector equation is calculated for all phi_test
        norm_vec = np.linalg.norm(Ghkl[np.newaxis, :] - kin[np.newaxis, :] + kout, axis=1)
        if (norm_vec < 0.01 * kin_norm).any(): # check whether any fulfills the criterion
            minarg = np.argmin(norm_vec)
            res_list.append([theta, phi_test[minarg]])
        Ghkl = np.dot(rot_mat,Ghkl)

print('Time was {0:1.2f} s'.format( (time.time()-start_time)))
# On my machine it takes 0.3s and res_list should be 
# [[1.0356115365192968, 1.578689775673263]]

您是否知道一种更快的计算方法，是从概念上通过求解完全不同的方程式还是通过使用我现有的方法使其更快？

Answer 1

存在依赖关系，因为 Ghkl 在每次迭代时更新并在下一次重复使用。相应的封闭形式可能难以追踪。因此，我将专注于提高最内层循环中其余代码的性能。

现在，我们正在计算 norm_vec，我认为可以使用下面列出的两种方法加快计算速度。

方法 #1 使用 Scipy's cdist -

from scipy.spatial.distance import cdist

norm_vec = cdist(kout,(kin-Ghkl)[None]) # Rest of code stays the same

方法 #2 使用 np.einsum -

sums = (Ghkl-kin)+kout
norm_vec_sq = np.einsum('ij,ij->i',sums,sums)
if (norm_vec_sq < (0.01 * kin_norm)**2 ).any():
    minarg = np.argmin(norm_vec_sq) # Rest of code stays the same

运行时测试 -

使用问题中列出的输入，我们得到这些结果：

In [91]: %timeit np.linalg.norm(Ghkl- kin + kout, axis=1)
10000 loops, best of 3: 31.1 µs per loop

In [92]: sums = (Ghkl-kin)+kout
    ...: norm_vec_sq = np.einsum('ij,ij->i',sums,sums)
    ...: 

In [93]: %timeit (Ghkl-kin)+kout                 # Approach2 - step1
100000 loops, best of 3: 7.09 µs per loop

In [94]: %timeit np.einsum('ij,ij->i',sums,sums) # Approach2 - step2
100000 loops, best of 3: 3.82 µs per loop

In [95]: %timeit cdist(kout,(kin-Ghkl)[None])    # Approach1
10000 loops, best of 3: 44.1 µs per loop

因此，方法 #1 没有显示出对这些输入大小的任何改进，但方法 #2 在计算 norm_vec.

时绕过 3x 加速

简短说明为什么 np.einsum 在这里是有益的： 那么，einsum 非常适合逐元素乘法和求和归约。我们正在利用这里问题的本质。 np.linalg.norm 给出了我们在输入的平方版本上沿第二个轴的求和。因此，einsum 对应项只是将相同的输入作为两个输入输入两次，从而处理平方，然后丢失第二个轴及其用字符串表示法表示的求和归约。它一次性完成这两项工作，这可能就是它在这里如此之快的原因。

如何使用 numpy 快速找到布拉格反射？

How to find Bragg reflexes fast with numpy?

python

performance

numpy

scientific-computing