Merge pull request #2960 from aml5600/fix/add-checks-false-initializa…

…tion kwarg to disable unit checking in `InitializationProblem`
SciML · Aug 18, 2024 · dc98d54 · dc98d54
2 parents 9938994 + 5e7f048
commit dc98d54
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 4 deletions.
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
@@ -782,6 +782,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         build_initializeprob = true,
         initialization_eqs = [],
         fully_determined = false,
+        check_units = true,
         kwargs...)
     eqs = equations(sys)
     dvs = unknowns(sys)
@@ -820,7 +821,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         end
         initializeprob = ModelingToolkit.InitializationProblem(
             sys, t, u0map, parammap; guesses, warn_initialize_determined,
-            initialization_eqs, eval_expression, eval_module, fully_determined)
+            initialization_eqs, eval_expression, eval_module, fully_determined, check_units)
         initializeprobmap = getu(initializeprob, unknowns(sys))
 
         zerovars = Dict(setdiff(unknowns(sys), keys(defaults(sys))) .=> 0.0)
@@ -1430,18 +1431,19 @@ function InitializationProblem{iip, specialize}(sys::AbstractODESystem,
         warn_initialize_determined = true,
         initialization_eqs = [],
         fully_determined = false,
+        check_units = true,
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `ODEProblem`")
     end
     if isempty(u0map) && get_initializesystem(sys) !== nothing
-        isys = get_initializesystem(sys; initialization_eqs)
+        isys = get_initializesystem(sys; initialization_eqs, check_units)
     elseif isempty(u0map) && get_initializesystem(sys) === nothing
         isys = structural_simplify(
-            generate_initializesystem(sys; initialization_eqs); fully_determined)
+            generate_initializesystem(sys; initialization_eqs, check_units); fully_determined)
     else
         isys = structural_simplify(
-            generate_initializesystem(sys; u0map, initialization_eqs); fully_determined)
+            generate_initializesystem(sys; u0map, initialization_eqs, check_units); fully_determined)
     end
 
     uninit = setdiff(unknowns(sys), [unknowns(isys); getfield.(observed(isys), :lhs)])

diff --git a/src/systems/nonlinear/initializesystem.jl b/src/systems/nonlinear/initializesystem.jl
@@ -10,6 +10,7 @@ function generate_initializesystem(sys::ODESystem;
         default_dd_value = 0.0,
         algebraic_only = false,
         initialization_eqs = [],
+        check_units = true,
         kwargs...)
     sts, eqs = unknowns(sys), equations(sys)
     idxs_diff = isdiffeq.(eqs)
@@ -102,6 +103,7 @@ function generate_initializesystem(sys::ODESystem;
         pars;
         defaults = merge(ModelingToolkit.defaults(sys), todict(u0), dd_guess),
         parameter_dependencies = parameter_dependencies(sys),
+        checks = check_units,
         name,
         kwargs...)
 

diff --git a/test/dq_units.jl b/test/dq_units.jl
@@ -185,3 +185,28 @@ end
 
 @named sys = ArrayParamTest(a = [1.0, 3.0]u"cm")
 @test ModelingToolkit.getdefault(sys.a) ≈ [0.01, 0.03]
+
+@testset "Initialization checks" begin
+    @mtkmodel PendulumUnits begin
+        @parameters begin
+            g, [unit = u"m/s^2"]
+            L, [unit = u"m"]
+        end
+        @variables begin
+            x(t), [unit = u"m"]
+            y(t), [state_priority = 10, unit = u"m"]
+            λ(t), [unit = u"s^-2"]
+        end
+        @equations begin
+            D(D(x)) ~ λ * x
+            D(D(y)) ~ λ * y - g
+            x^2 + y^2 ~ L^2
+        end
+    end
+    @mtkbuild pend = PendulumUnits()
+    u0 = [pend.x => 1.0, pend.y => 0.0]
+    p = [pend.g => 1.0, pend.L => 1.0]
+    guess = [pend.λ => 0.0]
+    @test prob = ODEProblem(
+        pend, u0, (0.0, 1.0), p; guesses = guess, check_units = false) isa Any
+end