Merge pull request #130 from SciML/multisize

Handle multisize undefs

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "LabelledArrays"
 uuid = "2ee39098-c373-598a-b85f-a56591580800"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "1.12.1"
+version = "1.12.2"
 
 [deps]
 ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2"

src/LabelledArrays.jl

@@ -64,8 +64,11 @@ function ArrayInterfaceCore.ismutable(::Type{<:LArray{T, N, D, Syms}}) where {T,
 end
 ArrayInterfaceCore.can_setindex(::Type{<:SLArray}) = false
 
+lenfun(x) = length(x)
+lenfun(::Symbol) = 1
 function ArrayInterfaceCore.undefmatrix(x::LArray{T, N, D, Syms}) where {T, N, D, Syms}
-    similar(x.__x, length(Syms), length(Syms))
+    n = sum(lenfun, Syms)
+    similar(x.__x, n, n)
 end
 
 function PreallocationTools.get_tmp(dc::PreallocationTools.DiffCache,

src/larray.jl

@@ -5,7 +5,7 @@ struct LArray{T, N, D <: AbstractArray{T, N}, Syms} <: DenseArray{T, N}
 end
 
 function LArray{T, N, D, Syms}(::UndefInitializer, n::Int64) where {T, N, D, Syms}
-    @assert length(Syms) == n
+    @assert sum(lenfun, Syms) == n
     LArray{T, N, D, Syms}(similar(D, n))
 end
 

test/larrays.jl

@@ -16,6 +16,10 @@ using LabelledArrays, Test, InteractiveUtils
     syms = (:a, :b, :c)
     @test typeof(typeof(x)(undef, 3)) == typeof(x)
 
+    z = LabelledArrays.LArray{Float64, 1, Vector{Float64}, (x = 1:4, y = 5:8)}(undef, 8)
+    @test length(z) == 8
+    @test size(LabelledArrays.ArrayInterfaceCore.undefmatrix(z)) == (8, 8)
+
     for (i, s) in enumerate(syms)
         @show i, s
         @test x[i] == x[s]