File size: 2,700 Bytes
f43af3c |
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 |
def ode_update_op(z0, dz, dt):
"""
General update operation for solving ODEs.
Args:
z0: Tensor or a list for Tensor whose shape is [..., dim]
State at t0.
dz: Tensor or a list for Tensor whose shape is [..., dim]
Differentiation of state.
dt: Tensor with shape [..., 1]
Equal to t1 - t0.
Returns:
"""
if isinstance(z0, list) or isinstance(z0, tuple):
return [item_z + dt * item_dz for item_z, item_dz in zip(z0, dz)]
else:
return z0 + dt * dz
def euler_step_method(diff_func, dt, z0):
"""
Euler method for solving ODEs.
Args:
diff_func: function(state)
Differential equation.
dt: Tensor with shape [..., 1]
Equal to t1 - t0.
z0: Tensor or a list for Tensor whose shape is [..., dim]
State at t0.
Returns:
Tensor or a list for Tensor whose shape is [..., dim], which is updated state.
"""
dz = diff_func(z0)
return ode_update_op(z0, dz, dt)
def rk2_step_method(diff_func, dt, z0):
"""
Second order Runge-Kutta method for solving ODEs.
Args:
diff_func: function(dt, state)
Differential equation.
dt: Tensor with shape [..., 1]
Equal to t1 - t0.
z0: Tensor or a list for Tensor whose shape is [..., dim]
State at t0.
Returns:
Tensor or a list for Tensor whose shape is [..., dim]
"""
# shape -> [..., dim]
k1 = diff_func(z0)
k2 = diff_func(ode_update_op(z0, k1, dt))
if isinstance(z0, list) or isinstance(z0, tuple):
return [item_z + (item_k1 + item_k2) * dt * 0.5 for item_z, item_k1, item_k2 in zip(z0, k1, k2)]
else:
return z0 + dt * (k1 + k2) * 0.5
def rk4_step_method(diff_func, dt, z0):
"""
Fourth order Runge-Kutta method for solving ODEs.
Args:
diff_func: function(dt, state)
Differential equation.
dt: Tensor with shape [..., 1]
Equal to t1 - t0.
z0: Tensor with shape [..., dim]
State at t0.
Returns:
Tensor with shape [..., dim], which is updated state.
"""
# shape -> [..., dim]
k1 = diff_func(z0)
k2 = diff_func(ode_update_op(z0, k1, dt / 2.0))
k3 = diff_func(ode_update_op(z0, k2, dt / 2.0))
k4 = diff_func(ode_update_op(z0, k3, dt))
if isinstance(z0, list) or isinstance(z0, tuple):
return [item_z + (item_k1 + 2.0 * item_k2 + 2.0 * item_k3 + item_k4) * dt / 6.0
for item_z, item_k1, item_k2, item_k3, item_k4 in zip(z0, k1, k2, k3, k4)]
else:
return z0 + dt * (k1 + k2 * 2.0 + k3 * 2.0 + k4) / 6.0
|