| import jax |
| import jax.numpy as jnp |
|
|
| @jax.jit |
| def lorenz63_step(state,t,params): |
| sigma, rho, beta = params |
| x, y, z = state |
| dxdt = sigma * (y - x) |
| dydt = x * (rho - z) - y |
| dzdt = x * y - beta * z |
| return jnp.stack([dxdt, dydt, dzdt],axis=-1) |
|
|
| def forward_euler_step(f,y,t,params,dt): |
| return y + dt * f(y, t, params) |
|
|
|
|
| def rk4_step(f,y,t,params,dt): |
| k1 = dt * f(y, t, params) |
| k2 = dt * f(y + 0.5 * k1, t + dt/2, params) |
| k3 = dt * f(y + 0.5 * k2, t + dt/2, params) |
| k4 = dt * f(y + k3, t + dt, params) |
| return y + (k1 + 2 * k2 + 2 * k3 + k4) / 6 |
|
|
| def integrate_lorenz(params,y0,t,dt): |
| def step(y_i,t_i): |
| y_next = rk4_step(lorenz63_step, y_i,t_i, params, dt) |
| return y_next, y_next |
| |
| _, ys = jax.lax.scan(step, y0, t[:-1]) |
| traj = jnp.concatenate([y0[None,...], ys], axis=0) |
| return traj |
