Finalize ODE codegen

Author: serenity4

src/DAECompiler.jl

@@ -10,6 +10,7 @@ module DAECompiler
     using Core.IR
     using SciMLBase
     using AutoHashEquals
+    using LinearAlgebra: LinearAlgebra
 
     include("utils.jl")
     include("intrinsics.jl")

src/problem_interface.jl

@@ -66,9 +66,9 @@ function ODECProblem(f, tspan::Tuple{Real, Real} = (0., 1.); guesses = nothing,
 end
 
 function DiffEqBase.get_concrete_problem(prob::ODECProblem, isadaptive; kwargs...)
-    odef = factory(Val(Settings(mode=prob.init === nothing ? ODE : ODENoInit)), prob.f)
+    (odef, n) = factory(Val(Settings(mode=prob.init === nothing ? ODE : ODENoInit)), prob.f)
 
-    u0 = zeros(length(prob.init))
+    u0 = zeros(n)
 
     if prob.init !== nothing
         for (which, val) in prob.init

src/transform/codegen/ode_factory.jl

@@ -5,22 +5,35 @@ Given an IR value `arg` that corresponds to `u` in SciML's ODE ABI, split it int
 the DAECompiler internal ABI.
 """
 function sciml_ode_split_u!(compact, line, arg, numstates)
-    nassgn = numstates[AssignedDiff]
     ntotalstates = numstates[AssignedDiff] + numstates[UnassignedDiff] + numstates[Algebraic] + numstates[AlgebraicDerivative]
 
-    u_mm = @insert_node_here compact line view(arg, 1:nassgn)::VectorViewType
-    u_unassgn = @insert_node_here compact line view(arg, (nassgn+1):(nassgn+numstates[UnassignedDiff]))::VectorViewType
-    alg = @insert_node_here compact line view(arg, (nassgn+numstates[UnassignedDiff]+1):ntotalstates)::VectorViewType
+    u_mm = @insert_node_here compact line view(arg,
+        1:numstates[AssignedDiff])::VectorViewType
+    u_unassgn = @insert_node_here compact line view(arg,
+        (numstates[AssignedDiff] + 1):(numstates[AssignedDiff] + numstates[UnassignedDiff]))::VectorViewType
+    alg = @insert_node_here compact line view(arg,
+        (numstates[AssignedDiff] + numstates[UnassignedDiff] + 1):(numstates[AssignedDiff] + numstates[UnassignedDiff] + numstates[Algebraic]))::VectorViewType
+    alg_derv = @insert_node_here compact line view(arg,
+        (numstates[AssignedDiff] + numstates[UnassignedDiff] + numstates[Algebraic] + 1):ntotalstates)::VectorViewType
+
+    return (u_mm, u_unassgn, alg, alg_derv)
+end
 
-    return (u_mm, u_unassgn, alg)
+function generate_ode_mass_matrix(nd, na)
+    n = nd + na
+    mass_matrix = zeros(Float64, n, n)
+    for i in 1:nd
+        mass_matrix[i, i] = 1.0
+    end
+    return mass_matrix
 end
 
-function make_odefunction(f)
-    ODEFunction(f)
+function make_odefunction(f, mass_matrix = LinearAlgebra.I, initf = nothing)
+    ODEFunction(f; mass_matrix, initialization_data = (initf === nothing ? nothing : initialization_data_ode(initf)))
 end
 
-function make_odefunction(f, initf)
-    ODEFunction(f; initialization_data = SciMLBase.OverrideInitData(NonlinearProblem((args...)->nothing, nothing, nothing), nothing, initf, nothing, nothing))
+function initialization_data_ode(initf)
+    return SciMLBase.OverrideInitData(NonlinearProblem((args...)->nothing, nothing, nothing), nothing, initf, nothing, nothing)
 end
 
 """
@@ -39,8 +52,8 @@ end
 ```
 
 """
-function ode_factory_gen(result::DAEIPOResult, ci::CodeInstance, key::TornCacheKey, world::UInt, init_key::Union{TornCacheKey, Nothing})
-    @ccall jl_safe_printf("$key\n"::Cstring)::Cvoid
+function ode_factory_gen(state::TransformationState, ci::CodeInstance, key::TornCacheKey, world::UInt, init_key::Union{TornCacheKey, Nothing})
+    result = state.result
     torn_ci = find_matching_ci(ci->isa(ci.owner, TornIRSpec) && ci.owner.key == key, ci.def, world)
     torn_ir = torn_ci.inferred
 
@@ -65,7 +78,7 @@ function ode_factory_gen(result::DAEIPOResult, ci::CodeInstance, key::TornCacheK
         sicm = ()
     end
 
-    odef_ci = rhs_finish!(result, ci, key, world, 1)
+    odef_ci = rhs_finish!(state, ci, key, world, 1)
 
     # Create a small opaque closure to adapt from SciML ABI to our own internal
     # ABI
@@ -76,7 +89,6 @@ function ode_factory_gen(result::DAEIPOResult, ci::CodeInstance, key::TornCacheK
     for var = 1:length(result.var_to_diff)
         kind = classify_var(result.var_to_diff, key, var)
         kind == nothing && continue
-        @ccall jl_safe_printf("$kind\n"::Cstring)::Cvoid
         numstates[kind] += 1
         (kind != AlgebraicDerivative) && push!(all_states, var)
     end
@@ -101,17 +113,17 @@ function ode_factory_gen(result::DAEIPOResult, ci::CodeInstance, key::TornCacheK
 
     # out_du_mm, out_eq, in_u_mm, in_u_unassgn, in_alg
     nassgn = numstates[AssignedDiff]
-    ntotalstates = numstates[AssignedDiff] + numstates[UnassignedDiff] + numstates[Algebraic]
+    ntotalstates = numstates[AssignedDiff] + numstates[UnassignedDiff] + numstates[Algebraic] + numstates[AlgebraicDerivative]
     out_du_mm = @insert_node_here oc_compact line view(du, 1:nassgn)::VectorViewType
     out_eq = @insert_node_here oc_compact line view(du, (nassgn+1):ntotalstates)::VectorViewType
 
-    (in_u_mm, in_u_unassgn, in_alg) = sciml_ode_split_u!(oc_compact, line, u, numstates)
+    (in_u_mm, in_u_unassgn, in_alg, in_alg_derv) = sciml_ode_split_u!(oc_compact, line, u, numstates)
 
     # Call DAECompiler-generated RHS with internal ABI
     oc_sicm = @insert_node_here oc_compact line getfield(self, 1)::Tuple
 
     # N.B: The ordering of arguments should match the ordering in the StateKind enum
-    @insert_node_here oc_compact line (:invoke)(odef_ci, oc_sicm, (), in_u_mm, in_u_unassgn, in_alg, out_du_mm, out_eq, t)::Nothing
+    @insert_node_here oc_compact line (:invoke)(odef_ci, oc_sicm, (), in_u_mm, in_u_unassgn, in_alg, in_alg_derv, out_du_mm, out_eq, t)::Nothing
 
     # Return
     @insert_node_here oc_compact line (return)::Union{}
@@ -129,14 +141,14 @@ function ode_factory_gen(result::DAEIPOResult, ci::CodeInstance, key::TornCacheK
 
     new_oc = @insert_node_here compact line (:new_opaque_closure)(argt, Union{}, Nothing, true, oc_source_method, sicm)::Core.OpaqueClosure true
 
-    if init_key !== nothing
-        initf = init_uncompress_gen!(compact, result, ci, init_key, key, world)
-        odef = @insert_node_here compact line make_odefunction(new_oc, initf)::ODEFunction true
-    else
-        odef = @insert_node_here compact line make_odefunction(new_oc)::ODEFunction true
-    end
+    nd = numstates[AssignedDiff] + numstates[UnassignedDiff]
+    na = numstates[Algebraic] + numstates[AlgebraicDerivative]
+    mass_matrix = na == 0 ? GlobalRef(LinearAlgebra, :I) : @insert_node_here compact line generate_ode_mass_matrix(nd, na)::Matrix{Float64}
+    initf = init_key !== nothing ? init_uncompress_gen!(compact, result, ci, init_key, key, world) : nothing
+    odef = @insert_node_here compact line make_odefunction(new_oc, mass_matrix, initf)::ODEFunction true
 
-    @insert_node_here compact line (return odef)::Core.OpaqueClosure true
+    odef_and_n = @insert_node_here compact line tuple(odef, nd + na)::Tuple true
+    @insert_node_here compact line (return odef_and_n)::Core.OpaqueClosure true
 
     ir_factory = Compiler.finish(compact)
 

test/basic.jl

@@ -20,6 +20,8 @@ end
 oneeq!()
 sol = solve(DAECProblem(oneeq!, (1,) .=> 1.), IDA())
 @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol.u[:, 1], exp.(sol.t)))
+sol = solve(ODECProblem(oneeq!, (1,) .=> 1.), Rodas5(autodiff=false))
+@test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol.u[:, 1], exp.(sol.t)))
 
 #= Initial Condition =#
 @noinline function oneeq_ic!()
@@ -33,6 +35,8 @@ oneeq_ic!()
 # TODO: Sundials is broken and doesn't respect the custom initialization (https://github.com/SciML/Sundials.jl/issues/469)
 sol = solve(DAECProblem(oneeq_ic!), DFBDF(autodiff=false))
 @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol.u[:, 1], exp.(sol.t)))
+sol = solve(ODECProblem(oneeq_ic!), Rodas5(autodiff=false))
+@test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol.u[:, 1], exp.(sol.t)))
 
 #= Pantelides =#
 function pantelides()
@@ -45,6 +49,8 @@ end
 pantelides()
 sol = solve(DAECProblem(pantelides, (1,) .=> 0.), DFBDF(autodiff=false))
 @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol.u[:, 1], sol.t))
+sol = solve(ODECProblem(pantelides, (1,) .=> 0.), Rodas5(autodiff=false))
+@test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol.u[:, 1], sol.t))
 
 #= Structural Singularity Removal =#
 function ssrm()
@@ -58,6 +64,8 @@ end
 ssrm()
 sol = solve(DAECProblem(ssrm, (1,) .=> 1.), DFBDF(autodiff=false))
 @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol.u[:, 1], exp.(0.5sol.t)))
+sol = solve(ODECProblem(ssrm, (1,) .=> 1.), Rodas5(autodiff=false))
+@test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol.u[:, 1], exp.(0.5sol.t)))
 
 #= Pantelides from init =#
 function pantelides_from_init()
@@ -66,4 +74,4 @@ function pantelides_from_init()
     initial!(x -ᵢ ddt(ddt(y)))
 end
 
-end
+end

test/ipo.jl

@@ -5,6 +5,7 @@ using DAECompiler
 using DAECompiler.Intrinsics
 using Sundials
 using SciMLBase
+using OrdinaryDiffEq
 
 #= Basic IPO: We need to read the incidence of the contained `-` =#
 @noinline function onecall!()
@@ -15,6 +16,8 @@ end
 onecall!()
 sol = solve(DAECProblem(onecall!, (1,) .=> 1.), IDA())
 @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[1, :], exp.(sol.t)))
+sol = solve(ODECProblem(onecall!, (1,) .=> 1.), Rodas5(autodiff=false))
+@test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[1, :], exp.(sol.t)))
 
 #= + Contained Equations =#
 function twocall!()
@@ -26,5 +29,8 @@ twocall!()
 sol = solve(DAECProblem(twocall!, (1, 2) .=> 1.), IDA())
 @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[1, :], exp.(sol.t)))
 @test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[2, :], exp.(sol.t)))
+sol = solve(ODECProblem(twocall!, (1, 2) .=> 1.), Rodas5(autodiff=false))
+@test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[1, :], exp.(sol.t)))
+@test all(map((x,y)->isapprox(x[], y, atol=1e-2), sol[2, :], exp.(sol.t)))
 
 end