Automatic symbolic function barrier generation

Specific Domains Modelling & Simulations
,
1

Hi,

I’m working with ModelingToolkit.jl and I have an issue with complex symbolic expressions, on which mtkcompile takes a lot of time.

I discovered that I can hide expressions behind a function barrier using @register_symbolic, and that prevents the simplification from happening. I tried a universal simpler approach like just defining the identity barrier f(x) = x but it seems that for structural_simplify to leave the expression alone, the arguments of the registered function cannot be complex expressions either.

So I need to wrap my expressions in functions that take only simple variables as arguments, then register those functions symbolically. The problem is that doing this manually for every expression is tedious and… I’m lazy. However, I am sure this could be automated, but my skills are not sufficient.

Here’s what I want to achieve:

Example usage with vector-variables:

using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D

function test()
   a = 6.0 
   @variables x(t)[1:2] y(t)[1:2] z(t)

   # Instead of writing this complex expression directly:
   eqz = z ~ sum(x)^2 + sum(y)/a - 2.0

   # I want to write:
   eqz = z ~ @symbolic_hide sum(x)^2 + sum(y)/a - 2.0
end

Expected behavior:

The macro (or function, I don’t even know what the best choice is since it can only be done at runtime) should automatically:

  1. Extract the variables from the expression (x, y, a)
  2. Determine which are symbolic variables and which are not
  3. Generate a unique function with proper type signatures
  4. Register it using @register_symbolic (this is the tricky part, since I want to call it inside a function)
  5. Return a call to that function

So the macro would generate something equivalent to:

function hidden_func(x::AbstractVector{T}, y::AbstractVector{T}, a::Float64) where T <: Number
    return sum(x)^2 + sum(y)/a - 2.0
end
@register_symbolic hidden_func(x::AbstractVector, y::AbstractVector, a::Float64)

eqz = z ~ hidden_func(collect(x), collect(y), a)

I tried some things with my basic knowledge and some llm, managed to get something working in scripts but never managed to make it work inside functions…

Any guidance would be much appreciated :slight_smile: