用于解决ODE Python系统的Runge-Kutta 4

问题描述

我为Runge-Kutta 4编写了用于求解ODE系统的代码。
它对于一维ODE很好用，但是当我尝试求解x'' + kx = 0时，尝试定义矢量函数时遇到了问题：

让u1 = x和u2 = x' = u1'，则系统如下所示：

u1' = u2
u2' = -k*u1

如果u = (u1,u2)和f(u,t) = (u2,-k*u1)，那么我们需要解决：

u' = f(u,t)

def f(u,t,omega=2):
    u,v = u
    return np.asarray([v,-omega**2*u])

我的整个代码是：

import numpy as np

def ode_RK4(f,X_0,dt,T):    
    N_t = int(round(T/dt))
    #  Create an array for the functions ui 
    u = np.zeros((len(X_0),N_t+1)) # Array u[j,:] corresponds to the j-solution
    t = np.linspace(0,N_t*dt,N_t + 1)
    # Initial conditions
    for j in range(len(X_0)):
        u[j,0] = X_0[j]
    # RK4
    for j in range(len(X_0)):
        for n in range(N_t):
            u1 = f(u[j,n] + 0.5*dt* f(u[j,n],t[n])[j],t[n] + 0.5*dt)[j]
            u2 = f(u[j,n] + 0.5*dt*u1,t[n] + 0.5*dt)[j]
            u3 = f(u[j,n] + dt*u2,t[n] + dt)[j]
            u[j,n+1] = u[j,n] + (1/6)*dt*( f(u[j,t[n])[j] + 2*u1 + 2*u2 + u3)
    
    return u,t

def demo_exp():
    import matplotlib.pyplot as plt
    
    def f(u,t):
        return np.asarray([u])

    u,t = ode_RK4(f,[1],0.1,1.5)
    
    plt.plot(t,u[0,:],"b*",np.exp(t),"r-")
    plt.show()
    
def demo_osci():
    import matplotlib.pyplot as plt
    
    def f(u,omega=2):
        # u,v = u Here I've got a problem
        return np.asarray([v,-omega**2*u])
    
    u,[2,0],2)
    
    for i in [1]:
        plt.plot(t,u[i,"b*")
    plt.show()
    

预先，谢谢。

解决方法

您走在正确的道路上，但是，当将诸如RK的时间积分方法应用于矢量值ODE时，实际上只需要矢量就可以完成与标量情况完全相同的事情。

因此，您跳过for j in range(len(X_0))循环和相关的索引编制，并确保将初始值作为向量（numpy数组）传递。

还清理了t的索引编制并将解决方案存储在列表中。

import numpy as np

def ode_RK4(f,X_0,dt,T):    
    N_t = int(round(T/dt))
    # Initial conditions
    usol = [X_0]
    u = np.copy(X_0)
    
    tt = np.linspace(0,N_t*dt,N_t + 1)
    # RK4
    for t in tt[:-1]:
        u1 = f(u + 0.5*dt* f(u,t),t + 0.5*dt)
        u2 = f(u + 0.5*dt*u1,t + 0.5*dt)
        u3 = f(u + dt*u2,t + dt)
        u = u + (1/6)*dt*( f(u,t) + 2*u1 + 2*u2 + u3)
        usol.append(u)
    return usol,tt

def demo_exp():
    import matplotlib.pyplot as plt
    
    def f(u,t):
        return np.asarray([u])

    u,t = ode_RK4(f,np.array([1]),0.1,1.5)
    
    plt.plot(t,u,"b*",t,np.exp(t),"r-")
    plt.show()
    
def demo_osci():
    import matplotlib.pyplot as plt
    
    def f(u,omega=2):
        u,v = u 
        return np.asarray([v,-omega**2*u])
    
    u,np.array([2,0]),2)
    
    u1 = [a[0] for a in u]
    
    for i in [1]:
        plt.plot(t,u1,"b*")
    plt.show()

,

模型是这样的： enter image description here

摘自Langtangen的《编程编程-Python》。