矢量化代码中 scipy fsolve 的速度

Speed of scipy fsolve in vectorised code

我有一个大小为 (254, 80) 的数组,我需要对其使用 scipy 的 fsolve。我发现在向量上使用 fsolve 的速度比在 for 循环中更快,但仅适用于长度不超过 100 个值的向量。在此之后,速度迅速下降并变得非常缓慢,有时甚至完全停止。

我目前正在遍历数组的一维并在较小的维度上使用矢量化 fsolve,但它仍然比我 expect/like 花费的时间更长。

有没有人对此有很好的解决方法或知道类似的函数可以很乐意处理更大尺寸的向量?或者如果我做错了什么......

这是当前代码:

for i in range(array.shape[0]):
    f = lambda y: a[i] - m[i]*y - md[i]*(( y**4 + 2*(y**2)*np.cos(Thetas[i,:]) )**0.25)
    ystar[i,:] = fsolve(f, y0[i]) 

(其余变量大小都差不多)

深入研究,我发现

这样的函数
f = lambda y: y*np.tanh(y) - a0/(m**2)

求解速度更快
f = lambda y: (m**2)y*np.tanh(y) - a0

其中 m 和 a0 是大型二维 np 数组。

谁能解释这是为什么?

谢谢, 雷切尔

虽然没有人回答,但我找到了一个解决方法,它避免了 fsolve 函数,而是使用了插值。幸运的是,初始猜测足够好,只需要几个 y 值。如果初始猜测知识不足,则此方法可能不合适。请注意,这仍然存在一些问题,但就我的目的而言,它表现良好......

ystar = np.empty((A,B))     # empty array for the solutions   
    num_ys = 20 #number of points to find where the solution is 
    y0_u = y0 #just so the calculated initial guess isn't overwritten
    for i in range(Thetas.shape[1]):
        ys = np.linspace(-.05,.2,num_ys)[:,None]*np.ones((num_ys,Thetas.shape[0])) + y0_u
        vals = (np.squeeze(eta) - np.squeeze(m)*ys*np.sqrt(g*np.tanh(ys**2*depth)) - np.squeeze(md)*np.sqrt(g*np.tanh(depth*np.sqrt(ys**4+2*(ys**2)*kB*np.cos(Thetas[:,i]+phi_bi)+kB**2)))*(( ys**4+2*(ys**2)*kB*np.cos(Thetas[:,i]+phi_bi)+kB**2 )**0.25))
        idxs_important = -1*(np.clip(np.vstack(((np.sign(vals[:-1]*vals[1:])-1),np.zeros((1,Thetas[:,i].size)))),-1,0) + np.clip(np.vstack((np.zeros((1,Thetas[:,i].size)),(np.sign(vals[:-1]*vals[1:]))-1)),-1,0))

        ys_chosen = idxs_important*ys
        ys_chosen[ys_chosen==0] = 10000 
        sorted_ys_idx = np.argsort(ys_chosen.T, axis = 1)
        sorted_ys = ((ys_chosen.T)[np.arange(np.shape(ys_chosen.T)[0])[:,np.newaxis],sorted_ys_idx]).T
        sorted_vals = (((vals*idxs_important).T)[np.arange(np.shape(vals.T)[0])[:,np.newaxis],sorted_ys_idx]).T
    #    interpolation bit 
        x_id = 0
        yposs = sorted_ys[:2,:]
        valposs = sorted_vals[:2,:]
        y = yposs[0,:] + (yposs[1,:] - yposs[0,:])*(x_id - valposs[0,:])/(valposs[1,:] - valposs[0,:])
        ystar[:,i] = np.squeeze(y)
        y0_u=ystar[:,i]