Implement the SciMLStructures interface for ArrayPartition

Project.toml

@@ -11,6 +11,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 SciMLPublic = "431bcebd-1456-4ced-9d72-93c2757fff0b"
+SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 
@@ -67,6 +68,7 @@ Random = "1"
 RecipesBase = "1.3.4"
 RecursiveArrayToolsShorthandConstructors = "1"
 ReverseDiff = "1.15"
+SciMLStructures = "1.10"
 SafeTestsets = "0.1"
 SciMLPublic = "1"
 SciMLTesting = "2.4"

src/RecursiveArrayTools.jl

@@ -14,6 +14,7 @@ module RecursiveArrayTools
         UnitUpperTriangular, UpperTriangular, ldiv!, mul!
     using RecipesBase: RecipesBase, @recipe
     using SciMLPublic: @public
+    using SciMLStructures: SciMLStructures, Tunable
     using StaticArraysCore: StaticArraysCore
     using SymbolicIndexingInterface: SymbolicIndexingInterface, ArraySymbolic,
         NotSymbolic, ParameterIndexingProxy, ParameterTimeseriesCollection,
@@ -200,6 +201,7 @@ module RecursiveArrayTools
     include("vector_of_array.jl")
     include("array_partition.jl")
     include("named_array_partition.jl")
+    include("scimlstructures.jl")
 
     function Base.show(io::IO, x::ArrayPartition)
         return invoke(show, Tuple{typeof(io), Any}, io, x)

src/scimlstructures.jl

@@ -0,0 +1,82 @@
+# SciMLStructures opts every numeric array into the interface and canonicalizes it
+# with `vec`. That is right for a contiguous array, but an `ArrayPartition` is
+# already an `AbstractVector`, so `vec` returns it unchanged and the buffer comes
+# back unflattened; the generic `repack` then errors, since an `ArrayPartition`
+# cannot be rebuilt by a trivial array constructor.
+#
+# The partitions are laid out one after another, which is the same stable ordering
+# `copyto!` and `Vector` use, so a caller can rely on the flat buffer matching
+# `collect(p)`.
+
+# `reduce(vcat, p.x)` rather than building from one partition: the partitions need
+# not share an array type or an element type, and concatenating promotes across both.
+# It also keeps a uniform static or device-backed partitioning in kind, where sizing
+# from `p.x[1]` would silently pick that partition's type for the whole buffer.
+#
+# The buffer must not alias `p`, since `canonicalize` reports `aliases = false`, and
+# `reduce(vcat, (x,))` returns `x` itself, so a lone partition is copied.
+function _flatten(p::ArrayPartition)
+    isempty(p.x) && throw(
+        ArgumentError("cannot canonicalize an `ArrayPartition` with no partitions")
+    )
+    length(p.x) == 1 && return copy(only(p.x))
+    return reduce(vcat, p.x)
+end
+
+SciMLStructures.isscimlstructure(::ArrayPartition) = true
+SciMLStructures.ismutablescimlstructure(::ArrayPartition) = true
+
+function SciMLStructures.canonicalize(::Tunable, p::ArrayPartition)
+    buffer = _flatten(p)
+    repack = let p = p
+        new_values -> SciMLStructures.replace(Tunable(), p, new_values)
+    end
+    return buffer, repack, false
+end
+
+function SciMLStructures.replace(
+        ::Tunable, p::ArrayPartition, new_values::AbstractArray
+    )
+    length(new_values) == length(p) || throw(
+        DimensionMismatch(
+            "expected $(length(p)) values to replace the tunable portion, got " *
+                "$(length(new_values))"
+        )
+    )
+    offsets = _offsets(p)
+    parts = ntuple(length(p.x)) do i
+        part = similar(p.x[i])
+        copyto!(part, view(new_values, (offsets[i] + 1):offsets[i + 1]))
+    end
+    return ArrayPartition(parts)
+end
+
+function SciMLStructures.replace!(
+        ::Tunable, p::ArrayPartition, new_values::AbstractArray
+    )
+    length(new_values) == length(p) || throw(
+        DimensionMismatch(
+            "expected $(length(p)) values to replace the tunable portion, got " *
+                "$(length(new_values))"
+        )
+    )
+    offsets = _offsets(p)
+    for i in eachindex(p.x)
+        copyto!(p.x[i], view(new_values, (offsets[i] + 1):offsets[i + 1]))
+    end
+    return nothing
+end
+
+# Cumulative partition boundaries, so partition `i` owns
+# `(offsets[i] + 1):offsets[i + 1]` of the flat buffer.
+function _offsets(p::ArrayPartition)
+    n = length(p.x)
+    offsets = ntuple(n + 1) do i
+        total = 0
+        for j in 1:(i - 1)
+            total += length(p.x[j])
+        end
+        total
+    end
+    return offsets
+end

