Search code examples
pythonodedifferential-equationsequation-solvingrunge-kutta

Runge-Kutta 4 method for complex differential equations in python


I am trying to solve the equations of motion of a complicated system using the RK4- method in python, but I am not very familiar with this method and need help in fixing my code. The differential equations are:

First part of equations:

First part of equations

next part of equations:

next part of equations

The code I have written is:

import numpy as np

p = 7.850
A = (0.01) ** 2
m1 = 1
m3 = 1
g = 9.81

r2 = 1

r3 = 1

r1 = 1

rcm = r1

r = 4 * r2

m2 = 1


def equations(t,y):
    theta, phi1, phi2, x1, w1, p1 = y

    f0 = x1
    f1 = w1
    f2 = p1  # Corrected the variable name to p1

    # Calculate f4 and f5 separately
    f4 = ((m1 * r2 * rcm * x1**2 * np.sin(theta + phi1)) - (m3 * r3 * (r-r2) * p1**2 * np.sin(phi2 - phi1)) - (g * m1 * r2 * np.sin(phi1)) + (g * m2 * ((r/2)-r2) * np.sin(phi1)) + (g * m3 * (r-r2) * np.sin(phi1)) - (m1 * r2 * rcm * np.cos(theta + phi1) * f3) + (m3 * r3 * (r-r2) * np.cos(phi2-phi1) * f5))/((m1 * r2**2))
    f5 = ((m3 * r3 * (r-r2) * np.sin(phi2 - phi1) * w1**2) + (g * m3 * r3 * np.sin(phi2)) + (m3 * r3 * (r-r2) * np.cos(phi2-phi1)))/(m3 * r3**2)

    # Calculate f3 using the previously calculated f4 and f5
    f3 = ((m1 * r2 * rcm * w1**2 * np.sin(theta + phi1)) + (g * m1 * rcm * np.sin(theta)) - (m1 * r2 * rcm * np.cos(theta + phi1) * f4))/(m1 * r2**2)
    theta1, phi1, phi2, x1, w1, p1 = y
    return np.array([f0, f1, f2, f3, f4, f5])


def RK4(t, y, h, f):
    theta, phi1, phi2, x1, w1, p1 = y
    # Runge Kutta standard calculations
    k1 = f(t, y)
    k2 = f(t + h/2, y + h/2 * k1)
    k3 = f(t + h/2, y + h/2 * k2)
    k4 = f(t + h, y + h * k3)

    return 1/6 * (k1 + 2 * k2 + 2 * k3 + k4)

def main():
    
     # Specify time interval and number of domain subintervals
    t0 = 0
    T = 100
    n = 10000
    # initial conditions
    y0 = [np.pi - np.arccos((2 - 1.5) / r2), np.arccos((2 - 1.5) / r2), np.arcsin(1.5/(r-r2)), 0, 0, 0]

    # Domain discretization
    t_n = np.linspace(t0, T, n)
    y_n = [np.array(y0)]

    # Step size
    h = (T - t0) / n

    while t_n[-1] < T:
        # Keep going until T is reached.
        # Store numerical solutions into y_n
        y_n.append(y_n[-1] + h * RK4(t_n[-1], y_n[-1], h, equations))
        t_n = np.append(t_n, t_n[-1] + h)

    print(y_n)

if __name__ == "__main__":
    main()

Running this code, i recieve the output:

[array([2.0943951 , 1.04719755, 0.52359878, 0.        , 0.        ,
       0.        ])]

This looks like it is returning the initial conditions meaning the system stays constant and that's not the outcome i am looking for. Any help and/ or alternative methods will be greatly appreciated as i would also like to simulate the system. Thank you in advance.


Solution

  • Use sympy or the CAS of your choice to get the many derivatives in the Euler-Lagrange system automatically correct.

    Most mechanical Lagrange functions in conservative systems are separable into kinetic and potential term

    L(x,Dx) = 1/2*Dx^T*M(x)*Dx - V(x),   D=d/dt the time derivative
    

    The dynamical equations can then be written as

    p = M(x)*Dx
    Dp = 1/2*grad_x(Dx^T*M(x)*Dx)-grad_x V(x)
    

    Using the impulse variable p leads to the Hamilton formalism with Dx = M(x)\p. The aim here however is to use the second derivatives, that is, the first derivatives as components of the state. To that aim take the time derivative of the first equation to get

    Dp = M1(x;Dx)*Dx + M(x)*DDx
    

    where M1(x;Dx) is the directional derivative of the matrix-valued function M(x) in direction Dx. To compare

    (M1(x;Dx)*Dx)[i] = M(x)[i,j;k]*Dx[j]*Dx[k]
    grad_x(Dx^T*M(x)*Dx)[i] = M(x)[j,k;i]*Dx[j]*D[k]
    

    with the semi-colon separating matrix position indices and derivative direction indices.

    Then

    M(x)*DDx = -grad_x V(x) + 1/2*grad_x(Dx^T*M(x)*Dx) - M1(x;Dx)*Dx
    

    is now a linear system for the second derivatives DDx that can be solved with methods of the linear algebra module.

    So if you can identify M(x) and V(x) from your given Lagrangian, you can step over the manual calculation of the derivatives and implement the ODE system systematically.