| use std::autodiff::autodiff; |
| |
| const N: usize = 32; |
| const xmin: f64 = 0.; |
| const xmax: f64 = 1.; |
| const ymin: f64 = 0.; |
| const ymax: f64 = 1.; |
| |
| #[inline(always)] |
| fn range(min: f64, max: f64, i: usize, N_var: usize) -> f64 { |
| (max - min) / (N_var as f64 - 1.) * i as f64 + min |
| } |
| |
| fn brusselator_f(x: f64, y: f64, t: f64) -> f64 { |
| let eq1 = (x - 0.3) * (x - 0.3) + (y - 0.6) * (y - 0.6) <= 0.1 * 0.1; |
| let eq2 = t >= 1.1; |
| if eq1 && eq2 { |
| 5.0 |
| } else { |
| 0.0 |
| } |
| } |
| |
| #[expect(unused)] |
| fn init_brusselator(u: &mut [f64], v: &mut [f64]) { |
| assert!(u.len() == N * N); |
| assert!(v.len() == N * N); |
| for i in 0..N { |
| for j in 0..N { |
| let x = range(xmin, xmax, i, N); |
| let y = range(ymin, ymax, j, N); |
| u[N * i + j] = 22.0 * (y * (1.0 - y)) * (y * (1.0 - y)).sqrt(); |
| v[N * i + j] = 27.0 * (x * (1.0 - x)) * (x * (1.0 - x)).sqrt(); |
| } |
| } |
| } |
| |
| #[no_mangle] |
| #[autodiff(dbrusselator_2d_loop, Reverse, Duplicated, Duplicated, Duplicated, Duplicated, Duplicated, Const)] |
| pub fn brusselator_2d_loop(d_u: &mut [f64;N*N], d_v: &mut [f64;N*N], u: &[f64;N*N], v: &[f64;N*N], p: &[f64;3], t: f64) { |
| let A = p[0]; |
| let B = p[1]; |
| let alpha = p[2]; |
| let dx = 1. / (N - 1) as f64; |
| let alpha = alpha / (dx * dx); |
| for i in 0..N { |
| for j in 0..N { |
| let x = range(xmin, xmax, i, N); |
| let y = range(ymin, ymax, j, N); |
| let ip1 = if i == N - 1 { i } else { i + 1 }; |
| let im1 = if i == 0 { i } else { i - 1 }; |
| let jp1 = if j == N - 1 { j } else { j + 1 }; |
| let jm1 = if j == 0 { j } else { j - 1 }; |
| let u2v = u[N * i + j] * u[N * i + j] * v[N * i + j]; |
| d_u[N * i + j] = alpha * (u[N * im1 + j] + u[N * ip1 + j] + u[N * i + jp1] + u[N * i + jm1] - 4. * u[N * i + j]) |
| + B + u2v - (A + 1.) * u[N * i + j] + brusselator_f(x, y, t); |
| d_v[N * i + j] = alpha * (v[N * im1 + j] + v[N * ip1 + j] + v[N * i + jp1] + v[N * i + jm1] - 4. * v[N * i + j]) |
| + A * u[N * i + j] - u2v; |
| } |
| } |
| } |
| |
| pub type StateType = [f64; 2 * N * N]; |
| |
| pub fn lorenz(x: &StateType, dxdt: &mut StateType, t: f64) { |
| let p = [3.4, 1., 10.]; |
| let (tmp1, tmp2) = dxdt.split_at_mut(N * N); |
| let mut dxdt1: [f64; N * N] = tmp1.try_into().unwrap(); |
| let mut dxdt2: [f64; N * N] = tmp2.try_into().unwrap(); |
| let (tmp1, tmp2) = x.split_at(N * N); |
| let u: [f64; N * N] = tmp1.try_into().unwrap(); |
| let v: [f64; N * N] = tmp2.try_into().unwrap(); |
| brusselator_2d_loop(&mut dxdt1, &mut dxdt2, &u, &v, &p, t); |
| } |
| |
