You can write a function that takes mtk variables as input and performs arbitrary calculations with them, also leading to optimised code once generated. The only caveat being that it must not contain branches, and if it needs to then you have to use ifelse instead. If you can’t do this, you can also register the function with @register_symbolic f(t, x, u, v...) which will stop symbolic tracing, but then you need to take care over the efficiency of your code.
1 Like