odetests.jl

load("ode.jl")
#load("winston.jl")

msdsys(t,y,m,b,k,Fext,xoff) = [0 1; -k./m b./m]*y + [0 0; sin(t.^2)./m k./m]*[Fext xoff]'
msdsys(t,y) = msdsys(t, y, 5, -0.5, 50, 0, 0)

function msd_analytic(t, m, b, k, x0)
    wn = sqrt(k./m)
    zeta = -b./m./2./wn
    wd = wn.*sqrt(1-zeta.^2)
    phi = atan(zeta./sqrt(1-zeta.^2))
    (t, x0.*exp(-zeta.*wn.*t)./sqrt(1-zeta.^2).*cos(wd.*t-phi))
end
msd_analytic(t, y) = msd_analytic(t, 5, -0.5, 50, y[1])

twinmsdsys(t,y,m,b1,k1,b2,k2,Fext) = [b1 -k1 b2-b1  k1-k2
                                      1   0      0      0
                                      0   0     b2    -k2
                                      0   0      1      0]*y + [1./m; 0; 0; 0]*Fext
                                      
                                      twinmsdsys(t, y) = twinmsdsys(t, y, 1, -0.5, 50, -0.005, 5000, 0)
                                      
twinmsd_analytic(t,m,b1,k1,b2,k2,x10,x20) = (t, msd_analytic(t, m, b1, k1, x10)[2] + 
                                                msd_analytic(t, m, b2, k2, x20)[2])
twinmsd_analytic(t,y) = twinmsd_analytic(t,1, -0.5, 50, -0.005, 5000, y[2], y[4])

type ODEResult{T1<:Real, T2<:Real, T3<:Real}
    solver::Function
    solver_fixed::Bool
    y0::Vector{T1}
    t::Vector{T2}
    y::Array{T3}
end

ODEResult(solver, fixed, y0, solution) = ODEResult(solver, fixed, y0, solution[1], solution[2])

function plot_ode_err(r::ODEResult, afcn::Function, a_state::Integer)
    p = FramedPlot()
    h_ = diff(r.t)
    h = r.solver_fixed ? string(h_[1]) : "varstep"
    y = squeeze(r.y[:,a_state])-afcn(r.t, r.y0)[2]
    yy = sprintf("%2e", max(y))
    setattr(p, "title", "$afcn: $(r.solver)@$h s\nmaxerr $yy")
    c = Curve(r.t, y)
    add(p, c)
    file(p, "err/$(afcn)_$(r.solver)_$(h).pdf")
end

function ode_runtest(odefun, y0, tmin, tmax, hs)
    variable_solvers = [ode45_dp
                        ode45_fb
                        ode45_ck
                        ode23
                        ]
    fixed_solvers = [ode4
                     ode4ms
                     ode4s_s
                     ode4s_kr
                     ]
    results = Array(ODEResult, length(variable_solvers) + length(fixed_solvers)*length(hs))
    i = 1
    for solver in variable_solvers
        print("$odefun $solver: ")
        @time results[i] = ODEResult(solver, false, y0, solver(odefun, [tmin; tmax], y0))
        i += 1
    end
    for solver in fixed_solvers, h in hs
        print("$odefun $solver@$h s: ")
        @time results[i] = ODEResult(solver, true, y0, solver(odefun, [tmin:h:tmax], y0))
        i += 1
    end
    results
end

function ode_test(odefun::Function, afcn::Function, a_state::Integer, tmin::Real, tmax::Real, hs, y0)
    results = ode_runtest(odefun, y0, tmin, tmax, hs)
    for r in results
        try
            plot_ode_err(r, afcn, a_state)
        end
    end
end

function ode_testall(tmax)
    ode_test(twinmsdsys, twinmsd_analytic, 2, 0, tmax, logspace(0, -3, 4), [0; 1; 0; 0.1])
    ode_test(msdsys, msd_analytic, 1, 0, tmax, logspace(0, -3, 4), [1; 0])
end
ode_testall() = ode_testall(10)