viva_math

-module(bench_ode).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/bench_ode.gleam").
-export([main/0]).

-if(?OTP_RELEASE >= 27).
-define(MODULEDOC(Str), -moduledoc(Str)).
-define(DOC(Str), -doc(Str)).
-else.
-define(MODULEDOC(Str), -compile([])).
-define(DOC(Str), -compile([])).
-endif.

?MODULEDOC(
    " Benchmarks for ODE solvers.\n"
    "\n"
    " Compares accuracy/cost of Euler, RK2, RK4, RKF45 and DOP54 on the\n"
    " canonical scaling problem `dx/dt = x` with `x(0) = 1`, target `x(1) = e`.\n"
).

-file("src/bench_ode.gleam", 48).
-spec bench_traj(
    binary(),
    fun((fun((float(), float()) -> float()), float(), float(), float()) -> float()),
    fun((float(), float()) -> float()),
    float(),
    integer(),
    float()
) -> nil.
bench_traj(Label, Method, F, Dt, Steps, Expected) ->
    T0 = erlang:monotonic_time(),
    Traj = viva_math@ode:integrate(Method, F, +0.0, 1.0, Dt, Steps),
    T1 = erlang:monotonic_time(),
    Last_pair@1 = case gleam@list:last(Traj) of
        {ok, Last_pair} -> Last_pair;
        _assert_fail ->
            erlang:error(#{gleam_error => let_assert,
                        message => <<"Pattern match failed, no pattern matched the value."/utf8>>,
                        file => <<?FILEPATH/utf8>>,
                        module => <<"bench_ode"/utf8>>,
                        function => <<"bench_traj"/utf8>>,
                        line => 59,
                        value => _assert_fail,
                        start => 1817,
                        'end' => 1859,
                        pattern_start => 1828,
                        pattern_end => 1841})
    end,
    Err = gleam@float:absolute_value(erlang:element(2, Last_pair@1) - Expected),
    gleam_stdlib:println(
        <<<<<<<<<<<<"  "/utf8, Label/binary>>/binary, " err = "/utf8>>/binary,
                        (erlang:float_to_binary(Err))/binary>>/binary,
                    "  time = "/utf8>>/binary,
                (erlang:integer_to_binary(T1 - T0))/binary>>/binary,
            "ns"/utf8>>
    ).

-file("src/bench_ode.gleam", 43).
-spec step_report(binary(), float(), float()) -> nil.
step_report(Label, Got, Expected) ->
    Err = gleam@float:absolute_value(Got - Expected),
    gleam_stdlib:println(
        <<<<<<"  "/utf8, Label/binary>>/binary, " err = "/utf8>>/binary,
            (erlang:float_to_binary(Err))/binary>>
    ).

-file("src/bench_ode.gleam", 12).
-spec main() -> nil.
main() ->
    gleam_stdlib:println(<<"\n=== viva_math ODE solver benchmarks ===\n"/utf8>>),
    F = fun(_, X) -> X end,
    gleam_stdlib:println(
        <<"Single step accuracy (dt = 0.1, expected = e^0.1):"/utf8>>
    ),
    Expected_short = 1.1051709180756477,
    step_report(
        <<"euler           "/utf8>>,
        viva_math@ode:euler(F, +0.0, 1.0, 0.1),
        Expected_short
    ),
    step_report(
        <<"rk2_midpoint    "/utf8>>,
        viva_math@ode:rk2_midpoint(F, +0.0, 1.0, 0.1),
        Expected_short
    ),
    step_report(
        <<"rk2_heun        "/utf8>>,
        viva_math@ode:rk2_heun(F, +0.0, 1.0, 0.1),
        Expected_short
    ),
    step_report(
        <<"rk4             "/utf8>>,
        viva_math@ode:rk4(F, +0.0, 1.0, 0.1),
        Expected_short
    ),
    {Rkf45_x, _} = viva_math@ode:rkf45(F, +0.0, 1.0, 0.1),
    step_report(<<"rkf45           "/utf8>>, Rkf45_x, Expected_short),
    {Dop54_x, _} = viva_math@ode:dop54(F, +0.0, 1.0, 0.1),
    step_report(<<"dop54           "/utf8>>, Dop54_x, Expected_short),
    gleam_stdlib:println(
        <<"\nFull-trajectory integration to t = 1.0, dt = 0.01:"/utf8>>
    ),
    Dt = 0.01,
    Steps = 100,
    bench_traj(
        <<"euler       "/utf8>>,
        fun viva_math@ode:euler/4,
        F,
        Dt,
        Steps,
        2.718281828459045
    ),
    bench_traj(
        <<"rk4         "/utf8>>,
        fun viva_math@ode:rk4/4,
        F,
        Dt,
        Steps,
        2.718281828459045
    ).