python - 如何使用 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 时速度提高了 3 倍。

简短说明为什么np.einsum在这里有用:嗯，einsum非常适合元素乘法和求和缩减。我们正在利用问题的本质。 np.linalg.norm 为我们提供了输入平方版本上沿第二个轴的求和。因此，einsum 的对应部分就是将相同的输入提供两次，从而处理平方，然后失去第二个轴及其求和减少，用其表示字符串表示法。它一次性完成了这两件事，这可能就是它如此快的原因。