馃悕馃悕馃悕
at main 2.1 kB view raw
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