handling of if-else-end conditionals in an MTK model (not requiring IfElse.jl)

[anandijain]

maybe related to #38, not sure.

if x < 0
    y = sqrt(-x)
else
    y = sqrt(x)
end

There are a collection of these type branches in my model (including nested). In MATLAB, the solver is able to handle this, but MTK currently cannot and modelingtoolkitize will fail.

I started working on a macro but I didn't know how to handle the original case all that well. It handles the very most basic case but nothing more.

function _toifelse(ex::Expr)
    args = map(_toifelse, ex.args)
    ex.head == :if ? Expr(:call, :ifelse, args...) : Expr(ex.head, args...)
end
_toifelse(x) = x
macro ifelse(ex)
    _toifelse(ex)
end

[anandijain]

macro to_ifelse(ex)
    esc(to_ifelse(ex))
end

function to_ifelse(ex)
    if ex isa Expr && ex.head in (:if, :elseif)
        :($(IfElse.ifelse)($(map(to_ifelse, vcat(ex.args, nothing)[1:3])...)))
    elseif ex isa Expr
        Expr(ex.head, map(to_ifelse, ex.args)...)
    else
        ex
    end
end

more robust

[hersle]

Was this fixed?

[ChrisRackauckas]

No, but it's more of a symbolics issue. We'd need an alternative tracer since a dispatch-based tracer cannot handle if, at least in its current form. If if was made dispatchable in Julia then that would be another solution... but that would be a weird new feature 😅

[pepijndevos]

Here is a macro that can convert nested if statements to ifelse calls by doing some kind of static single assignment transform. Not sure if there is a good place where this should be contributed.

ffirst(x) = first(first(x))

opmap = Dict(
:(+=) => :(+),
:(-=) => :(-),
:(*=) => :(*),
:(/=) => :(/),
:(%=) => :(%),
:(&=) => :(&),
:(|=) => :(|),
)

ifelsexpr(cond, lhs, rhs) = :($lhs = $ifelse($cond, $rhs, $(Expr(:isdefined, lhs)) ? $lhs : NaN))

function assignments(cond, expr, res, mod)
    if expr isa LineNumberNode
        push!(res, expr)
    elseif !(expr isa Expr)
        push!(res, :($ifelse($cond, $expr, NaN)))
    elseif expr.head == :if || expr.head == :elseif
        tcond = :($(expr.args[1]) & $cond)
        assignments(tcond, expr.args[2], res, mod)
        if length(expr.args) == 3
            fcond = :($(expr.args[1]) & !($cond))
            assignments(fcond, expr.args[3], res, mod)
        end
    elseif expr.head == :call && expr.args[1] == ifelse
        tcond = :($(expr.args[2]) & $cond)
        assignments(tcond, expr.args[3], res, mod)
        if length(expr.args) == 4
            fcond = :($(expr.args[2]) & !($cond))
            assignments(fcond, expr.args[4], res, mod)
        end
    elseif expr.head == :block
        for e in expr.args
            assignments(cond, e, res, mod)
        end
    elseif expr.head == :(=)
        lhs, rhs = expr.args
        push!(res, ifelsexpr(cond, lhs, rhs))
    elseif expr.head in keys(opmap)
        lhs, rhs = expr.args
        push!(res, ifelsexpr(cond, lhs, Expr(:call, opmap[expr.head], lhs , rhs)))
    else
        @warn("Only if and assignement supported: $expr")
        push!(res, :(if$cond; $expr; end))
    end
    return res
end

function ssa_(expr, mod)
    expr = macroexpand(mod, expr)
    @assert expr.head == :if || expr.head == :elseif
    res = Expr(:block)
    cond = expr.args[1]
    assignments(cond, expr.args[2], res.args, mod)
    if length(expr.args) == 3
        assignments(:(!($cond)), expr.args[3], res.args, mod)
    end
    res
end

macro ssa(expr)
    esc(ssa_(expr, __module__))
end

[ChrisRackauckas]

Yeah it's interesting but it doesn't solve this because the tracing still requires that the code is hitting dispatchable ifelse functions. Users would have to use this macro for our code to work. We would need to for example setup the abstract interpreter to do this transform before dispatching, or write a more complex tracer which @MasonProtter looked into before. Otherwise it indeed requires user intervention. That said, this macro is a nice thing to help users do that manual intervention so if they find this issue, they should do it

[pepijndevos]

DAECompiler is an alternative approach that indeed uses custom interpreters to run Julia IR rather than tracing symbolics. The above macro was from an experiment to use JuliaSimCompiler.

[Uroc327]

This would be a great feature. The documentation says Generally, a code which is compatible with forward-mode automatic differentiation is compatible with modelingtoolkitize. To live up to that statement, I think automatic conversion of branches to symbolics-compatible ifelse is necessary.