Merge pull request #3448 from vyudu/fix_float

fix: propagate `tofloat`, `use_union` to `better_varmap_to_vars`

Author: ChrisRackauckas

src/systems/diffeqs/abstractodesystem.jl

@@ -1504,5 +1504,6 @@ function InitializationProblem{iip, specialize}(sys::AbstractSystem,
     else
         NonlinearLeastSquaresProblem
     end
-    TProb(isys, u0map, parammap; kwargs..., build_initializeprob = false, is_initializeprob = true)
+    TProb(isys, u0map, parammap; kwargs...,
+        build_initializeprob = false, is_initializeprob = true)
 end

src/systems/discrete_system/discrete_system.jl

@@ -302,7 +302,6 @@ function SciMLBase.DiscreteProblem(
         parammap = SciMLBase.NullParameters();
         eval_module = @__MODULE__,
         eval_expression = false,
-        use_union = false,
         kwargs...
 )
     if !iscomplete(sys)

src/systems/discrete_system/implicit_discrete_system.jl

@@ -321,7 +321,6 @@ function SciMLBase.ImplicitDiscreteProblem(
         parammap = SciMLBase.NullParameters();
         eval_module = @__MODULE__,
         eval_expression = false,
-        use_union = false,
         kwargs...
 )
     if !iscomplete(sys)

src/systems/jumps/jumpsystem.jl

@@ -383,7 +383,6 @@ end
 ```julia
 DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan,
                            parammap = DiffEqBase.NullParameters;
-                           use_union = true,
                            kwargs...)
 ```
 
@@ -403,7 +402,6 @@ dprob = DiscreteProblem(complete(js), u₀map, tspan, parammap)
 """
 function DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothing},
         parammap = DiffEqBase.NullParameters();
-        use_union = true,
         eval_expression = false,
         eval_module = @__MODULE__,
         kwargs...)
@@ -416,7 +414,7 @@ function DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan::Union{Tuple,
     end
 
     _f, u0, p = process_SciMLProblem(EmptySciMLFunction, sys, u0map, parammap;
-        t = tspan === nothing ? nothing : tspan[1], use_union, tofloat = false, check_length = false, build_initializeprob = false)
+        t = tspan === nothing ? nothing : tspan[1], tofloat = false, check_length = false, build_initializeprob = false)
     f = DiffEqBase.DISCRETE_INPLACE_DEFAULT
 
     observedfun = ObservedFunctionCache(
@@ -451,14 +449,13 @@ struct DiscreteProblemExpr{iip} end
 
 function DiscreteProblemExpr{iip}(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothing},
         parammap = DiffEqBase.NullParameters();
-        use_union = true,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed `JumpSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `DiscreteProblemExpr`")
     end
 
     _, u0, p = process_SciMLProblem(EmptySciMLFunction, sys, u0map, parammap;
-        t = tspan === nothing ? nothing : tspan[1], use_union, tofloat = false, check_length = false)
+        t = tspan === nothing ? nothing : tspan[1], tofloat = false, check_length = false)
     # identity function to make syms works
     quote
         f = DiffEqBase.DISCRETE_INPLACE_DEFAULT
@@ -475,7 +472,6 @@ end
 ```julia
 DiffEqBase.ODEProblem(sys::JumpSystem, u0map, tspan,
                            parammap = DiffEqBase.NullParameters;
-                           use_union = true,
                            kwargs...)
 ```
 
@@ -497,7 +493,6 @@ oprob = ODEProblem(complete(js), u₀map, tspan, parammap)
 """
 function DiffEqBase.ODEProblem(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothing},
         parammap = DiffEqBase.NullParameters();
-        use_union = false,
         eval_expression = false,
         eval_module = @__MODULE__,
         kwargs...)
@@ -517,7 +512,7 @@ function DiffEqBase.ODEProblem(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothi
             build_initializeprob = false, kwargs...)
     else
         _, u0, p = process_SciMLProblem(EmptySciMLFunction, sys, u0map, parammap;
-            t = tspan === nothing ? nothing : tspan[1], use_union, tofloat = false,
+            t = tspan === nothing ? nothing : tspan[1], tofloat = false,
             check_length = false)
         f = (du, u, p, t) -> (du .= 0; nothing)
         observedfun = ObservedFunctionCache(sys; eval_expression, eval_module,

src/systems/optimization/optimizationsystem.jl

@@ -295,7 +295,6 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
         cons_sparse = false, checkbounds = false,
         linenumbers = true, parallel = SerialForm(),
         eval_expression = false, eval_module = @__MODULE__,
-        use_union = false,
         checks = true,
         kwargs...) where {iip}
     if !iscomplete(sys)
