python - 使用 fsolve 求解数组或函数列表的最快方法

Question

我有下面的工作功能。我有一个函数，可以从中计算一阶和二阶导数。然后我需要找到一阶导数为零而第二个导数为负的 theta 值。我必须为大量点计算这个。点数等于K1和K2的长度。使用 sympy 我计算一阶和二阶导数。我目前迭代所有导数并为每个导数求解方程。有没有更快的方法来做到这一点，一旦 K1 和 K2 的长度增加 > 1000 ，这对我的应用程序来说需要很长时间。

import numpy as np
import sympy as sp
from scipy.optimize import fsolve
from sympy.utilities.lambdify import lambdify

def get_cpd(K1, K2):
       '''
        *args:
            K1, 1D numpy array: mode I stress intensity factors
            K2, 1D numpy array: mode II stress intensity factor
        *Note:
            K1 and K2 should have the same length
       '''
        # Define symbols
        r, theta = sp.symbols("r theta")

        # Shear stress intensity
        sif_shear = 1/2*sp.cos(theta/2)*(K1*sp.sin(theta)+K2*(3*sp.cos(theta)-1))

        # Determine the first and second derivative w.r.t. theta
        first_derivative = sp.diff(sif_shear, theta)
        second_derivative = sp.diff(first_derivative, theta)

        cpd_lst = []
        for first, second in zip(first_derivative, second_derivative):
            # Lambdify function such that it can evaluate an array of points
            func1 = sp.lambdify(theta, first, "numpy")
            func2 = sp.lambdify(theta, second, "numpy")

            # initialize array from -π/2 to π/2, this is used for the initial guesses of the solver
            x = np.linspace(-np.pi/2, np.pi/2, num=50)

            # Solve the first derivative for all initial guesses to find possible propagation angles
            y1 = fsolve(func1, x)

            # Evaluate the second derivative in the roots of the first derivative
            y2 = func2(y1)
                   
            # Get roots of first derivative between -π/2 to π/2 
            # and where second derivative is negative
            y1 = np.round(y1, 4)
            y1 = y1[(y1 > -np.pi/2) & (y1 < np.pi/2) & (y2 < 0)]

            # get unique roots
            cpd = np.unique(y1)

            cpd_lst.append(cpd)

        return cpd_lst

输入示例:

K1 = np.random.rand(10000,)
K2 = np.random.rand(10000,)
get_cpd(K1, K2)

Answer 1

最好的办法是在符号参数方面尽可能多地尝试符号化地处理方程。可以得到一个解析解，例如first_derivative但你需要稍微改造一下。在这里，我将 sin/cos 重写为 exp，然后使用替换 exp(I*theta/2) = sqrt(z)得到 z 的三次多项式:

In [150]: K1, K2 = symbols('K1, K2', real=True)                                                                                                                                                             

In [151]: theta = Symbol('theta', real=True)                                                                                                                                                                

In [152]: sif_shear = S.Half*sp.cos(theta/2)*(K1*sin(theta)+K2*(3*cos(theta)-1))                                                                                                                            

In [153]: eq = diff(sif_shear, theta)                                                                                                                                                                       

In [154]: eq                                                                                                                                                                                                
Out[154]: 
  ⎛K₁⋅sin(θ)   K₂⋅(3⋅cos(θ) - 1)⎞    ⎛θ⎞                                   
  ⎜───────── + ─────────────────⎟⋅sin⎜─⎟                                   
  ⎝    2               2        ⎠    ⎝2⎠   ⎛K₁⋅cos(θ)   3⋅K₂⋅sin(θ)⎞    ⎛θ⎞
- ────────────────────────────────────── + ⎜───────── - ───────────⎟⋅cos⎜─⎟
                    2                      ⎝    2            2     ⎠    ⎝2⎠

In [155]: eqz = fraction(cancel(eq.rewrite(exp).subs(exp(I*theta/2), sqrt(z))))[0].collect(z)                                                                                                               

In [156]: eqz                                                                                                                                                                                               
Out[156]: 
                 3                    2                            
3⋅K₁ - 9⋅ⅈ⋅K₂ + z ⋅(3⋅K₁ + 9⋅ⅈ⋅K₂) + z ⋅(K₁ + ⅈ⋅K₂) + z⋅(K₁ - ⅈ⋅K₂)

现在 sympy 可以解决这个问题( roots(eqz, z) )，但是三次方的通用公式非常复杂，所以这可能不是最好的方法。给定 K1 的特定浮点值和 K2尽管 sympy 可以轻松获得 nroots 的根源否则你可以使用 numpy 的 roots功能。

In [157]: eqzp = eqz.subs({K1:0.2, K2:0.5})                                                                                                                                                                 

In [158]: eqzp                                                                                                                                                                                              
Out[158]: 
 3                  2                                              
z ⋅(0.6 + 4.5⋅ⅈ) + z ⋅(0.2 + 0.5⋅ⅈ) + z⋅(0.2 - 0.5⋅ⅈ) + 0.6 - 4.5⋅ⅈ

In [159]: Poly(eqzp, z).nroots()                                                                                                                                                                            
Out[159]: [-0.617215947987055 + 0.786793793538333⋅ⅈ, -0.491339121039621 - 0.870968350823388⋅ⅈ, 0.993562347047054 + 0.113286638798883⋅ⅈ]

In [163]: coeffs = [complex(c) for c in Poly(eqzp, z).all_coeffs()]                                                                                                                                            

In [164]: np.roots(coeffs)                                                                                                                                                                                  
Out[164]: 
array([ 0.99356235+0.11328664j, -0.61721595+0.78679379j,
       -0.49133912-0.87096835j])

无论哪种方式，这都会为 z 提供 3 个可能的值这是 exp(I*theta)所以你可以得到 theta (模 2*pi ):

In [167]: r1, r2, r3 = Poly(eqzp, z).nroots()                                                                                                                                                               

In [168]: get_theta = lambda r: acos((r + r.conjugate())/2)                                                                                                                                                 

In [169]: get_theta(r1)                                                                                                                                                                                     
Out[169]: 2.23599562043958

In [170]: get_theta(r2)                                                                                                                                                                                     
Out[170]: 2.08442292239622

In [171]: get_theta(r3)                                                                                                                                                                                     
Out[171]: 0.113530366549989

我们所做的转换意味着 +-这些值可以是原始方程的解，因此我们可以通过代入例如:

In [178]: eq.subs({K1:0.2, K2:0.5}).subs(theta, get_theta(r1))                                                                                                                                              
Out[178]: -5.55111512312578e-17

In [179]: eq.subs({K1:0.2, K2:0.5}).subs(theta, get_theta(r2))                                                                                                                                              
Out[179]: -0.124767626702216

In [180]: eq.subs({K1:0.2, K2:0.5}).subs(theta, -get_theta(r2))                                                                                                                                             
Out[180]: 5.55111512312578e-17