diff --git a/Project.toml b/Project.toml index cf07790..cac7537 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "BlockSparseArrays" uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4" authors = ["ITensor developers and contributors"] -version = "0.7.14" +version = "0.7.15" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/abstractblocksparsearray/abstractblocksparsearray.jl b/src/abstractblocksparsearray/abstractblocksparsearray.jl index f82e2ca..3d688e4 100644 --- a/src/abstractblocksparsearray/abstractblocksparsearray.jl +++ b/src/abstractblocksparsearray/abstractblocksparsearray.jl @@ -63,6 +63,19 @@ function Base.setindex!(a::AbstractBlockSparseArray{<:Any,0}, value, ::Block{0}) return a end +# Custom `_convert` works around the issue that +# `convert(::Type{<:Diagonal}, ::AbstractMatrix)` isnt' defined +# in Julia v1.10 (https://github.com/JuliaLang/julia/pull/48895, +# https://github.com/JuliaLang/julia/pull/52487). +# TODO: Delete once we drop support for Julia v1.10. +_convert(::Type{T}, a::AbstractArray) where {T} = convert(T, a) +using LinearAlgebra: LinearAlgebra, Diagonal, diag, isdiag +_construct(T::Type{<:Diagonal}, a::AbstractMatrix) = T(diag(a)) +function _convert(T::Type{<:Diagonal}, a::AbstractMatrix) + LinearAlgebra.checksquare(a) + return isdiag(a) ? _construct(T, a) : throw(InexactError(:convert, T, a)) +end + function Base.setindex!( a::AbstractBlockSparseArray{<:Any,N}, value, I::Vararg{Block{1},N} ) where {N} @@ -74,7 +87,12 @@ function Base.setindex!( ), ) end - blocks(a)[Int.(I)...] = value + # Custom `_convert` works around the issue that + # `convert(::Type{<:Diagonal}, ::AbstractMatrix)` isnt' defined + # in Julia v1.10 (https://github.com/JuliaLang/julia/pull/48895, + # https://github.com/JuliaLang/julia/pull/52487). + # TODO: Delete once we drop support for Julia v1.10. + blocks(a)[Int.(I)...] = _convert(blocktype(a), value) return a end diff --git a/src/abstractblocksparsearray/arraylayouts.jl b/src/abstractblocksparsearray/arraylayouts.jl index 684fe64..4476438 100644 --- a/src/abstractblocksparsearray/arraylayouts.jl +++ b/src/abstractblocksparsearray/arraylayouts.jl @@ -23,13 +23,37 @@ function ArrayLayouts.MemoryLayout( end function Base.similar( - mul::MulAdd{<:BlockLayout{<:SparseLayout},<:BlockLayout{<:SparseLayout},<:Any,<:Any,A,B}, + mul::MulAdd{ + <:BlockLayout{<:SparseLayout,BlockLayoutA}, + <:BlockLayout{<:SparseLayout,BlockLayoutB}, + LayoutC, + T, + A, + B, + C, + }, elt::Type, axes, -) where {A,B} - # TODO: Use something like `Base.promote_op(*, A, B)` to determine the output block type. - output_blocktype = similartype(blocktype(A), Type{elt}, Tuple{blockaxistype.(axes)...}) - return similar(BlockSparseArray{elt,length(axes),output_blocktype}, axes) +) where {BlockLayoutA,BlockLayoutB,LayoutC,T,A,B,C} + + # TODO: Consider using this instead: + # ```julia + # blockmultype = MulAdd{BlockLayoutA,BlockLayoutB,LayoutC,T,blocktype(A),blocktype(B),C} + # output_blocktype = Base.promote_op( + # similar, blockmultype, Type{elt}, Tuple{eltype.(eachblockaxis.(axes))...} + # ) + # ``` + # The issue is that it in some cases it seems to lose some information about the block types. + + # TODO: Maybe this should be: + # output_blocktype = Base.promote_op( + # mul!, blocktype(mul.A), blocktype(mul.B), blocktype(mul.C), typeof(mul.α), typeof(mul.β) + # ) + + output_blocktype = Base.promote_op(*, blocktype(mul.A), blocktype(mul.B)) + output_blocktype′ = + !isconcretetype(output_blocktype) ? AbstractMatrix{elt} : output_blocktype + return similar(BlockSparseArray{elt,length(axes),output_blocktype′}, axes) end # Materialize a SubArray view. diff --git a/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl index 9a4a269..c7d2888 100644 --- a/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl +++ b/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl @@ -231,8 +231,11 @@ end function blocksparse_similar(a, elt::Type, axes::Tuple) ndims = length(axes) - blockt = similartype(blocktype(a), Type{elt}, Tuple{blockaxistype.(axes)...}) - return BlockSparseArray{elt,ndims,blockt}(undef, axes) + # TODO: Define a version of `similartype` that catches the case + # where the output isn't concrete and returns an `AbstractArray`. + blockt = Base.promote_op(similar, blocktype(a), Type{elt}, Tuple{blockaxistype.(axes)...}) + blockt′ = !isconcretetype(blockt) ? AbstractArray{elt,ndims} : blockt + return BlockSparseArray{elt,ndims,blockt′}(undef, axes) end @interface ::AbstractBlockSparseArrayInterface function Base.similar( a::AbstractArray, elt::Type, axes::Tuple{Vararg{Int}} @@ -283,13 +286,11 @@ function Base.similar( elt::Type, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}, ) - # TODO: Use `@interface interface(a) similar(...)`. return @interface interface(a) similar(a, elt, axes) end # Fixes ambiguity error. function Base.similar(a::AnyAbstractBlockSparseArray, elt::Type, axes::Tuple{}) - # TODO: Use `@interface interface(a) similar(...)`. return @interface interface(a) similar(a, elt, axes) end @@ -301,7 +302,6 @@ function Base.similar( AbstractBlockedUnitRange{<:Integer},Vararg{AbstractBlockedUnitRange{<:Integer}} }, ) - # TODO: Use `@interface interface(a) similar(...)`. return @interface interface(a) similar(a, elt, axes) end @@ -311,7 +311,6 @@ function Base.similar( elt::Type, axes::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}, ) - # TODO: Use `@interface interface(a) similar(...)`. return @interface interface(a) similar(a, elt, axes) end @@ -321,9 +320,17 @@ function Base.similar( elt::Type, axes::Tuple{AbstractBlockedUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}, ) - # TODO: Use `@interface interface(a) similar(...)`. return @interface interface(a) similar(a, elt, axes) end +function Base.similar(a::AnyAbstractBlockSparseArray, elt::Type) + return @interface interface(a) similar(a, elt, axes(a)) +end +function Base.similar( + a::AnyAbstractBlockSparseArray, + axes::Tuple{AbstractBlockedUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}, +) + return @interface interface(a) similar(a, eltype(a), axes) +end # Fixes ambiguity errors with BlockArrays. function Base.similar( @@ -343,7 +350,6 @@ end function Base.similar( a::AnyAbstractBlockSparseArray, elt::Type, axes::Tuple{Base.OneTo,Vararg{Base.OneTo}} ) - # TODO: Use `@interface interface(a) similar(...)`. return @interface interface(a) similar(a, elt, axes) end diff --git a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index b4ce403..9492bac 100644 --- a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -42,6 +42,14 @@ function eachblockstoredindex(a::AbstractArray) return Block.(Tuple.(eachstoredindex(blocks(a)))) end +function SparseArraysBase.isstored(a::AbstractArray, I1::Block{1}, Irest::Block{1}...) + I = (I1, Irest...) + return isstored(blocks(a), Int.(I)...) +end +function SparseArraysBase.isstored(a::AbstractArray{<:Any,N}, I::Block{N}) where {N} + return isstored(a, Tuple(I)...) +end + using DiagonalArrays: diagindices # Block version of `DiagonalArrays.diagindices`. function blockdiagindices(a::AbstractArray) diff --git a/src/blocksparsearrayinterface/getunstoredblock.jl b/src/blocksparsearrayinterface/getunstoredblock.jl index d221286..fcb0760 100644 --- a/src/blocksparsearrayinterface/getunstoredblock.jl +++ b/src/blocksparsearrayinterface/getunstoredblock.jl @@ -11,6 +11,7 @@ end ax = ntuple(N) do d return only(axes(f.axes[d][Block(I[d])])) end + !isconcretetype(A) && return zero!(similar(Array{eltype(A),N}, ax)) return zero!(similar(A, ax)) end @inline function (f::GetUnstoredBlock)( diff --git a/src/blocksparsearrayinterface/map.jl b/src/blocksparsearrayinterface/map.jl index 60a6bc8..3798254 100644 --- a/src/blocksparsearrayinterface/map.jl +++ b/src/blocksparsearrayinterface/map.jl @@ -18,6 +18,38 @@ function union_eachblockstoredindex(as::AbstractArray...) return ∪(map(eachblockstoredindex, as)...) end +# Get a view of a block assuming it is stored. +function viewblock_stored(a::AbstractArray{<:Any,N}, I::Vararg{Block{1},N}) where {N} + return blocks(a)[Int.(I)...] +end +function viewblock_stored(a::AbstractArray{<:Any,N}, I::Block{N}) where {N} + return viewblock_stored(a, Tuple(I)...) +end + +using FillArrays: Zeros +# Get a view of a block if it is stored, otherwise return a lazy zeros. +function viewblock_or_zeros(a::AbstractArray{<:Any,N}, I::Vararg{Block{1},N}) where {N} + if isstored(a, I...) + return viewblock_stored(a, I...) + else + block_ax = map((ax, i) -> eachblockaxis(ax)[Int(i)], axes(a), I) + return Zeros{eltype(a)}(block_ax) + end +end +function viewblock_or_zeros(a::AbstractArray{<:Any,N}, I::Block{N}) where {N} + return viewblock_or_zeros(a, Tuple(I)...) +end + +function map_block!(f, a_dest::AbstractArray, I::Block, a_srcs::AbstractArray...) + a_srcs_I = map(a_src -> viewblock_or_zeros(a_src, I), a_srcs) + if isstored(a_dest, I) + a_dest[I] .= f.(a_srcs_I...) + else + a_dest[I] = Broadcast.broadcast_preserving_zero_d(f, a_srcs_I...) + end + return a_dest +end + function map_blockwise!(f, a_dest::AbstractArray, a_srcs::AbstractArray...) # TODO: This assumes element types are numbers, generalize this logic. f_preserves_zeros = f(zero.(eltype.(a_srcs))...) == zero(eltype(a_dest)) @@ -27,22 +59,7 @@ function map_blockwise!(f, a_dest::AbstractArray, a_srcs::AbstractArray...) BlockRange(a_dest) end for I in Is - # TODO: Use: - # block_dest = @view a_dest[I] - # or: - # block_dest = @view! a_dest[I] - block_dest = blocks_maybe_single(a_dest)[Int.(Tuple(I))...] - # TODO: Use: - # block_srcs = map(a_src -> @view(a_src[I]), a_srcs) - block_srcs = map(a_srcs) do a_src - return blocks_maybe_single(a_src)[Int.(Tuple(I))...] - end - # TODO: Use `map!!` to handle immutable blocks. - map!(f, block_dest, block_srcs...) - # Replace the entire block, handles initializing new blocks - # or if blocks are immutable. - # TODO: Use `a_dest[I] = block_dest`. - blocks(a_dest)[Int.(Tuple(I))...] = block_dest + map_block!(f, a_dest, I, a_srcs...) end return a_dest end @@ -151,8 +168,12 @@ end function map_stored_blocks(f, a::AbstractArray) block_stored_indices = collect(eachblockstoredindex(a)) if isempty(block_stored_indices) + eltype_a′ = Base.promote_op(f, eltype(a)) blocktype_a′ = Base.promote_op(f, blocktype(a)) - return BlockSparseArray{eltype(blocktype_a′),ndims(a),blocktype_a′}(undef, axes(a)) + eltype_a′′ = !isconcretetype(eltype_a′) ? Any : eltype_a′ + blocktype_a′′ = + !isconcretetype(blocktype_a′) ? AbstractArray{eltype_a′′,ndims(a)} : blocktype_a′ + return BlockSparseArray{eltype_a′′,ndims(a),blocktype_a′′}(undef, axes(a)) end stored_blocks = map(B -> f(@view!(a[B])), block_stored_indices) blocktype_a′ = eltype(stored_blocks) diff --git a/src/factorizations/svd.jl b/src/factorizations/svd.jl index b57b8eb..19e5c6c 100644 --- a/src/factorizations/svd.jl +++ b/src/factorizations/svd.jl @@ -22,10 +22,21 @@ function MatrixAlgebraKit.default_svd_algorithm( return BlockPermutedDiagonalAlgorithm(alg) end +function output_type( + ::typeof(svd_compact!), + A::Type{<:AbstractMatrix{T}}, + Alg::Type{<:MatrixAlgebraKit.AbstractAlgorithm}, +) where {T} + USVᴴ = Base.promote_op(svd_compact!, A, Alg) + !isconcretetype(USVᴴ) && + return Tuple{AbstractMatrix{T},AbstractMatrix{realtype(T)},AbstractMatrix{T}} + return USVᴴ +end + function similar_output( ::typeof(svd_compact!), A, S_axes, alg::MatrixAlgebraKit.AbstractAlgorithm ) - BU, BS, BVᴴ = fieldtypes(Base.promote_op(svd_compact!, blocktype(A), typeof(alg.alg))) + BU, BS, BVᴴ = fieldtypes(output_type(svd_compact!, blocktype(A), typeof(alg.alg))) U = similar(A, BlockType(BU), (axes(A, 1), S_axes[1])) S = similar(A, BlockType(BS), S_axes) Vᴴ = similar(A, BlockType(BVᴴ), (S_axes[2], axes(A, 2))) diff --git a/test/Project.toml b/test/Project.toml index 8c52cd7..536c8dd 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -4,6 +4,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a" BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" BlockSparseArrays = "2c9a651f-6452-4ace-a6ac-809f4280fbb4" +DerivableInterfaces = "6c5e35bf-e59e-4898-b73c-732dcc4ba65f" DiagonalArrays = "74fd4be6-21e2-4f6f-823a-4360d37c7a77" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" @@ -25,6 +26,7 @@ Aqua = "0.8" ArrayLayouts = "1" BlockArrays = "1" BlockSparseArrays = "0.7" +DerivableInterfaces = "0.5" DiagonalArrays = "0.3" GPUArraysCore = "0.2" JLArrays = "0.2" diff --git a/test/test_abstract_blocktype.jl b/test/test_abstract_blocktype.jl new file mode 100644 index 0000000..2107e25 --- /dev/null +++ b/test/test_abstract_blocktype.jl @@ -0,0 +1,44 @@ +using Adapt: adapt +using BlockArrays: Block +using BlockSparseArrays: BlockSparseMatrix, blockstoredlength +using JLArrays: JLArray +using SparseArraysBase: storedlength +using Test: @test, @test_broken, @testset + +elts = (Float32, Float64, ComplexF32) +arrayts = (Array, JLArray) +@testset "Abstract block type (arraytype=$arrayt, eltype=$elt)" for arrayt in arrayts, + elt in elts + + dev = adapt(arrayt) + a = BlockSparseMatrix{elt,AbstractMatrix{elt}}(undef, [2, 3], [2, 3]) + @test sprint(show, MIME"text/plain"(), a) isa String + @test iszero(storedlength(a)) + @test iszero(blockstoredlength(a)) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + a[Block(2, 2)] = dev(randn(elt, 3, 3)) + @test !iszero(a[Block(1, 1)]) + @test a[Block(1, 1)] isa arrayt{elt,2} + @test !iszero(a[Block(2, 2)]) + @test a[Block(2, 2)] isa arrayt{elt,2} + @test iszero(a[Block(2, 1)]) + @test a[Block(2, 1)] isa Matrix{elt} + @test iszero(a[Block(1, 2)]) + @test a[Block(1, 2)] isa Matrix{elt} + + b = copy(a) + @test Array(b) ≈ Array(a) + + b = a + a + @test Array(b) ≈ Array(a) + Array(a) + + b = 3a + @test Array(b) ≈ 3Array(a) + + if arrayt === Array + b = a * a + @test Array(b) ≈ Array(a) * Array(a) + else + @test_broken a * a + end +end diff --git a/test/test_basics.jl b/test/test_basics.jl index cde6f2e..a464fa3 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -46,33 +46,13 @@ using TestExtras: @constinferred using TypeParameterAccessors: TypeParameterAccessors, Position include("TestBlockSparseArraysUtils.jl") -@testset "similartype_unchecked" begin - @test @constinferred(similartype_unchecked(Array{Float32}, NTuple{2,Int})) === - Matrix{Float32} - @test @constinferred(similartype_unchecked(Array{Float32}, NTuple{2,Base.OneTo{Int}})) === - Matrix{Float32} - if VERSION < v"1.11-" - # Not type stable in Julia 1.10. - @test similartype_unchecked(AbstractArray{Float32}, NTuple{2,Int}) === Matrix{Float32} - @test similartype_unchecked(JLArray{Float32}, NTuple{2,Int}) === JLMatrix{Float32} - @test similartype_unchecked(JLArray{Float32}, NTuple{2,Base.OneTo{Int}}) === - JLMatrix{Float32} - else - @test @constinferred(similartype_unchecked(AbstractArray{Float32}, NTuple{2,Int})) === - Matrix{Float32} - @test @constinferred(similartype_unchecked(JLArray{Float32}, NTuple{2,Int})) === - JLMatrix{Float32} - @test @constinferred( - similartype_unchecked(JLArray{Float32}, NTuple{2,Base.OneTo{Int}}) - ) === JLMatrix{Float32} - end -end - +elts = (Float32, Float64, ComplexF32) arrayts = (Array, JLArray) -@testset "BlockSparseArrays (arraytype=$arrayt, eltype=$elt)" for arrayt in arrayts, - elt in (Float32, Float64, Complex{Float32}, Complex{Float64}) +@testset "BlockSparseArrays basics (arraytype=$arrayt, eltype=$elt)" for arrayt in arrayts, + elt in elts + + dev = adapt(arrayt) - dev(a) = adapt(arrayt, a) @testset "Broken" begin # TODO: Fix this and turn it into a proper test. a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) @@ -411,774 +391,6 @@ arrayts = (Array, JLArray) @test a_jl[Block(1, 2)] isa JLMatrix{elt} @test adapt(Array, a_jl[Block(1, 2)]) == a_12 end - @testset "Tensor algebra" begin - a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = dev(randn(elt, size(a[b]))) - end - @test eltype(a) == elt - @test blockstoredlength(a) == 2 - @test storedlength(a) == 2 * 4 + 3 * 3 - - # TODO: Broken on GPU. - if arrayt ≠ Array - a = dev(BlockSparseArray{elt}(undef, [2, 3], [3, 4])) - @test_broken a[Block(1, 2)] .= 2 - end - - # TODO: Broken on GPU. - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - a[Block(1, 2)] .= 2 - @test eltype(a) == elt - @test all(==(2), a[Block(1, 2)]) - @test iszero(a[Block(1, 1)]) - @test iszero(a[Block(2, 1)]) - @test iszero(a[Block(2, 2)]) - @test blockstoredlength(a) == 1 - @test storedlength(a) == 2 * 4 - - # TODO: Broken on GPU. - if arrayt ≠ Array - a = dev(BlockSparseArray{elt}(undef, [2, 3], [3, 4])) - @test_broken a[Block(1, 2)] .= 0 - end - - # TODO: Broken on GPU. - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - a[Block(1, 2)] .= 0 - @test eltype(a) == elt - @test iszero(a[Block(1, 1)]) - @test iszero(a[Block(2, 1)]) - @test iszero(a[Block(1, 2)]) - @test iszero(a[Block(2, 2)]) - @test blockstoredlength(a) == 1 - @test storedlength(a) == 2 * 4 - - # Test similar on broadcasted expressions. - a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) - bc = Broadcast.broadcasted(+, a, a) - a′ = similar(bc, Float32) - @test a′ isa BlockSparseArray{Float32} - @test blocktype(a′) <: arrayt{Float32,2} - @test axes(a) == (blockedrange([2, 3]), blockedrange([3, 4])) - - # Test similar on broadcasted expressions with axes specified. - a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) - bc = Broadcast.broadcasted(+, a, a) - a′ = similar( - bc, Float32, (blockedrange([2, 4]), blockedrange([2, 5]), blockedrange([2, 2])) - ) - @test a′ isa BlockSparseArray{Float32} - @test blocktype(a′) <: arrayt{Float32,3} - @test axes(a′) == (blockedrange([2, 4]), blockedrange([2, 5]), blockedrange([2, 2])) - - a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = dev(randn(elt, size(a[b]))) - end - b = similar(a, complex(elt)) - @test eltype(b) == complex(eltype(a)) - @test iszero(b) - @test blockstoredlength(b) == 0 - @test storedlength(b) == 0 - @test size(b) == size(a) - @test blocksize(b) == blocksize(a) - - # Regression test for https://github.com/ITensor/BlockSparseArrays.jl/issues/98 - a = dev(BlockSparseArray{elt}(undef)) - b = similar(a, Float64, (Base.OneTo(1),)) - @test b isa BlockSparseVector{Float64} - @test size(b) == (1,) - @test blocksize(b) == (1,) - - a = dev(BlockSparseArray{elt}(undef, [2, 3], [3, 4])) - b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]] - c = @view b[Block(1, 1)] - @test iszero(a) - @test iszero(storedlength(a)) - @test iszero(b) - @test iszero(storedlength(b)) - # TODO: Broken on GPU. - @test iszero(c) broken = arrayt ≠ Array - @test iszero(storedlength(c)) - @allowscalar a[5, 7] = 1 - @test !iszero(a) - @test storedlength(a) == 3 * 4 - @test !iszero(b) - @test storedlength(b) == 3 * 4 - # TODO: Broken on GPU. - @test !iszero(c) broken = arrayt ≠ Array - @test storedlength(c) == 3 * 4 - d = @view a[1:4, 1:6] - @test iszero(d) - @test storedlength(d) == 2 * 3 - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = copy(a) - b[1, 1] = 11 - @test b[1, 1] == 11 - @test a[1, 1] ≠ 11 - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = copy(a) - b .*= 2 - @test b ≈ 2a - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = copy(a) - b ./= 2 - @test b ≈ a / 2 - - a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = dev(randn(elt, size(a[b]))) - end - b = 2 * a - @allowscalar @test Array(b) ≈ 2 * Array(a) - @test eltype(b) == elt - @test blockstoredlength(b) == 2 - @test storedlength(b) == 2 * 4 + 3 * 3 - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = (2 + 3im) * a - @test Array(b) ≈ (2 + 3im) * Array(a) - @test eltype(b) == complex(elt) - @test blockstoredlength(b) == 2 - @test storedlength(b) == 2 * 4 + 3 * 3 - - a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = dev(randn(elt, size(a[b]))) - end - b = a + a - @allowscalar @test Array(b) ≈ 2 * Array(a) - @test eltype(b) == elt - @test blockstoredlength(b) == 2 - @test storedlength(b) == 2 * 4 + 3 * 3 - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - x = BlockSparseArray{elt}(undef, ([3, 4], [2, 3])) - @views for b in [Block(1, 2), Block(2, 1)] - x[b] = randn(elt, size(x[b])) - end - b = a .+ a .+ 3 .* PermutedDimsArray(x, (2, 1)) - @test Array(b) ≈ 2 * Array(a) + 3 * permutedims(Array(x), (2, 1)) - @test eltype(b) == elt - @test blockstoredlength(b) == 2 - @test storedlength(b) == 2 * 4 + 3 * 3 - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = permutedims(a, (2, 1)) - @test Array(b) ≈ permutedims(Array(a), (2, 1)) - @test eltype(b) == elt - @test blockstoredlength(b) == 2 - @test storedlength(b) == 2 * 4 + 3 * 3 - - a = dev(BlockSparseArray{elt}(undef, [1, 1, 1], [1, 2, 3], [2, 2, 1], [1, 2, 1])) - a[Block(3, 2, 2, 3)] = dev(randn(elt, 1, 2, 2, 1)) - perm = (2, 3, 4, 1) - for b in (PermutedDimsArray(a, perm), permutedims(a, perm)) - @test @allowscalar(Array(b)) == permutedims(Array(a), perm) - @test issetequal(eachblockstoredindex(b), [Block(2, 2, 3, 3)]) - @test @allowscalar b[Block(2, 2, 3, 3)] == permutedims(a[Block(3, 2, 2, 3)], perm) - end - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = map(x -> 2x, a) - @test Array(b) ≈ 2 * Array(a) - @test eltype(b) == elt - @test size(b) == size(a) - @test blocksize(b) == (2, 2) - @test blockstoredlength(b) == 2 - @test storedlength(b) == 2 * 4 + 3 * 3 - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = a[[Block(2), Block(1)], [Block(2), Block(1)]] - @test b[Block(1, 1)] == a[Block(2, 2)] - @test b[Block(1, 2)] == a[Block(2, 1)] - @test b[Block(2, 1)] == a[Block(1, 2)] - @test b[Block(2, 2)] == a[Block(1, 1)] - @test size(b) == size(a) - @test blocksize(b) == (2, 2) - @test storedlength(b) == storedlength(a) - @test blockstoredlength(b) == 2 - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = a[Block(1):Block(2), Block(1):Block(2)] - @test b == a - @test size(b) == size(a) - @test blocksize(b) == (2, 2) - @test storedlength(b) == storedlength(a) - @test blockstoredlength(b) == 2 - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = a[Block(1):Block(1), Block(1):Block(2)] - @test b == Array(a)[1:2, 1:end] - @test b[Block(1, 1)] == a[Block(1, 1)] - @test b[Block(1, 2)] == a[Block(1, 2)] - @test size(b) == (2, 7) - @test blocksize(b) == (1, 2) - @test storedlength(b) == storedlength(a[Block(1, 2)]) - @test blockstoredlength(b) == 1 - - a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = dev(randn(elt, size(a[b]))) - end - I = 2:4 - I′, J′ = falses.(size(a)) - I′[I] .= true - J′[I] .= true - for b in (a[I, I], @view(a[I, I]), a[I′, J′], @view(a[I′, J′])) - @allowscalar @test b == Array(a)[2:4, 2:4] - @test size(b) == (3, 3) - @test blocksize(b) == (2, 2) - @test storedlength(b) == 1 * 1 + 2 * 2 - @test blockstoredlength(b) == 2 - for f in (getindex, view) - # TODO: Broken on GPU. - @allowscalar begin - @test size(f(b, Block(1, 1))) == (1, 2) - @test size(f(b, Block(2, 1))) == (2, 2) - @test size(f(b, Block(1, 2))) == (1, 1) - @test size(f(b, Block(2, 2))) == (2, 1) - @test f(b, Block(1, 1)) == a[Block(1, 1)[2:2, 2:3]] - @test f(b, Block(2, 1)) == a[Block(2, 1)[1:2, 2:3]] - @test f(b, Block(1, 2)) == a[Block(1, 2)[2:2, 1:1]] - @test f(b, Block(2, 2)) == a[Block(2, 2)[1:2, 1:1]] - end - end - end - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = a[Block(2, 1)[1:2, 2:3]] - @test b == Array(a)[3:4, 2:3] - @test size(b) == (2, 2) - @test blocksize(b) == (1, 1) - @test storedlength(b) == 2 * 2 - @test blockstoredlength(b) == 1 - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = PermutedDimsArray(a, (2, 1)) - @test blockstoredlength(b) == 2 - @test Array(b) == permutedims(Array(a), (2, 1)) - c = 2 * b - @test blockstoredlength(c) == 2 - @test Array(c) == 2 * permutedims(Array(a), (2, 1)) - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = a' - @test blockstoredlength(b) == 2 - @test Array(b) == Array(a)' - c = 2 * b - @test blockstoredlength(c) == 2 - @test Array(c) == 2 * Array(a)' - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = transpose(a) - @test blockstoredlength(b) == 2 - @test Array(b) == transpose(Array(a)) - c = 2 * b - @test blockstoredlength(c) == 2 - @test Array(c) == 2 * transpose(Array(a)) - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = a[Block(1), Block(1):Block(2)] - @test size(b) == (2, 7) - @test blocksize(b) == (1, 2) - @test b[Block(1, 1)] == a[Block(1, 1)] - @test b[Block(1, 2)] == a[Block(1, 2)] - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = copy(a) - x = randn(elt, size(@view(a[Block(2, 2)]))) - b[Block(2), Block(2)] = x - @test b[Block(2, 2)] == x - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = copy(a) - b[Block(1, 1)] .= 1 - @test b[Block(1, 1)] == trues(blocksizes(b)[1, 1]) - - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - b = @view a[Block(2, 2)] - @test size(b) == (3, 4) - for i in parentindices(b) - @test i isa Base.OneTo{Int} - end - @test parentindices(b)[1] == 1:3 - @test parentindices(b)[2] == 1:4 - - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - b = @view a[Block(2, 2)[1:2, 2:2]] - @test size(b) == (2, 1) - for i in parentindices(b) - @test i isa UnitRange{Int} - end - @test parentindices(b)[1] == 1:2 - @test parentindices(b)[2] == 2:2 - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - x = randn(elt, 1, 2) - @view(a[Block(2, 2)])[1:1, 1:2] = x - @test a[Block(2, 2)][1:1, 1:2] == x - @test @view(a[Block(2, 2)])[1:1, 1:2] == x - @test a[3:3, 4:5] == x - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - x = randn(elt, 1, 2) - @views a[Block(2, 2)][1:1, 1:2] = x - @test a[Block(2, 2)][1:1, 1:2] == x - @test @view(a[Block(2, 2)])[1:1, 1:2] == x - @test a[3:3, 4:5] == x - - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - b = @views a[Block(2, 2)][1:2, 2:3] - @test b isa SubArray{<:Any,<:Any,<:BlockView} - for i in parentindices(b) - @test i isa UnitRange{Int} - end - x = randn(elt, 2, 2) - b .= x - @test a[Block(2, 2)[1:2, 2:3]] == x - @test a[Block(2, 2)[1:2, 2:3]] == b - @test blockstoredlength(a) == 1 - - # Noncontiguous slicing. - a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - a[Block(1, 1)] = dev(randn(elt, 2, 2)) - a[Block(2, 2)] = dev(randn(elt, 3, 3)) - I = ([3, 5], [2, 4]) - @test Array(a[I...]) == Array(a)[I...] - - # Noncontiguous slicing. - a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - a[Block(1, 1)] = dev(randn(elt, 2, 2)) - a[Block(2, 2)] = dev(randn(elt, 3, 3)) - I = (:, [2, 4]) - @test Array(a[I...]) == Array(a)[I...] - - a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) - @views for b in [Block(1, 1), Block(2, 2)] - a[b] = randn(elt, size(a[b])) - end - for I in (Block.(1:2), [Block(1), Block(2)]) - b = @view a[I, I] - for I in CartesianIndices(a) - @test b[I] == a[I] - end - for block in BlockRange(a) - @test b[block] == a[block] - end - end - - a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - @views for b in [Block(1, 1), Block(2, 2)] - # TODO: Use `blocksizes(a)[Int.(Tuple(b))...]` once available. - a[b] = dev(randn(elt, size(a[b]))) - end - for I in ([Block(2), Block(1)],) - b = @view a[I, I] - @test b[Block(1, 1)] == a[Block(2, 2)] - @test b[Block(2, 1)] == a[Block(1, 2)] - @test b[Block(1, 2)] == a[Block(2, 1)] - @test b[Block(2, 2)] == a[Block(1, 1)] - @allowscalar begin - @test b[1, 1] == a[3, 3] - @test b[4, 4] == a[1, 1] - b[4, 4] = 44 - @test b[4, 4] == 44 - end - end - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - b = a[Block(2):Block(2), Block(1):Block(2)] - @test blockstoredlength(b) == 1 - @test b == Array(a)[3:5, 1:end] - - a = BlockSparseArray{elt}(undef, ([2, 3, 4], [2, 3, 4])) - # TODO: Define `block_diagindices`. - @views for b in [Block(1, 1), Block(2, 2), Block(3, 3)] - a[b] = randn(elt, size(a[b])) - end - for (I1, I2) in ( - (mortar([Block(2)[2:3], Block(3)[1:3]]), mortar([Block(2)[2:3], Block(3)[2:3]])), - ([Block(2)[2:3], Block(3)[1:3]], [Block(2)[2:3], Block(3)[2:3]]), - ) - for b in (a[I1, I2], @view(a[I1, I2])) - # TODO: Rename `blockstoredlength`. - @test blockstoredlength(b) == 2 - @test b[Block(1, 1)] == a[Block(2, 2)[2:3, 2:3]] - @test b[Block(2, 2)] == a[Block(3, 3)[1:3, 2:3]] - end - end - - a = dev(BlockSparseArray{elt}(undef, ([3, 3], [3, 3]))) - # TODO: Define `block_diagindices`. - @views for b in [Block(1, 1), Block(2, 2)] - a[b] = dev(randn(elt, size(a[b]))) - end - I = mortar([Block(1)[1:2], Block(2)[1:2]]) - b = a[:, I] - @test b[Block(1, 1)] == a[Block(1, 1)][:, 1:2] - @test b[Block(2, 1)] == a[Block(2, 1)][:, 1:2] - @test b[Block(1, 2)] == a[Block(1, 2)][:, 1:2] - @test b[Block(2, 2)] == a[Block(2, 2)][:, 1:2] - @test blocklengths.(axes(b)) == ([3, 3], [2, 2]) - # TODO: Rename `blockstoredlength`. - @test blocksize(b) == (2, 2) - @test blockstoredlength(b) == 2 - - a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) - @views for b in [Block(1, 2), Block(2, 1)] - a[b] = randn(elt, size(a[b])) - end - @test isassigned(a, 1, 1) - @test isassigned(a, 1, 1, 1) - @test !isassigned(a, 1, 1, 2) - @test isassigned(a, 5, 7) - @test isassigned(a, 5, 7, 1) - @test !isassigned(a, 5, 7, 2) - @test !isassigned(a, 0, 1) - @test !isassigned(a, 5, 8) - @test isassigned(a, Block(1), Block(1)) - @test isassigned(a, Block(2), Block(2)) - @test !isassigned(a, Block(1), Block(0)) - @test !isassigned(a, Block(3), Block(2)) - @test isassigned(a, Block(1, 1)) - @test isassigned(a, Block(2, 2)) - @test !isassigned(a, Block(1, 0)) - @test !isassigned(a, Block(3, 2)) - @test isassigned(a, Block(1)[1], Block(1)[1]) - @test isassigned(a, Block(2)[3], Block(2)[4]) - @test !isassigned(a, Block(1)[0], Block(1)[1]) - @test !isassigned(a, Block(2)[3], Block(2)[5]) - @test !isassigned(a, Block(1)[1], Block(0)[1]) - @test !isassigned(a, Block(3)[3], Block(2)[4]) - - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - @test iszero(a) - @test iszero(blockstoredlength(a)) - fill!(a, 0) - @test iszero(a) - @test iszero(blockstoredlength(a)) - fill!(a, 2) - @test !iszero(a) - @test all(==(2), a) - @test blockstoredlength(a) == 4 - fill!(a, 0) - @test iszero(a) - @test iszero(blockstoredlength(a)) - - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - @test iszero(a) - @test iszero(blockstoredlength(a)) - zero!(a) - @test iszero(a) - @test iszero(blockstoredlength(a)) - fill!(a, 2) - @test !iszero(a) - @test all(==(2), a) - @test blockstoredlength(a) == 4 - zero!(a) - @test iszero(a) - @test iszero(blockstoredlength(a)) - - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - @test iszero(a) - @test iszero(blockstoredlength(a)) - a .= 0 - @test iszero(a) - @test iszero(blockstoredlength(a)) - a .= 2 - @test !iszero(a) - @test all(==(2), a) - @test blockstoredlength(a) == 4 - a .= 0 - @test iszero(a) - @test iszero(blockstoredlength(a)) - - # TODO: Broken on GPU. - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - for I in (Block.(1:2), [Block(1), Block(2)]) - b = @view a[I, I] - x = randn(elt, 3, 4) - b[Block(2, 2)] = x - # These outputs a block of zeros, - # for some reason the block - # is not getting set. - # I think the issue is that: - # ```julia - # @view(@view(a[I, I]))[Block(1, 1)] - # ``` - # creates a doubly-wrapped SubArray - # instead of flattening down to a - # single SubArray wrapper. - @test a[Block(2, 2)] == x - @test b[Block(2, 2)] == x - end - - function f1() - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]] - x = randn(elt, 3, 4) - b[Block(1, 1)] .= x - return (; a, b, x) - end - function f2() - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]] - x = randn(elt, 3, 4) - b[Block(1, 1)] = x - return (; a, b, x) - end - for abx in (f1(), f2()) - (; a, b, x) = abx - @test b isa SubArray{<:Any,<:Any,<:BlockSparseArray} - @test blockstoredlength(b) == 1 - @test b[Block(1, 1)] == x - @test @view(b[Block(1, 1)]) isa Matrix{elt} - for blck in [Block(2, 1), Block(1, 2), Block(2, 2)] - @test iszero(b[blck]) - end - @test blockstoredlength(a) == 1 - @test a[Block(2, 2)] == x - for blck in [Block(1, 1), Block(2, 1), Block(1, 2)] - @test iszero(a[blck]) - end - @test_throws DimensionMismatch b[Block(1, 1)] .= randn(2, 3) - end - - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - b = @views a[[Block(2), Block(1)], [Block(2), Block(1)]][Block(2, 1)] - @test iszero(b) - @test size(b) == (2, 4) - x = randn(elt, 2, 4) - b .= x - @test b == x - @test a[Block(1, 2)] == x - @test blockstoredlength(a) == 1 - - a = BlockSparseArray{elt}(undef, [4, 3, 2], [4, 3, 2]) - @views for B in [Block(1, 1), Block(2, 2), Block(3, 3)] - a[B] = randn(elt, size(a[B])) - end - b = @view a[[Block(3), Block(2), Block(1)], [Block(3), Block(2), Block(1)]] - @test b isa SubArray{<:Any,<:Any,<:BlockSparseArray} - c = @view b[4:8, 4:8] - @test c isa SubArray{<:Any,<:Any,<:BlockSparseArray} - @test size(c) == (5, 5) - @test blockstoredlength(c) == 2 - @test blocksize(c) == (2, 2) - @test blocklengths.(axes(c)) == ([2, 3], [2, 3]) - @test size(c[Block(1, 1)]) == (2, 2) - @test c[Block(1, 1)] == a[Block(2, 2)[2:3, 2:3]] - @test size(c[Block(2, 2)]) == (3, 3) - @test c[Block(2, 2)] == a[Block(1, 1)[1:3, 1:3]] - @test size(c[Block(2, 1)]) == (3, 2) - @test iszero(c[Block(2, 1)]) - @test size(c[Block(1, 2)]) == (2, 3) - @test iszero(c[Block(1, 2)]) - - x = randn(elt, 3, 3) - c[Block(2, 2)] = x - @test c[Block(2, 2)] == x - @test a[Block(1, 1)[1:3, 1:3]] == x - - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]] - for index in parentindices(@view(b[Block(1, 1)])) - @test index isa Base.OneTo{Int} - end - - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) - a[Block(1, 1)] = randn(elt, 2, 3) - b = @view a[Block(1, 1)[1:2, 1:1]] - @test b isa SubArray{elt,2,Matrix{elt}} - for i in parentindices(b) - @test i isa UnitRange{Int} - end - - a = BlockSparseArray{elt}(undef, [2, 2, 2, 2], [2, 2, 2, 2]) - @views for I in [Block(1, 1), Block(2, 2), Block(3, 3), Block(4, 4)] - a[I] = randn(elt, size(a[I])) - end - for I in (blockedrange([4, 4]), BlockedVector(Block.(1:4), [2, 2])) - b = @view a[I, I] - @test copy(b) == a - @test blocksize(b) == (2, 2) - @test blocklengths.(axes(b)) == ([4, 4], [4, 4]) - # TODO: Fix in Julia 1.11 (https://github.com/ITensor/ITensors.jl/pull/1539). - if VERSION < v"1.11-" - @test b[Block(1, 1)] == a[Block.(1:2), Block.(1:2)] - @test b[Block(2, 1)] == a[Block.(3:4), Block.(1:2)] - @test b[Block(1, 2)] == a[Block.(1:2), Block.(3:4)] - @test b[Block(2, 2)] == a[Block.(3:4), Block.(3:4)] - end - c = @view b[Block(2, 2)] - @test blocksize(c) == (1, 1) - @test c == a[Block.(3:4), Block.(3:4)] - end - - a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) - a[Block(1, 1)] = randn(elt, 2, 2) - a[Block(2, 2)] = randn(elt, 3, 3) - for I in (mortar([Block(1)[2:2], Block(2)[2:3]]), [Block(1)[2:2], Block(2)[2:3]]) - b = @view a[:, I] - @test b == Array(a)[:, [2, 4, 5]] - end - - # Merge and permute blocks. - a = BlockSparseArray{elt}(undef, [2, 2, 2, 2], [2, 2, 2, 2]) - @views for I in [Block(1, 1), Block(2, 2), Block(3, 3), Block(4, 4)] - a[I] = randn(elt, size(a[I])) - end - for I in ( - BlockVector([Block(4), Block(3), Block(2), Block(1)], [2, 2]), - BlockedVector([Block(4), Block(3), Block(2), Block(1)], [2, 2]), - ) - b = @view a[I, I] - J = [Block(4), Block(3), Block(2), Block(1)] - @test b == a[J, J] - @test copy(b) == a[J, J] - @test blocksize(b) == (2, 2) - @test blocklengths.(axes(b)) == ([4, 4], [4, 4]) - @test b[Block(1, 1)] == Array(a)[[7, 8, 5, 6], [7, 8, 5, 6]] - c = @views b[Block(1, 1)][2:3, 2:3] - @test c == Array(a)[[8, 5], [8, 5]] - @test copy(c) == Array(a)[[8, 5], [8, 5]] - c = @view b[Block(1, 1)[2:3, 2:3]] - @test c == Array(a)[[8, 5], [8, 5]] - @test copy(c) == Array(a)[[8, 5], [8, 5]] - end - - # TODO: Add more tests of this, it may - # only be working accidentally. - a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) - a[Block(1, 1)] = randn(elt, 2, 2) - a[Block(2, 2)] = randn(elt, 3, 3) - @test a[2:4, 4] == Array(a)[2:4, 4] - # TODO: Fix this. - @test_broken a[4, 2:4] == Array(a)[4, 2:4] - end - @testset "view!" begin - for blk in ((Block(2, 2),), (Block(2), Block(2))) - a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) - b = view!(a, blk...) - x = randn(elt, 3, 3) - b .= x - @test b == x - @test a[blk...] == x - @test @view(a[blk...]) == x - @test view!(a, blk...) == x - @test @view!(a[blk...]) == x - end - for blk in ((Block(2, 2),), (Block(2), Block(2))) - a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) - b = @view! a[blk...] - x = randn(elt, 3, 3) - b .= x - @test b == x - @test a[blk...] == x - @test @view(a[blk...]) == x - @test view!(a, blk...) == x - @test @view!(a[blk...]) == x - end - for blk in ((Block(2, 2)[2:3, 1:2],), (Block(2)[2:3], Block(2)[1:2])) - a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) - b = view!(a, blk...) - x = randn(elt, 2, 2) - b .= x - @test b == x - @test a[blk...] == x - @test @view(a[blk...]) == x - @test view!(a, blk...) == x - @test @view!(a[blk...]) == x - end - for blk in ((Block(2, 2)[2:3, 1:2],), (Block(2)[2:3], Block(2)[1:2])) - a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) - b = @view! a[blk...] - x = randn(elt, 2, 2) - b .= x - @test b == x - @test a[blk...] == x - @test @view(a[blk...]) == x - @test view!(a, blk...) == x - @test @view!(a[blk...]) == x - end - # 0-dim case - # Regression test for https://github.com/ITensor/BlockSparseArrays.jl/issues/148 - for I in ((), (Block(),)) - a = dev(BlockSparseArray{elt}(undef)) - @test !isstored(a) - @test iszero(blockstoredlength(a)) - @test isempty(eachblockstoredindex(a)) - @test iszero(a) - b = @view! a[I...] - @test isstored(a) - @test isone(blockstoredlength(a)) - @test issetequal(eachblockstoredindex(a), [Block()]) - @test iszero(adapt(Array)(a)) - @test b isa arrayt{elt,0} - @test size(b) == () - # Converting to `Array` works around a bug in `iszero(JLArray{Float64}(undef))`. - @test iszero(adapt(Array)(b)) - end - end @testset "LinearAlgebra" begin a1 = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) a1[Block(1, 1)] = dev(randn(elt, size(@view(a1[Block(1, 1)])))) diff --git a/test/test_factorizations.jl b/test/test_factorizations.jl index b934f7e..f162c81 100644 --- a/test/test_factorizations.jl +++ b/test/test_factorizations.jl @@ -1,3 +1,4 @@ +using Adapt: adapt using BlockArrays: Block, BlockedMatrix, BlockedVector, blocks, mortar using BlockSparseArrays: BlockSparseArrays, @@ -473,3 +474,25 @@ end @test sort(diagview(D[Block(1, 1)]); by=abs, rev=true) ≈ D1[1:1] @test sort(diagview(D[Block(2, 2)]); by=abs, rev=true) ≈ D2[1:2] end + +@testset "Abstract block type" begin + arrayt = Array + elt = Float32 + dev = adapt(arrayt) + + a = BlockSparseMatrix{elt,AbstractMatrix{elt}}(undef, ([2, 3], [2, 3])) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + a[Block(2, 2)] = dev(randn(elt, 3, 3)) + @test_broken eig_full(a) + @test_broken eigh_full(a) + @test_broken svd_compact(a) + @test_broken svd_full(a) + @test_broken left_orth(a) + @test_broken right_orth(a) + @test_broken left_polar(a) + @test_broken right_polar(a) + @test_broken qr_compact(a) + @test_broken qr_full(a) + @test_broken lq_compact(a) + @test_broken lq_full(a) +end diff --git a/test/test_map.jl b/test/test_map.jl new file mode 100644 index 0000000..6eaf907 --- /dev/null +++ b/test/test_map.jl @@ -0,0 +1,733 @@ +using Adapt: adapt +using BlockArrays: Block, BlockRange, blockedrange, blocksize, blocksizes, mortar +using BlockSparseArrays: + BlockSparseArray, + BlockSparseVector, + BlockVector, + BlockView, + BlockedVector, + blocklengths, + blockstoredlength, + blocktype, + eachblockstoredindex +using DerivableInterfaces: zero! +using GPUArraysCore: @allowscalar +using JLArrays: JLArray +using SparseArraysBase: storedlength +using Test: @test, @test_broken, @test_throws, @testset + +elts = (Float32, Float64, ComplexF32) +arrayts = (Array, JLArray) +@testset "map and slicing (arraytype=$arrayt, eltype=$elt)" for arrayt in arrayts, + elt in elts + + dev = adapt(arrayt) + + a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = dev(randn(elt, size(a[b]))) + end + @test eltype(a) == elt + @test blockstoredlength(a) == 2 + @test storedlength(a) == 2 * 4 + 3 * 3 + + # TODO: Broken on GPU. + if arrayt ≠ Array + a = dev(BlockSparseArray{elt}(undef, [2, 3], [3, 4])) + @test_broken a[Block(1, 2)] .= 2 + end + + # TODO: Broken on GPU. + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + a[Block(1, 2)] .= 2 + @test eltype(a) == elt + @test all(==(2), a[Block(1, 2)]) + @test iszero(a[Block(1, 1)]) + @test iszero(a[Block(2, 1)]) + @test iszero(a[Block(2, 2)]) + @test blockstoredlength(a) == 1 + @test storedlength(a) == 2 * 4 + + # TODO: Broken on GPU. + if arrayt ≠ Array + a = dev(BlockSparseArray{elt}(undef, [2, 3], [3, 4])) + @test_broken a[Block(1, 2)] .= 0 + end + + # TODO: Broken on GPU. + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + a[Block(1, 2)] .= 0 + @test eltype(a) == elt + @test iszero(a[Block(1, 1)]) + @test iszero(a[Block(2, 1)]) + @test iszero(a[Block(1, 2)]) + @test iszero(a[Block(2, 2)]) + @test blockstoredlength(a) == 1 + @test storedlength(a) == 2 * 4 + + # Test similar on broadcasted expressions. + a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) + bc = Broadcast.broadcasted(+, a, a) + a′ = similar(bc, Float32) + @test a′ isa BlockSparseArray{Float32} + @test blocktype(a′) <: arrayt{Float32,2} + @test axes(a) == (blockedrange([2, 3]), blockedrange([3, 4])) + + # Test similar on broadcasted expressions with axes specified. + a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) + bc = Broadcast.broadcasted(+, a, a) + a′ = similar( + bc, Float32, (blockedrange([2, 4]), blockedrange([2, 5]), blockedrange([2, 2])) + ) + @test a′ isa BlockSparseArray{Float32} + @test blocktype(a′) <: arrayt{Float32,3} + @test axes(a′) == (blockedrange([2, 4]), blockedrange([2, 5]), blockedrange([2, 2])) + + a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = dev(randn(elt, size(a[b]))) + end + b = similar(a, complex(elt)) + @test eltype(b) == complex(eltype(a)) + @test iszero(b) + @test blockstoredlength(b) == 0 + @test storedlength(b) == 0 + @test size(b) == size(a) + @test blocksize(b) == blocksize(a) + + # Regression test for https://github.com/ITensor/BlockSparseArrays.jl/issues/98 + a = dev(BlockSparseArray{elt}(undef)) + b = similar(a, Float64, (Base.OneTo(1),)) + @test b isa BlockSparseVector{Float64} + @test size(b) == (1,) + @test blocksize(b) == (1,) + + a = dev(BlockSparseArray{elt}(undef, [2, 3], [3, 4])) + b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]] + c = @view b[Block(1, 1)] + @test iszero(a) + @test iszero(storedlength(a)) + @test iszero(b) + @test iszero(storedlength(b)) + # TODO: Broken on GPU. + @test iszero(c) broken = arrayt ≠ Array + @test iszero(storedlength(c)) + @allowscalar a[5, 7] = 1 + @test !iszero(a) + @test storedlength(a) == 3 * 4 + @test !iszero(b) + @test storedlength(b) == 3 * 4 + # TODO: Broken on GPU. + @test !iszero(c) broken = arrayt ≠ Array + @test storedlength(c) == 3 * 4 + d = @view a[1:4, 1:6] + @test iszero(d) + @test storedlength(d) == 2 * 3 + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = copy(a) + b[1, 1] = 11 + @test b[1, 1] == 11 + @test a[1, 1] ≠ 11 + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = copy(a) + b .*= 2 + @test b ≈ 2a + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = copy(a) + b ./= 2 + @test b ≈ a / 2 + + a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = dev(randn(elt, size(a[b]))) + end + b = 2 * a + @allowscalar @test Array(b) ≈ 2 * Array(a) + @test eltype(b) == elt + @test blockstoredlength(b) == 2 + @test storedlength(b) == 2 * 4 + 3 * 3 + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = (2 + 3im) * a + @test Array(b) ≈ (2 + 3im) * Array(a) + @test eltype(b) == complex(elt) + @test blockstoredlength(b) == 2 + @test storedlength(b) == 2 * 4 + 3 * 3 + + a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = dev(randn(elt, size(a[b]))) + end + b = a + a + @allowscalar @test Array(b) ≈ 2 * Array(a) + @test eltype(b) == elt + @test blockstoredlength(b) == 2 + @test storedlength(b) == 2 * 4 + 3 * 3 + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + x = BlockSparseArray{elt}(undef, ([3, 4], [2, 3])) + @views for b in [Block(1, 2), Block(2, 1)] + x[b] = randn(elt, size(x[b])) + end + b = a .+ a .+ 3 .* PermutedDimsArray(x, (2, 1)) + @test Array(b) ≈ 2 * Array(a) + 3 * permutedims(Array(x), (2, 1)) + @test eltype(b) == elt + @test blockstoredlength(b) == 2 + @test storedlength(b) == 2 * 4 + 3 * 3 + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = permutedims(a, (2, 1)) + @test Array(b) ≈ permutedims(Array(a), (2, 1)) + @test eltype(b) == elt + @test blockstoredlength(b) == 2 + @test storedlength(b) == 2 * 4 + 3 * 3 + + a = dev(BlockSparseArray{elt}(undef, [1, 1, 1], [1, 2, 3], [2, 2, 1], [1, 2, 1])) + a[Block(3, 2, 2, 3)] = dev(randn(elt, 1, 2, 2, 1)) + perm = (2, 3, 4, 1) + for b in (PermutedDimsArray(a, perm), permutedims(a, perm)) + @test @allowscalar(Array(b)) == permutedims(Array(a), perm) + @test issetequal(eachblockstoredindex(b), [Block(2, 2, 3, 3)]) + @test @allowscalar b[Block(2, 2, 3, 3)] == permutedims(a[Block(3, 2, 2, 3)], perm) + end + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = map(x -> 2x, a) + @test Array(b) ≈ 2 * Array(a) + @test eltype(b) == elt + @test size(b) == size(a) + @test blocksize(b) == (2, 2) + @test blockstoredlength(b) == 2 + @test storedlength(b) == 2 * 4 + 3 * 3 + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = a[[Block(2), Block(1)], [Block(2), Block(1)]] + @test b[Block(1, 1)] == a[Block(2, 2)] + @test b[Block(1, 2)] == a[Block(2, 1)] + @test b[Block(2, 1)] == a[Block(1, 2)] + @test b[Block(2, 2)] == a[Block(1, 1)] + @test size(b) == size(a) + @test blocksize(b) == (2, 2) + @test storedlength(b) == storedlength(a) + @test blockstoredlength(b) == 2 + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = a[Block(1):Block(2), Block(1):Block(2)] + @test b == a + @test size(b) == size(a) + @test blocksize(b) == (2, 2) + @test storedlength(b) == storedlength(a) + @test blockstoredlength(b) == 2 + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = a[Block(1):Block(1), Block(1):Block(2)] + @test b == Array(a)[1:2, 1:end] + @test b[Block(1, 1)] == a[Block(1, 1)] + @test b[Block(1, 2)] == a[Block(1, 2)] + @test size(b) == (2, 7) + @test blocksize(b) == (1, 2) + @test storedlength(b) == storedlength(a[Block(1, 2)]) + @test blockstoredlength(b) == 1 + + a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = dev(randn(elt, size(a[b]))) + end + I = 2:4 + I′, J′ = falses.(size(a)) + I′[I] .= true + J′[I] .= true + for b in (a[I, I], @view(a[I, I]), a[I′, J′], @view(a[I′, J′])) + @allowscalar @test b == Array(a)[2:4, 2:4] + @test size(b) == (3, 3) + @test blocksize(b) == (2, 2) + @test storedlength(b) == 1 * 1 + 2 * 2 + @test blockstoredlength(b) == 2 + for f in (getindex, view) + # TODO: Broken on GPU. + @allowscalar begin + @test size(f(b, Block(1, 1))) == (1, 2) + @test size(f(b, Block(2, 1))) == (2, 2) + @test size(f(b, Block(1, 2))) == (1, 1) + @test size(f(b, Block(2, 2))) == (2, 1) + @test f(b, Block(1, 1)) == a[Block(1, 1)[2:2, 2:3]] + @test f(b, Block(2, 1)) == a[Block(2, 1)[1:2, 2:3]] + @test f(b, Block(1, 2)) == a[Block(1, 2)[2:2, 1:1]] + @test f(b, Block(2, 2)) == a[Block(2, 2)[1:2, 1:1]] + end + end + end + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = a[Block(2, 1)[1:2, 2:3]] + @test b == Array(a)[3:4, 2:3] + @test size(b) == (2, 2) + @test blocksize(b) == (1, 1) + @test storedlength(b) == 2 * 2 + @test blockstoredlength(b) == 1 + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = PermutedDimsArray(a, (2, 1)) + @test blockstoredlength(b) == 2 + @test Array(b) == permutedims(Array(a), (2, 1)) + c = 2 * b + @test blockstoredlength(c) == 2 + @test Array(c) == 2 * permutedims(Array(a), (2, 1)) + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = a' + @test blockstoredlength(b) == 2 + @test Array(b) == Array(a)' + c = 2 * b + @test blockstoredlength(c) == 2 + @test Array(c) == 2 * Array(a)' + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = transpose(a) + @test blockstoredlength(b) == 2 + @test Array(b) == transpose(Array(a)) + c = 2 * b + @test blockstoredlength(c) == 2 + @test Array(c) == 2 * transpose(Array(a)) + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = a[Block(1), Block(1):Block(2)] + @test size(b) == (2, 7) + @test blocksize(b) == (1, 2) + @test b[Block(1, 1)] == a[Block(1, 1)] + @test b[Block(1, 2)] == a[Block(1, 2)] + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = copy(a) + x = randn(elt, size(@view(a[Block(2, 2)]))) + b[Block(2), Block(2)] = x + @test b[Block(2, 2)] == x + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = copy(a) + b[Block(1, 1)] .= 1 + @test b[Block(1, 1)] == trues(blocksizes(b)[1, 1]) + + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + b = @view a[Block(2, 2)] + @test size(b) == (3, 4) + for i in parentindices(b) + @test i isa Base.OneTo{Int} + end + @test parentindices(b)[1] == 1:3 + @test parentindices(b)[2] == 1:4 + + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + b = @view a[Block(2, 2)[1:2, 2:2]] + @test size(b) == (2, 1) + for i in parentindices(b) + @test i isa UnitRange{Int} + end + @test parentindices(b)[1] == 1:2 + @test parentindices(b)[2] == 2:2 + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + x = randn(elt, 1, 2) + @view(a[Block(2, 2)])[1:1, 1:2] = x + @test a[Block(2, 2)][1:1, 1:2] == x + @test @view(a[Block(2, 2)])[1:1, 1:2] == x + @test a[3:3, 4:5] == x + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + x = randn(elt, 1, 2) + @views a[Block(2, 2)][1:1, 1:2] = x + @test a[Block(2, 2)][1:1, 1:2] == x + @test @view(a[Block(2, 2)])[1:1, 1:2] == x + @test a[3:3, 4:5] == x + + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + b = @views a[Block(2, 2)][1:2, 2:3] + @test b isa SubArray{<:Any,<:Any,<:BlockView} + for i in parentindices(b) + @test i isa UnitRange{Int} + end + x = randn(elt, 2, 2) + b .= x + @test a[Block(2, 2)[1:2, 2:3]] == x + @test a[Block(2, 2)[1:2, 2:3]] == b + @test blockstoredlength(a) == 1 + + # Noncontiguous slicing. + a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + a[Block(2, 2)] = dev(randn(elt, 3, 3)) + I = ([3, 5], [2, 4]) + @test Array(a[I...]) == Array(a)[I...] + + # Noncontiguous slicing. + a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + a[Block(2, 2)] = dev(randn(elt, 3, 3)) + I = (:, [2, 4]) + if arrayt === Array + @test Array(a[I...]) == Array(a)[I...] + else + # TODO: Broken on GPU, fix this. + @test_broken a[I...] + end + + a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) + @views for b in [Block(1, 1), Block(2, 2)] + a[b] = randn(elt, size(a[b])) + end + for I in (Block.(1:2), [Block(1), Block(2)]) + b = @view a[I, I] + for I in CartesianIndices(a) + @test b[I] == a[I] + end + for block in BlockRange(a) + @test b[block] == a[block] + end + end + + a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + @views for b in [Block(1, 1), Block(2, 2)] + # TODO: Use `blocksizes(a)[Int.(Tuple(b))...]` once available. + a[b] = dev(randn(elt, size(a[b]))) + end + for I in ([Block(2), Block(1)],) + b = @view a[I, I] + @test b[Block(1, 1)] == a[Block(2, 2)] + @test b[Block(2, 1)] == a[Block(1, 2)] + @test b[Block(1, 2)] == a[Block(2, 1)] + @test b[Block(2, 2)] == a[Block(1, 1)] + @allowscalar begin + @test b[1, 1] == a[3, 3] + @test b[4, 4] == a[1, 1] + b[4, 4] = 44 + @test b[4, 4] == 44 + end + end + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + b = a[Block(2):Block(2), Block(1):Block(2)] + @test blockstoredlength(b) == 1 + @test b == Array(a)[3:5, 1:end] + + a = BlockSparseArray{elt}(undef, ([2, 3, 4], [2, 3, 4])) + # TODO: Define `block_diagindices`. + @views for b in [Block(1, 1), Block(2, 2), Block(3, 3)] + a[b] = randn(elt, size(a[b])) + end + for (I1, I2) in ( + (mortar([Block(2)[2:3], Block(3)[1:3]]), mortar([Block(2)[2:3], Block(3)[2:3]])), + ([Block(2)[2:3], Block(3)[1:3]], [Block(2)[2:3], Block(3)[2:3]]), + ) + for b in (a[I1, I2], @view(a[I1, I2])) + # TODO: Rename `blockstoredlength`. + @test blockstoredlength(b) == 2 + @test b[Block(1, 1)] == a[Block(2, 2)[2:3, 2:3]] + @test b[Block(2, 2)] == a[Block(3, 3)[1:3, 2:3]] + end + end + + a = dev(BlockSparseArray{elt}(undef, ([3, 3], [3, 3]))) + # TODO: Define `block_diagindices`. + @views for b in [Block(1, 1), Block(2, 2)] + a[b] = dev(randn(elt, size(a[b]))) + end + I = mortar([Block(1)[1:2], Block(2)[1:2]]) + b = a[:, I] + @test b[Block(1, 1)] == a[Block(1, 1)][:, 1:2] + @test b[Block(2, 1)] == a[Block(2, 1)][:, 1:2] + @test b[Block(1, 2)] == a[Block(1, 2)][:, 1:2] + @test b[Block(2, 2)] == a[Block(2, 2)][:, 1:2] + @test blocklengths.(axes(b)) == ([3, 3], [2, 2]) + # TODO: Rename `blockstoredlength`. + @test blocksize(b) == (2, 2) + @test blockstoredlength(b) == 2 + + a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end + @test isassigned(a, 1, 1) + @test isassigned(a, 1, 1, 1) + @test !isassigned(a, 1, 1, 2) + @test isassigned(a, 5, 7) + @test isassigned(a, 5, 7, 1) + @test !isassigned(a, 5, 7, 2) + @test !isassigned(a, 0, 1) + @test !isassigned(a, 5, 8) + @test isassigned(a, Block(1), Block(1)) + @test isassigned(a, Block(2), Block(2)) + @test !isassigned(a, Block(1), Block(0)) + @test !isassigned(a, Block(3), Block(2)) + @test isassigned(a, Block(1, 1)) + @test isassigned(a, Block(2, 2)) + @test !isassigned(a, Block(1, 0)) + @test !isassigned(a, Block(3, 2)) + @test isassigned(a, Block(1)[1], Block(1)[1]) + @test isassigned(a, Block(2)[3], Block(2)[4]) + @test !isassigned(a, Block(1)[0], Block(1)[1]) + @test !isassigned(a, Block(2)[3], Block(2)[5]) + @test !isassigned(a, Block(1)[1], Block(0)[1]) + @test !isassigned(a, Block(3)[3], Block(2)[4]) + + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + @test iszero(a) + @test iszero(blockstoredlength(a)) + fill!(a, 0) + @test iszero(a) + @test iszero(blockstoredlength(a)) + fill!(a, 2) + @test !iszero(a) + @test all(==(2), a) + @test blockstoredlength(a) == 4 + fill!(a, 0) + @test iszero(a) + @test iszero(blockstoredlength(a)) + + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + @test iszero(a) + @test iszero(blockstoredlength(a)) + zero!(a) + @test iszero(a) + @test iszero(blockstoredlength(a)) + fill!(a, 2) + @test !iszero(a) + @test all(==(2), a) + @test blockstoredlength(a) == 4 + zero!(a) + @test iszero(a) + @test iszero(blockstoredlength(a)) + + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + @test iszero(a) + @test iszero(blockstoredlength(a)) + a .= 0 + @test iszero(a) + @test iszero(blockstoredlength(a)) + a .= 2 + @test !iszero(a) + @test all(==(2), a) + @test blockstoredlength(a) == 4 + a .= 0 + @test iszero(a) + @test iszero(blockstoredlength(a)) + + # TODO: Broken on GPU. + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + for I in (Block.(1:2), [Block(1), Block(2)]) + b = @view a[I, I] + x = randn(elt, 3, 4) + b[Block(2, 2)] = x + # These outputs a block of zeros, + # for some reason the block + # is not getting set. + # I think the issue is that: + # ```julia + # @view(@view(a[I, I]))[Block(1, 1)] + # ``` + # creates a doubly-wrapped SubArray + # instead of flattening down to a + # single SubArray wrapper. + @test a[Block(2, 2)] == x + @test b[Block(2, 2)] == x + end + + function f1() + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]] + x = randn(elt, 3, 4) + b[Block(1, 1)] .= x + return (; a, b, x) + end + function f2() + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]] + x = randn(elt, 3, 4) + b[Block(1, 1)] = x + return (; a, b, x) + end + for abx in (f1(), f2()) + (; a, b, x) = abx + @test b isa SubArray{<:Any,<:Any,<:BlockSparseArray} + @test blockstoredlength(b) == 1 + @test b[Block(1, 1)] == x + @test @view(b[Block(1, 1)]) isa Matrix{elt} + for blck in [Block(2, 1), Block(1, 2), Block(2, 2)] + @test iszero(b[blck]) + end + @test blockstoredlength(a) == 1 + @test a[Block(2, 2)] == x + for blck in [Block(1, 1), Block(2, 1), Block(1, 2)] + @test iszero(a[blck]) + end + @test_throws DimensionMismatch b[Block(1, 1)] .= randn(2, 3) + end + + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + b = @views a[[Block(2), Block(1)], [Block(2), Block(1)]][Block(2, 1)] + @test iszero(b) + @test size(b) == (2, 4) + x = randn(elt, 2, 4) + b .= x + @test b == x + @test a[Block(1, 2)] == x + @test blockstoredlength(a) == 1 + + a = BlockSparseArray{elt}(undef, [4, 3, 2], [4, 3, 2]) + @views for B in [Block(1, 1), Block(2, 2), Block(3, 3)] + a[B] = randn(elt, size(a[B])) + end + b = @view a[[Block(3), Block(2), Block(1)], [Block(3), Block(2), Block(1)]] + @test b isa SubArray{<:Any,<:Any,<:BlockSparseArray} + c = @view b[4:8, 4:8] + @test c isa SubArray{<:Any,<:Any,<:BlockSparseArray} + @test size(c) == (5, 5) + @test blockstoredlength(c) == 2 + @test blocksize(c) == (2, 2) + @test blocklengths.(axes(c)) == ([2, 3], [2, 3]) + @test size(c[Block(1, 1)]) == (2, 2) + @test c[Block(1, 1)] == a[Block(2, 2)[2:3, 2:3]] + @test size(c[Block(2, 2)]) == (3, 3) + @test c[Block(2, 2)] == a[Block(1, 1)[1:3, 1:3]] + @test size(c[Block(2, 1)]) == (3, 2) + @test iszero(c[Block(2, 1)]) + @test size(c[Block(1, 2)]) == (2, 3) + @test iszero(c[Block(1, 2)]) + + x = randn(elt, 3, 3) + c[Block(2, 2)] = x + @test c[Block(2, 2)] == x + @test a[Block(1, 1)[1:3, 1:3]] == x + + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]] + for index in parentindices(@view(b[Block(1, 1)])) + @test index isa Base.OneTo{Int} + end + + a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + a[Block(1, 1)] = randn(elt, 2, 3) + b = @view a[Block(1, 1)[1:2, 1:1]] + @test b isa SubArray{elt,2,Matrix{elt}} + for i in parentindices(b) + @test i isa UnitRange{Int} + end + + a = BlockSparseArray{elt}(undef, [2, 2, 2, 2], [2, 2, 2, 2]) + @views for I in [Block(1, 1), Block(2, 2), Block(3, 3), Block(4, 4)] + a[I] = randn(elt, size(a[I])) + end + for I in (blockedrange([4, 4]), BlockedVector(Block.(1:4), [2, 2])) + b = @view a[I, I] + @test copy(b) == a + @test blocksize(b) == (2, 2) + @test blocklengths.(axes(b)) == ([4, 4], [4, 4]) + # TODO: Fix in Julia 1.11 (https://github.com/ITensor/ITensors.jl/pull/1539). + if VERSION < v"1.11-" + @test b[Block(1, 1)] == a[Block.(1:2), Block.(1:2)] + @test b[Block(2, 1)] == a[Block.(3:4), Block.(1:2)] + @test b[Block(1, 2)] == a[Block.(1:2), Block.(3:4)] + @test b[Block(2, 2)] == a[Block.(3:4), Block.(3:4)] + end + c = @view b[Block(2, 2)] + @test blocksize(c) == (1, 1) + @test c == a[Block.(3:4), Block.(3:4)] + end + + a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) + a[Block(1, 1)] = randn(elt, 2, 2) + a[Block(2, 2)] = randn(elt, 3, 3) + for I in (mortar([Block(1)[2:2], Block(2)[2:3]]), [Block(1)[2:2], Block(2)[2:3]]) + b = @view a[:, I] + @test b == Array(a)[:, [2, 4, 5]] + end + + # Merge and permute blocks. + a = BlockSparseArray{elt}(undef, [2, 2, 2, 2], [2, 2, 2, 2]) + @views for I in [Block(1, 1), Block(2, 2), Block(3, 3), Block(4, 4)] + a[I] = randn(elt, size(a[I])) + end + for I in ( + BlockVector([Block(4), Block(3), Block(2), Block(1)], [2, 2]), + BlockedVector([Block(4), Block(3), Block(2), Block(1)], [2, 2]), + ) + b = @view a[I, I] + J = [Block(4), Block(3), Block(2), Block(1)] + @test b == a[J, J] + @test copy(b) == a[J, J] + @test blocksize(b) == (2, 2) + @test blocklengths.(axes(b)) == ([4, 4], [4, 4]) + @test b[Block(1, 1)] == Array(a)[[7, 8, 5, 6], [7, 8, 5, 6]] + c = @views b[Block(1, 1)][2:3, 2:3] + @test c == Array(a)[[8, 5], [8, 5]] + @test copy(c) == Array(a)[[8, 5], [8, 5]] + c = @view b[Block(1, 1)[2:3, 2:3]] + @test c == Array(a)[[8, 5], [8, 5]] + @test copy(c) == Array(a)[[8, 5], [8, 5]] + end + + # TODO: Add more tests of this, it may + # only be working accidentally. + a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) + a[Block(1, 1)] = randn(elt, 2, 2) + a[Block(2, 2)] = randn(elt, 3, 3) + @test a[2:4, 4] == Array(a)[2:4, 4] + # TODO: Fix this. + @test_broken a[4, 2:4] == Array(a)[4, 2:4] +end diff --git a/test/test_similartype_unchecked.jl b/test/test_similartype_unchecked.jl new file mode 100644 index 0000000..884ab0d --- /dev/null +++ b/test/test_similartype_unchecked.jl @@ -0,0 +1,26 @@ +using BlockSparseArrays: similartype_unchecked +using JLArrays: JLArray, JLMatrix +using Test: @test, @testset +using TestExtras: @constinferred + +@testset "similartype_unchecked" begin + @test @constinferred(similartype_unchecked(Array{Float32}, NTuple{2,Int})) === + Matrix{Float32} + @test @constinferred(similartype_unchecked(Array{Float32}, NTuple{2,Base.OneTo{Int}})) === + Matrix{Float32} + if VERSION < v"1.11-" + # Not type stable in Julia 1.10. + @test similartype_unchecked(AbstractArray{Float32}, NTuple{2,Int}) === Matrix{Float32} + @test similartype_unchecked(JLArray{Float32}, NTuple{2,Int}) === JLMatrix{Float32} + @test similartype_unchecked(JLArray{Float32}, NTuple{2,Base.OneTo{Int}}) === + JLMatrix{Float32} + else + @test @constinferred(similartype_unchecked(AbstractArray{Float32}, NTuple{2,Int})) === + Matrix{Float32} + @test @constinferred(similartype_unchecked(JLArray{Float32}, NTuple{2,Int})) === + JLMatrix{Float32} + @test @constinferred( + similartype_unchecked(JLArray{Float32}, NTuple{2,Base.OneTo{Int}}) + ) === JLMatrix{Float32} + end +end diff --git a/test/test_view_bang.jl b/test/test_view_bang.jl new file mode 100644 index 0000000..62eb09d --- /dev/null +++ b/test/test_view_bang.jl @@ -0,0 +1,76 @@ +using Adapt: adapt +using BlockArrays: Block +using BlockSparseArrays: + BlockSparseArray, @view!, blockstoredlength, eachblockstoredindex, view! +using JLArrays: JLArray +using SparseArraysBase: isstored +using Test: @test, @testset + +elts = (Float32, Float64, ComplexF32) +arrayts = (Array, JLArray) +@testset "view! (arraytype=$arrayt, eltype=$elt)" for arrayt in arrayts, elt in elts + dev = adapt(arrayt) + + for blk in ((Block(2, 2),), (Block(2), Block(2))) + a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) + b = view!(a, blk...) + x = randn(elt, 3, 3) + b .= x + @test b == x + @test a[blk...] == x + @test @view(a[blk...]) == x + @test view!(a, blk...) == x + @test @view!(a[blk...]) == x + end + for blk in ((Block(2, 2),), (Block(2), Block(2))) + a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) + b = @view! a[blk...] + x = randn(elt, 3, 3) + b .= x + @test b == x + @test a[blk...] == x + @test @view(a[blk...]) == x + @test view!(a, blk...) == x + @test @view!(a[blk...]) == x + end + for blk in ((Block(2, 2)[2:3, 1:2],), (Block(2)[2:3], Block(2)[1:2])) + a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) + b = view!(a, blk...) + x = randn(elt, 2, 2) + b .= x + @test b == x + @test a[blk...] == x + @test @view(a[blk...]) == x + @test view!(a, blk...) == x + @test @view!(a[blk...]) == x + end + for blk in ((Block(2, 2)[2:3, 1:2],), (Block(2)[2:3], Block(2)[1:2])) + a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) + b = @view! a[blk...] + x = randn(elt, 2, 2) + b .= x + @test b == x + @test a[blk...] == x + @test @view(a[blk...]) == x + @test view!(a, blk...) == x + @test @view!(a[blk...]) == x + end + # 0-dim case + # Regression test for https://github.com/ITensor/BlockSparseArrays.jl/issues/148 + for I in ((), (Block(),)) + a = dev(BlockSparseArray{elt}(undef)) + @test !isstored(a) + @test iszero(blockstoredlength(a)) + @test isempty(eachblockstoredindex(a)) + @test iszero(a) + b = @view! a[I...] + @test isstored(a) + @test isone(blockstoredlength(a)) + @test issetequal(eachblockstoredindex(a), [Block()]) + @test iszero(adapt(Array)(a)) + @test b isa arrayt{elt,0} + @test size(b) == () + # Converting to `Array` works around a bug in `iszero(JLArray{Float64}(undef))`. + @test iszero(adapt(Array)(b)) + end +end