diff --git a/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl b/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl index e87963f362..cdd7014cdd 100644 --- a/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl +++ b/lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl @@ -396,8 +396,9 @@ function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::OD M = integrator.f.mass_matrix M isa UniformScaling && return update_coefficients!(M, u, p, t) - algebraic_vars = [all(iszero, x) for x in eachcol(M)] - algebraic_eqs = [all(iszero, x) for x in eachrow(M)] + algebraic_vars = vec(all(iszero, M, dims = 1)) + algebraic_eqs = vec(all(iszero, M, dims = 2)) + (iszero(algebraic_vars) || iszero(algebraic_eqs)) && return tmp = get_tmp_cache(integrator)[1] @@ -456,7 +457,7 @@ function _initialize_dae!(integrator::OrdinaryDiffEqCore.ODEIntegrator, prob::OD nlsolve = default_nlsolve(alg.nlsolve, isinplace, u, nlprob, isAD) nlsol = solve(nlprob, nlsolve; abstol = alg.abstol, reltol = integrator.opts.reltol) - alg_u .= nlsol + alg_u .= nlsol.u recursivecopy!(integrator.uprev, integrator.u) if alg_extrapolates(integrator.alg) diff --git a/test/gpu/Project.toml b/test/gpu/Project.toml index e0e8db295f..fdf60e309f 100644 --- a/test/gpu/Project.toml +++ b/test/gpu/Project.toml @@ -6,6 +6,8 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" +OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8" +OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" [compat] diff --git a/test/gpu/simple_dae.jl b/test/gpu/simple_dae.jl new file mode 100644 index 0000000000..b8a73f0fb3 --- /dev/null +++ b/test/gpu/simple_dae.jl @@ -0,0 +1,65 @@ +using OrdinaryDiffEqRosenbrock +using OrdinaryDiffEqNonlinearSolve +using CUDA +using LinearAlgebra +using Adapt +using SparseArrays +using Test + +#= +du[1] = -u[1] +du[2] = -0.5*u[2] + 0 = u[1] + u[2] - u[3] + 0 = -u[1] + u[2] - u[4] +=# + +function dae!(du, u, p, t) + mul!(du, p, u) +end + +p = [-1 0 0 0 + 1 -0.5 0 0 + 1 1 -1 0 + -1 1 0 -1] + +# mass_matrix = [1 0 0 0 +# 0 1 0 0 +# 0 0 0 0 +# 0 0 0 0] +mass_matrix = Diagonal([1, 1, 0, 0]) +jac_prototype = sparse(map(x -> iszero(x) ? 0.0 : 1.0, p)) + +u0 = [1.0, 1.0, 0.5, 0.5] # force init +odef = ODEFunction(dae!, mass_matrix = mass_matrix, jac_prototype = jac_prototype) + +tspan = (0.0, 5.0) +prob = ODEProblem(odef, u0, tspan, p) +sol = solve(prob, Rodas5P()) + +# gpu version +mass_matrix_d = adapt(CuArray, mass_matrix) + +# TODO: jac_prototype fails +# jac_prototype_d = adapt(CuArray, jac_prototype) +# jac_prototype_d = CUDA.CUSPARSE.CuSparseMatrixCSR(jac_prototype) +jac_prototype_d = nothing + +u0_d = adapt(CuArray, u0) +p_d = adapt(CuArray, p) +odef_d = ODEFunction(dae!, mass_matrix = mass_matrix_d, jac_prototype = jac_prototype_d) +prob_d = ODEProblem(odef_d, u0_d, tspan, p_d) +sol_d = solve(prob_d, Rodas5P()) + +@testset "Test constraints in GPU sol" begin + for t in sol_d.t + u = Vector(sol_d(t)) + @test isapprox(u[1] + u[2], u[3]; atol = 1e-6) + @test isapprox(-u[1] + u[2], u[4]; atol = 1e-6) + end +end + +@testset "Compare GPU to CPU solution" begin + for t in tspan[begin]:0.1:tspan[end] + @test Vector(sol_d(t)) ≈ sol(t) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 09c0af2144..48f1ca0b17 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -205,6 +205,7 @@ end @time @safetestset "Linear Exponential GPU" include("gpu/linear_exp.jl") @time @safetestset "Reaction-Diffusion Stiff Solver GPU" include("gpu/reaction_diffusion_stiff.jl") @time @safetestset "Scalar indexing bug bypass" include("gpu/hermite_test.jl") + @time @safetestset "simple dae on GPU" include("gpu/simple_dae.jl") end if !is_APPVEYOR && GROUP == "QA"