pendulum/methods.py at master · kencx/pendulum · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import numpy as np
from math import pi


class Derivative():

    @staticmethod
    def pend(state, t, l, g):

        '''
        Derivative of the state of a simple pendulum

        Args:
        state (array): State function
        t (int): Time
        l (float): Length of string
        g (float): Gravitational acceleration


        Output: Derivative of input state (array)
        '''

        x_dot = state[1] # dx/dt = v
        v_dot = -(g/l)*np.sin(state[0]) #dv/dt = -g/l * sin(x)
        return np.array([x_dot, v_dot])


    @staticmethod
    def double_pend(state, t, l1, l2, m1, m2, g):

        '''
        Derivate of the state of a double pendulum

        Args:
        state (array): State function
        t (int): Time
        l1, l2 (float): Length of string 1 and string 2
        m1, m2 (float): Mass of mass 1 and mass 2
        g (float): Gravitational acceleration

        Output: Derivative of input state
        '''

        delta = state[1]-state[0]
        g0 = state[2]
        g1 = state[3]
        g2 = (m2*l1*state[2]*state[2]*np.sin(delta)*np.cos(delta) \
        + m2*g*np.sin(state[1])*np.cos(delta) + m2*l2*state[3]*state[3]*np.sin(delta) \
        - (m1+m2)*g*np.sin(state[0])) / ((m1+m2)*l1 - m2*l1*np.cos(delta)*np.cos(delta))
        g3 = ((-m2*l2*state[3]*state[3]*np.sin(delta)*np.cos(delta)) \
        + (m1+m2) * (g*np.sin(state[0])*np.cos(delta) - l1*state[2]*state[2]*np.sin(delta) - g*np.sin(state[1]))) \
        / ((m1+m2)*l2 - m2*l2*np.cos(delta)*np.cos(delta))
        return np.array([g0, g1, g2, g3])


def rk4(f, t, dt, derivs, *args):

    '''
    Runge-Kutta 4 method

    Args:
    f (array): State function at time i
    t (int): Time
    dt (float): Timestep
    derivs (func): Derivate function
    args (float): Additional arguments required for derivs function


    Outputs: State function at time i+1
    '''

    # find a better way to run the function when different derivs function is passed

    if derivs == Derivative.pend:
        k1 = derivs(f, t, args[0], args[1]) * dt
        k2 = derivs(f + 0.5*k1, t + 0.5*dt, args[0], args[1]) * dt
        k3 = derivs(f + 0.5*k2, t + 0.5*dt, args[0], args[1]) * dt
        k4 = derivs(f + k3, t + dt, args[0], args[1]) * dt
        return f + (k1 + 2*k2 + 2*k3 + k4)/6

    if derivs == Derivative.double_pend:
        k1 = derivs(f, t, args[0], args[1], args[2], args[3], args[4]) * dt
        k2 = derivs(f + 0.5*k1, t + 0.5*dt, args[0], args[1], args[2], args[3], args[4]) * dt
        k3 = derivs(f + 0.5*k2, t + 0.5*dt, args[0], args[1], args[2], args[3], args[4]) * dt
        k4 = derivs(f + k3, t + dt, args[0], args[1], args[2], args[3], args[4]) * dt
        return f + (k1 + 2*k2 + 2*k3 + k4)/6


def wrapping(point):
    '''
    Restricts angles to -pi < x < pi

    Args:
    point (float): Point

    Output: Point within -pi < x < pi
    '''
    # change this to remove if, elif as it is unnecessary
    if point > pi:
        while point > pi:
            point -= 2*pi

    elif point < -pi:
        while point < -pi:
            point += 2*pi

    return point