shape-runtime 0.3.0

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

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

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

function vec_add_scaled(a: Array<number>, b: Array<number>, s: number) -> Array<number> {
    let scaled: Array<number> = vec_scale(b, s);
    vec_add(a, scaled)
}

/// Euler integrator for scalar ODE
pub fn euler(f: (number, number) => number, y0: number, t_start: number, t_end: number, dt: number) {
    let steps: int = floor((t_end - t_start) / dt) as int;
    let mut t: number = t_start;
    let mut y: number = y0;
    let mut results: Array<{t: number, y: number}> = [];

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

    results
}

/// RK4 integrator for scalar ODE
pub fn rk4(f: (number, number) => number, y0: number, t_start: number, t_end: number, dt: number) {
    let steps: int = floor((t_end - t_start) / dt) as int;
    let mut t: number = t_start;
    let mut y: number = y0;
    let mut results: Array<{t: number, y: number}> = [];

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

        let k1: number = f(t, y);
        let k2: number = f(t + dt / 2.0, y + k1 * dt / 2.0);
        let k3: number = f(t + dt / 2.0, y + k2 * dt / 2.0);
        let k4: number = 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: (number, Array<number>) => Array<number>, y0_vec: Array<number>, t_start: number, t_end: number, dt: number) {
    let steps: int = floor((t_end - t_start) / dt) as int;
    let mut t: number = t_start;
    let mut y: Array<number> = y0_vec;
    // NOTE: results seeded with first record rather than `[]` because v0.3
    // inference cannot match Array<{t:number, y:Array<number>}> through push.
    let mut results = [{ t: t, y: y }];
    let mut first: bool = true;

    for i in range(0, steps + 1) {
        if !first {
            results.push({ t: t, y: y });
        }
        first = false;
        let dy: Array<number> = f(t, y);
        y = vec_add_scaled(y, dy, dt);
        t = t + dt;
    }

    results
}

/// RK4 integrator for vector systems
pub fn rk4_system(f: (number, Array<number>) => Array<number>, y0_vec: Array<number>, t_start: number, t_end: number, dt: number) {
    let steps: int = floor((t_end - t_start) / dt) as int;
    let mut t: number = t_start;
    let mut y: Array<number> = y0_vec;
    // NOTE: results seeded per v0.3 inference gap (see euler_system).
    let mut results = [{ t: t, y: y }];
    let mut first: bool = true;

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

        let k1: Array<number> = f(t, y);
        let ys2: Array<number> = vec_add_scaled(y, k1, dt / 2.0);
        let k2: Array<number> = f(t + dt / 2.0, ys2);
        let ys3: Array<number> = vec_add_scaled(y, k2, dt / 2.0);
        let k3: Array<number> = f(t + dt / 2.0, ys3);
        let ys4: Array<number> = vec_add_scaled(y, k3, dt);
        let k4: Array<number> = f(t + dt, ys4);

        let sk2: Array<number> = vec_scale(k2, 2.0);
        let sk3: Array<number> = vec_scale(k3, 2.0);
        let half1: Array<number> = vec_add(k1, sk2);
        let half2: Array<number> = vec_add(sk3, k4);
        let step: Array<number> = vec_add(half1, half2);
        y = vec_add_scaled(y, step, dt / 6.0);
        t = t + dt;
    }

    results
}

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

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

