Merge pull request #546 from rokke-git/vectors

adds vector syntax for `ArrayPartition`s and `VectorOfArray`s

Author: ChrisRackauckas

README.md

@@ -32,6 +32,12 @@ vA = VectorOfArray(a)
 vB = VectorOfArray(b)
 
 vA .* vB # Now all standard array stuff works!
+
+# you can also create it directly with a vector-like syntax:
+c = VA[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
+d = VA[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
+
+c .* d
 ```
 
 ### ArrayPartition
@@ -44,15 +50,15 @@ pB = ArrayPartition(b)
 
 pA .* pB # Now all standard array stuff works!
 
-# or do:
+# or using the vector syntax:
 x0 = rand(3, 3)
 v0 = rand(3, 3)
 a0 = rand(3, 3)
-u0 = ArrayPartition(x0, v0, a0)
-u0.x[1] == x0 # true
+u0 = AP[x0, v0, a0]
+u0.x[1] === x0 # true
 
 u0 .+= 1
-u0.x[2] == v0 # still true
+u0.x[2] === v0 # still true
 
 # do some calculations creating a new partitioned array
 unew = u0 * 10

src/RecursiveArrayTools.jl

@@ -133,14 +133,14 @@ module RecursiveArrayTools
     Base.convert(T::Type{<:GPUArraysCore.AnyGPUArray}, VA::AbstractVectorOfArray) = stack(VA.u)
     (T::Type{<:GPUArraysCore.AnyGPUArray})(VA::AbstractVectorOfArray) = T(Array(VA))
 
-    export VectorOfArray, DiffEqArray, AbstractVectorOfArray, AbstractDiffEqArray,
+    export VectorOfArray, VA, DiffEqArray, AbstractVectorOfArray, AbstractDiffEqArray,
         AllObserved, vecarr_to_vectors, tuples
 
     export recursivecopy, recursivecopy!, recursivefill!, vecvecapply, copyat_or_push!,
         vecvec_to_mat, recursive_one, recursive_mean, recursive_bottom_eltype,
         recursive_unitless_bottom_eltype, recursive_unitless_eltype
 
-    export ArrayPartition, NamedArrayPartition
+    export ArrayPartition, AP, NamedArrayPartition
 
     include("precompilation.jl")
 

src/array_partition.jl

@@ -722,3 +722,28 @@ end
 function Adapt.adapt_structure(to, ap::ArrayPartition)
     return ArrayPartition(map(x -> Adapt.adapt(to, x), ap.x)...)
 end
+
+"""
+```julia
+AP[ matrices, ]
+```
+
+Create an `ArrayPartition` using vector syntax. Equivalent to `ArrayPartition(matrices)`, but looks nicer with nesting.
+
+# Examples:
+
+Simple examples:
+```julia
+ArrayPartition([1,2,3], [1 2;3 4]) == AP[[1,2,3], [1 2;3 4]] # true
+AP[1u"m/s^2", 1u"m/s", 1u"m"]
+```
+
+With an ODEProblem:
+```julia
+func(u, p, t) = AP[5u.x[1], u.x[2]./2]
+ODEProblem(func, AP[ [1.,2.,3.], [1. 2.;3. 4.] ], (0, 1)) |> solve
+```
+
+"""
+struct AP end
+Base.getindex(::Type{AP}, xs...) = ArrayPartition(xs...)

src/vector_of_array.jl

@@ -1655,3 +1655,31 @@ end
 end
 unpack_args_voa(i, args::Tuple{Any}) = (unpack_voa(args[1], i),)
 unpack_args_voa(::Any, args::Tuple{}) = ()
+
+"""
+```julia
+VA[ matrices, ]
+```
+
+Create an `VectorOfArray` using vector syntax. Equivalent to `VectorOfArray([matrices])`, but looks nicer with nesting.
+
+# Simple example:
+```julia
+VectorOfArray([[1,2,3], [1 2;3 4]]) == VA[[1,2,3], [1 2;3 4]] # true
+```
+
+# All the layers:
+```julia
+nested = VA[
+    fill(1, 2, 3),
+    VA[
+        VA[8, [1, 2, 3], [1 2;3 4], VA[1, 2, 3]],
+        fill(2, 3, 4),
+        VA[3ones(3), zeros(3)],
+    ],
+]
+```
+
+"""
+struct VA end
+Base.getindex(::Type{VA}, xs...) = VectorOfArray(collect(xs))

test/basic_indexing.jl

@@ -6,6 +6,9 @@ testa = cat(recs..., dims = 2)
 testva = VectorOfArray(recs)
 @test maximum(testva) == maximum(maximum.(recs))
 
+testva = VA[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
+@test maximum(testva) == maximum(maximum.(recs))
+
 # broadcast with array
 X = rand(3, 3)
 mulX = sqrt.(abs.(testva .* X))
@@ -161,15 +164,15 @@ diffeq = DiffEqArray(recs, t)
 @test diffeq[:, (end - 1):end].t == t[(length(t) - 1):length(t)]
 
 # Test views of heterogeneous arrays (issue #453)
-f = VectorOfArray([[1.0], [2.0, 3.0]])
+f = VA[[1.0], [2.0, 3.0]]
 @test length(view(f, :, 1)) == 1
 @test length(view(f, :, 2)) == 2
 @test view(f, :, 1) == [1.0]
 @test view(f, :, 2) == [2.0, 3.0]
 @test collect(view(f, :, 1)) == f[:, 1]
 @test collect(view(f, :, 2)) == f[:, 2]
 
-f2 = VectorOfArray([[1.0, 2.0], [3.0]])
+f2 = VA[[1.0, 2.0], [3.0]]
 @test length(view(f2, :, 1)) == 2
 @test length(view(f2, :, 2)) == 1
 @test view(f2, :, 1) == [1.0, 2.0]
@@ -178,7 +181,7 @@ f2 = VectorOfArray([[1.0, 2.0], [3.0]])
 @test collect(view(f2, :, 2)) == f2[:, 2]
 
 # Test `end` with ragged arrays
-ragged = VectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0]])
+ragged = VA[[1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0]]
 @test ragged[end, 1] == 2.0
 @test ragged[end, 2] == 5.0
 @test ragged[end, 3] == 9.0
@@ -192,7 +195,7 @@ ragged = VectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0]])
 @test ragged[:, 2:end] == VectorOfArray(ragged.u[2:end])
 @test ragged[:, (end - 1):end] == VectorOfArray(ragged.u[(end - 1):end])
 
-ragged2 = VectorOfArray([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0], [7.0, 8.0, 9.0]])
+ragged2 = VA[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0], [7.0, 8.0, 9.0]]
 @test ragged2[end, 1] == 4.0
 @test ragged2[end, 2] == 6.0
 @test ragged2[end, 3] == 9.0
@@ -228,7 +231,7 @@ ragged_range_idx = 1:lastindex(ragged, 1)
 @test identity.(ragged_range_idx) === ragged_range_idx
 
 # Broadcasting of heterogeneous arrays (issue #454)
-u = VectorOfArray([[1.0], [2.0, 3.0]])
+u = VA[[1.0], [2.0, 3.0]]
 @test length(view(u, :, 1)) == 1
 @test length(view(u, :, 2)) == 2
 # broadcast assignment into selected column (last index Int)
@@ -293,7 +296,7 @@ u[1:2, 1, [1, 3], 2] .= [1.0 3.0; 2.0 4.0]
 @test u[end, 1, end] == u.u[end][end, 1, end]
 
 # Test that views can be modified
-f3 = VectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]])
+f3 = VA[[1.0, 2.0], [3.0, 4.0, 5.0]]
 v = view(f3, :, 2)
 @test length(v) == 3
 v[1] = 10.0
@@ -384,7 +387,7 @@ mulX .= sqrt.(abs.(testva .* testvb))
 @test mulX == ref
 
 # https://github.com/SciML/RecursiveArrayTools.jl/issues/49
-a = ArrayPartition(1:5, 1:6)
+a = AP[1:5, 1:6]
 a[1:8]
 a[[1, 3, 8]]
 

test/downstream/downstream_events.jl

@@ -1,9 +1,9 @@
 using OrdinaryDiffEq, StaticArrays, RecursiveArrayTools
-u0 = ArrayPartition(SVector{1}(50.0), SVector{1}(0.0))
+u0 = AP[SVector{1}(50.0), SVector{1}(0.0)]
 tspan = (0.0, 15.0)
 
 function f(u, p, t)
-    return ArrayPartition(SVector{1}(u[2]), SVector{1}(-9.81))
+    return AP[SVector{1}(u[2]), SVector{1}(-9.81)]
 end
 
 prob = ODEProblem(f, u0, tspan)
@@ -13,7 +13,7 @@ function condition(u, t, integrator) # Event when event_f(u,t,k) == 0
 end
 
 affect! = nothingf = affect_neg! = function (integrator)
-    return integrator.u = ArrayPartition(SVector{1}(integrator.u[1]), SVector{1}(-integrator.u[2]))
+    return integrator.u = AP[SVector{1}(integrator.u[1]), SVector{1}(-integrator.u[2])]
 end
 
 callback = ContinuousCallback(condition, affect!, affect_neg!, interp_points = 100)

test/downstream/odesolve.jl

@@ -4,7 +4,7 @@ function lorenz(du, u, p, t)
     du[2] = u[1] * (28.0 - u[3]) - u[2]
     return du[3] = u[1] * u[2] - (8 / 3) * u[3]
 end
-u0 = ArrayPartition([1.0, 0.0], [0.0])
+u0 = AP[[1.0, 0.0], [0.0]]
 @test ArrayInterface.zeromatrix(u0) isa Matrix
 tspan = (0.0, 100.0)
 prob = ODEProblem(lorenz, u0, tspan)
@@ -25,26 +25,26 @@ function mymodel(F, vars)
     return
 end
 # To show that the function works
-F = ArrayPartition([0.0 0.0; 0.0 0.0], [0.0 0.0; 0.0 0.0])
-u0 = ArrayPartition([0.1 1.2; 0.1 1.2], [0.1 1.2; 0.1 1.2])
+F = AP[[0.0 0.0; 0.0 0.0], [0.0 0.0; 0.0 0.0]]
+u0 = AP[[0.1 1.2; 0.1 1.2], [0.1 1.2; 0.1 1.2]]
 result = mymodel(F, u0)
 nlsolve(mymodel, u0)
 
 # Nested ArrayPartition solves
 
 function dyn(u, p, t)
     return ArrayPartition(
-        ArrayPartition(zeros(1), [0.0]),
-        ArrayPartition(zeros(1), [0.0])
+        AP[zeros(1), [0.0]],
+        AP[zeros(1), [0.0]]
     )
 end
 
 @test solve(
     ODEProblem(
         dyn,
         ArrayPartition(
-            ArrayPartition(zeros(1), [-1.0]),
-            ArrayPartition(zeros(1), [0.75])
+            AP[zeros(1), [-1.0]],
+            AP[zeros(1), [0.75]]
         ),
         (0.0, 1.0)
     ),
@@ -55,8 +55,8 @@ end
     ODEProblem(
         dyn,
         ArrayPartition(
-            ArrayPartition(zeros(1), [-1.0]),
-            ArrayPartition(zeros(1), [0.75])
+            AP[zeros(1), [-1.0]],
+            AP[zeros(1), [0.75]]
         ),
         (0.0, 1.0)
     ),