shape-runtime 0.2.0

/// @module std::core::ode
/// ODE Integrators
///
/// Basic Euler and RK4 integrators for scalar and vector systems.

function vec_add(a, b) {
    let mut out = [];
    for i in range(0, a.len()) {
        out.push(a[i] + b[i]);
    }
    out
}

function vec_scale(a, s) {
    let mut out = [];
    for i in range(0, a.len()) {
        out.push(a[i] * s);
    }
    out
}

function vec_add_scaled(a, b, s) {
    vec_add(a, vec_scale(b, s))
}

/// Euler integrator for scalar ODE
pub fn euler(f, y0, t_start, t_end, dt) {
    let steps = floor((t_end - t_start) / dt);
    let mut t = t_start;
    let mut y = y0;
    let mut results = [];

    for i in range(0, steps + 1) {
        results.push({ t: t, y: y });
        y = y + f(t, y) * dt;
        t = t + dt;
    }

    results
}

/// RK4 integrator for scalar ODE
pub fn rk4(f, y0, t_start, t_end, dt) {
    let steps = floor((t_end - t_start) / dt);
    let mut t = t_start;
    let mut y = y0;
    let mut results = [];

    for i in range(0, steps + 1) {
        results.push({ t: t, y: y });

        let k1 = f(t, y);
        let k2 = f(t + dt / 2.0, y + k1 * dt / 2.0);
        let k3 = f(t + dt / 2.0, y + k2 * dt / 2.0);
        let k4 = f(t + dt, y + k3 * dt);

        y = y + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4);
        t = t + dt;
    }

    results
}

/// Euler integrator for vector systems
pub fn euler_system(f, y0_vec, t_start, t_end, dt) {
    let steps = floor((t_end - t_start) / dt);
    let mut t = t_start;
    let mut y = y0_vec;
    let mut results = [];

    for i in range(0, steps + 1) {
        results.push({ t: t, y: y });
        let dy = f(t, y);
        y = vec_add_scaled(y, dy, dt);
        t = t + dt;
    }

    results
}

/// RK4 integrator for vector systems
pub fn rk4_system(f, y0_vec, t_start, t_end, dt) {
    let steps = floor((t_end - t_start) / dt);
    let mut t = t_start;
    let mut y = y0_vec;
    let mut results = [];

    for i in range(0, steps + 1) {
        results.push({ t: t, y: y });

        let k1 = f(t, y);
        let k2 = f(t + dt / 2.0, vec_add_scaled(y, k1, dt / 2.0));
        let k3 = f(t + dt / 2.0, vec_add_scaled(y, k2, dt / 2.0));
        let k4 = f(t + dt, vec_add_scaled(y, k3, dt));

        let step = vec_add(vec_add(k1, vec_scale(k2, 2.0)), vec_add(vec_scale(k3, 2.0), k4));
        y = vec_add_scaled(y, step, dt / 6.0);
        t = t + dt;
    }

    results
}

// ===== Adaptive Step-Size Integrators =====

function vec_sub(a, b) {
    let mut out = [];
    for i in range(0, a.len()) {
        out.push(a[i] - b[i]);
    }
    out
}

function vec_norm(a) {
    let mut s = 0.0;
    for i in range(0, a.len()) {
        s = s + a[i] * a[i];
    }
    sqrt(s)
}

/// RK45 adaptive integrator for scalar ODE (Dormand-Prince method)
///
/// Automatically adjusts step size to maintain the requested tolerance.
/// Uses an embedded 4th/5th order pair for error estimation.
///
/// @param f - derivative function (t, y) => dy/dt
/// @param y0 - initial value
/// @param t_start - start time
/// @param t_end - end time
/// @param tol - error tolerance (default 1e-6)
/// @param dt_init - initial step size (default (t_end - t_start) / 100)
/// @param dt_min - minimum step size (default 1e-12)
/// @param dt_max - maximum step size (default (t_end - t_start) / 4)
/// @returns array of { t, y } records at each accepted step
pub fn rk45(f, y0, t_start, t_end, tol = 0.000001, dt_init = 0.0, dt_min = 0.000000000001, dt_max = 0.0) {
    let safety = 0.9;
    let max_factor = 5.0;
    let min_factor = 0.2;

    let mut h = dt_init;
    if h == 0.0 {
        h = (t_end - t_start) / 100.0;
    }
    let mut h_max = dt_max;
    if h_max == 0.0 {
        h_max = (t_end - t_start) / 4.0;
    }

    let mut t = t_start;
    let mut y = y0;
    let mut results = [];
    results.push({ t: t, y: y });

    let max_steps = 100000;
    let mut step_count = 0;

    while t < t_end && step_count < max_steps {
        // Clamp step to not overshoot t_end
        if t + h > t_end {
            h = t_end - t;
        }
        if h < dt_min {
            h = dt_min;
        }

        // Dormand-Prince stages
        let k1 = f(t, y);
        let k2 = f(t + h / 5.0, y + h * k1 / 5.0);
        let k3 = f(t + 3.0 * h / 10.0, y + h * (3.0 * k1 / 40.0 + 9.0 * k2 / 40.0));
        let k4 = f(t + 4.0 * h / 5.0, y + h * (44.0 * k1 / 45.0 - 56.0 * k2 / 15.0 + 32.0 * k3 / 9.0));
        let k5 = f(t + 8.0 * h / 9.0, y + h * (19372.0 * k1 / 6561.0 - 25360.0 * k2 / 2187.0 + 64448.0 * k3 / 6561.0 - 212.0 * k4 / 729.0));
        let k6 = f(t + h, y + h * (9017.0 * k1 / 3168.0 - 355.0 * k2 / 33.0 + 46732.0 * k3 / 5247.0 + 49.0 * k4 / 176.0 - 5103.0 * k5 / 18656.0));

        // 5th order solution (for advancing)
        let y5 = y + h * (35.0 * k1 / 384.0 + 500.0 * k3 / 1113.0 + 125.0 * k4 / 192.0 - 2187.0 * k5 / 6784.0 + 11.0 * k6 / 84.0);

        // 4th order solution (for error estimate)
        let k7 = f(t + h, y5);
        let y4 = y + h * (5179.0 * k1 / 57600.0 + 7571.0 * k3 / 16695.0 + 393.0 * k4 / 640.0 - 92097.0 * k5 / 339200.0 + 187.0 * k6 / 2100.0 + k7 / 40.0);

        // Error estimate
        let mut err = abs(y5 - y4);
        if err < 0.000000000000001 {
            err = 0.000000000000001;
        }

        if err <= tol {
            // Accept step
            t = t + h;
            y = y5;
            results.push({ t: t, y: y });
        }

        // Adjust step size
        let mut factor = safety * pow(tol / err, 0.2);
        if factor > max_factor {
            factor = max_factor;
        }
        if factor < min_factor {
            factor = min_factor;
        }
        h = h * factor;
        if h > h_max {
            h = h_max;
        }

        step_count = step_count + 1;
    }

    results
}