function vec_norm(a: Array<number>) -> number {
    let mut s: number = 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: (number, number) => number, y0: number, t_start: number, t_end: number, tol: number = 0.000001, dt_init: number = 0.0, dt_min: number = 0.000000000001, dt_max: number = 0.0) {
    let safety: number = 0.9;
    let max_factor: number = 5.0;
    let min_factor: number = 0.2;

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

    let mut t: number = t_start;
    let mut y: number = y0;
    let mut results: Array<{t: number, y: number}> = [];
    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: number = f(t, y);
        let k2: number = f(t + h / 5.0, y + h * k1 / 5.0);
        let k3: number = f(t + 3.0 * h / 10.0, y + h * (3.0 * k1 / 40.0 + 9.0 * k2 / 40.0));
        let k4: number = 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: number = 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: number = 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: number = 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: number = f(t + h, y5);
        let y4: number = 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: number = 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 pow_factor: number = pow(tol / err, 0.2);
        let mut factor: number = safety * pow_factor;
        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: (number, Array<number>) => Array<number>, y0_vec: Array<number>, t_start: number, t_end: number, tol: number = 0.000001, dt_init: number = 0.0, dt_min: number = 0.000000000001, dt_max: number = 0.0) {
    let safety: number = 0.9;
    let max_factor: number = 5.0;
    let min_factor: number = 0.2;

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

    let mut t: number = t_start;
    let mut y: Array<number> = y0_vec;
    // NOTE: results seeded per v0.3 inference gap (see euler_system).
    let mut results = [{ 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: Array<number> = f(t, y);
        let yk2: Array<number> = vec_add_scaled(y, k1, h / 5.0);
        let k2: Array<number> = f(t + h / 5.0, yk2);

        let y3a: Array<number> = vec_add_scaled(y, k1, 3.0 * h / 40.0);
        let y3b: Array<number> = vec_scale(k2, 9.0 * h / 40.0);
        let y3: Array<number> = vec_add(y3a, y3b);
        let k3: Array<number> = f(t + 3.0 * h / 10.0, y3);

        let y4a: Array<number> = vec_add_scaled(y, k1, 44.0 * h / 45.0);
        let y4b: Array<number> = vec_scale(k2, -56.0 * h / 15.0);
        let y4c: Array<number> = vec_scale(k3, 32.0 * h / 9.0);
        let y4_stage: Array<number> = vec_add(vec_add(y4a, y4b), y4c);
        let k4: Array<number> = f(t + 4.0 * h / 5.0, y4_stage);

        let y5a: Array<number> = vec_add_scaled(y, k1, 19372.0 * h / 6561.0);
        let y5b: Array<number> = vec_scale(k2, -25360.0 * h / 2187.0);
        let y5c: Array<number> = vec_scale(k3, 64448.0 * h / 6561.0);
        let y5d: Array<number> = vec_scale(k4, -212.0 * h / 729.0);
        let y5_stage: Array<number> = vec_add(vec_add(vec_add(y5a, y5b), y5c), y5d);
        let k5: Array<number> = f(t + 8.0 * h / 9.0, y5_stage);

        let y6a: Array<number> = vec_add_scaled(y, k1, 9017.0 * h / 3168.0);
        let y6b: Array<number> = vec_scale(k2, -355.0 * h / 33.0);
        let y6c: Array<number> = vec_scale(k3, 46732.0 * h / 5247.0);
        let y6d: Array<number> = vec_scale(k4, 49.0 * h / 176.0);
        let y6e: Array<number> = vec_scale(k5, -5103.0 * h / 18656.0);
        let y6: Array<number> = vec_add(vec_add(vec_add(vec_add(y6a, y6b), y6c), y6d), y6e);
        let k6: Array<number> = f(t + h, y6);

        // 5th order solution
        let yn_a: Array<number> = vec_add_scaled(y, k1, 35.0 * h / 384.0);
        let yn_b: Array<number> = vec_scale(k3, 500.0 * h / 1113.0);
        let yn_c: Array<number> = vec_scale(k4, 125.0 * h / 192.0);
        let yn_d: Array<number> = vec_scale(k5, -2187.0 * h / 6784.0);
        let yn_e: Array<number> = vec_scale(k6, 11.0 * h / 84.0);
        let y_next: Array<number> = vec_add(vec_add(vec_add(vec_add(yn_a, yn_b), yn_c), yn_d), yn_e);

        // 4th order solution (for error)
        let k7: Array<number> = f(t + h, y_next);
        let y4s_a: Array<number> = vec_add_scaled(y, k1, 5179.0 * h / 57600.0);
        let y4s_b: Array<number> = vec_scale(k3, 7571.0 * h / 16695.0);
        let y4s_c: Array<number> = vec_scale(k4, 393.0 * h / 640.0);
        let y4s_d: Array<number> = vec_scale(k5, -92097.0 * h / 339200.0);
        let y4s_e: Array<number> = vec_scale(k6, 187.0 * h / 2100.0);
        let y4s_f: Array<number> = vec_scale(k7, h / 40.0);
        let y4_sol: Array<number> = vec_add(vec_add(vec_add(vec_add(vec_add(y4s_a, y4s_b), y4s_c), y4s_d), y4s_e), y4s_f);

        // Error estimate (vector norm)
        let err_vec: Array<number> = vec_sub(y_next, y4_sol);
        let mut err: number = 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 pow_factor: number = pow(tol / err, 0.2);
        let mut factor: number = safety * pow_factor;
        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
}