Np.cross производит неправильные результаты, ищет рабочую альтернативу

Question

Я переписываю код анализа для временных рядов Molecular Dynamics. Из-за огромного количества временных шагов (150 000 для каждого прогона моделирования), которые необходимо проанализировать, очень важно, чтобы мой код был как можно быстрее.

Старый код очень медленный (на самом деле это требует от 300 до 500 раз больше времени по сравнению с моим), потому что он был написан для анализа нескольких тысяч файлов PDB, а не пучка, полная различных симуляций (около 60) , каждый из которых имеет 150 000 временных шагов. Я знаю, что C или Fortran был бы швейцарским армейским ножом в этом случае, но мой опыт работы с c был .....

Поэтому я стараюсь использовать как можно больше numpy/scipy для моего кода на Python. Поскольку у меня есть лицензия на ускоренное распространение anaconda с mkl, это действительно значительное ускорение.

Теперь у меня возникла проблема, и я надеюсь, что смогу объяснить это так, как вы понимаете, что я имею в виду.

У меня есть три массива, каждый из которых имеет форму (n, 3, 20). В первой строке находятся все остатки моего пептида, обычно около 23-31. Во второй строке находятся координаты в порядке xyz, а в третьей строке - некоторые конкретные временные шаги.

Теперь я вычисляю кручение для каждого остатка на каждом временном шаге. мой код для случая массивов с формой (п, 3,1) его:

def fast_torsion(d1, d2, d3): 
    tt = dot(d1, np.cross(d2, d3)) 
    tb = dot(d1, d1) * dot(d2, d2) 
    torsion = np.zeros([len(d1), 1]) 
    for i in xrange(len(d1)): 
     if tb[i] != 0: 
      torsion[i] = tt[i]/tb[i] 
    return torsion 

Теперь я пытался использовать один и тот же код для массивов с расширенной третьей осью, но функция перекрестного произведения производит неправильные значения по сравнению с оригинальным медленным кодом, который использует цикл for. Я пробовал этот код с моими большими массивами, это примерно в 10-20 раз быстрее, чем решение для цикла, и примерно в 200 раз быстрее, чем старый код.

То, что я пытаюсь сделать, это то, что np.cross() вычисляет только произведение перекрестка по второй оси (xyz) и выполняет итерацию по двум другим осям. В случае с короткой третьей осью он работает нормально, но с большими массивами он работает только в первый раз. Я также пробовал настройки оси, но у меня не было никаких шансов.

Я также могу использовать Cython или numba, если это единственное решение для моей проблемы.

P.S. Извините за мой английский. Надеюсь, вы все поймете.

Answer 1

np.cross имеет axisa, axisb и axisc именованных аргументов, чтобы выбрать, где в аргументах ввода и вывода являются векторы для кросса умножается. Я думаю, что вы хотите использовать:

np.cross(d2, d3, axisa=1, axisb=1, axisc=1) 

Если вы не включаете axisc=1, результат умножения будет в конце выходного массива.

Кроме того, вы можете избежать явно зацикливания над torsion массивом, выполнив:

torsion = np.zeros((len(d1), 1) 
idx = (tb !=0) 
torsion[idx] = tt[idx]/tb[idx]