test/Core/scimlstructures_test.jl

@@ -0,0 +1,130 @@
+using RecursiveArrayTools, SciMLStructures, StaticArrays, Test
+using SciMLStructures: Tunable, Constants, Caches, Discrete, Initials, Input,
+    canonicalize, hasportion, isscimlstructure, ismutablescimlstructure
+
+# SciMLStructures canonicalizes a generic array with `vec`, which returns an
+# `ArrayPartition` unchanged since it is already an `AbstractVector`, and its generic
+# `repack` cannot rebuild one. The extension implements the interface properly.
+@testset "SciMLStructures interface for ArrayPartition" begin
+    p = ArrayPartition([1.0, 2.0], [3.0, 4.0, 5.0])
+
+    @testset "traits" begin
+        @test isscimlstructure(p)
+        @test ismutablescimlstructure(p)
+        @test hasportion(Tunable(), p)
+    end
+
+    @testset "canonicalize flattens in the documented order" begin
+        buffer, repack, aliases = canonicalize(Tunable(), p)
+        @test buffer isa AbstractVector
+        @test !(buffer isa ArrayPartition)
+        @test length(buffer) == length(p)
+        # The partitions are laid out one after another, matching `collect`.
+        @test buffer == collect(p)
+        @test buffer == [1.0, 2.0, 3.0, 4.0, 5.0]
+        # The partitions are separate arrays, so the flat buffer has to be a copy.
+        @test aliases == false
+        buffer[1] = -1.0
+        @test p.x[1][1] == 1.0
+    end
+
+    @testset "repack rebuilds the partitioning" begin
+        _, repack, _ = canonicalize(Tunable(), p)
+        back = repack([9.0, 8.0, 7.0, 6.0, 5.0])
+        @test back isa ArrayPartition
+        @test map(length, back.x) == map(length, p.x)
+        @test collect(back) == [9.0, 8.0, 7.0, 6.0, 5.0]
+        # The original is untouched.
+        @test collect(p) == [1.0, 2.0, 3.0, 4.0, 5.0]
+    end
+
+    @testset "replace matches repack" begin
+        _, repack, _ = canonicalize(Tunable(), p)
+        new_values = [9.0, 8.0, 7.0, 6.0, 5.0]
+        @test collect(SciMLStructures.replace(Tunable(), p, new_values)) ==
+            collect(repack(new_values))
+    end
+
+    @testset "replace! mutates in place" begin
+        q = ArrayPartition([0.0, 0.0], [0.0, 0.0, 0.0])
+        parts = q.x
+        @test SciMLStructures.replace!(Tunable(), q, [1.0, 2.0, 3.0, 4.0, 5.0]) === nothing
+        @test collect(q) == [1.0, 2.0, 3.0, 4.0, 5.0]
+        # The same partition arrays were written into, not replaced.
+        @test q.x[1] === parts[1]
+        @test q.x[2] === parts[2]
+    end
+
+    @testset "wrong lengths are rejected" begin
+        @test_throws DimensionMismatch SciMLStructures.replace(Tunable(), p, [1.0])
+        @test_throws DimensionMismatch SciMLStructures.replace!(
+            Tunable(), ArrayPartition([0.0], [0.0]), [1.0]
+        )
+    end
+
+    @testset "other portions are absent" begin
+        for portion in (Constants(), Caches(), Discrete(), Initials(), Input())
+            @test !hasportion(portion, p)
+            @test canonicalize(portion, p) == (nothing, nothing, nothing)
+        end
+    end
+
+    @testset "uneven and single partitions" begin
+        for q in (
+                ArrayPartition([1.0]),
+                ArrayPartition([1.0], [2.0, 3.0], [4.0, 5.0, 6.0]),
+                ArrayPartition([1.0, 2.0], Float64[], [3.0]),
+            )
+            buffer, repack, _ = canonicalize(Tunable(), q)
+            @test buffer == collect(q)
+            back = repack(collect(Float64, 1:length(q)))
+            @test map(length, back.x) == map(length, q.x)
+            @test collect(back) == collect(Float64, 1:length(q))
+        end
+    end
+
+    @testset "the partition array type is preserved" begin
+        # Concatenating keeps a uniformly static partitioning static rather than
+        # dropping it to a `Vector`, and `replace` rebuilds each partition in kind.
+        q = ArrayPartition(MVector(1.0, 2.0), MVector(3.0, 4.0))
+        buffer, repack, _ = canonicalize(Tunable(), q)
+        @test buffer == [1.0, 2.0, 3.0, 4.0]
+        @test buffer isa MVector
+        back = repack([5.0, 6.0, 7.0, 8.0])
+        @test back.x[1] isa MVector
+        @test collect(back) == [5.0, 6.0, 7.0, 8.0]
+    end
+
+    # The partitions need not share an array type or an element type, so the buffer
+    # has to promote across them rather than take its type from one of them.
+    @testset "partitions of differing types" begin
+        mixed_eltype = ArrayPartition([1, 2], [3.0, 4.5])
+        buffer, repack, _ = canonicalize(Tunable(), mixed_eltype)
+        # Sizing from the first partition would give a `Vector{Int}` and truncate.
+        @test eltype(buffer) === Float64
+        @test buffer == [1.0, 2.0, 3.0, 4.5]
+
+        mixed_array = ArrayPartition([1.0, 2.0], MVector(3.0, 4.0))
+        buffer2, _, _ = canonicalize(Tunable(), mixed_array)
+        @test buffer2 == [1.0, 2.0, 3.0, 4.0]
+        @test length(buffer2) == 4
+    end
+
+    # `reduce(vcat, (x,))` returns `x` itself, so the single-partition case has to
+    # copy or the buffer would alias `p` despite `aliases` being reported `false`.
+    @testset "a single partition still yields an independent buffer" begin
+        q = ArrayPartition([1.0, 2.0, 3.0])
+        buffer, _, aliases = canonicalize(Tunable(), q)
+        @test aliases == false
+        @test buffer !== q.x[1]
+        buffer[1] = -1.0
+        @test q.x[1][1] == 1.0
+    end
+
+    @testset "integer partitions" begin
+        q = ArrayPartition([1, 2], [3, 4])
+        buffer, repack, _ = canonicalize(Tunable(), q)
+        @test buffer == [1, 2, 3, 4]
+        @test collect(repack([5, 6, 7, 8])) == [5, 6, 7, 8]
+    end
+end

test/runtests.jl

@@ -23,6 +23,7 @@ run_tests(;
         @time @safetestset "Table traits" include("Core/tabletraits.jl")
         @time @safetestset "StaticArrays Tests" include("Core/copy_static_array_test.jl")
         @time @safetestset "Linear Algebra Tests" include("Core/linalg.jl")
+        @time @safetestset "SciMLStructures Tests" include("Core/scimlstructures_test.jl")
         return @time @safetestset "Measurement Tests" include("Core/measurements.jl")
     end,
     groups = Dict(