/// RK45 adaptive integrator for vector systems (Dormand-Prince method)
///
/// @param f - derivative function (t, y_vec) => dy_vec/dt
/// @param y0_vec - initial state vector
/// @param t_start - start time
/// @param t_end - end time
/// @param tol - error tolerance (default 1e-6)
/// @param dt_init - initial step size (default (t_end - t_start) / 100)
/// @param dt_min - minimum step size (default 1e-12)
/// @param dt_max - maximum step size (default (t_end - t_start) / 4)
/// @returns array of { t, y } records at each accepted step
pub fn rk45_system(f, y0_vec, t_start, t_end, tol = 0.000001, dt_init = 0.0, dt_min = 0.000000000001, dt_max = 0.0) {
    let safety = 0.9;
    let max_factor = 5.0;
    let min_factor = 0.2;

    let mut h = dt_init;
    if h == 0.0 {
        h = (t_end - t_start) / 100.0;
    }
    let mut h_max = dt_max;
    if h_max == 0.0 {
        h_max = (t_end - t_start) / 4.0;
    }

    let mut t = t_start;
    let mut y = y0_vec;
    let mut results = [];
    results.push({ t: t, y: y });

    let max_steps = 100000;
    let mut step_count = 0;

    while t < t_end && step_count < max_steps {
        if t + h > t_end {
            h = t_end - t;
        }
        if h < dt_min {
            h = dt_min;
        }

        // Dormand-Prince stages (vector version)
        let k1 = f(t, y);
        let k2 = f(t + h / 5.0, vec_add_scaled(y, k1, h / 5.0));

        let y3 = vec_add(vec_add_scaled(y, k1, 3.0 * h / 40.0), vec_scale(k2, 9.0 * h / 40.0));
        let k3 = f(t + 3.0 * h / 10.0, y3);

        let y4_stage = vec_add(vec_add(vec_add_scaled(y, k1, 44.0 * h / 45.0), vec_scale(k2, -56.0 * h / 15.0)), vec_scale(k3, 32.0 * h / 9.0));
        let k4 = f(t + 4.0 * h / 5.0, y4_stage);

        let y5_stage = vec_add(vec_add(vec_add(vec_add_scaled(y, k1, 19372.0 * h / 6561.0), vec_scale(k2, -25360.0 * h / 2187.0)), vec_scale(k3, 64448.0 * h / 6561.0)), vec_scale(k4, -212.0 * h / 729.0));
        let k5 = f(t + 8.0 * h / 9.0, y5_stage);

        let y6 = vec_add(vec_add(vec_add(vec_add(vec_add_scaled(y, k1, 9017.0 * h / 3168.0), vec_scale(k2, -355.0 * h / 33.0)), vec_scale(k3, 46732.0 * h / 5247.0)), vec_scale(k4, 49.0 * h / 176.0)), vec_scale(k5, -5103.0 * h / 18656.0));
        let k6 = f(t + h, y6);

        // 5th order solution
        let y_next = vec_add(vec_add(vec_add(vec_add(vec_add_scaled(y, k1, 35.0 * h / 384.0), vec_scale(k3, 500.0 * h / 1113.0)), vec_scale(k4, 125.0 * h / 192.0)), vec_scale(k5, -2187.0 * h / 6784.0)), vec_scale(k6, 11.0 * h / 84.0));

        // 4th order solution (for error)
        let k7 = f(t + h, y_next);
        let y4_sol = vec_add(vec_add(vec_add(vec_add(vec_add(vec_add_scaled(y, k1, 5179.0 * h / 57600.0), vec_scale(k3, 7571.0 * h / 16695.0)), vec_scale(k4, 393.0 * h / 640.0)), vec_scale(k5, -92097.0 * h / 339200.0)), vec_scale(k6, 187.0 * h / 2100.0)), vec_scale(k7, h / 40.0));

        // Error estimate (vector norm)
        let err_vec = vec_sub(y_next, y4_sol);
        let mut err = vec_norm(err_vec);
        if err < 0.000000000000001 {
            err = 0.000000000000001;
        }

        if err <= tol {
            t = t + h;
            y = y_next;
            results.push({ t: t, y: y });
        }

        let mut factor = safety * pow(tol / err, 0.2);
        if factor > max_factor {
            factor = max_factor;
        }
        if factor < min_factor {
            factor = min_factor;
        }
        h = h * factor;
        if h > h_max {
            h = h_max;
        }

        step_count = step_count + 1;
    }

    results
}