使用scipy.odeint解决预定义函数的非线性ODE

Question

I want to solve a non linear ordinary differential equation of the form

Theta2 = (C + j(Theta2))**-1 * (f(t) – g(Theta1) -h(Theta0))

Where f(), g(), h(), and j() are functions already defined that take Theta2, Theta1, Theta0 or t as an input. Theta2 and Theta1 are the second and first derivative of Theta0 with time t.

I have been solving the equation without the j(Theta2) term using the SciPy.odeint function using the following code:

from scipy.integrate import odeint

def ODE():
    def g(Theta, t):
        Theta0 = Theta[0]
        Theta1 = Theta[1]
        Theta2 = (1/C)*( f(t) - g(Theta1)  - h(Theta0))

        return Theta1, Theta2
    init = 0, 0 # Initial conditions on theta0 and theta1 (velocity) at t=0
    sol=odeint(g, init, t)
    A = sol[:,1]
    B = sol[:,0]
    return(A, B)

Answer 1

The equation could be re-written as:

          F(t, theta, theta')
theta'' = -------------------
            a + b*theta''

where a and b are constants, and F corresponds to (f(t) – g(Theta1) -h(Theta0)).

It is a second order polynomial function of theta'', with 2 solutions (considering b!=0 and a^2 + 4*b*F>0) :

theta'' = -( sqrt(a^2 + 4*b*F) +/- a )/(2*b)

This new equation is of the form y' = f(t, y) which could be solved using regular ODE solver.

Here is an example using solve_ivp which is the replacement for odeint :

import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

a = 20
b = 1

def f(t, y, dydt):
    return t + y**2

def ode_function_plus(t, Y):
    y = Y[0]
    dydt = Y[1]

    d2y_dt2 = -(np.sqrt(a**2 + 4*b*f(t, y, dydt)) + a )/(2*b)
    return [dydt, d2y_dt2]

def ode_function_minus(t, Y):
    y = Y[0]
    dydt = Y[1]

    d2y_dt2 = -(np.sqrt(a**2 + 4*b*f(t, y, dydt)) - a )/(2*b)
    return [dydt, d2y_dt2]

# Solve
t_span = [0, 4]
Y0 = [10, 1]
sol_plus = solve_ivp(ode_function_plus, t_span, Y0)
sol_minus = solve_ivp(ode_function_minus, t_span, Y0)
print(sol_plus.message)

# Graph
plt.plot(sol_plus.t, sol_plus.y[0, :], label='solution +a');
plt.plot(sol_minus.t, sol_minus.y[0, :], label='solution -a');
plt.xlabel('time'); plt.ylabel('y'); plt.legend();