RK4求解器Python代码

  • 该类代码仿照MATLAB中的ode45方法的需求进行编写
  • 可以进行继承，对被积函数与截止函数（积分截止条件）进行重构。
  • 该类的输入为时间步长，积分初值以及积分时间范围，返回积分时间与积分结果
  • 下面是代码
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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
import numpy as np


class RK4Solver:
    # 该类为写的任意函数的求解器。任意函数为func。
    # func的输入为(t, y)，输出为dydt，y和dydt均为numpy.array类型，其中y尽量为行向量
    # break_func为终止条件，输入为t, y，输出为(value, isterminal, direction)
    # (value, isterminal, direction)中：
    # value: 待检测的数值
    # isterminal: 是否终止，1为终止，0为继续
    # direction: 过0点的方向，1为从负到正，-1为从正到负
    # 输入参数解析：
    def __init__(self, dt: float, y0, tspan: list):
        # dt:       时间步长
        # y0:       积分初值
        # tspan:    积分上下限
        self.dt = dt
        # 如果输入错了，纠正一下
        if type(y0) is list:
            k = len(y0)
            y11 = np.zeros([1, k])
            for w in range(k):
                y11[0, w] = y0[w]
            y0 = y11
        self.y0 = y0
        self.tspan = tspan
        # 进行数据初始化
        # 时间序列
        self.t_list = np.arange(min(self.tspan), max(self.tspan), self.dt)
        # 时间序列长度
        self.M = np.shape(self.t_list)[0]
        # 空间维度
        self.N = max(np.shape(y0))
        # 定义时间为行，空间为列
        self.y_list = np.zeros([self.M, self.N])
        # 对定义好的初值先塞进去
        # 为了兼容性，这里用for循环对其进行赋值
        for w in range(self.N):
            self.y_list[0, w] = self.y0[0, w]
        # 定义截至监测list
        self.value_list = np.zeros([self.M, 1])
        # 记录截止的点
        self.break_point = self.M

    def func(self, t, y):
        # 写的一个三变量的demo，其方程为dz/dt = z
        # 先将函数取出来
        x0 = y[0, 0]
        y0 = y[0, 1]
        z0 = y[0, 2]
        # 此处省略x个计算步骤
        # 定义返回的微分表达式
        dydt = np.zeros([1, 3])
        dydt[0, 0] = x0
        dydt[0, 1] = y0
        dydt[0, 2] = z0
        return dydt

    def break_func(self, t, y):
        # 该函数为截至函数的demo
        # 输入为t, y
        # 输出为(value, isterminal, direction)
        # value: 待检测的数值
        # isterminal: 是否终止，1为终止，0为继续
        # direction: 过0点的方向，1为从负到正，-1为从正到负
        value = y[0] - 10
        isterminal = 1
        direction = 1
        return value, isterminal, direction

    def RK4step(self, t, y):
        # 进行单步的RK4积分
        # 为了后续的兼容，输入输出应该全是numpy.array变量
        if type(y) is list:
            raise ValueError('The function return must be numpy.array!')
        # 进行numpy.array的一维转换
        yi = np.zeros([1, self.N])
        for i in range(self.N):
            yi[0, i] = y[i]
        y = yi
        # 进行正常的RK4积分步骤
        dydt1 = self.func(t, y)

        dydt2 = self.func(t + self.dt / 2, y + dydt1 / 2 * self.dt)

        dydt3 = self.func(t + self.dt / 2, y + dydt2 / 2 * self.dt)

        dydt4 = self.func(t + self.dt, y + dydt3 * self.dt)

        dydt = (dydt1 + 2 * dydt2 + 2 * dydt3 + dydt4) / 6

        return dydt

    def calculate(self):
        # 由于初始化已经在init中，这里直接开始写demo
        # 直接进行for循环吧
        # 由于初值在init中率先赋值给y_list了，所以跳过0序列
        for i in range(1, self.M):
            dydti = self.RK4step(self.t_list[i - 1], self.y_list[i - 1, :])
            yi = self.y_list[i, :] + dydti
            self.y_list[i] = yi
            # 下面进行截至判定
            if i == 1:
                # 如果是第一次循环，需要先对判定第一个判定一下
                valuei0, _, _ = self.break_func(self.t_list[i - 1], self.y_list[i - 1, :])
                self.value_list[i - 1] = valuei0
            value, isterminal, direction = self.break_func(self.t_list[i], self.y_list[i, :])
            self.value_list[i] = value
            # 进行终止判定
            if isterminal == 1:
                # 下面是终止判定
                checkvalue = self.value_list[i] * self.value_list[i - 1]
                if checkvalue <= 0:
                    # 也就是变号的时候
                    if direction > 0:
                        # 从负数到正数
                        if self.value_list[i - 1] <= 0 and self.value_list[i] >= 0:
                            self.break_point = i
                            break
                    else:
                        # 从正到负
                        if self.value_list[i - 1] >= 0 and self.value_list[i] <= 0:
                            self.break_point = i
                            break
                else:
                    continue
        return self.y_list

    def run(self):
        # 定义最后的运行函数
        y_list = self.calculate()
        t_list = self.t_list[:self.break_point]
        return t_list, y_list[:self.break_point, :]


if __main__ == '__main__':
    # 这是一个测试的demo
    Y = RK4Solver(0.1, [1, 1, 1], [0, 10])
    t, y = Y.run()