single_stage_solver

PartitionedSolvers/src/PartitionedSolvers.jl

@@ -14,5 +14,6 @@ include("wrappers.jl")
 include("smoothers.jl")
 include("amg.jl")
 include("nonlinear_solvers.jl")
+include("ode_solvers.jl")
 
 end # module

PartitionedSolvers/src/interfaces.jl

@@ -328,7 +328,12 @@ function set(p::NonlinearProblem;kwargs...)
     else
         attrs = attributes(p)
     end
-    NonlinearProblem(st,x,b,A,p.workspace,attrs)
+    if hasproperty(data,:workspace)
+        ws = data.workspace
+    else
+        ws = workspace(p)
+    end
+    NonlinearProblem(st,x,b,A,ws,attrs)
 end
 
 function update(p::NonlinearProblem;kwargs...)

PartitionedSolvers/src/ode_solvers.jl

@@ -0,0 +1,90 @@
+
+function nonlinear_stage_problem_statement(p)
+    (;scheme,coeffs,ode,x1,x0,dt) = workspace(p)
+    x = solution(p)
+    r = residual(p)
+    j = jacobian(p)
+    x1 = scheme(x1,x0,x,dt)
+    ode = update(ode,residual=r,jacobian=j,solution=x1,coefficients=coeffs(dt))
+    r = residual(ode)
+    j = jacobian(ode)
+    p = update(p,residual=r,jacobian=j)
+end
+
+function nonlinear_stage_problem(scheme,coeffs,x1,x0,x,ode,dt)
+    workspace = (;scheme,coeffs,x1,x0,ode,dt)
+    nonlinear_problem(
+        nonlinear_stage_problem_statement,
+        x,residual(ode),jacobian(ode),workspace)
+end
+
+function update_nonlinear_stage_problem(x1,x0,p,dt)
+    (;scheme,coeffs,ode,dt) = workspace(p)
+    update(p,workspace=(;scheme,coeffs,x1,x0,ode,dt))
+end
+
+function single_stage_solver(scheme,coeffs,ode;
+        dt = (interval(ode)[end]-interval(ode)[1])/10,
+        solver = default_solver,
+    )
+    t,u,v = solution(ode)
+    t = interval(ode)[1]
+    ode = update(ode,solution=(t,u,v),coefficients=coeffs(dt))
+    x = copy(u)
+    x0 = (t,u,v)
+    x1 = map(copy,x0) # TODO Too much copies?
+    p = nonlinear_stage_problem(scheme,coeffs,x1,x0,x,ode,dt)
+    s = solver(p)
+    workspace = (;s,dt,scheme,p,x1)
+    ode_solver(single_stage_solver_update,single_stage_solver_step,ode,workspace)
+end
+
+function single_stage_solver_update(workspace,ode)
+    (;s,dt,scheme,p,x1) = workspace
+    t,u,v = solution(ode)
+    x = solution(s)
+    x0 = (t,u,v)
+    p = update_nonlinear_stage_problem(x1,x0,p,dt)
+    s = update(s,problem=p)
+    workspace = (;s,dt,scheme,p,x1)
+end
+
+function single_stage_solver_step(workspace,ode,phase=:start)
+    (;s,dt,scheme,p,x1) = workspace
+    t,u,v = solution(ode)
+    if phase === :start
+        t = first(interval(ode))
+        phase = :advance
+    end
+    s = solve(s)
+    x = solution(s)
+    x0 = solution(ode)
+    x0 = scheme(x0,x0,x,dt)
+    t, = x0
+    ode = update(ode,solution=x0)
+    tend = last(interval(ode))
+    if t > tend || t≈tend
+        phase = :stop
+    end
+    workspace = (;s,dt,scheme,p,x1)
+    workspace = single_stage_solver_update(workspace,ode)
+    workspace,ode,phase
+end
+
+function backward_euler_scheme!((t1,u1,v1),(t0,u0,v0),x,dt)
+    t1 = t0 + dt
+    v1 .= (x .- u0) ./ dt
+    u1 .= x
+    (t1,u1,v1)
+end
+
+function backward_euler_coefficients(dt)
+    (1.0,1.0/dt)
+end
+
+function backward_euler(ode;kwargs...)
+    coeffs = backward_euler_coefficients
+    scheme = backward_euler_scheme!
+    single_stage_solver(scheme,coeffs,ode;kwargs...)
+end
+

