Как я могу ускорить миллионы вычислений небольших собственных значений?

Question

Как я могу ускорить миллионы вычислений собственных значений с помощью Python?

Я хочу вычислить главные напряжения многих (миллионов) 3d-тензоров напряжений.

[[sig_xx, sig_xy, sig_zx], [sig_xy, sig_yy, sig_yz], --> [sig_1, sig_2, sig_3] [sig_zx, sig_yz, sig_zz]]

Мое текущее решение использует apply -функцию pandas.DataFrame. Для каждой строки в кадре данных отдельная проблема собственных значений решается с помощью функции numpy.linalg.eigvalsh.

Для тензоров напряжений 1_000_000 моя apply -функция занимает около 85 с.
Многопроцессорному решению требуется 39 с на 4 (старых) ядрах.

На 40-ядерном станке это выглядит лучше:

• обычное время расчета: 50,66 с
• параллельное время расчета: 3,94 с

Как я могу ускорить расчет?

ввод

   mid_sig_xx  mid_sig_yy  mid_sig_zz  mid_sig_xy  mid_sig_yz  mid_sig_zx       foo
0    0.452220    0.097544    0.713128    0.704898    0.568947    0.205623  0.377253
1    0.813111    0.789671    0.607063    0.334711    0.526826    0.490880  0.517193
2    0.655822    0.784863    0.910415    0.167779    0.888867    0.737995  0.656025
3    0.558476    0.492004    0.779018    0.109135    0.966178    0.311049  0.779220
4    0.639169    0.724089    0.520573    0.424256    0.909775    0.562966  0.244805

выход

      sig_1     sig_2     sig_3
0  1.410364  0.399706 -0.547178
1  1.638306  0.468745  0.102793
2  2.059439  0.543883 -0.252221
3  1.695600  0.491881 -0.357983
4  1.915084  0.288536 -0.319789

пример кода

# -*- coding: utf-8 -*-

from joblib import Parallel, delayed
import multiprocessing
import numpy as np
import numpy.linalg as nl
import pandas as pd
import time

def apply_parallel(df, func, prefix="", n_groups=None):
    if n_groups is None:
        n_groups = multiprocessing.cpu_count() * 2
    split_idx = len(df)//n_groups + 1
    df_grouped = [df[i*split_idx:(i+1)*split_idx] for i in range(n_groups)]
    retLst = Parallel(n_jobs=multiprocessing.cpu_count())(delayed(func)(group, prefix) for group in df_grouped)
    return pd.concat(retLst)

def calc_sig_principal(df, prefix="", inplace=False):
    """calc principal components of a second order tensor"""

    def sig123(sig):
        """solve eigen value problem"""
        sig_xx, sig_yy, sig_zz, sig_xy, sig_yz, sig_zx = sig
        x = np.array([[sig_xx, sig_xy, sig_zx],
                      [sig_xy, sig_yy, sig_yz],
                      [sig_zx, sig_yz, sig_zz]])
        return sorted(nl.eigvalsh(x), reverse=True)

    sigs =  df[["{}sig_xx".format(prefix),
               "{}sig_yy".format(prefix),
               "{}sig_zz".format(prefix),
               "{}sig_xy".format(prefix),
               "{}sig_yz".format(prefix),
               "{}sig_zx".format(prefix)]].apply(sig123, axis=1)
    return pd.DataFrame(np.vstack(sigs), columns=["sig_1", "sig_2", "sig_3"], index=df.index)


df = pd.DataFrame(np.random.random((1000000, 7)), columns=["mid_sig_xx", "mid_sig_yy", "mid_sig_zz", "mid_sig_xy", "mid_sig_yz", "mid_sig_zx", "foo"])

print('parallel version: ')
t0 = time.time()
df_res = apply_parallel(df, calc_sig_principal, prefix="mid_")
t1 = time.time()
print("parallel calculation time: {:.2f}s".format(t1-t0))
# print(df_res.head())

print('regular version: ')
t0 = time.time()
df_res2 = calc_sig_principal(df, prefix="mid_")
t1 = time.time()
print("regular calculation time: {:.2f}s".format(t1-t0))
# print(df_res2.head())

assert len(df_res) == len(df_res2)
assert np.isclose((df_res-df_res2).values.sum(), 0), "{}: sum should be zero!".format((df_res-df_res2).values.sum())

Спасибо