Merge pull request #490 from SciML/autoopt

Changes to make auto-optimize work

Author: ChrisRackauckas

Project.toml

@@ -12,6 +12,7 @@ DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 GeneralizedGenerated = "6b9d7cbe-bcb9-11e9-073f-15a7a543e2eb"
+LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/ModelingToolkit.jl

@@ -1,7 +1,7 @@
 module ModelingToolkit
 
 using DiffEqBase, Distributed
-using StaticArrays, LinearAlgebra, SparseArrays
+using StaticArrays, LinearAlgebra, SparseArrays, LabelledArrays
 using Latexify, Unitful, ArrayInterface
 using MacroTools
 using UnPack: @unpack

src/build_function.jl

@@ -24,7 +24,7 @@ i.e., f(u,p,args...) for the out-of-place and scalar functions and
 ```julia
 build_function(ex, args...;
                conv = simplified_expr, expression = Val{true},
-               checkbounds = false, constructor=nothing,
+               checkbounds = false,
                linenumbers = false, target = JuliaTarget())
 ```
 
@@ -46,8 +46,6 @@ Keyword Arguments:
 
 - `checkbounds`: For whether to enable bounds checking inside of the generated
   function. Defaults to false, meaning that `@inbounds` is applied.
-- `constructor`: Allows for an arbitrary constructor function to be passed in
-  for handling expressions of "weird" types. Defaults to nothing.
 - `linenumbers`: Determines whether the generated function expression retains
   the line numbers. Defaults to true.
 - `target`: The output target of the compilation process. Possible options are:
@@ -104,7 +102,7 @@ end
 # Scalar output
 function _build_function(target::JuliaTarget, op::Operation, args...;
                          conv = simplified_expr, expression = Val{true},
-                         checkbounds = false, constructor=nothing,
+                         checkbounds = false,
                          linenumbers = true, headerfun=addheader)
 
     argnames = [gensym(:MTKArg) for i in 1:length(args)]
@@ -165,7 +163,7 @@ end
 
 function _build_function(target::JuliaTarget, rhss, args...;
                          conv = simplified_expr, expression = Val{true},
-                         checkbounds = false, constructor=nothing,
+                         checkbounds = false,
                          linenumbers = false, multithread=nothing,
                          headerfun=addheader, outputidxs=nothing,
                          skipzeros = false, parallel=SerialForm())
@@ -323,41 +321,39 @@ function _build_function(target::JuliaTarget, rhss, args...;
 
     if rhss isa Matrix
         arr_sys_expr = build_expr(:vcat, [build_expr(:row,[conv(rhs) for rhs ∈ rhss[i,:]]) for i in 1:size(rhss,1)])
-        # : x because ??? what to do in the general case?
-        _constructor = constructor === nothing ? :($(first(argnames)) isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.SMatrix{$(size(rhss)...)} :  x->(out = similar(typeof($(fargs.args[1])),$(size(rhss)...)); out .= x)) : constructor
     elseif typeof(rhss) <: Array && !(typeof(rhss) <: Vector)
         vector_form = build_expr(:vect, [conv(rhs) for rhs ∈ rhss])
         arr_sys_expr = :(reshape($vector_form,$(size(rhss)...)))
-        _constructor = constructor === nothing ? :($(first(argnames)) isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.SArray{$(size(rhss)...)} :  x->(out = similar(typeof($(fargs.args[1])),$(size(rhss)...)); out .= x)) : constructor
     elseif rhss isa SparseMatrixCSC
         vector_form = build_expr(:vect, [conv(rhs) for rhs ∈ nonzeros(rhss)])
         arr_sys_expr = :(SparseMatrixCSC{eltype($(first(argnames))),Int}($(size(rhss)...), $(rhss.colptr), $(rhss.rowval), $vector_form))
-        # Static and sparse? Probably not a combo that will actually be hit, but give a default anyways
-        _constructor = constructor === nothing ? :($(first(argnames)) isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.SMatrix{$(size(rhss)...)} : x->x) : constructor
     else # Vector
         arr_sys_expr = build_expr(:vect, [conv(rhs) for rhs ∈ rhss])
-        # Handle vector constructor separately using `typeof(u)` to support things like LabelledArrays
-        _constructor = constructor === nothing ? :($(first(argnames)) isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.similar_type(typeof($(fargs.args[1])), eltype(X)) : x->convert(typeof($(fargs.args[1])),x)) : constructor
     end
 
+	xname = gensym(:MTK)
+
+	arr_sys_expr = (typeof(rhss) <: Vector || typeof(rhss) <: Matrix) && !(eltype(rhss) <: AbstractArray) ? quote
+		if typeof($(fargs.args[1])) <: Union{ModelingToolkit.StaticArrays.SArray,ModelingToolkit.LabelledArrays.SLArray}
+			$xname = ModelingToolkit.StaticArrays.@SArray $arr_sys_expr
+			convert(typeof($(fargs.args[1])),$xname)
+		else
+			$xname = $arr_sys_expr
+			if !(typeof($(fargs.args[1])) <: Array)
+				convert(typeof($(fargs.args[1])),$xname)
+			else
+				$xname
+			end
+		end
+	end : arr_sys_expr
+
     let_expr = Expr(:let, var_eqs, tuple_sys_expr)
     arr_let_expr = Expr(:let, var_eqs, arr_sys_expr)
     bounds_block = checkbounds ? let_expr : :(@inbounds begin $let_expr end)
-    arr_bounds_block = checkbounds ? arr_let_expr : :(@inbounds begin $arr_let_expr end)
+    oop_bounds_block = checkbounds ? arr_let_expr : :(@inbounds begin $arr_let_expr end)
     ip_bounds_block = checkbounds ? ip_let_expr : :(@inbounds begin $ip_let_expr end)
 
-    oop_body_block = :(
-      # If u is a weird non-StaticArray type and we want a sparse matrix, just do the optimized sparse anyways
-      if $(fargs.args[1]) isa Array || (!(typeof($(fargs.args[1])) <: StaticArray) && $(rhss isa SparseMatrixCSC))
-          return $arr_bounds_block
-      else
-          X = $bounds_block
-          construct = $_constructor
-          return construct(X)
-      end
-    )
-
-    oop_ex = headerfun(oop_body_block, fargs, false)
+    oop_ex = headerfun(oop_bounds_block, fargs, false)
     iip_ex = headerfun(ip_bounds_block, fargs, true; X=X)
 
     if !linenumbers