ODE model with preallocated MVectors and still compatible with AD/stiff solvers?

I am writing an ODE model that describes the positions of a system of N interacting particles living in a 2D box. My ODE state vector u is an N x 2 array of numbers. Each call to the model updates a pre-allocated N x 2 array named F_arr, which I preallocated using DiffCache from PreallocationTools.jl. I am able to solve the model using non-stiff and stiff solvers (e.g. Tsit5 and Rodas5), and it’s compatible with dual numbers/autodifferentiation.

Now I’m wondering if I can improve the performance through using MVectors. Since the position of each particle is of fixed length (2), I was trying to come up with a way to store each particles position as an MVector. To do this, I changed the state variable u to be a vector (of length N) of MVectors (of length 2(. This approach worked when solving the model using Tsit5, but it I ran into issues when using stiff solvers such as Rodas 5 because apparently the stiff solvers expect u to be an N-dimensional array of numbers.

Any suggestions for how to incorporate MVectors into my model to maximize (or at least improve) performance, while also being compatible with stiff solvers and Dual numbers?

It would probably help if you post the exact error that you got.

Arrays of arrays are not generally supported. But even if they were, it wouldn’t be any faster. Arrays of SVectors are equivalent in memory to just an array, except worse for linear algebra since it requires a reinterpret to a vector. Arrays of MVectors are simply worse in memory than using an array. So this won’t be an acceleration, it’ll be the opposite.

If you want to do this, you could always do mi = MVector{2}(u[i],u[i+1]) and use that. If you do it locally in f then, under certain escape analysis rules, it will even compile out so it will not allocate and can be a nice tool. But, keeping that in the u is not going to be performance helpful.