`respecialize` always promotes `Float16` and `Float32` to `Float64`

[bradcarman]

Using MTK v10.21.0

All Int types are preserved, but Float types are not. Take the following MWE...

using ModelingToolkit
using ModelingToolkit: t_nounits as t
using Test

function ConstantGeneric(;name, k::T) where T

    pars = @parameters begin
        k::T = k
    end

    systems = @named begin
        output = Blocks.RealOutput()
    end
    eqs = [
        output.u ~ k
    ]

    return ODESystem(eqs, t, [], pars; name, systems)
end

function TypeTest(T1::Type, T2::Type; name)
    vars = @variables begin
        output(t)
    end
    systems = @named begin
        c1 = ConstantGeneric(k = T1(1))
        c2 = ConstantGeneric(k = T2(2))
    end
    equations = [
        output ~ ifelse(sin(t) > 0, c1.output.u, c2.output.u)
    ]
    ODESystem(equations, t; name, systems)
end

function run_test(T1, T2)
    @mtkcompile sys = TypeTest(T1, T2)
    sys = respecialize(sys)
    prob = ODEProblem(sys, Any[], (0,1))
    sol = solve(prob)

    @test typeof(sol(0.0; idxs=sys.c1.output.u)) == T1
    @test typeof(sol(0.0; idxs=sys.c2.output.u)) == T2 
    @test typeof(sol(0.0; idxs=sys.output)) == promote_type(T1,T2)
end

run_test(Int16, Int16) # Test Passed
run_test(Int16, Int32) # Test Passed
run_test(Float16, Float16) # ERROR
run_test(Float32, Float32) # ERROR
run_test(Float64, Float64) # Test Passed
run_test(Int16, Float16) # ERROR
#=
Test Failed at c:\Work\Packages\BoeingLLFLibrary\scripts\type_testing.jl:45
  Expression: typeof(sol(0.0; idxs = sys.c2.output.u)) == T2
   Evaluated: Float64 == Float16
=#

[AayushSabharwal]

There's a couple things to note:

• respecialize does nothing here. It only operates on nonnumeric parameters, and says so in the docstring.
• Even if respecialize operated on numeric parameters, it would do nothing here since all parameters have concrete symtypes (Int16, Float32 etc).

Now onto the crux of the issue - the floating point types not working. What happens here is that any parameter declared with @parameters is considered tunable. Now, integer types can't be tunable (they're not differentiable) so MTK puts them in the Constants portion. The floating point types, however, can be tunable so they're grouped together. We can see this here:

julia> @mtkcompile sys = TypeTest(Int16, Int32)
julia> prob = ODEProblem(sys, [], (0.0, 1.0))
julia> prob.p
MTKParameters{StaticArraysCore.SizedVector{0, Float64, Vector{Float64}}, Vector{Float64}, Tuple{}, Tuple{Vector{Int32}, Vector{Int16}}, Tuple{}, Tuple{}}(Float64[], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], (), (Int32[2], Int16[1]), (), ())

Note how there are Int32 and Int32 vectors holding one element each.

julia> @mtkcompile sys = TypeTest(Float16, Float32)
julia> prob = ODEProblem(sys, [], (0.0, 1.0))
julia> prob.p
MTKParameters{Vector{Float64}, Vector{Float64}, Tuple{}, Tuple{}, Tuple{}, Tuple{}}([1.0, 2.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], (), (), (), ())

Now note how the parameters are just in the first Vector{Float64}. This is Float64 and not Float32 (the result of promote_type(Float16, Float32)) is because tunables have the same eltype as that of prob.u0.

To make this work the way you intend, the parameters in ConstantGeneric must be declared with the tunable = false metadata to ensure they're considered part of the Constants portion:

julia> function ConstantGeneric(;name, k::T) where T

           pars = @parameters begin
               k::T = k, [tunable = false]
           end

           systems = @named begin
               output = Blocks.RealOutput()
           end
           eqs = [
               output.u ~ k
           ]

           return ODESystem(eqs, t, [], pars; name, systems)
       end

MTKv10 has a shorthand for this: the @constants macro

julia> function ConstantGeneric(;name, k::T) where T

           pars = @constants begin
               k::T = k
           end

           systems = @named begin
               output = Blocks.RealOutput()
           end
           eqs = [
               output.u ~ k
           ]

           return ODESystem(eqs, t, [], pars; name, systems)
       end

With this change:

julia> @mtkcompile sys = TypeTest(Float16, Float32)
julia> prob = ODEProblem(sys, [], (0.0, 1.0))
julia> prob.p
MTKParameters{StaticArraysCore.SizedVector{0, Float64, Vector{Float64}}, Vector{Float64}, Tuple{}, Tuple{Vector{Float16}, Vector{Float32}}, Tuple{}, Tuple{}}(Float64[], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], (), (Float16[1.0], Float32[2.0]), (), ())

The parameters are stored as intended.