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

Question

对于x射线衍射，需要找到Laue方程的解

G|u hkl-k|u in+| k|u in |*（sin（θ）cos（phi），sin（θ）sin（phi），cos（θ））=0

其中G_hkl是一个给定的三维向量，k_in可以选择为（0,0,1），θ和φ是满足方程的自由参数。在一个典型的实验中，gkuhkl绕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使用^{}-

from scipy.spatial.distance import cdist

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

应用程序#2使用^{}-

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在计算3x时加快了速度。你知道吗

为什么np.einsum在这里是有益的简短说明：好吧，einsum对于元素级乘法和和归约非常有用。我们正在利用这个问题的本质。np.linalg.norm在输入的平方版本上沿第二个轴求和。因此，einsum对应的方法就是将两个输入的两倍输入到同一个输入中，这样就可以处理平方运算，然后丢失第二个轴的和归约，即用字符串表示法表示的和归约。这两个都是一次性的，这可能就是为什么它在这里这么快。你知道吗