如何加快数百万次小特征值的计算？

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]]

我当前的解决方案使用pandas.DataFrame的apply-Function。 对于数据框中的每一行，都可以通过numpy.linalg.eigvalsh函数解决单独的特征值问题。

对于1_000_000应力张量，我的apply函数大约需要85s。
多处理解决方案在4个（旧）内核上需要39s。

在40核机器上，它看起来更好：

• 常规计算时间：50.66s
• 并行计算时间：3.94s

如何加快计算速度？

输入

   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())

谢谢