馃悕馃悕馃悕
1
2# ODE solvers
3
4def rk4_step(get_derivative, take_step, position_0):
5 direction_0 = get_derivative(position_0)
6 position_1 = take_step(position_0, direction_0, 0.5)
7 direction_1 = get_derivative(position_1)
8 position_2 = take_step(position_0, direction_1, 0.5)
9 direction_2 = get_derivative(position_2)
10 position_3 = take_step(position_0, direction_2, 1)
11 direction_3 = get_derivative(position_3)
12 final_direction = (direction_0 + 2 * (direction_1 + direction_2) + direction_3) / 6
13 return take_step(position_0, final_direction, 1)
14
15def euler_step(get_derivative, take_step, position_0):
16 return take_step(position_0, get_derivative(position_0), 1)
17
18# From old diffusion notebook:
19
20def _euler_step(get_derivative, take_step, position_0):
21 return take_step(position_0, get_derivative(position_0), 0, 1)
22
23def _heun_step(get_derivative, take_step, position_0):
24 direction_0 = get_derivative(position_0)
25 position_1 = take_step(position_0, direction_0, 0, 1)
26 direction_1 = get_derivative(position_1, 1)
27 final_direction = tuple((a + b) / 2 for (a,b) in zip(direction_0, direction_1))
28 return take_step(position_0, final_direction, 0, 1)
29
30def _rk2_step(get_derivative, take_step, position_0):
31 direction_0 = get_derivative(position_0)
32 position_1 = take_step(position_0, direction_0, 0, 0.5)
33 final_direction = get_derivative(position_1, 0.5)
34 return take_step(position_0, final_direction, 0, 1)
35
36def _rk4_step(get_derivative, take_step, position_0):
37 direction_0 = get_derivative(position_0)
38 position_1 = take_step(position_0, direction_0, 0, 0.5)
39 direction_1 = get_derivative(position_1, 0.5)
40 position_2 = take_step(position_0, direction_1, 0, 0.5)
41 direction_2 = get_derivative(position_2, 0.5)
42 position_3 = take_step(position_0, direction_2, 0, 1)
43 direction_3 = get_derivative(position_3, 1)
44 final_direction = tuple((a + 2 * (b + c) + d) / 6 for (a,b,c,d) in zip(direction_0, direction_1, direction_2, direction_3))
45 return take_step(position_0, final_direction, 0, 1)
46