@@ -338,10 +337,10 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
     elseif has_index_cache(sys) && get_index_cache(sys) !== nothing
         p = MTKParameters(sys, parammap, u0map)
     else
-        p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false, use_union)
+        p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false)
     end
-    lb = varmap_to_vars(dvs .=> lb, dvs; defaults = defs, tofloat = false, use_union)
-    ub = varmap_to_vars(dvs .=> ub, dvs; defaults = defs, tofloat = false, use_union)
+    lb = varmap_to_vars(dvs .=> lb, dvs; defaults = defs, tofloat = false)
+    ub = varmap_to_vars(dvs .=> ub, dvs; defaults = defs, tofloat = false)
 
     if !isnothing(lb) && all(lb .== -Inf) && !isnothing(ub) && all(ub .== Inf)
         lb = nothing
@@ -538,7 +537,6 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0map,
         checkbounds = false,
         linenumbers = false, parallel = SerialForm(),
         eval_expression = false, eval_module = @__MODULE__,
-        use_union = false,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed `OptimizationSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `OptimizationProblemExpr`")
@@ -578,10 +576,10 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0map,
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
         p = MTKParameters(sys, parammap, u0map)
     else
-        p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false, use_union)
+        p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false)
     end
-    lb = varmap_to_vars(dvs .=> lb, dvs; defaults = defs, tofloat = false, use_union)
-    ub = varmap_to_vars(dvs .=> ub, dvs; defaults = defs, tofloat = false, use_union)
+    lb = varmap_to_vars(dvs .=> lb, dvs; defaults = defs, tofloat = false)
+    ub = varmap_to_vars(dvs .=> ub, dvs; defaults = defs, tofloat = false)
 
     if !isnothing(lb) && all(lb .== -Inf) && !isnothing(ub) && all(ub .== Inf)
         lb = nothing

src/systems/parameter_buffer.jl

@@ -27,7 +27,7 @@ This requires that `complete` has been called on the system (usually via
 the default behavior).
 """
 function MTKParameters(
-        sys::AbstractSystem, p, u0 = Dict(); tofloat = false, use_union = false,
+        sys::AbstractSystem, p, u0 = Dict(); tofloat = false,
         t0 = nothing, substitution_limit = 1000)
     ic = if has_index_cache(sys) && get_index_cache(sys) !== nothing
         get_index_cache(sys)

src/systems/problem_utils.jl

@@ -330,17 +330,17 @@ struct MissingGuessError <: Exception
     vals::Vector{Any}
 end
 
-function Base.showerror(io::IO, err::MissingGuessError) 
-    println(io, 
-            """
-            Cyclic guesses detected in the system. Symbolic values were found for the following variables/parameters in the map: \
-            """)
+function Base.showerror(io::IO, err::MissingGuessError)
+    println(io,
+        """
+        Cyclic guesses detected in the system. Symbolic values were found for the following variables/parameters in the map: \
+        """)
     for (sym, val) in zip(err.syms, err.vals)
         println(io, "$sym  => $val")
     end
     println(io,
-            """
-            In order to resolve this, please provide additional numeric guesses so that the chain can be resolved to assign numeric values to each variable.            """)
+        """
+        In order to resolve this, please provide additional numeric guesses so that the chain can be resolved to assign numeric values to each variable.            """)
 end
 
 """
@@ -351,20 +351,20 @@ in `varmap`. Does not perform symbolic substitution in the values of `varmap`.
 
 Keyword arguments:
 - `tofloat`: Convert values to floating point numbers using `float`.
-- `use_union`: Use a `Union`-typed array if the values have heterogeneous types.
 - `container_type`: The type of container to use for the values.
 - `toterm`: The `toterm` method to use for converting symbolics.
 - `promotetoconcrete`: whether the promote to a concrete buffer (respecting
-  `tofloat` and `use_union`). Defaults to `container_type <: AbstractArray`.
+  `tofloat`). Defaults to `container_type <: AbstractArray`.
 - `check`: Error if any variables in `vars` do not have a mapping in `varmap`. Uses
   [`missingvars`](@ref) to perform the check.
 - `allow_symbolic` allows the returned array to contain symbolic values. If this is `true`,
   `promotetoconcrete` is set to `false`.
 - `is_initializeprob, guesses`: Used to determine whether the system is missing guesses.
 """
 function better_varmap_to_vars(varmap::AbstractDict, vars::Vector;
-        tofloat = true, use_union = true, container_type = Array,
-        toterm = default_toterm, promotetoconcrete = nothing, check = true, allow_symbolic = false, is_initializeprob = false)
+        tofloat = true, container_type = Array,
+        toterm = default_toterm, promotetoconcrete = nothing, check = true,
+        allow_symbolic = false, is_initializeprob = false)
     isempty(vars) && return nothing
 
     if check
@@ -382,8 +382,8 @@ function better_varmap_to_vars(varmap::AbstractDict, vars::Vector;
         end
 
         if !isempty(missingsyms)
-            is_initializeprob ? throw(MissingGuessError(missingsyms, missingvals)) : 
-                throw(UnexpectedSymbolicValueInVarmap(missingsyms[1], missingvals[1]))
+            is_initializeprob ? throw(MissingGuessError(missingsyms, missingvals)) :
+            throw(UnexpectedSymbolicValueInVarmap(missingsyms[1], missingvals[1]))
         end
     end
 
@@ -393,7 +393,7 @@ function better_varmap_to_vars(varmap::AbstractDict, vars::Vector;
 
     promotetoconcrete === nothing && (promotetoconcrete = container_type <: AbstractArray)
     if promotetoconcrete && !allow_symbolic
-        vals = promote_to_concrete(vals; tofloat = tofloat, use_union = use_union)
+        vals = promote_to_concrete(vals; tofloat = tofloat, use_union = false)
     end
 
     if isempty(vals)
@@ -731,8 +731,7 @@ Keyword arguments:
 - `fully_determined`: Override whether the initialization system is fully determined.
 - `check_initialization_units`: Enable or disable unit checks when constructing the
   initialization problem.
-- `tofloat`, `use_union`, `is_initializeprob`: Passed to [`better_varmap_to_vars`](@ref) for building `u0` (and
-  possibly `p`).
+- `tofloat`, `is_initializeprob`: Passed to [`better_varmap_to_vars`](@ref) for building `u0` (and possibly `p`).
 - `u0_constructor`: A function to apply to the `u0` value returned from `better_varmap_to_vars`
   to construct the final `u0` value.
 - `du0map`: A map of derivatives to values. See `implicit_dae`.
@@ -762,7 +761,7 @@ function process_SciMLProblem(
         implicit_dae = false, t = nothing, guesses = AnyDict(),
         warn_initialize_determined = true, initialization_eqs = [],
         eval_expression = false, eval_module = @__MODULE__, fully_determined = nothing,
-        check_initialization_units = false, tofloat = true, use_union = false,
+        check_initialization_units = false, tofloat = true,
         u0_constructor = identity, du0map = nothing, check_length = true,
         symbolic_u0 = false, warn_cyclic_dependency = false,
         circular_dependency_max_cycle_length = length(all_symbols(sys)),
@@ -841,7 +840,7 @@ function process_SciMLProblem(
     evaluate_varmap!(op, dvs; limit = substitution_limit)
 
     u0 = better_varmap_to_vars(
-        op, dvs; tofloat = true, use_union = false,
+        op, dvs; tofloat,
         container_type = u0Type, allow_symbolic = symbolic_u0, is_initializeprob)
 
     if u0 !== nothing
@@ -867,15 +866,15 @@ function process_SciMLProblem(
     if is_split(sys)
         p = MTKParameters(sys, op)
     else
-        p = better_varmap_to_vars(op, ps; tofloat, use_union, container_type = pType)
+        p = better_varmap_to_vars(op, ps; tofloat, container_type = pType)
     end
 
     if implicit_dae && du0map !== nothing
         ddvs = map(Differential(iv), dvs)
         du0map = to_varmap(du0map, ddvs)
         merge!(op, du0map)
         du0 = varmap_to_vars(op, ddvs; toterm = identity,
-            tofloat = true)
+            tofloat)
         kwargs = merge(kwargs, (; ddvs))
     else
         du0 = nothing
@@ -944,8 +943,8 @@ function get_u0_p(sys,
         u0map,
         parammap = nothing;
         t0 = nothing,
-        use_union = true,
         tofloat = true,
+        use_union = true,
         symbolic_u0 = false)
     dvs = unknowns(sys)
     ps = parameters(sys; initial_parameters = true)
@@ -985,7 +984,7 @@ function get_u0_p(sys,
     if symbolic_u0
         u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false, use_union = false)
     else
-        u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = true, use_union)
+        u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat, use_union)
     end
     p = varmap_to_vars(parammap, ps; defaults = defs, tofloat, use_union)
     p = p === nothing ? SciMLBase.NullParameters() : p

src/variables.jl

@@ -186,15 +186,14 @@ function varmap_to_vars(varmap, varlist; defaults = Dict(), check = true,
 
     vals = if eltype(varmap) <: Pair # `varmap` is a dict or an array of pairs
         varmap = todict(varmap)
-        _varmap_to_vars(varmap, varlist; defaults = defaults, check = check,
-            toterm = toterm)
+        _varmap_to_vars(varmap, varlist; defaults, check, toterm)
     else # plain array-like initialization
         varmap
     end
 
     promotetoconcrete === nothing && (promotetoconcrete = container_type <: AbstractArray)
     if promotetoconcrete
-        vals = promote_to_concrete(vals; tofloat = tofloat, use_union = use_union)
+        vals = promote_to_concrete(vals; tofloat, use_union)
     end
 
     if isempty(vals)

test/initial_values.jl

@@ -208,12 +208,14 @@ end
            x^2 + y^2 ~ 1]
     @mtkbuild pend = ODESystem(eqs, t)
 
-    @test_throws ModelingToolkit.MissingGuessError ODEProblem(pend, [x => 1], (0, 1), [g => 1], guesses = [y => λ, λ => y + 1])
+    @test_throws ModelingToolkit.MissingGuessError ODEProblem(
+        pend, [x => 1], (0, 1), [g => 1], guesses = [y => λ, λ => y + 1])
     ODEProblem(pend, [x => 1], (0, 1), [g => 1], guesses = [y => λ, λ => 0.5])
 
     # Throw multiple if multiple are missing
     @variables a(t) b(t) c(t) d(t) e(t)
     eqs = [D(a) ~ b, D(b) ~ c, D(c) ~ d, D(d) ~ e, D(e) ~ 1]
     @mtkbuild sys = ODESystem(eqs, t)
-    @test_throws ["a(t)", "c(t)"] ODEProblem(sys, [e => 2, a => b, b => a + 1, c => d, d => c + 1], (0, 1))
+    @test_throws ["a(t)", "c(t)"] ODEProblem(
+        sys, [e => 2, a => b, b => a + 1, c => d, d => c + 1], (0, 1))
 end

test/jumpsystem.jl

@@ -533,3 +533,21 @@ let
     Xsamp /= Nsims
     @test abs(Xsamp - Xf(0.2, p) < 0.05 * Xf(0.2, p))
 end
+
+@testset "JumpProcess simulation should be Int64 valued (#3446)" begin
+    @parameters p d
+    @variables X(t)
+    rate1 = p
+    rate2 = X * d
+    affect1 = [X ~ X + 1]
+    affect2 = [X ~ X - 1]
+    j1 = ConstantRateJump(rate1, affect1)
+    j2 = ConstantRateJump(rate2, affect2)
+
+    # Works.
+    @mtkbuild js = JumpSystem([j1, j2], t, [X], [p, d])
+    dprob = DiscreteProblem(js, [X => 15], (0.0, 10.0), [p => 2.0, d => 0.5])
+    jprob = JumpProblem(js, dprob, Direct())
+    sol = solve(jprob, SSAStepper())
+    @test eltype(sol[X]) === Int64
+end

test/odesystem.jl

@@ -745,7 +745,7 @@ let
     # No longer supported, Tuple used instead
     # pmap = Pair{Any, Union{Int, Float64}}[k1 => 1, k2 => 1.0]
     # tspan = (0.0, 1.0)
-    # prob = ODEProblem(sys, u0map, tspan, pmap, use_union = true)
+    # prob = ODEProblem(sys, u0map, tspan, pmap)
     # @test eltype(prob.p) === Union{Float64, Int}
 end
 
@@ -1208,7 +1208,7 @@ end
     sys = structural_simplify(ODESystem([D(x) ~ P], t, [x], [P]; name = :sys))
 
     function x_at_1(P)
-        prob = ODEProblem(sys, [x => P], (0.0, 1.0), [sys.P => P], use_union = false)
+        prob = ODEProblem(sys, [x => P], (0.0, 1.0), [sys.P => P])
         return solve(prob, Tsit5())(1.0)
     end
 

test/split_parameters.jl

@@ -123,9 +123,8 @@ sol = solve(prob, ImplicitEuler());
 # ------------------------ Mixed Type Conserved
 
 prob = ODEProblem(
-    sys, [], tspan, []; tofloat = false, use_union = true, build_initializeprob = false)
+    sys, [], tspan, []; tofloat = false, build_initializeprob = false)
 
-@test prob.p isa Vector{Union{Float64, Int64}}
 sol = solve(prob, ImplicitEuler());
 @test sol.retcode == ReturnCode.Success