PartitionedSolvers/test/ode_solvers_tests.jl

@@ -0,0 +1,83 @@
+module ODESolversTests
+
+import PartitionedSolvers as PS
+using Test
+import PartitionedArrays as PA
+
+function mock_ode_1(u)
+    r = zeros(1)
+    j = PA.sparse_matrix([1],[1],[0.0],1,1)
+    v = 0*u
+    x = (0,u,v)
+    ts = (0,10)
+    coeffs = (1.,1.)
+    workspace = nothing
+    PS.ode_problem(x,r,j,ts,coeffs,workspace) do ode
+        (t,u2,v2) = PS.solution(ode)
+        du,dv = PS.coefficients(ode)
+        r = PS.residual(ode)
+        j = PS.jacobian(ode)
+        if r !== nothing
+            r .= v2 .- 32
+            ode = PS.update(ode,residual = r)
+        end
+        if j !== nothing
+            j .= dv
+            ode = PS.update(ode,jacobian = j)
+        end
+        ode
+    end |> PS.update
+end
+
+function mock_ode_2(u)
+    r = zeros(1)
+    j = PA.sparse_matrix([1],[1],[0.0],1,1)
+    v = 0*u
+    x = (0,u,v)
+    ts = (1,10)
+    coeffs = (1.,1.)
+    workspace = nothing
+    PS.ode_problem(x,r,j,ts,coeffs,workspace) do ode
+        (t,u2,v2) = PS.solution(ode)
+        du,dv = PS.coefficients(ode)
+        r = PS.residual(ode)
+        j = PS.jacobian(ode)
+        if r !== nothing
+            r .=  v2 .+ (u2 ./ t) .- 3*t
+            #r .=  v2 .+ u2
+            ode = PS.update(ode,residual = r)
+        end
+        if j !== nothing
+            j .=  (du ./ t) .+ dv
+            #j .=  du .+ dv
+            ode = PS.update(ode,jacobian = j)
+        end
+        ode
+    end |> PS.update
+end
+
+u = [0.0]
+p = mock_ode_1(u)
+solver = p -> PS.newton_raphson(p;verbose=false)
+s = PS.backward_euler(p;dt=1,solver)
+for s in PS.history(s)
+    t,u,v = PS.solution(s)
+    @show PS.solution(s)
+end
+t,u,v = PS.solution(s)
+@test u == [320.]
+@test v == [32.]
+
+u = [1.0]
+p = mock_ode_2(u)
+solver = p -> PS.newton_raphson(p;verbose=false)
+s = PS.backward_euler(p;dt=0.001,solver)
+s = PS.solve(s)
+t,u,v = PS.solution(s)
+tol = 1.0e-4
+@test abs(u[1] - 100)/100 < tol
+@test abs(v[1] - 20)/20 < tol
+@show PS.solution(s)
+
+
+end # module

PartitionedSolvers/test/runtests.jl

@@ -8,6 +8,7 @@ using Test
     @testset "interfaces" begin include("interfaces_tests.jl") end
     @testset "wrappers" begin include("wrappers_tests.jl") end
     @testset "nonlinear_solvers" begin include("nonlinear_solvers_tests.jl") end
+    @testset "ode_solvers" begin include("ode_solvers_tests.jl") end
     @testset "smoothers" begin include("smoothers_tests.jl") end
     @testset "amg" begin include("amg_tests.jl") end
 end