使用 numpy.piecewise 的问题
Problems using numpy.piecewise
1。核心问题与疑问
我将在下面提供一个可执行示例，但让我先向您介绍一下问题。
我正在使用 scipy.integrate 中的 solve_ivp 来解决初始值问题 (see documentation)。事实上，我必须两次调用求解器，一次向前积分，一次向后积分。 （我将不得不不必要地深入我的具体问题来解释为什么这是必要的，但请相信我——它是！）
sol0 = solve_ivp(rhs,[0,-1e8],y0,rtol=10e-12,atol=10e-12,dense_output=True)
sol1 = solve_ivp(rhs,[0, 1e8],y0,rtol=10e-12,atol=10e-12,dense_output=True)
这里rhs是初值问题y(t) = rhs(t,y)的右边函数。在我的例子中，y 有六个组成部分 y[0] 到 y[5]。 y0=y(0) 是初始条件。 [0,±1e8]是各自的积分范围，时间上一前一后。 rtol 和 atol 是公差。
重要的是，你看到我标记了 dense_output=True，这意味着求解器不仅 return 数值网格上的解，而且还作为插值函数 sol0.sol(t) 和 sol1.sol(t).
我现在的主要目标是定义一个分段函数，比如 sol(t)，它为 t<0 取值 sol0.sol(t) 和值 sol1.sol(t)对于 t>=0。所以主要问题是：我该怎么做？
我认为 numpy.piecewise 应该是我执行此操作的首选工具。但是我在使用它时遇到了问题，正如您将在下面看到的，我向您展示了到目前为止我尝试过的方法。
2。示例代码
下框中的代码解决了我的例子的初值问题。大部分代码都是rhs函数的定义，具体细节对题来说并不重要
import numpy as np
from scipy.integrate import solve_ivp
# aux definitions and constants
sin=np.sin; cos=np.cos; tan=np.tan; sqrt=np.sqrt; pi=np.pi;
c = 299792458
Gm = 5.655090674872875e26
# define right hand side function of initial value problem, y'(t) = rhs(t,y)
def rhs(t,y):
p,e,i,Om,om,f = y
sinf=np.sin(f); cosf=np.cos(f); Q=sqrt(p/Gm); opecf=1+e*cosf;
R = Gm**2/(c**2*p**3)*opecf**2*(3*(e**2 + 1) + 2*e*cosf - 4*e**2*cosf**2)
S = Gm**2/(c**2*p**3)*4*opecf**3*e*sinf
rhs = np.zeros(6)
rhs[0] = 2*sqrt(p**3/Gm)/opecf*S
rhs[1] = Q*(sinf*R + (2*cosf + e*(1 + cosf**2))/opecf*S)
rhs[2] = 0
rhs[3] = 0
rhs[4] = Q/e*(-cosf*R + (2 + e*cosf)/opecf*sinf*S)
rhs[5] = sqrt(Gm/p**3)*opecf**2 + Q/e*(cosf*R - (2 + e*cosf)/opecf*sinf*S)
return rhs
# define initial values, y0
y0=[3.3578528933149297e13,0.8846,2.34921,3.98284,1.15715,0]
# integrate twice from t = 0, once backward in time (sol0) and once forward in time (sol1)
sol0 = solve_ivp(rhs,[0,-1e8],y0,rtol=10e-12,atol=10e-12,dense_output=True)
sol1 = solve_ivp(rhs,[0, 1e8],y0,rtol=10e-12,atol=10e-12,dense_output=True)
解决方案功能可以从这里分别通过 sol0.sol 和 sol1.sol 解决。例如，让我们绘制第 4 个分量：
from matplotlib import pyplot as plt
t0 = np.linspace(-1,0,500)*1e8
t1 = np.linspace( 0,1,500)*1e8
plt.plot(t0,sol0.sol(t0)[4])
plt.plot(t1,sol1.sol(t1)[4])
plt.title('plot 1')
plt.show()
3。尝试构建分段函数失败
3.1 直接从 sol0.sol 和 sol1.sol
中构建向量值分段函数
def sol(t): return np.piecewise(t,[t<0,t>=0],[sol0.sol,sol1.sol])
t = np.linspace(-1,1,1000)*1e8
print(sol(t))
这会导致在.../numpy/lib/function_base.py 的第 628 行中出现以下分段错误：
TypeError: NumPy boolean array indexing assignment requires a 0 or 1-dimensional input, input has 2 dimensions
我不确定，但我确实认为这是因为以下原因：在 documentation of piecewise 中它提到了第三个参数：
funclistlist of callables, f(x,*args,**kw), or scalars
[...]. It should take a 1d array as input and give an 1d array or a scalar value as output. [...].
我想问题是，我的解决方案有六个组成部分。因此，在时间网格上评估输出将是一个二维数组。有人可以确认这确实是问题所在吗？因为我认为这确实大大限制了 piecewise 的用处。
3.2 尝试相同的方法，但只针对一个组件（例如第 4 个组件）
def sol4(t): return np.piecewise(t,[t<0,t>=0],[sol0.sol(t)[4],sol1.sol(t)[4]])
t = np.linspace(-1,1,1000)*1e8
print(sol4(t))
这会导致与上面相同的文件的第 624 行出现此错误：
ValueError: NumPy boolean array indexing assignment cannot assign 1000 input values to the 500 output values where the mask is true
与之前的错误相反，不幸的是，到目前为止我不知道为什么它不起作用。
3.3 类似的尝试，但是首先为第 4 个组件定义函数
def sol40(t): return sol0.sol(t)[4]
def sol41(t): return sol1.sol(t)[4]
def sol4(t): return np.piecewise(t,[t<0,t>=0],[sol40,sol41])
t = np.linspace(-1,1,1000)
plt.plot(t,sol4(t))
plt.title('plot 2')
plt.show()
现在这不会导致错误，我可以生成一个图，但是这个图看起来并不像它应该的那样。它应该看起来像上面的情节 1。同样在这里，我到目前为止不知道发生了什么。
感谢您的帮助！
您可以查看 numpy.piecewise 源代码。这个功能没有什么特别的，所以我建议手动完成所有事情。
def sol(t):
ans = np.empty((6, len(t)))
ans[:, t<0] = sol0.sol(t[t<0])
ans[:, t>=0] = sol1.sol(t[t>=0])
return ans
关于你失败的尝试。是的，piecewise 排除函数 return 一维数组。您的第二次尝试失败，因为文档说 funclist 参数应该是函数列表或 scalars 但您发送了数组列表。与它甚至可以与数组一起使用的文档相反，您只应使用与 t < 0 和 t >= 0 大小相同的数组，例如：
def sol4(t): return np.piecewise(t,[t<0,t>=0],[sol0.sol(t[t<0])[4],sol1.sol(t[t>=0])[4]])
1。核心问题与疑问
我将在下面提供一个可执行示例，但让我先向您介绍一下问题。
我正在使用 scipy.integrate 中的 solve_ivp 来解决初始值问题 (see documentation)。事实上，我必须两次调用求解器，一次向前积分，一次向后积分。 （我将不得不不必要地深入我的具体问题来解释为什么这是必要的，但请相信我——它是！）
sol0 = solve_ivp(rhs,[0,-1e8],y0,rtol=10e-12,atol=10e-12,dense_output=True)
sol1 = solve_ivp(rhs,[0, 1e8],y0,rtol=10e-12,atol=10e-12,dense_output=True)
这里rhs是初值问题y(t) = rhs(t,y)的右边函数。在我的例子中，y 有六个组成部分 y[0] 到 y[5]。 y0=y(0) 是初始条件。 [0,±1e8]是各自的积分范围，时间上一前一后。 rtol 和 atol 是公差。
重要的是，你看到我标记了 dense_output=True，这意味着求解器不仅 return 数值网格上的解，而且还作为插值函数 sol0.sol(t) 和 sol1.sol(t).
我现在的主要目标是定义一个分段函数，比如 sol(t)，它为 t<0 取值 sol0.sol(t) 和值 sol1.sol(t)对于 t>=0。所以主要问题是：我该怎么做？
我认为 numpy.piecewise 应该是我执行此操作的首选工具。但是我在使用它时遇到了问题，正如您将在下面看到的，我向您展示了到目前为止我尝试过的方法。
2。示例代码
下框中的代码解决了我的例子的初值问题。大部分代码都是rhs函数的定义，具体细节对题来说并不重要
import numpy as np
from scipy.integrate import solve_ivp
# aux definitions and constants
sin=np.sin; cos=np.cos; tan=np.tan; sqrt=np.sqrt; pi=np.pi;
c = 299792458
Gm = 5.655090674872875e26
# define right hand side function of initial value problem, y'(t) = rhs(t,y)
def rhs(t,y):
p,e,i,Om,om,f = y
sinf=np.sin(f); cosf=np.cos(f); Q=sqrt(p/Gm); opecf=1+e*cosf;
R = Gm**2/(c**2*p**3)*opecf**2*(3*(e**2 + 1) + 2*e*cosf - 4*e**2*cosf**2)
S = Gm**2/(c**2*p**3)*4*opecf**3*e*sinf
rhs = np.zeros(6)
rhs[0] = 2*sqrt(p**3/Gm)/opecf*S
rhs[1] = Q*(sinf*R + (2*cosf + e*(1 + cosf**2))/opecf*S)
rhs[2] = 0
rhs[3] = 0
rhs[4] = Q/e*(-cosf*R + (2 + e*cosf)/opecf*sinf*S)
rhs[5] = sqrt(Gm/p**3)*opecf**2 + Q/e*(cosf*R - (2 + e*cosf)/opecf*sinf*S)
return rhs
# define initial values, y0
y0=[3.3578528933149297e13,0.8846,2.34921,3.98284,1.15715,0]
# integrate twice from t = 0, once backward in time (sol0) and once forward in time (sol1)
sol0 = solve_ivp(rhs,[0,-1e8],y0,rtol=10e-12,atol=10e-12,dense_output=True)
sol1 = solve_ivp(rhs,[0, 1e8],y0,rtol=10e-12,atol=10e-12,dense_output=True)
解决方案功能可以从这里分别通过 sol0.sol 和 sol1.sol 解决。例如，让我们绘制第 4 个分量：
from matplotlib import pyplot as plt
t0 = np.linspace(-1,0,500)*1e8
t1 = np.linspace( 0,1,500)*1e8
plt.plot(t0,sol0.sol(t0)[4])
plt.plot(t1,sol1.sol(t1)[4])
plt.title('plot 1')
plt.show()
3。尝试构建分段函数失败
3.1 直接从 sol0.sol 和 sol1.sol
def sol(t): return np.piecewise(t,[t<0,t>=0],[sol0.sol,sol1.sol])
t = np.linspace(-1,1,1000)*1e8
print(sol(t))
这会导致在.../numpy/lib/function_base.py 的第 628 行中出现以下分段错误：
TypeError: NumPy boolean array indexing assignment requires a 0 or 1-dimensional input, input has 2 dimensions
我不确定，但我确实认为这是因为以下原因：在 documentation of piecewise 中它提到了第三个参数：
funclistlist of callables, f(x,*args,**kw), or scalars
[...]. It should take a 1d array as input and give an 1d array or a scalar value as output. [...].
我想问题是，我的解决方案有六个组成部分。因此，在时间网格上评估输出将是一个二维数组。有人可以确认这确实是问题所在吗？因为我认为这确实大大限制了 piecewise 的用处。
3.2 尝试相同的方法，但只针对一个组件（例如第 4 个组件）
def sol4(t): return np.piecewise(t,[t<0,t>=0],[sol0.sol(t)[4],sol1.sol(t)[4]])
t = np.linspace(-1,1,1000)*1e8
print(sol4(t))
这会导致与上面相同的文件的第 624 行出现此错误：
ValueError: NumPy boolean array indexing assignment cannot assign 1000 input values to the 500 output values where the mask is true
与之前的错误相反，不幸的是，到目前为止我不知道为什么它不起作用。
3.3 类似的尝试，但是首先为第 4 个组件定义函数
def sol40(t): return sol0.sol(t)[4]
def sol41(t): return sol1.sol(t)[4]
def sol4(t): return np.piecewise(t,[t<0,t>=0],[sol40,sol41])
t = np.linspace(-1,1,1000)
plt.plot(t,sol4(t))
plt.title('plot 2')
plt.show()
现在这不会导致错误，我可以生成一个图，但是这个图看起来并不像它应该的那样。它应该看起来像上面的情节 1。同样在这里，我到目前为止不知道发生了什么。
感谢您的帮助！
您可以查看 numpy.piecewise 源代码。这个功能没有什么特别的，所以我建议手动完成所有事情。
def sol(t):
ans = np.empty((6, len(t)))
ans[:, t<0] = sol0.sol(t[t<0])
ans[:, t>=0] = sol1.sol(t[t>=0])
return ans
关于你失败的尝试。是的，piecewise 排除函数 return 一维数组。您的第二次尝试失败，因为文档说 funclist 参数应该是函数列表或 scalars 但您发送了数组列表。与它甚至可以与数组一起使用的文档相反，您只应使用与 t < 0 和 t >= 0 大小相同的数组，例如：
def sol4(t): return np.piecewise(t,[t<0,t>=0],[sol0.sol(t[t<0])[4],sol1.sol(t[t>=0])[4]])