SciML / RecursiveArrayTools.jl

Tools for easily handling objects like arrays of arrays and deeper nestings in scientific machine learning (SciML) and other applications
https://docs.sciml.ai/RecursiveArrayTools/stable/
Other
212 stars 57 forks source link

Specialize `Base.similar` for `VectorOfArray` with multidimensional parent #359

Closed jlchan closed 7 months ago

jlchan commented 7 months ago

This PR specializes Base.similar for VectorOfArray with a homogeneous (e.g., all elements are the same type) multidimensional parent array (introduced in https://github.com/SciML/RecursiveArrayTools.jl/pull/357). This ensures that similar(VectorOfArray(array)).u is similar to array.u where array is a homogeneous multi-dimensional array.

This PR also adds Base.parent for VectorOfArray (parent(vec) just returns vec.u)

jlchan commented 7 months ago

@ranocha I think with these changes, we can deal with multidimensional arrays of SVector

For example, the following works for me under this PR:

using StaticArrays
using RecursiveArrayTools
using OrdinaryDiffEq

function rhs!(duu::VectorOfArray, uu::VectorOfArray, p, t)
    du = parent(duu)
    u = parent(uu)
    du .= u
end

u = fill(SVector{2}(ones(2)), 2, 3)
ode = ODEProblem(rhs!, VectorOfArray(u), (0.0, 1.0))
sol = solve(ode, Tsit5())
ChrisRackauckas commented 7 months ago

Add that as a regression test in the downstream portion (https://github.com/SciML/RecursiveArrayTools.jl/blob/master/test/upstream.jl which is for some reason named upstream we should fix that and move it to the other test group https://github.com/SciML/RecursiveArrayTools.jl/blob/master/test/runtests.jl#L63

ranocha commented 7 months ago

@ranocha I think with these changes, we can deal with multidimensional arrays of SVector

Nice 👍