Added single_stage_solver

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

@@ -1,66 +1,90 @@
 
-function nonlinear_stage_problem(f,ode0,t,x0,coeffs)
-    workspace = nothing
-    nonlinear_problem(x0,residual(ode0),jacobian(ode0),workspace) do p
-        x = solution(p)
-        r = residual(p)
-        j = jacobian(p)
-        ode = update(ode0,coefficients=coeffs,residual=r,jacobian=j,solution=(t,f(x)...))
-        r = residual(ode)
-        j = jacobian(ode)
-        p = update(p,residual=r,jacobian=j)
-    end
+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 backward_euler_update(workspace,ode)
-    (;s,x,dt,ex) = workspace
-    t, = solution(ode)
+function single_stage_solver_update(workspace,ode)
+    (;s,dt,scheme,p,x1) = workspace
+    t,u,v = solution(ode)
     x = solution(s)
-    coeffs = coefficients(ode)
-    p = nonlinear_stage_problem(ex,ode,t,x,coeffs)
-    p = update(p,solution=x)
+    x0 = (t,u,v)
+    p = update_nonlinear_stage_problem(x1,x0,p,dt)
     s = update(s,problem=p)
-    workspace = (;s,x,dt,ex)
+    workspace = (;s,dt,scheme,p,x1)
 end
 
-function backward_euler_step(workspace,ode,phase=:start)
-    (;s,x,dt,ex) = workspace
+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)
-    t += dt
-    u,v = ex(x)
-    u .= x
-    ode = update(ode,solution=(t,u,v))
+    x0 = solution(ode)
+    x0 = scheme(x0,x0,x,dt)
+    t, = x0
+    ode = update(ode,solution=x0)
     tend = last(interval(ode))
-    if t >= tend
+    if t > tend || t≈tend
         phase = :stop
     end
-    workspace = (;s,x,dt,ex)
-    workspace = backward_euler_update(workspace,ode)
+    workspace = (;s,dt,scheme,p,x1)
+    workspace = single_stage_solver_update(workspace,ode)
     workspace,ode,phase
 end
 
-function backward_euler(ode;
-        dt = (interval(ode)[end]-interval(ode)[1])/10,
-        solver = default_solver,
-    )
-    @assert uses_mutable_types(ode)
-    coeffs = (1.0,1/dt)
-    t,u,v = solution(ode)
-    t = interval(ode)[1]
-    x = copy(u)
-    function ex(x)
-        v .= (x .- u) ./ dt
-        (x,v)
-    end
-    p = nonlinear_stage_problem(ex,ode,t,x,coeffs)
-    s = solver(p)
-    workspace = (;s,x,dt,ex)
-    ode_solver(backward_euler_update,backward_euler_step,ode,workspace)
+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

@@ -4,7 +4,7 @@ import PartitionedSolvers as PS
 using Test
 import PartitionedArrays as PA
 
-function mock_ode(u)
+function mock_ode_1(u)
     r = zeros(1)
     j = PA.sparse_matrix([1],[1],[0.0],1,1)
     v = 0*u
@@ -18,26 +18,66 @@ function mock_ode(u)
         r = PS.residual(ode)
         j = PS.jacobian(ode)
         if r !== nothing
-            r .= 2 .* u2.^2 .+ v2 .- 4*t .+ 1
+            r .= v2 .- 32
             ode = PS.update(ode,residual = r)
         end
         if j !== nothing
-            j .= 4 .* u2 .* du .+ dv
+            j .= dv
             ode = PS.update(ode,jacobian = j)
         end
         ode
     end |> PS.update
 end
 
-u = [2.0]
-p = mock_ode(u)
-s = PS.backward_euler(p)
+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)
-    @test v[1] != 0
 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