From 49e7df639c76c4ea7d67639d2612253abe49164e Mon Sep 17 00:00:00 2001 From: mtfishman <7855256+mtfishman@users.noreply.github.com> Date: Thu, 2 Oct 2025 00:47:40 +0000 Subject: [PATCH 1/2] Format .jl files (Runic) --- Project.toml | 2 +- docs/make.jl | 28 +- docs/make_index.jl | 16 +- docs/make_readme.jl | 16 +- examples/README.jl | 2 +- .../BlockSparseArraysTensorAlgebraExt.jl | 84 +- .../BlockSparseArraysTensorProductsExt.jl | 8 +- .../BlockArraysExtensions.jl | 622 ++++---- src/BlockArraysExtensions/blockedunitrange.jl | 422 ++--- src/BlockArraysExtensions/blockrange.jl | 26 +- .../BlockArraysSparseArraysBaseExt.jl | 4 +- src/BlockSparseArrays.jl | 12 +- .../abstractblocksparsearray.jl | 156 +- .../abstractblocksparsematrix.jl | 114 +- .../abstractblocksparsevector.jl | 2 +- src/abstractblocksparsearray/arraylayouts.jl | 104 +- src/abstractblocksparsearray/broadcast.jl | 90 +- src/abstractblocksparsearray/cat.jl | 2 +- src/abstractblocksparsearray/linearalgebra.jl | 218 +-- src/abstractblocksparsearray/map.jl | 90 +- .../sparsearrayinterface.jl | 28 +- .../unblockedsubarray.jl | 90 +- src/abstractblocksparsearray/views.jl | 468 +++--- .../wrappedabstractblocksparsearray.jl | 452 +++--- src/blocksparsearray/blockdiagonalarray.jl | 24 +- src/blocksparsearray/blocksparsearray.jl | 374 ++--- src/blocksparsearrayinterface/arraylayouts.jl | 56 +- .../blocksparsearrayinterface.jl | 548 +++---- src/blocksparsearrayinterface/broadcast.jl | 84 +- src/blocksparsearrayinterface/cat.jl | 18 +- .../getunstoredblock.jl | 54 +- .../linearalgebra.jl | 16 +- src/blocksparsearrayinterface/map.jl | 230 +-- src/blocksparsearrayinterface/views.jl | 2 +- src/factorizations/eig.jl | 182 +-- src/factorizations/lq.jl | 206 +-- src/factorizations/orthnull.jl | 256 +-- src/factorizations/polar.jl | 74 +- src/factorizations/qr.jl | 214 +-- src/factorizations/svd.jl | 340 ++-- src/factorizations/truncation.jl | 114 +- src/factorizations/utility.jl | 96 +- test/TestBlockSparseArraysUtils.jl | 12 +- test/runtests.jl | 80 +- test/test_abstract_blocktype.jl | 220 +-- test/test_aqua.jl | 4 +- test/test_basics.jl | 1008 ++++++------ test/test_blockrange.jl | 20 +- test/test_exports.jl | 22 +- test/test_factorizations.jl | 788 ++++----- test/test_genericblockindex.jl | 478 +++--- test/test_issues.jl | 26 +- test/test_map.jl | 1408 ++++++++--------- test/test_similartype_unchecked.jl | 38 +- test/test_tensoralgebraext.jl | 114 +- test/test_view_bang.jl | 128 +- 56 files changed, 5152 insertions(+), 5138 deletions(-) diff --git a/Project.toml b/Project.toml index 338cbc82..eb569a34 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.10.6" +version = "0.10.7" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/docs/make.jl b/docs/make.jl index 49b91db6..4fc83b0e 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -2,26 +2,26 @@ using BlockSparseArrays: BlockSparseArrays using Documenter: Documenter, DocMeta, deploydocs, makedocs DocMeta.setdocmeta!( - BlockSparseArrays, :DocTestSetup, quote - using BlockSparseArrays - using LinearAlgebra: Diagonal - end; recursive=true + BlockSparseArrays, :DocTestSetup, quote + using BlockSparseArrays + using LinearAlgebra: Diagonal + end; recursive = true ) include("make_index.jl") makedocs(; - modules=[BlockSparseArrays], - authors="ITensor developers and contributors", - sitename="BlockSparseArrays.jl", - format=Documenter.HTML(; - canonical="https://ITensor.github.io/BlockSparseArrays.jl", - edit_link="main", - assets=["assets/favicon.ico", "assets/extras.css"], - ), - pages=["Home" => "index.md", "Reference" => "reference.md"], + modules = [BlockSparseArrays], + authors = "ITensor developers and contributors", + sitename = "BlockSparseArrays.jl", + format = Documenter.HTML(; + canonical = "https://ITensor.github.io/BlockSparseArrays.jl", + edit_link = "main", + assets = ["assets/favicon.ico", "assets/extras.css"], + ), + pages = ["Home" => "index.md", "Reference" => "reference.md"], ) deploydocs(; - repo="github.com/ITensor/BlockSparseArrays.jl", devbranch="main", push_preview=true + repo = "github.com/ITensor/BlockSparseArrays.jl", devbranch = "main", push_preview = true ) diff --git a/docs/make_index.jl b/docs/make_index.jl index 97ca8dc0..0b2d95d0 100644 --- a/docs/make_index.jl +++ b/docs/make_index.jl @@ -2,20 +2,20 @@ using Literate: Literate using BlockSparseArrays: BlockSparseArrays function ccq_logo(content) - include_ccq_logo = """ + include_ccq_logo = """ ```@raw html Flatiron Center for Computational Quantum Physics logo. Flatiron Center for Computational Quantum Physics logo. ``` """ - content = replace(content, "{CCQ_LOGO}" => include_ccq_logo) - return content + content = replace(content, "{CCQ_LOGO}" => include_ccq_logo) + return content end Literate.markdown( - joinpath(pkgdir(BlockSparseArrays), "examples", "README.jl"), - joinpath(pkgdir(BlockSparseArrays), "docs", "src"); - flavor=Literate.DocumenterFlavor(), - name="index", - postprocess=ccq_logo, + joinpath(pkgdir(BlockSparseArrays), "examples", "README.jl"), + joinpath(pkgdir(BlockSparseArrays), "docs", "src"); + flavor = Literate.DocumenterFlavor(), + name = "index", + postprocess = ccq_logo, ) diff --git a/docs/make_readme.jl b/docs/make_readme.jl index 097d78be..ad2e38a0 100644 --- a/docs/make_readme.jl +++ b/docs/make_readme.jl @@ -2,20 +2,20 @@ using Literate: Literate using BlockSparseArrays: BlockSparseArrays function ccq_logo(content) - include_ccq_logo = """ + include_ccq_logo = """ Flatiron Center for Computational Quantum Physics logo. """ - content = replace(content, "{CCQ_LOGO}" => include_ccq_logo) - return content + content = replace(content, "{CCQ_LOGO}" => include_ccq_logo) + return content end Literate.markdown( - joinpath(pkgdir(BlockSparseArrays), "examples", "README.jl"), - joinpath(pkgdir(BlockSparseArrays)); - flavor=Literate.CommonMarkFlavor(), - name="README", - postprocess=ccq_logo, + joinpath(pkgdir(BlockSparseArrays), "examples", "README.jl"), + joinpath(pkgdir(BlockSparseArrays)); + flavor = Literate.CommonMarkFlavor(), + name = "README", + postprocess = ccq_logo, ) diff --git a/examples/README.jl b/examples/README.jl index c83b5dec..f97e3f44 100644 --- a/examples/README.jl +++ b/examples/README.jl @@ -1,5 +1,5 @@ # # BlockSparseArrays.jl -# +# # [![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://itensor.github.io/BlockSparseArrays.jl/stable/) # [![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://itensor.github.io/BlockSparseArrays.jl/dev/) # [![Build Status](https://github.com/ITensor/BlockSparseArrays.jl/actions/workflows/Tests.yml/badge.svg?branch=main)](https://github.com/ITensor/BlockSparseArrays.jl/actions/workflows/Tests.yml?query=branch%3Amain) diff --git a/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl b/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl index 7912ef51..8a224f3a 100644 --- a/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl +++ b/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl @@ -2,17 +2,17 @@ module BlockSparseArraysTensorAlgebraExt using BlockSparseArrays: AbstractBlockSparseArray, blockreshape using TensorAlgebra: - TensorAlgebra, - BlockedTrivialPermutation, - BlockedTuple, - FusionStyle, - ReshapeFusion, - fuseaxes + TensorAlgebra, + BlockedTrivialPermutation, + BlockedTuple, + FusionStyle, + ReshapeFusion, + fuseaxes struct BlockReshapeFusion <: FusionStyle end function TensorAlgebra.FusionStyle(::Type{<:AbstractBlockSparseArray}) - return BlockReshapeFusion() + return BlockReshapeFusion() end using BlockArrays: Block, blocklength, blocks @@ -20,46 +20,46 @@ using BlockSparseArrays: blocksparse using SparseArraysBase: eachstoredindex using TensorAlgebra: TensorAlgebra, matricize, unmatricize function TensorAlgebra.matricize( - ::BlockReshapeFusion, a::AbstractArray, biperm::BlockedTrivialPermutation{2} -) - ax = fuseaxes(axes(a), biperm) - reshaped_blocks_a = reshape(blocks(a), map(blocklength, ax)) - key(I) = Block(Tuple(I)) - value(I) = matricize(reshaped_blocks_a[I], biperm) - Is = eachstoredindex(reshaped_blocks_a) - bs = if isempty(Is) - # Catch empty case and make sure the type is constrained properly. - # This seems to only be necessary in Julia versions below v1.11, - # try removing it when we drop support for those versions. - keytype = Base.promote_op(key, eltype(Is)) - valtype = Base.promote_op(value, eltype(Is)) - valtype′ = !isconcretetype(valtype) ? AbstractMatrix{eltype(a)} : valtype - Dict{keytype,valtype′}() - else - Dict(key(I) => value(I) for I in Is) - end - return blocksparse(bs, ax) + ::BlockReshapeFusion, a::AbstractArray, biperm::BlockedTrivialPermutation{2} + ) + ax = fuseaxes(axes(a), biperm) + reshaped_blocks_a = reshape(blocks(a), map(blocklength, ax)) + key(I) = Block(Tuple(I)) + value(I) = matricize(reshaped_blocks_a[I], biperm) + Is = eachstoredindex(reshaped_blocks_a) + bs = if isempty(Is) + # Catch empty case and make sure the type is constrained properly. + # This seems to only be necessary in Julia versions below v1.11, + # try removing it when we drop support for those versions. + keytype = Base.promote_op(key, eltype(Is)) + valtype = Base.promote_op(value, eltype(Is)) + valtype′ = !isconcretetype(valtype) ? AbstractMatrix{eltype(a)} : valtype + Dict{keytype, valtype′}() + else + Dict(key(I) => value(I) for I in Is) + end + return blocksparse(bs, ax) end using BlockArrays: blocklengths function TensorAlgebra.unmatricize( - ::BlockReshapeFusion, - m::AbstractMatrix, - blocked_ax::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}}, -) - ax = Tuple(blocked_ax) - reshaped_blocks_m = reshape(blocks(m), map(blocklength, ax)) - function f(I) - block_axes_I = BlockedTuple( - map(ntuple(identity, length(ax))) do i - return Base.axes1(ax[i][Block(I[i])]) - end, - blocklengths(blocked_ax), + ::BlockReshapeFusion, + m::AbstractMatrix, + blocked_ax::BlockedTuple{2, <:Any, <:Tuple{Vararg{AbstractUnitRange}}}, ) - return unmatricize(reshaped_blocks_m[I], block_axes_I) - end - bs = Dict(Block(Tuple(I)) => f(I) for I in eachstoredindex(reshaped_blocks_m)) - return blocksparse(bs, ax) + ax = Tuple(blocked_ax) + reshaped_blocks_m = reshape(blocks(m), map(blocklength, ax)) + function f(I) + block_axes_I = BlockedTuple( + map(ntuple(identity, length(ax))) do i + return Base.axes1(ax[i][Block(I[i])]) + end, + blocklengths(blocked_ax), + ) + return unmatricize(reshaped_blocks_m[I], block_axes_I) + end + bs = Dict(Block(Tuple(I)) => f(I) for I in eachstoredindex(reshaped_blocks_m)) + return blocksparse(bs, ax) end end diff --git a/ext/BlockSparseArraysTensorProductsExt/BlockSparseArraysTensorProductsExt.jl b/ext/BlockSparseArraysTensorProductsExt/BlockSparseArraysTensorProductsExt.jl index d3733d89..82f46fc2 100644 --- a/ext/BlockSparseArraysTensorProductsExt/BlockSparseArraysTensorProductsExt.jl +++ b/ext/BlockSparseArraysTensorProductsExt/BlockSparseArraysTensorProductsExt.jl @@ -5,10 +5,10 @@ using TensorProducts: TensorProducts, tensor_product # TODO: Dispatch on `FusionStyle` to allow different kinds of products, # for example to allow merging common symmetry sectors. function TensorProducts.tensor_product(a1::BlockUnitRange, a2::BlockUnitRange) - new_blockaxes = vec( - map(splat(tensor_product), Iterators.product(eachblockaxis(a1), eachblockaxis(a2))) - ) - return blockrange(new_blockaxes) + new_blockaxes = vec( + map(splat(tensor_product), Iterators.product(eachblockaxis(a1), eachblockaxis(a2))) + ) + return blockrange(new_blockaxes) end end diff --git a/src/BlockArraysExtensions/BlockArraysExtensions.jl b/src/BlockArraysExtensions/BlockArraysExtensions.jl index ff977f39..b7ac41e3 100644 --- a/src/BlockArraysExtensions/BlockArraysExtensions.jl +++ b/src/BlockArraysExtensions/BlockArraysExtensions.jl @@ -1,114 +1,114 @@ using ArrayLayouts: ArrayLayouts, MemoryLayout, sub_materialize using BlockArrays: - BlockArrays, - AbstractBlockArray, - AbstractBlockVector, - Block, - BlockIndex, - BlockIndexRange, - BlockRange, - BlockSlice, - BlockVector, - BlockedOneTo, - BlockedUnitRange, - BlockedVector, - block, - blockaxes, - blockedrange, - blockindex, - blocks, - findblock, - findblockindex + BlockArrays, + AbstractBlockArray, + AbstractBlockVector, + Block, + BlockIndex, + BlockIndexRange, + BlockRange, + BlockSlice, + BlockVector, + BlockedOneTo, + BlockedUnitRange, + BlockedVector, + block, + blockaxes, + blockedrange, + blockindex, + blocks, + findblock, + findblockindex using Dictionaries: Dictionary, Indices using SparseArraysBase: - SparseArraysBase, - eachstoredindex, - getunstoredindex, - isstored, - setunstoredindex!, - storedlength + SparseArraysBase, + eachstoredindex, + getunstoredindex, + isstored, + setunstoredindex!, + storedlength -function view!(a::AbstractArray{<:Any,N}, index::Block{N}) where {N} - return view!(a, Tuple(index)...) +function view!(a::AbstractArray{<:Any, N}, index::Block{N}) where {N} + return view!(a, Tuple(index)...) end -function view!(a::AbstractArray{<:Any,N}, index::Vararg{Block{1},N}) where {N} - blocks(a)[Int.(index)...] = blocks(a)[Int.(index)...] - return blocks(a)[Int.(index)...] +function view!(a::AbstractArray{<:Any, N}, index::Vararg{Block{1}, N}) where {N} + blocks(a)[Int.(index)...] = blocks(a)[Int.(index)...] + return blocks(a)[Int.(index)...] end # Fix ambiguity error. -function view!(a::AbstractArray{<:Any,0}) - blocks(a)[] = blocks(a)[] - return blocks(a)[] +function view!(a::AbstractArray{<:Any, 0}) + blocks(a)[] = blocks(a)[] + return blocks(a)[] end -function view!(a::AbstractArray{<:Any,N}, index::BlockIndexRange{N}) where {N} - # TODO: Is there a better code pattern for this? - indices = ntuple(N) do dim - return Tuple(Block(index))[dim][index.indices[dim]] - end - return view!(a, indices...) +function view!(a::AbstractArray{<:Any, N}, index::BlockIndexRange{N}) where {N} + # TODO: Is there a better code pattern for this? + indices = ntuple(N) do dim + return Tuple(Block(index))[dim][index.indices[dim]] + end + return view!(a, indices...) end -function view!(a::AbstractArray{<:Any,N}, index::Vararg{BlockIndexRange{1},N}) where {N} - b = view!(a, Block.(index)...) - r = map(index -> only(index.indices), index) - return @view b[r...] +function view!(a::AbstractArray{<:Any, N}, index::Vararg{BlockIndexRange{1}, N}) where {N} + b = view!(a, Block.(index)...) + r = map(index -> only(index.indices), index) + return @view b[r...] end using MacroTools: @capture is_getindex_expr(expr::Expr) = (expr.head === :ref) is_getindex_expr(x) = false macro view!(expr) - if !is_getindex_expr(expr) - error("@view must be used with getindex syntax (as `@view! a[i,j,...]`)") - end - @capture(expr, array_[indices__]) - return :(view!($(esc(array)), $(esc.(indices)...))) + if !is_getindex_expr(expr) + error("@view must be used with getindex syntax (as `@view! a[i,j,...]`)") + end + @capture(expr, array_[indices__]) + return :(view!($(esc(array)), $(esc.(indices)...))) end # A return type for `blocks(array)` when `array` isn't blocked. # Represents a vector with just that single block. -struct SingleBlockView{N,Array<:AbstractArray{<:Any,N}} <: AbstractArray{Array,N} - array::Array +struct SingleBlockView{N, Array <: AbstractArray{<:Any, N}} <: AbstractArray{Array, N} + array::Array end Base.parent(a::SingleBlockView) = a.array Base.size(a::SingleBlockView) = ntuple(Returns(1), ndims(a)) blocks_maybe_single(a) = blocks(a) blocks_maybe_single(a::Array) = SingleBlockView(a) -function Base.getindex(a::SingleBlockView{N}, index::Vararg{Int,N}) where {N} - @assert all(isone, index) - return parent(a) +function Base.getindex(a::SingleBlockView{N}, index::Vararg{Int, N}) where {N} + @assert all(isone, index) + return parent(a) end # A wrapper around a potentially blocked array that is not blocked. -struct NonBlockedArray{T,N,Array<:AbstractArray{T,N}} <: AbstractArray{T,N} - array::Array +struct NonBlockedArray{T, N, Array <: AbstractArray{T, N}} <: AbstractArray{T, N} + array::Array end Base.parent(a::NonBlockedArray) = a.array Base.size(a::NonBlockedArray) = size(parent(a)) -Base.getindex(a::NonBlockedArray{<:Any,N}, I::Vararg{Integer,N}) where {N} = parent(a)[I...] +Base.getindex(a::NonBlockedArray{<:Any, N}, I::Vararg{Integer, N}) where {N} = parent(a)[I...] # Views of `NonBlockedArray`/`NonBlockedVector` are eager. # This fixes an issue in Julia 1.11 where reindexing defaults to using views. # TODO: Maybe reconsider this design, and allows views to work in slicing. Base.view(a::NonBlockedArray, I...) = a[I...] BlockArrays.blocks(a::NonBlockedArray) = SingleBlockView(parent(a)) -const NonBlockedVector{T,Array} = NonBlockedArray{T,1,Array} +const NonBlockedVector{T, Array} = NonBlockedArray{T, 1, Array} NonBlockedVector(array::AbstractVector) = NonBlockedArray(array) # BlockIndices works around an issue that the indices of BlockSlice # are restricted to AbstractUnitRange{Int}. -struct BlockIndices{B,T<:Integer,I<:AbstractVector{T}} <: AbstractVector{T} - blocks::B - indices::I +struct BlockIndices{B, T <: Integer, I <: AbstractVector{T}} <: AbstractVector{T} + blocks::B + indices::I end for f in (:axes, :unsafe_indices, :axes1, :first, :last, :size, :length, :unsafe_length) - @eval Base.$f(S::BlockIndices) = Base.$f(S.indices) + @eval Base.$f(S::BlockIndices) = Base.$f(S.indices) end Base.getindex(S::BlockIndices, i::Integer) = getindex(S.indices, i) # TODO: Move this to a `BlockArraysExtensions` library. function blockedunitrange_getindices(a::AbstractBlockedUnitRange, indices::BlockIndices) - # TODO: Is this a good definition? It ignores `indices.indices`. - return a[indices.blocks] + # TODO: Is this a good definition? It ignores `indices.indices`. + return a[indices.blocks] end # Generalization of to `BlockArrays._blockslice`: @@ -119,28 +119,28 @@ end # TODO: Ideally this would be handled in BlockArrays.jl # once slicing like `A[Block(1)[[1, 2]]]` is supported. function _blockslice(x, y::AbstractUnitRange) - return BlockSlice(x, y) + return BlockSlice(x, y) end function _blockslice(x, y::AbstractVector) - return BlockIndices(x, y) + return BlockIndices(x, y) end # TODO: Constrain the type of `BlockIndices` more, this seems # to assume that `S.blocks` is a list of blocks as opposed to # a flat list of block indices like the definition below. function Base.getindex(S::BlockIndices, i::BlockSlice{<:Block{1}}) - # TODO: Check that `i.indices` is consistent with `S.indices`. - # It seems like this isn't handling the case where `i` is a - # subslice of a block correctly (i.e. it ignores `i.indices`). - @assert length(S.indices[Block(i)]) == length(i.indices) - return _blockslice(S.blocks[Int(Block(i))], S.indices[Block(i)]) + # TODO: Check that `i.indices` is consistent with `S.indices`. + # It seems like this isn't handling the case where `i` is a + # subslice of a block correctly (i.e. it ignores `i.indices`). + @assert length(S.indices[Block(i)]) == length(i.indices) + return _blockslice(S.blocks[Int(Block(i))], S.indices[Block(i)]) end function Base.getindex( - S::BlockIndices{<:AbstractBlockVector{<:BlockIndex{1}}}, i::BlockSlice{<:Block{1}} -) - @assert length(S.indices[Block(i)]) == length(i.indices) - return _blockslice(S.blocks[Block(i)], S.indices[Block(i)]) + S::BlockIndices{<:AbstractBlockVector{<:BlockIndex{1}}}, i::BlockSlice{<:Block{1}} + ) + @assert length(S.indices[Block(i)]) == length(i.indices) + return _blockslice(S.blocks[Block(i)], S.indices[Block(i)]) end # This is used in slicing like: @@ -148,17 +148,17 @@ end # I = BlockedVector([Block(4), Block(3), Block(2), Block(1)], [2, 2]) # a[I, I] function Base.getindex( - S::BlockIndices{<:AbstractBlockVector{<:Block{1}}}, i::BlockSlice{<:Block{1}} -) - # TODO: Check for conistency of indices. - # Wrapping the indices in `NonBlockedVector` reinterprets the blocked indices - # as a single block, since the result shouldn't be blocked. - return NonBlockedVector(BlockIndices(S.blocks[Block(i)], S.indices[Block(i)])) + S::BlockIndices{<:AbstractBlockVector{<:Block{1}}}, i::BlockSlice{<:Block{1}} + ) + # TODO: Check for conistency of indices. + # Wrapping the indices in `NonBlockedVector` reinterprets the blocked indices + # as a single block, since the result shouldn't be blocked. + return NonBlockedVector(BlockIndices(S.blocks[Block(i)], S.indices[Block(i)])) end function Base.getindex( - S::BlockIndices{<:BlockedVector{<:Block{1},<:BlockRange{1}}}, i::BlockSlice{<:Block{1}} -) - return i + S::BlockIndices{<:BlockedVector{<:Block{1}, <:BlockRange{1}}}, i::BlockSlice{<:Block{1}} + ) + return i end # Views of `BlockIndices` are eager. # This fixes an issue in Julia 1.11 where reindexing defaults to using views. @@ -174,41 +174,41 @@ Base.view(S::BlockIndices, i) = S[i] # This is similar to the definition: # @interface interface(a) to_indices(a, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}}) function Base.getindex( - a::NonBlockedVector{<:Integer,<:BlockIndices}, I::UnitRange{<:Integer} -) - ax = only(axes(parent(a).indices)) - brs = to_blockindices(ax, I) - inds = blockedunitrange_getindices(ax, I) - return NonBlockedVector(parent(a)[BlockSlice(brs, inds)]) + a::NonBlockedVector{<:Integer, <:BlockIndices}, I::UnitRange{<:Integer} + ) + ax = only(axes(parent(a).indices)) + brs = to_blockindices(ax, I) + inds = blockedunitrange_getindices(ax, I) + return NonBlockedVector(parent(a)[BlockSlice(brs, inds)]) end function Base.getindex(S::BlockIndices, i::BlockSlice{<:BlockRange{1}}) - # TODO: Check that `i.indices` is consistent with `S.indices`. - # TODO: Turn this into a `blockedunitrange_getindices` definition. - subblocks = S.blocks[Int.(i.block)] - subindices = mortar( - map(1:length(i.block)) do I - r = blocks(i.indices)[I] - return S.indices[first(r)]:S.indices[last(r)] - end, - ) - return BlockIndices(subblocks, subindices) + # TODO: Check that `i.indices` is consistent with `S.indices`. + # TODO: Turn this into a `blockedunitrange_getindices` definition. + subblocks = S.blocks[Int.(i.block)] + subindices = mortar( + map(1:length(i.block)) do I + r = blocks(i.indices)[I] + return S.indices[first(r)]:S.indices[last(r)] + end, + ) + return BlockIndices(subblocks, subindices) end # Used when performing slices like: # @views a[[Block(2), Block(1)]][2:4, 2:4] function Base.getindex(S::BlockIndices, i::BlockSlice{<:BlockVector{<:BlockIndex{1}}}) - subblocks = mortar( - map(blocks(i.block)) do br - return S.blocks[Int(Block(br))][only(br.indices)] - end, - ) - subindices = mortar( - map(blocks(i.block)) do br - S.indices[br] - end, - ) - return BlockIndices(subblocks, subindices) + subblocks = mortar( + map(blocks(i.block)) do br + return S.blocks[Int(Block(br))][only(br.indices)] + end, + ) + subindices = mortar( + map(blocks(i.block)) do br + S.indices[br] + end, + ) + return BlockIndices(subblocks, subindices) end # Similar to the definition of `BlockArrays.BlockSlices`: @@ -217,29 +217,29 @@ end # ``` # but includes `BlockIndices`, where the blocks aren't contiguous. const BlockSliceCollection = Union{ - Base.Slice, - BlockSlice{<:Block{1}}, - BlockSlice{<:BlockRange{1}}, - BlockIndices{<:Vector{<:Block{1}}}, + Base.Slice, + BlockSlice{<:Block{1}}, + BlockSlice{<:BlockRange{1}}, + BlockIndices{<:Vector{<:Block{1}}}, } const BlockIndexRangeSlice = BlockSlice{ - <:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}} + <:BlockVector{<:BlockIndex{1}, <:Vector{<:BlockIndexRange{1}}}, } const BlockIndexRangeSlices = BlockIndices{ - <:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}} + <:BlockVector{<:BlockIndex{1}, <:Vector{<:BlockIndexRange{1}}}, } const BlockIndexVectorSlices = BlockIndices{ - <:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexVector}} + <:BlockVector{<:BlockIndex{1}, <:Vector{<:BlockIndexVector}}, } const GenericBlockIndexVectorSlices = BlockIndices{ - <:BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector}} + <:BlockVector{<:GenericBlockIndex{1}, <:Vector{<:BlockIndexVector}}, } const SubBlockSliceCollection = Union{ - Base.Slice, - BlockIndexRangeSlice, - BlockIndexRangeSlices, - BlockIndexVectorSlices, - GenericBlockIndexVectorSlices, + Base.Slice, + BlockIndexRangeSlice, + BlockIndexRangeSlices, + BlockIndexVectorSlices, + GenericBlockIndexVectorSlices, } # TODO: This is type piracy. This is used in `reindex` when making @@ -272,12 +272,12 @@ const SubBlockSliceCollection = Union{ # Block(2)[2] # ``` function Base.getindex( - a::BlockVector{<:BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}}, - I::BlockSlice{<:Block{1}}, -) - # Check that the block slice corresponds to the correct block. - @assert I.indices == only(axes(a))[Block(I)] - return blocks(a)[Int(Block(I))] + a::BlockVector{<:BlockIndex{1}, <:AbstractVector{<:BlockIndexRange{1}}}, + I::BlockSlice{<:Block{1}}, + ) + # Check that the block slice corresponds to the correct block. + @assert I.indices == only(axes(a))[Block(I)] + return blocks(a)[Int(Block(I))] end # TODO: Use `Tuple` conversion once @@ -285,103 +285,103 @@ end block_to_cartesianindex(b::Block) = CartesianIndex(b.n) function blocks_to_cartesianindices(i::Indices{<:Block}) - return block_to_cartesianindex.(i) + return block_to_cartesianindex.(i) end function blocks_to_cartesianindices(d::Dictionary{<:Block}) - return Dictionary(blocks_to_cartesianindices(eachindex(d)), d) + return Dictionary(blocks_to_cartesianindices(eachindex(d)), d) end -function blockreshape(a::AbstractArray, dims::Tuple{Vector{Int},Vararg{Vector{Int}}}) - return blockreshape(a, blockedrange.(dims)) +function blockreshape(a::AbstractArray, dims::Tuple{Vector{Int}, Vararg{Vector{Int}}}) + return blockreshape(a, blockedrange.(dims)) end function blockreshape(a::AbstractArray, dim1::Vector{Int}, dim_rest::Vararg{Vector{Int}}) - return blockreshape(a, (dim1, dim_rest...)) + return blockreshape(a, (dim1, dim_rest...)) end # Fix ambiguity error. function blockreshape(a::AbstractArray) - return blockreshape(a, ()) + return blockreshape(a, ()) end tuple_oneto(n) = ntuple(identity, n) function _blockreshape(a::AbstractArray, axes::Tuple{Vararg{AbstractUnitRange}}) - reshaped_blocks_a = reshape(blocks(a), blocklength.(axes)) - function f(I) - block_axes_I = map(ntuple(identity, length(axes))) do i - return Base.axes1(axes[i][Block(I[i])]) + reshaped_blocks_a = reshape(blocks(a), blocklength.(axes)) + function f(I) + block_axes_I = map(ntuple(identity, length(axes))) do i + return Base.axes1(axes[i][Block(I[i])]) + end + # TODO: Better converter here. + return reshape(reshaped_blocks_a[I], block_axes_I) + end + bs = Dict(Block(Tuple(I)) => f(I) for I in eachstoredindex(reshaped_blocks_a)) + if !isconcretetype(keytype(bs)) || !isconcretetype(valtype(bs)) + # This branch only seems to be required in Julia 1.10, not Julia 1.11. + # TODO: Remove this branch once Julia 1.10 support is dropped. + bs = Dict{Block{length(axes), Int}, AbstractArray{eltype(a), length(axes)}}() end - # TODO: Better converter here. - return reshape(reshaped_blocks_a[I], block_axes_I) - end - bs = Dict(Block(Tuple(I)) => f(I) for I in eachstoredindex(reshaped_blocks_a)) - if !isconcretetype(keytype(bs)) || !isconcretetype(valtype(bs)) - # This branch only seems to be required in Julia 1.10, not Julia 1.11. - # TODO: Remove this branch once Julia 1.10 support is dropped. - bs = Dict{Block{length(axes),Int},AbstractArray{eltype(a),length(axes)}}() - end - return blocksparse(bs, axes) + return blocksparse(bs, axes) end function blockreshape( - a::AbstractArray, axes::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}} -) - return _blockreshape(a, axes) + a::AbstractArray, axes::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}} + ) + return _blockreshape(a, axes) end # Fix ambiguity error. function blockreshape(a::AbstractArray, axes::Tuple{}) - return _blockreshape(a, axes) + return _blockreshape(a, axes) end function blockreshape( - a::AbstractArray, axis1::AbstractUnitRange, axes_rest::Vararg{AbstractUnitRange} -) - return blockreshape(a, (axis1, axes_rest...)) + a::AbstractArray, axis1::AbstractUnitRange, axes_rest::Vararg{AbstractUnitRange} + ) + return blockreshape(a, (axis1, axes_rest...)) end function cartesianindices(axes::Tuple, b::Block) - return CartesianIndices(ntuple(dim -> axes[dim][Tuple(b)[dim]], length(axes))) + return CartesianIndices(ntuple(dim -> axes[dim][Tuple(b)[dim]], length(axes))) end # Get the range within a block. function blockindexrange(axis::AbstractUnitRange, r::AbstractUnitRange) - bi1 = findblockindex(axis, first(r)) - bi2 = findblockindex(axis, last(r)) - b = block(bi1) - # Range must fall within a single block. - @assert b == block(bi2) - i1 = blockindex(bi1) - i2 = blockindex(bi2) - return b[i1:i2] + bi1 = findblockindex(axis, first(r)) + bi2 = findblockindex(axis, last(r)) + b = block(bi1) + # Range must fall within a single block. + @assert b == block(bi2) + i1 = blockindex(bi1) + i2 = blockindex(bi2) + return b[i1:i2] end function blockindexrange( - axes::Tuple{Vararg{AbstractUnitRange,N}}, I::CartesianIndices{N} -) where {N} - brs = blockindexrange.(axes, I.indices) - b = Block(block.(brs)) - rs = map(br -> only(br.indices), brs) - return b[rs...] + axes::Tuple{Vararg{AbstractUnitRange, N}}, I::CartesianIndices{N} + ) where {N} + brs = blockindexrange.(axes, I.indices) + b = Block(block.(brs)) + rs = map(br -> only(br.indices), brs) + return b[rs...] end function blockindexrange(a::AbstractArray, I::CartesianIndices) - return blockindexrange(axes(a), I) + return blockindexrange(axes(a), I) end # Get the blocks the range spans across. function blockrange(axis::AbstractUnitRange, r::UnitRange) - return findblock(axis, first(r)):findblock(axis, last(r)) + return findblock(axis, first(r)):findblock(axis, last(r)) end # Occurs when slicing with `a[2:4, 2:4]`. function blockrange(axis::BlockedOneTo{<:Integer}, r::BlockedUnitRange{<:Integer}) - # TODO: Check the blocks are commensurate. - return findblock(axis, first(r)):findblock(axis, last(r)) + # TODO: Check the blocks are commensurate. + return findblock(axis, first(r)):findblock(axis, last(r)) end function blockrange(axis::AbstractUnitRange, r::Int) - ## return findblock(axis, r) - return error("Slicing with integer values isn't supported.") + ## return findblock(axis, r) + return error("Slicing with integer values isn't supported.") end # This handles changing the blocking, for example: @@ -390,9 +390,9 @@ end # a[I, I] # TODO: Generalize to `AbstractBlockedUnitRange`. function blockrange(axis::BlockedOneTo{<:Integer}, r::BlockedOneTo{<:Integer}) - # TODO: Probably this is incorrect and should be something like: - # return findblock(axis, first(r)):findblock(axis, last(r)) - return only(blockaxes(r)) + # TODO: Probably this is incorrect and should be something like: + # return findblock(axis, first(r)):findblock(axis, last(r)) + return only(blockaxes(r)) end # This handles block merging: @@ -403,101 +403,101 @@ end # I = BlockVector([Block(4), Block(3), Block(2), Block(1)], [2, 2]) # a[I, I] function blockrange(axis::AbstractUnitRange, r::AbstractBlockVector{<:Block{1}}) - for b in r - @assert b ∈ blockaxes(axis, 1) - end - return only(blockaxes(r)) + for b in r + @assert b ∈ blockaxes(axis, 1) + end + return only(blockaxes(r)) end function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Block{1}}) - for b in r - @assert b ∈ blockaxes(axis, 1) - end - return r + for b in r + @assert b ∈ blockaxes(axis, 1) + end + return r end using BlockArrays: BlockSlice function blockrange(axis::AbstractUnitRange, r::BlockSlice) - return blockrange(axis, r.block) + return blockrange(axis, r.block) end function blockrange(a::AbstractUnitRange, r::BlockIndices) - return blockrange(a, r.blocks) + return blockrange(a, r.blocks) end function blockrange(axis::AbstractUnitRange, r::Block{1}) - return r:r + return r:r end function blockrange(axis::AbstractUnitRange, r::BlockIndexRange) - return Block(r):Block(r) + return Block(r):Block(r) end function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:BlockIndexRange{1}}) - return error("Slicing not implemented for range of type `$(typeof(r))`.") + return error("Slicing not implemented for range of type `$(typeof(r))`.") end function blockrange( - axis::AbstractUnitRange, - r::BlockVector{<:BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}}, -) - return map(Block, blocks(r)) + axis::AbstractUnitRange, + r::BlockVector{<:BlockIndex{1}, <:AbstractVector{<:BlockIndexRange{1}}}, + ) + return map(Block, blocks(r)) end # This handles slicing with `:`/`Colon()`. function blockrange(axis::AbstractUnitRange, r::Base.Slice) - # TODO: Maybe use `BlockRange`, but that doesn't output - # the same thing. - return only(blockaxes(axis)) + # TODO: Maybe use `BlockRange`, but that doesn't output + # the same thing. + return only(blockaxes(axis)) end function blockrange(axis::AbstractUnitRange, r::NonBlockedVector) - return Block.(Base.OneTo(1)) + return Block.(Base.OneTo(1)) end function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Integer}) - return Block.(Base.OneTo(1)) + return Block.(Base.OneTo(1)) end function blockrange(axis::AbstractUnitRange, r::BlockIndexVector) - return Block(r):Block(r) + return Block(r):Block(r) end function blockrange( - axis::AbstractUnitRange, - r::BlockVector{<:BlockIndex{1},<:AbstractVector{<:BlockIndexVector}}, -) - return map(Block, blocks(r)) + axis::AbstractUnitRange, + r::BlockVector{<:BlockIndex{1}, <:AbstractVector{<:BlockIndexVector}}, + ) + return map(Block, blocks(r)) end function blockrange( - axis::AbstractUnitRange, - r::BlockVector{<:GenericBlockIndex{1},<:AbstractVector{<:BlockIndexVector}}, -) - return map(Block, blocks(r)) + axis::AbstractUnitRange, + r::BlockVector{<:GenericBlockIndex{1}, <:AbstractVector{<:BlockIndexVector}}, + ) + return map(Block, blocks(r)) end function blockrange(axis::AbstractUnitRange, r) - return error("Slicing not implemented for range of type `$(typeof(r))`.") + return error("Slicing not implemented for range of type `$(typeof(r))`.") end # This takes a range of indices `indices` of array `a` # and maps it to the range of indices within block `block`. function blockindices(a::AbstractArray, block::Block, indices::Tuple) - return blockindices(axes(a), block, indices) + return blockindices(axes(a), block, indices) end function blockindices(axes::Tuple, block::Block, indices::Tuple) - return blockindices.(axes, Tuple(block), indices) + return blockindices.(axes, Tuple(block), indices) end function blockindices(axis::AbstractUnitRange, block::Block, indices::AbstractUnitRange) - indices_within_block = intersect(indices, axis[block]) - if iszero(length(indices_within_block)) - # Falls outside of block - return 1:0 - end - return only(blockindexrange(axis, indices_within_block).indices) + indices_within_block = intersect(indices, axis[block]) + if iszero(length(indices_within_block)) + # Falls outside of block + return 1:0 + end + return only(blockindexrange(axis, indices_within_block).indices) end # This catches the case of `Vector{<:Block{1}}`. @@ -505,146 +505,146 @@ end # by the version with `indices::AbstractUnitRange`. # TODO: This should get fixed in a better way inside of `BlockArrays`. function blockindices( - axis::AbstractUnitRange, block::Block, indices::AbstractVector{<:Block{1}} -) - if block ∉ indices - # Falls outside of block - return 1:0 - end - return Base.OneTo(length(axis[block])) + axis::AbstractUnitRange, block::Block, indices::AbstractVector{<:Block{1}} + ) + if block ∉ indices + # Falls outside of block + return 1:0 + end + return Base.OneTo(length(axis[block])) end function blockindices(a::AbstractUnitRange, b::Block, r::BlockIndices) - return blockindices(a, b, r.blocks) + return blockindices(a, b, r.blocks) end function blockindices( - a::AbstractUnitRange, - b::Block, - r::BlockVector{<:BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}}, -) - # TODO: Change to iterate over `BlockRange(r)` - # once https://github.com/JuliaArrays/BlockArrays.jl/issues/404 - # is fixed. - for bl in blocks(r) - if b == Block(bl) - return only(bl.indices) + a::AbstractUnitRange, + b::Block, + r::BlockVector{<:BlockIndex{1}, <:AbstractVector{<:BlockIndexRange{1}}}, + ) + # TODO: Change to iterate over `BlockRange(r)` + # once https://github.com/JuliaArrays/BlockArrays.jl/issues/404 + # is fixed. + for bl in blocks(r) + if b == Block(bl) + return only(bl.indices) + end end - end - return error("Block not found.") + return error("Block not found.") end function cartesianindices(a::AbstractArray, b::Block) - return cartesianindices(axes(a), b) + return cartesianindices(axes(a), b) end # Output which blocks of `axis` are contained within the unit range `range`. # The start and end points must match. function findblocks(axis::AbstractUnitRange, range::AbstractUnitRange) - # TODO: Add a test that the start and end points of the ranges match. - return findblock(axis, first(range)):findblock(axis, last(range)) + # TODO: Add a test that the start and end points of the ranges match. + return findblock(axis, first(range)):findblock(axis, last(range)) end _block(indices) = block(indices) _block(indices::CartesianIndices) = Block(ntuple(Returns(1), ndims(indices))) function combine_axes(as::Vararg{Tuple}) - @assert allequal(length.(as)) - ndims = length(first(as)) - return ntuple(ndims) do dim - dim_axes = map(a -> a[dim], as) - return reduce(BlockArrays.combine_blockaxes, dim_axes) - end + @assert allequal(length.(as)) + ndims = length(first(as)) + return ntuple(ndims) do dim + dim_axes = map(a -> a[dim], as) + return reduce(BlockArrays.combine_blockaxes, dim_axes) + end end # Returns `BlockRange` # Convert the block of the axes to blocks of the subaxes. function subblocks(axes::Tuple, subaxes::Tuple, block::Block) - @assert length(axes) == length(subaxes) - return BlockRange( - ntuple(length(axes)) do dim - findblocks(subaxes[dim], axes[dim][Tuple(block)[dim]]) - end, - ) + @assert length(axes) == length(subaxes) + return BlockRange( + ntuple(length(axes)) do dim + findblocks(subaxes[dim], axes[dim][Tuple(block)[dim]]) + end, + ) end # Returns `Vector{<:Block}` function subblocks(axes::Tuple, subaxes::Tuple, blocks) - return mapreduce(vcat, blocks; init=eltype(blocks)[]) do block - return vec(subblocks(axes, subaxes, block)) - end + return mapreduce(vcat, blocks; init = eltype(blocks)[]) do block + return vec(subblocks(axes, subaxes, block)) + end end # Returns `Vector{<:CartesianIndices}` function blocked_cartesianindices(axes::Tuple, subaxes::Tuple, blocks) - return map(subblocks(axes, subaxes, blocks)) do block - return cartesianindices(subaxes, block) - end + return map(subblocks(axes, subaxes, blocks)) do block + return cartesianindices(subaxes, block) + end end # Represents a view of a block of a blocked array. -struct BlockView{T,N,Array<:AbstractArray{T,N}} <: AbstractArray{T,N} - array::Array - block::Tuple{Vararg{Block{1,Int},N}} +struct BlockView{T, N, Array <: AbstractArray{T, N}} <: AbstractArray{T, N} + array::Array + block::Tuple{Vararg{Block{1, Int}, N}} end Base.parent(a::BlockView) = a.array function Base.axes(a::BlockView) - # TODO: Try to avoid conversion to `Base.OneTo{Int}`, or just convert - # the element type to `Int` with `Int.(...)`. - # When the axes of `parent(a)` are `GradedOneTo`, the block is `LabelledUnitRange`, - # which has element type `LabelledInteger`. That causes conversion problems - # in some generic Base Julia code, for example when printing `BlockView`. - return ntuple(ndims(a)) do dim - return Base.OneTo{Int}(only(axes(axes(parent(a), dim)[a.block[dim]]))) - end + # TODO: Try to avoid conversion to `Base.OneTo{Int}`, or just convert + # the element type to `Int` with `Int.(...)`. + # When the axes of `parent(a)` are `GradedOneTo`, the block is `LabelledUnitRange`, + # which has element type `LabelledInteger`. That causes conversion problems + # in some generic Base Julia code, for example when printing `BlockView`. + return ntuple(ndims(a)) do dim + return Base.OneTo{Int}(only(axes(axes(parent(a), dim)[a.block[dim]]))) + end end function Base.size(a::BlockView) - return length.(axes(a)) + return length.(axes(a)) end -function Base.getindex(a::BlockView{<:Any,N}, index::Vararg{Int,N}) where {N} - return blocks(parent(a))[Int.(a.block)...][index...] +function Base.getindex(a::BlockView{<:Any, N}, index::Vararg{Int, N}) where {N} + return blocks(parent(a))[Int.(a.block)...][index...] end -function Base.setindex!(a::BlockView{<:Any,N}, value, index::Vararg{Int,N}) where {N} - b = @view! parent(a)[a.block...] - b[index...] = value - return a +function Base.setindex!(a::BlockView{<:Any, N}, value, index::Vararg{Int, N}) where {N} + b = @view! parent(a)[a.block...] + b[index...] = value + return a end function Base.fill!(a::BlockView, value) - b = @view! parent(a)[a.block...] - fill!(b, value) + b = @view! parent(a)[a.block...] + return fill!(b, value) end using Base.Broadcast: AbstractArrayStyle, Broadcasted, broadcasted materialize_blockviews(x) = x materialize_blockviews(a::BlockView) = blocks(parent(a))[Int.(a.block)...] function materialize_blockviews(bc::Broadcasted) - return broadcasted(bc.f, map(materialize_blockviews, bc.args)...) + return broadcasted(bc.f, map(materialize_blockviews, bc.args)...) end function Base.copyto!(a::BlockView, bc::Broadcasted) - b = @view! parent(a)[a.block...] - bc′ = materialize_blockviews(bc) - copyto!(b, bc′) - return a + b = @view! parent(a)[a.block...] + bc′ = materialize_blockviews(bc) + copyto!(b, bc′) + return a end function Base.copyto!(a::BlockView, bc::Broadcasted{<:AbstractArrayStyle{0}}) - b = @view! parent(a)[a.block...] - copyto!(b, bc) - return a + b = @view! parent(a)[a.block...] + copyto!(b, bc) + return a end function Base.copyto!(a::BlockView, src::AbstractArray) - b = @view! parent(a)[a.block...] - copyto!(b, src) - return a + b = @view! parent(a)[a.block...] + copyto!(b, src) + return a end function SparseArraysBase.storedlength(a::BlockView) - # TODO: Store whether or not the block is stored already as - # a Bool in `BlockView`. - I = CartesianIndex(Int.(a.block)) - # TODO: Use `eachblockstoredindex`. - if I ∈ eachstoredindex(blocks(parent(a))) - return storedlength(blocks(parent(a))[I]) - end - return 0 + # TODO: Store whether or not the block is stored already as + # a Bool in `BlockView`. + I = CartesianIndex(Int.(a.block)) + # TODO: Use `eachblockstoredindex`. + if I ∈ eachstoredindex(blocks(parent(a))) + return storedlength(blocks(parent(a))[I]) + end + return 0 end ## # Allow more fine-grained control: @@ -658,7 +658,7 @@ end ## return sub_materialize(MemoryLayout(a), a) ## end function ArrayLayouts.sub_materialize(a::BlockView) - return blocks(parent(a))[Int.(a.block)...] + return blocks(parent(a))[Int.(a.block)...] end # SVD additions @@ -670,18 +670,18 @@ using BlockArrays: BlockedMatrix # Here, we hijack this system to determine if there is any structure we can exploit # default: SVD is most efficient with BlockedArray function eigencopy_oftype(A::AbstractBlockArray, S) - return BlockedMatrix{S}(A) + return BlockedMatrix{S}(A) end -function svd!(A::BlockedMatrix; full::Bool=false, alg::Algorithm=default_svd_alg(A)) - F = svd!(parent(A); full, alg) +function svd!(A::BlockedMatrix; full::Bool = false, alg::Algorithm = default_svd_alg(A)) + F = svd!(parent(A); full, alg) - # restore block pattern - m = length(F.S) - bax1, bax2, bax3 = axes(A, 1), blockedrange([m]), axes(A, 2) + # restore block pattern + m = length(F.S) + bax1, bax2, bax3 = axes(A, 1), blockedrange([m]), axes(A, 2) - u = BlockedArray(F.U, (bax1, bax2)) - s = BlockedVector(F.S, (bax2,)) - vt = BlockedArray(F.Vt, (bax2, bax3)) - return SVD(u, s, vt) + u = BlockedArray(F.U, (bax1, bax2)) + s = BlockedVector(F.S, (bax2,)) + vt = BlockedArray(F.Vt, (bax2, bax3)) + return SVD(u, s, vt) end diff --git a/src/BlockArraysExtensions/blockedunitrange.jl b/src/BlockArraysExtensions/blockedunitrange.jl index 230591d0..6ad7080f 100644 --- a/src/BlockArraysExtensions/blockedunitrange.jl +++ b/src/BlockArraysExtensions/blockedunitrange.jl @@ -1,25 +1,25 @@ using BlockArrays: - BlockArrays, - AbstractBlockedUnitRange, - AbstractBlockVector, - Block, - BlockIndex, - BlockIndexRange, - BlockRange, - BlockSlice, - BlockVector, - block, - blockedrange, - blockfirsts, - blockindex, - blocklengths, - findblock, - findblockindex, - mortar + BlockArrays, + AbstractBlockedUnitRange, + AbstractBlockVector, + Block, + BlockIndex, + BlockIndexRange, + BlockRange, + BlockSlice, + BlockVector, + block, + blockedrange, + blockfirsts, + blockindex, + blocklengths, + findblock, + findblockindex, + mortar # Get the axes of each block of a block array. function eachblockaxes(a::AbstractArray) - return map(axes, blocks(a)) + return map(axes, blocks(a)) end axis(a::AbstractVector) = axes(a, 1) @@ -27,28 +27,28 @@ axis(a::AbstractVector) = axes(a, 1) # Get the axis of each block of a blocked unit # range. function eachblockaxis(a::AbstractVector) - return map(axis, blocks(a)) + return map(axis, blocks(a)) end function blockaxistype(a::AbstractVector) - return eltype(eachblockaxis(a)) + return eltype(eachblockaxis(a)) end # Take a collection of axes and mortar them # into a single blocked axis. function mortar_axis(axs) - return blockrange(axs) + return blockrange(axs) end function mortar_axis(axs::Vector{<:Base.OneTo{<:Integer}}) - return blockedrange(length.(axs)) + return blockedrange(length.(axs)) end # Custom `BlockedUnitRange` constructor that takes a unit range # and a set of block lengths, similar to `BlockArray(::AbstractArray, blocklengths...)`. function blockedunitrange(a::AbstractUnitRange, blocklengths) - blocklengths_shifted = copy(blocklengths) - blocklengths_shifted[1] += (first(a) - 1) - blocklasts = cumsum(blocklengths_shifted) - return BlockArrays._BlockedUnitRange(first(a), blocklasts) + blocklengths_shifted = copy(blocklengths) + blocklengths_shifted[1] += (first(a) - 1) + blocklasts = cumsum(blocklengths_shifted) + return BlockArrays._BlockedUnitRange(first(a), blocklasts) end # TODO: Move this to a `BlockArraysExtensions` library. @@ -56,8 +56,8 @@ end # block of the value `k`, while this finds the block of the index `k`. # This could make use of the `BlockIndices` object, i.e. `block(BlockIndices(a)[index])`. function blockedunitrange_findblock(a::AbstractBlockedUnitRange, index::Integer) - @boundscheck index in 1:length(a) || throw(BoundsError(a, index)) - return @inbounds findblock(a, index + first(a) - 1) + @boundscheck index in 1:length(a) || throw(BoundsError(a, index)) + return @inbounds findblock(a, index + first(a) - 1) end # TODO: Move this to a `BlockArraysExtensions` library. @@ -65,12 +65,12 @@ end # block index of the value `k`, while this finds the block index of the index `k`. # This could make use of the `BlockIndices` object, i.e. `BlockIndices(a)[index]`. function blockedunitrange_findblockindex(a::AbstractBlockedUnitRange, index::Integer) - @boundscheck index in 1:length(a) || throw(BoundsError()) - return @inbounds findblockindex(a, index + first(a) - 1) + @boundscheck index in 1:length(a) || throw(BoundsError()) + return @inbounds findblockindex(a, index + first(a) - 1) end function blockedunitrange_getindices(a::AbstractUnitRange, indices) - return a[indices] + return a[indices] end # TODO: Move this to a `BlockArraysExtensions` library. @@ -79,26 +79,26 @@ end # `blocked_getindex`. See the discussion here: # https://github.com/JuliaArrays/BlockArrays.jl/issues/347 function blockedunitrange_getindices( - a::AbstractBlockedUnitRange, indices::AbstractUnitRange{<:Integer} -) - first_blockindex = blockedunitrange_findblockindex(a, first(indices)) - last_blockindex = blockedunitrange_findblockindex(a, last(indices)) - first_block = block(first_blockindex) - last_block = block(last_blockindex) - blocklengths = if first_block == last_block - [length(indices)] - else - map(first_block:last_block) do block - if block == first_block - return length(a[first_block]) - blockindex(first_blockindex) + 1 - end - if block == last_block - return blockindex(last_blockindex) - end - return length(a[block]) + a::AbstractBlockedUnitRange, indices::AbstractUnitRange{<:Integer} + ) + first_blockindex = blockedunitrange_findblockindex(a, first(indices)) + last_blockindex = blockedunitrange_findblockindex(a, last(indices)) + first_block = block(first_blockindex) + last_block = block(last_blockindex) + blocklengths = if first_block == last_block + [length(indices)] + else + map(first_block:last_block) do block + if block == first_block + return length(a[first_block]) - blockindex(first_blockindex) + 1 + end + if block == last_block + return blockindex(last_blockindex) + end + return length(a[block]) + end end - end - return blockedunitrange(indices .+ (first(a) - 1), blocklengths) + return blockedunitrange(indices .+ (first(a) - 1), blocklengths) end # TODO: Make sure this handles block labels (AbstractGradedUnitRange) correctly. @@ -109,117 +109,117 @@ end # return blockedrange(blocklengths) # ``` function blockedunitrange_getindices( - a::AbstractBlockedUnitRange, indices::AbstractBlockVector{<:Block{1}} -) - blks = map(bs -> mortar(map(b -> a[b], bs)), blocks(indices)) - # We pass `length.(blks)` to `mortar` in order - # to pass block labels to the axes of the output, - # if they exist. This makes it so that - # `only(axes(a[indices])) isa `GradedUnitRange` - # if `a isa `GradedUnitRange`, for example. - # Note there is a more specialized definition: - # ```julia - # function blockedunitrange_getindices( - # a::AbstractGradedUnitRange, indices::AbstractBlockVector{<:Block{1}} - # ) - # ``` - # that does a better job of preserving labels, since `length` - # may drop labels for certain block types. - return mortar(blks, length.(blks)) + a::AbstractBlockedUnitRange, indices::AbstractBlockVector{<:Block{1}} + ) + blks = map(bs -> mortar(map(b -> a[b], bs)), blocks(indices)) + # We pass `length.(blks)` to `mortar` in order + # to pass block labels to the axes of the output, + # if they exist. This makes it so that + # `only(axes(a[indices])) isa `GradedUnitRange` + # if `a isa `GradedUnitRange`, for example. + # Note there is a more specialized definition: + # ```julia + # function blockedunitrange_getindices( + # a::AbstractGradedUnitRange, indices::AbstractBlockVector{<:Block{1}} + # ) + # ``` + # that does a better job of preserving labels, since `length` + # may drop labels for certain block types. + return mortar(blks, length.(blks)) end # TODO: Move this to a `BlockArraysExtensions` library. function blockedunitrange_getindices(a::AbstractBlockedUnitRange, indices::BlockIndexRange) - return a[block(indices)][only(indices.indices)] + return a[block(indices)][only(indices.indices)] end # TODO: Move this to a `BlockArraysExtensions` library. function blockedunitrange_getindices(a::AbstractBlockedUnitRange, indices::BlockSlice) - # TODO: Is this a good definition? It ignores `indices.indices`. - return a[indices.block] + # TODO: Is this a good definition? It ignores `indices.indices`. + return a[indices.block] end # TODO: Move this to a `BlockArraysExtensions` library. function blockedunitrange_getindices( - a::AbstractBlockedUnitRange, indices::AbstractVector{<:Integer} -) - return map(index -> a[index], indices) + a::AbstractBlockedUnitRange, indices::AbstractVector{<:Integer} + ) + return map(index -> a[index], indices) end # TODO: Move this to a `BlockArraysExtensions` library. # TODO: Make a special definition for `BlockedVector{<:Block{1}}` in order # to merge blocks. function blockedunitrange_getindices( - a::AbstractBlockedUnitRange, indices::AbstractVector{<:Union{Block{1},BlockIndexRange{1}}} -) - # Without converting `indices` to `Vector`, - # mapping `indices` outputs a `BlockVector` - # which is harder to reason about. - blocks = map(index -> a[index], Vector(indices)) - return mortar(blocks, (mortar_axis(map(axis, blocks)),)) + a::AbstractBlockedUnitRange, indices::AbstractVector{<:Union{Block{1}, BlockIndexRange{1}}} + ) + # Without converting `indices` to `Vector`, + # mapping `indices` outputs a `BlockVector` + # which is harder to reason about. + blocks = map(index -> a[index], Vector(indices)) + return mortar(blocks, (mortar_axis(map(axis, blocks)),)) end # TODO: Move this to a `BlockArraysExtensions` library. function blockedunitrange_getindices(a::AbstractBlockedUnitRange, indices::Block{1}) - return a[indices] + return a[indices] end function blockedunitrange_getindices( - a::AbstractBlockedUnitRange, - indices::BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}, -) - return mortar(map(b -> a[b], blocks(indices))) + a::AbstractBlockedUnitRange, + indices::BlockVector{<:BlockIndex{1}, <:Vector{<:BlockIndexRange{1}}}, + ) + return mortar(map(b -> a[b], blocks(indices))) end function blockedunitrange_getindices( - a::AbstractBlockedUnitRange, indices::AbstractVector{Bool} -) - blocked_indices = BlockedVector(indices, axes(a)) - bs = map(Base.OneTo(blocklength(blocked_indices))) do b - binds = blocked_indices[Block(b)] - bstart = blockfirsts(only(axes(blocked_indices)))[b] - return findall(binds) .+ (bstart - 1) - end - return mortar(filter(!isempty, bs)) + a::AbstractBlockedUnitRange, indices::AbstractVector{Bool} + ) + blocked_indices = BlockedVector(indices, axes(a)) + bs = map(Base.OneTo(blocklength(blocked_indices))) do b + binds = blocked_indices[Block(b)] + bstart = blockfirsts(only(axes(blocked_indices)))[b] + return findall(binds) .+ (bstart - 1) + end + return mortar(filter(!isempty, bs)) end # TODO: Move this to a `BlockArraysExtensions` library. function blockedunitrange_getindices(a::AbstractBlockedUnitRange, indices) - return error("Not implemented.") + return error("Not implemented.") end # The blocks of the corresponding slice. _blocks(a::AbstractUnitRange, indices) = error("Not implemented") function _blocks(a::AbstractUnitRange, indices::AbstractUnitRange) - return findblock(a, first(indices)):findblock(a, last(indices)) + return findblock(a, first(indices)):findblock(a, last(indices)) end function _blocks(a::AbstractUnitRange, indices::BlockRange) - return indices + return indices end # Slice `a` by `I`, returning a: # `BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}` # with the `BlockIndex{1}` corresponding to each value of `I`. function to_blockindices(a::AbstractBlockedUnitRange{<:Integer}, I::UnitRange{<:Integer}) - return mortar( - map(blocks(blockedunitrange_getindices(a, I))) do r - bi_first = findblockindex(a, first(r)) - bi_last = findblockindex(a, last(r)) - @assert block(bi_first) == block(bi_last) - return block(bi_first)[blockindex(bi_first):blockindex(bi_last)] - end, - ) + return mortar( + map(blocks(blockedunitrange_getindices(a, I))) do r + bi_first = findblockindex(a, first(r)) + bi_last = findblockindex(a, last(r)) + @assert block(bi_first) == block(bi_last) + return block(bi_first)[blockindex(bi_first):blockindex(bi_last)] + end, + ) end -struct GenericBlockIndex{N,TI<:Tuple{Vararg{Integer,N}},Tα<:Tuple{Vararg{Any,N}}} - I::TI - α::Tα +struct GenericBlockIndex{N, TI <: Tuple{Vararg{Integer, N}}, Tα <: Tuple{Vararg{Any, N}}} + I::TI + α::Tα end -@inline function GenericBlockIndex(a::NTuple{N,Block{1}}, b::Tuple) where {N} - return GenericBlockIndex(Int.(a), b) +@inline function GenericBlockIndex(a::NTuple{N, Block{1}}, b::Tuple) where {N} + return GenericBlockIndex(Int.(a), b) end @inline function GenericBlockIndex(::Tuple{}, b::Tuple{}) - return GenericBlockIndex{0,Tuple{},Tuple{}}((), ()) + return GenericBlockIndex{0, Tuple{}, Tuple{}}((), ()) end @inline GenericBlockIndex(a::Integer, b) = GenericBlockIndex((a,), (b,)) @inline GenericBlockIndex(a::Tuple, b) = GenericBlockIndex(a, (b,)) @@ -228,70 +228,70 @@ end @inline GenericBlockIndex(a::Block, b::Tuple) = GenericBlockIndex(a.n, b) @inline GenericBlockIndex(a::Block, b) = GenericBlockIndex(a, (b,)) @inline function GenericBlockIndex( - I::Tuple{Vararg{Integer,N}}, α::Tuple{Vararg{Any,M}} -) where {M,N} - M <= N || throw(ArgumentError("number of indices must not exceed the number of blocks")) - α2 = ntuple(k -> k <= M ? α[k] : 1, N) - GenericBlockIndex(I, α2) + I::Tuple{Vararg{Integer, N}}, α::Tuple{Vararg{Any, M}} + ) where {M, N} + M <= N || throw(ArgumentError("number of indices must not exceed the number of blocks")) + α2 = ntuple(k -> k <= M ? α[k] : 1, N) + return GenericBlockIndex(I, α2) end BlockArrays.block(b::GenericBlockIndex) = Block(b.I...) BlockArrays.blockindex(b::GenericBlockIndex{1}) = b.α[1] -function GenericBlockIndex(indcs::Tuple{Vararg{GenericBlockIndex{1},N}}) where {N} - GenericBlockIndex(block.(indcs), blockindex.(indcs)) +function GenericBlockIndex(indcs::Tuple{Vararg{GenericBlockIndex{1}, N}}) where {N} + return GenericBlockIndex(block.(indcs), blockindex.(indcs)) end function Base.checkindex( - ::Type{Bool}, axis::AbstractBlockedUnitRange, ind::GenericBlockIndex{1} -) - return checkindex(Bool, axis, block(ind)) && - checkbounds(Bool, axis[block(ind)], blockindex(ind)) + ::Type{Bool}, axis::AbstractBlockedUnitRange, ind::GenericBlockIndex{1} + ) + return checkindex(Bool, axis, block(ind)) && + checkbounds(Bool, axis[block(ind)], blockindex(ind)) end Base.to_index(i::GenericBlockIndex) = i function print_tuple_elements(io::IO, @nospecialize(t)) - if !isempty(t) - print(io, t[1]) - for n in t[2:end] - print(io, ", ", n) + if !isempty(t) + print(io, t[1]) + for n in t[2:end] + print(io, ", ", n) + end end - end - return nothing + return nothing end function Base.show(io::IO, B::GenericBlockIndex) - show(io, Block(B.I...)) - print(io, "[") - print_tuple_elements(io, B.α) - print(io, "]") - return nothing + show(io, Block(B.I...)) + print(io, "[") + print_tuple_elements(io, B.α) + print(io, "]") + return nothing end # https://github.com/JuliaArrays/BlockArrays.jl/blob/v1.6.3/src/views.jl#L31-L32 _maybetail(::Tuple{}) = () _maybetail(t::Tuple) = Base.tail(t) -@inline function Base.to_indices(A, inds, I::Tuple{GenericBlockIndex{1},Vararg{Any}}) - return (inds[1][I[1]], to_indices(A, _maybetail(inds), Base.tail(I))...) +@inline function Base.to_indices(A, inds, I::Tuple{GenericBlockIndex{1}, Vararg{Any}}) + return (inds[1][I[1]], to_indices(A, _maybetail(inds), Base.tail(I))...) end using Base: @propagate_inbounds @propagate_inbounds function Base.getindex(b::AbstractVector, K::GenericBlockIndex{1}) - return b[Block(K.I[1])][K.α[1]] + return b[Block(K.I[1])][K.α[1]] end @propagate_inbounds function Base.getindex( - b::AbstractArray{T,N}, K::GenericBlockIndex{N} -) where {T,N} - return b[block(K)][K.α...] + b::AbstractArray{T, N}, K::GenericBlockIndex{N} + ) where {T, N} + return b[block(K)][K.α...] end @propagate_inbounds function Base.getindex( - b::AbstractArray, K::GenericBlockIndex{1}, J::GenericBlockIndex{1}... -) - return b[GenericBlockIndex(tuple(K, J...))] + b::AbstractArray, K::GenericBlockIndex{1}, J::GenericBlockIndex{1}... + ) + return b[GenericBlockIndex(tuple(K, J...))] end # TODO: Delete this once `BlockArrays.BlockIndex` is generalized. @inline function Base.to_indices( - A, inds, I::Tuple{AbstractVector{<:GenericBlockIndex{1}},Vararg{Any}} -) - return (unblock(A, inds, I), to_indices(A, _maybetail(inds), Base.tail(I))...) + A, inds, I::Tuple{AbstractVector{<:GenericBlockIndex{1}}, Vararg{Any}} + ) + return (unblock(A, inds, I), to_indices(A, _maybetail(inds), Base.tail(I))...) end # This is a specialization of `BlockArrays.unblock`: @@ -300,44 +300,44 @@ end # BlockArrays.jl. # TODO: Ideally this would be defined in BlockArrays.jl once the slicing # there is made more generic. -function BlockArrays.unblock(A, inds, I::Tuple{GenericBlockIndex{1},Vararg{Any}}) - B = first(I) - return _blockslice(B, inds[1][B]) +function BlockArrays.unblock(A, inds, I::Tuple{GenericBlockIndex{1}, Vararg{Any}}) + B = first(I) + return _blockslice(B, inds[1][B]) end # Work around the fact that it is type piracy to define # `Base.getindex(a::Block, b...)`. -_getindex(a::Block{N}, b::Vararg{Any,N}) where {N} = GenericBlockIndex(a, b) -_getindex(a::Block{N}, b::Vararg{Integer,N}) where {N} = a[b...] -_getindex(a::Block{N}, b::Vararg{AbstractUnitRange{<:Integer},N}) where {N} = a[b...] -_getindex(a::Block{N}, b::Vararg{AbstractVector,N}) where {N} = BlockIndexVector(a, b) +_getindex(a::Block{N}, b::Vararg{Any, N}) where {N} = GenericBlockIndex(a, b) +_getindex(a::Block{N}, b::Vararg{Integer, N}) where {N} = a[b...] +_getindex(a::Block{N}, b::Vararg{AbstractUnitRange{<:Integer}, N}) where {N} = a[b...] +_getindex(a::Block{N}, b::Vararg{AbstractVector, N}) where {N} = BlockIndexVector(a, b) # Fix ambiguity. _getindex(a::Block{0}) = a[] -struct BlockIndexVector{N,BT,I<:NTuple{N,AbstractVector},TB<:Integer} <: AbstractArray{BT,N} - block::Block{N,TB} - indices::I - function BlockIndexVector{N,BT}( - block::Block{N,TB}, indices::I - ) where {N,BT,I<:NTuple{N,AbstractVector},TB<:Integer} - return new{N,BT,I,TB}(block, indices) - end +struct BlockIndexVector{N, BT, I <: NTuple{N, AbstractVector}, TB <: Integer} <: AbstractArray{BT, N} + block::Block{N, TB} + indices::I + function BlockIndexVector{N, BT}( + block::Block{N, TB}, indices::I + ) where {N, BT, I <: NTuple{N, AbstractVector}, TB <: Integer} + return new{N, BT, I, TB}(block, indices) + end end -function BlockIndexVector{1,BT}(block::Block{1}, indices::AbstractVector) where {BT} - return BlockIndexVector{1,BT}(block, (indices,)) +function BlockIndexVector{1, BT}(block::Block{1}, indices::AbstractVector) where {BT} + return BlockIndexVector{1, BT}(block, (indices,)) end function BlockIndexVector( - block::Block{N,TB}, indices::NTuple{N,AbstractVector} -) where {N,TB<:Integer} - BT = Base.promote_op(_getindex, typeof(block), eltype.(indices)...) - return BlockIndexVector{N,BT}(block, indices) + block::Block{N, TB}, indices::NTuple{N, AbstractVector} + ) where {N, TB <: Integer} + BT = Base.promote_op(_getindex, typeof(block), eltype.(indices)...) + return BlockIndexVector{N, BT}(block, indices) end function BlockIndexVector(block::Block{1}, indices::AbstractVector) - return BlockIndexVector(block, (indices,)) + return BlockIndexVector(block, (indices,)) end Base.size(a::BlockIndexVector) = length.(a.indices) -function Base.getindex(a::BlockIndexVector{N}, I::Vararg{Integer,N}) where {N} - return _getindex(Block(a), getindex.(a.indices, I)...) +function Base.getindex(a::BlockIndexVector{N}, I::Vararg{Integer, N}) where {N} + return _getindex(Block(a), getindex.(a.indices, I)...) end BlockArrays.block(b::BlockIndexVector) = b.block BlockArrays.Block(b::BlockIndexVector) = b.block @@ -346,47 +346,53 @@ Base.copy(a::BlockIndexVector) = BlockIndexVector(a.block, copy.(a.indices)) # Copied from BlockArrays.BlockIndexRange. function Base.show(io::IO, B::BlockIndexVector) - show(io, Block(B)) - print(io, "[") - print_tuple_elements(io, B.indices) - print(io, "]") + show(io, Block(B)) + print(io, "[") + print_tuple_elements(io, B.indices) + return print(io, "]") end Base.show(io::IO, ::MIME"text/plain", B::BlockIndexVector) = show(io, B) function Base.getindex(b::AbstractBlockedUnitRange, Kkr::BlockIndexVector{1}) - return b[block(Kkr)][Kkr.indices...] + return b[block(Kkr)][Kkr.indices...] end using ArrayLayouts: LayoutArray -@propagate_inbounds Base.getindex(b::AbstractArray{T,N}, K::BlockIndexVector{N}) where {T,N} = b[block( - K -)][K.indices...] -@propagate_inbounds Base.getindex(b::LayoutArray{T,N}, K::BlockIndexVector{N}) where {T,N} = b[block( - K -)][K.indices...] -@propagate_inbounds Base.getindex(b::LayoutArray{T,1}, K::BlockIndexVector{1}) where {T} = b[block( - K -)][K.indices...] +@propagate_inbounds Base.getindex(b::AbstractArray{T, N}, K::BlockIndexVector{N}) where {T, N} = b[ + block( + K + ), +][K.indices...] +@propagate_inbounds Base.getindex(b::LayoutArray{T, N}, K::BlockIndexVector{N}) where {T, N} = b[ + block( + K + ), +][K.indices...] +@propagate_inbounds Base.getindex(b::LayoutArray{T, 1}, K::BlockIndexVector{1}) where {T} = b[ + block( + K + ), +][K.indices...] function blockedunitrange_getindices( - a::AbstractBlockedUnitRange, - indices::BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexVector{1}}}, -) - blks = map(b -> a[b], blocks(indices)) - # Preserve any extra structure in the axes, like a - # Kronecker structure, symmetry sectors, etc. - ax = mortar_axis(map(b -> axis(a[b]), blocks(indices))) - return mortar(blks, (ax,)) + a::AbstractBlockedUnitRange, + indices::BlockVector{<:BlockIndex{1}, <:Vector{<:BlockIndexVector{1}}}, + ) + blks = map(b -> a[b], blocks(indices)) + # Preserve any extra structure in the axes, like a + # Kronecker structure, symmetry sectors, etc. + ax = mortar_axis(map(b -> axis(a[b]), blocks(indices))) + return mortar(blks, (ax,)) end function blockedunitrange_getindices( - a::AbstractBlockedUnitRange, - indices::BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector{1}}}, -) - blks = map(b -> a[b], blocks(indices)) - # Preserve any extra structure in the axes, like a - # Kronecker structure, symmetry sectors, etc. - ax = mortar_axis(map(b -> axis(a[b]), blocks(indices))) - return mortar(blks, (ax,)) + a::AbstractBlockedUnitRange, + indices::BlockVector{<:GenericBlockIndex{1}, <:Vector{<:BlockIndexVector{1}}}, + ) + blks = map(b -> a[b], blocks(indices)) + # Preserve any extra structure in the axes, like a + # Kronecker structure, symmetry sectors, etc. + ax = mortar_axis(map(b -> axis(a[b]), blocks(indices))) + return mortar(blks, (ax,)) end # This is a specialization of `BlockArrays.unblock`: @@ -395,18 +401,18 @@ end # BlockArrays.jl. # TODO: Ideally this would be defined in BlockArrays.jl once the slicing # there is made more generic. -function BlockArrays.unblock(A, inds, I::Tuple{BlockIndexVector{1},Vararg{Any}}) - B = first(I) - return _blockslice(B, inds[1][B]) +function BlockArrays.unblock(A, inds, I::Tuple{BlockIndexVector{1}, Vararg{Any}}) + B = first(I) + return _blockslice(B, inds[1][B]) end function to_blockindices(a::AbstractBlockedUnitRange{<:Integer}, I::AbstractArray{Bool}) - I_blocks = blocks(BlockedVector(I, blocklengths(a))) - I′_blocks = map(eachindex(I_blocks)) do b - I_b = findall(I_blocks[b]) - return BlockIndexVector(Block(b), I_b) - end - return mortar(filter(!isempty, I′_blocks)) + I_blocks = blocks(BlockedVector(I, blocklengths(a))) + I′_blocks = map(eachindex(I_blocks)) do b + I_b = findall(I_blocks[b]) + return BlockIndexVector(Block(b), I_b) + end + return mortar(filter(!isempty, I′_blocks)) end # This handles non-blocked slices. diff --git a/src/BlockArraysExtensions/blockrange.jl b/src/BlockArraysExtensions/blockrange.jl index ea991e13..21f8a089 100644 --- a/src/BlockArraysExtensions/blockrange.jl +++ b/src/BlockArraysExtensions/blockrange.jl @@ -1,38 +1,38 @@ using BlockArrays: BlockArrays, AbstractBlockedUnitRange, Block, blockedrange, blocklasts -struct BlockUnitRange{T,B,CS,R<:AbstractBlockedUnitRange{T,CS}} <: - AbstractBlockedUnitRange{T,CS} - r::R - eachblockaxis::B +struct BlockUnitRange{T, B, CS, R <: AbstractBlockedUnitRange{T, CS}} <: + AbstractBlockedUnitRange{T, CS} + r::R + eachblockaxis::B end function blockrange(eachblockaxis) - return BlockUnitRange(blockedrange(length.(eachblockaxis)), eachblockaxis) + return BlockUnitRange(blockedrange(length.(eachblockaxis)), eachblockaxis) end function blockrange(first::Integer, eachblockaxis) - return BlockUnitRange(blockedrange(first, length.(eachblockaxis)), eachblockaxis) + return BlockUnitRange(blockedrange(first, length.(eachblockaxis)), eachblockaxis) end Base.first(r::BlockUnitRange) = first(r.r) Base.last(r::BlockUnitRange) = last(r.r) BlockArrays.blocklasts(r::BlockUnitRange) = blocklasts(r.r) eachblockaxis(r::BlockUnitRange) = r.eachblockaxis function Base.getindex(r::BlockUnitRange, I::Block{1}) - return eachblockaxis(r)[Int(I)] .+ (first(r.r[I]) - 1) + return eachblockaxis(r)[Int(I)] .+ (first(r.r[I]) - 1) end function Base.getindex(r::BlockUnitRange, I::BlockRange{1}) - return blockrange(first(r), eachblockaxis(r)[Int.(I)]) + return blockrange(first(r), eachblockaxis(r)[Int.(I)]) end Base.axes(r::BlockUnitRange) = (blockrange(eachblockaxis(r)),) Base.axes1(r::BlockUnitRange) = blockrange(eachblockaxis(r)) using BlockArrays: BlockedOneTo -const BlockOneTo{T<:Integer,B,CS,R<:BlockedOneTo{T,CS}} = BlockUnitRange{T,B,CS,R} +const BlockOneTo{T <: Integer, B, CS, R <: BlockedOneTo{T, CS}} = BlockUnitRange{T, B, CS, R} Base.axes(S::Base.Slice{<:BlockOneTo}) = (S.indices,) Base.axes1(S::Base.Slice{<:BlockOneTo}) = S.indices Base.unsafe_indices(S::Base.Slice{<:BlockOneTo}) = (S.indices,) function BlockArrays.combine_blockaxes(r1::BlockUnitRange, r2::BlockUnitRange) - if eachblockaxis(r1) ≠ eachblockaxis(r2) - return throw(ArgumentError("BlockUnitRanges must have the same block axes")) - end - return r1 + if eachblockaxis(r1) ≠ eachblockaxis(r2) + return throw(ArgumentError("BlockUnitRanges must have the same block axes")) + end + return r1 end diff --git a/src/BlockArraysSparseArraysBaseExt/BlockArraysSparseArraysBaseExt.jl b/src/BlockArraysSparseArraysBaseExt/BlockArraysSparseArraysBaseExt.jl index 7d104de2..e543e5dd 100644 --- a/src/BlockArraysSparseArraysBaseExt/BlockArraysSparseArraysBaseExt.jl +++ b/src/BlockArraysSparseArraysBaseExt/BlockArraysSparseArraysBaseExt.jl @@ -2,10 +2,10 @@ using BlockArrays: AbstractBlockArray, BlocksView using SparseArraysBase: SparseArraysBase, storedlength function SparseArraysBase.storedlength(a::AbstractBlockArray) - return sum(b -> storedlength(b), blocks(a); init=zero(Int)) + return sum(b -> storedlength(b), blocks(a); init = zero(Int)) end # TODO: Handle `BlocksView` wrapping a sparse array? function SparseArraysBase.eachstoredindex(a::BlocksView) - return CartesianIndices(a) + return CartesianIndices(a) end diff --git a/src/BlockSparseArrays.jl b/src/BlockSparseArrays.jl index 8949bad7..56498147 100644 --- a/src/BlockSparseArrays.jl +++ b/src/BlockSparseArrays.jl @@ -1,12 +1,12 @@ module BlockSparseArrays export BlockSparseArray, - BlockSparseMatrix, - BlockSparseVector, - blockstoredlength, - eachblockstoredindex, - eachstoredblock, - sparsemortar + BlockSparseMatrix, + BlockSparseVector, + blockstoredlength, + eachblockstoredindex, + eachstoredblock, + sparsemortar # possible upstream contributions include("BlockArraysExtensions/blockedunitrange.jl") diff --git a/src/abstractblocksparsearray/abstractblocksparsearray.jl b/src/abstractblocksparsearray/abstractblocksparsearray.jl index 5d5805cd..5d4c7f34 100644 --- a/src/abstractblocksparsearray/abstractblocksparsearray.jl +++ b/src/abstractblocksparsearray/abstractblocksparsearray.jl @@ -1,7 +1,7 @@ using BlockArrays: - BlockArrays, AbstractBlockArray, Block, BlockIndex, BlockedUnitRange, blocks + BlockArrays, AbstractBlockArray, Block, BlockIndex, BlockedUnitRange, blocks -abstract type AbstractBlockSparseArray{T,N} <: AbstractBlockArray{T,N} end +abstract type AbstractBlockSparseArray{T, N} <: AbstractBlockArray{T, N} end ## Base `AbstractArray` interface @@ -11,20 +11,20 @@ Base.axes(::AbstractBlockSparseArray) = error("Not implemented") # TODO: Decide what a good default is. blockstype(arraytype::Type{<:AbstractBlockSparseArray}) = SparseArrayDOK{AbstractArray} function blockstype(arraytype::Type{<:AbstractBlockSparseArray{T}}) where {T} - return SparseArrayDOK{AbstractArray{T}} + return SparseArrayDOK{AbstractArray{T}} end -function blockstype(arraytype::Type{<:AbstractBlockSparseArray{T,N}}) where {T,N} - return SparseArrayDOK{AbstractArray{T,N},N} +function blockstype(arraytype::Type{<:AbstractBlockSparseArray{T, N}}) where {T, N} + return SparseArrayDOK{AbstractArray{T, N}, N} end # Specialized in order to fix ambiguity error with `BlockArrays`. -function Base.getindex(a::AbstractBlockSparseArray{<:Any,N}, I::Vararg{Int,N}) where {N} - return @interface interface(a) getindex(a, I...) +function Base.getindex(a::AbstractBlockSparseArray{<:Any, N}, I::Vararg{Int, N}) where {N} + return @interface interface(a) getindex(a, I...) end # Specialized in order to fix ambiguity error with `BlockArrays`. -function Base.getindex(a::AbstractBlockSparseArray{<:Any,0}) - return @interface interface(a) getindex(a) +function Base.getindex(a::AbstractBlockSparseArray{<:Any, 0}) + return @interface interface(a) getindex(a) end ## # Fix ambiguity error with `BlockArrays`. @@ -45,22 +45,22 @@ end # Specialized in order to fix ambiguity error with `BlockArrays`. function Base.setindex!( - a::AbstractBlockSparseArray{<:Any,N}, value, I::Vararg{Int,N} -) where {N} - @interface interface(a) setindex!(a, value, I...) - return a + a::AbstractBlockSparseArray{<:Any, N}, value, I::Vararg{Int, N} + ) where {N} + @interface interface(a) setindex!(a, value, I...) + return a end # Fix ambiguity error. -function Base.setindex!(a::AbstractBlockSparseArray{<:Any,0}, value) - @interface interface(a) setindex!(a, value) - return a +function Base.setindex!(a::AbstractBlockSparseArray{<:Any, 0}, value) + @interface interface(a) setindex!(a, value) + return a end # Catch zero-dimensional case to avoid scalar indexing. -function Base.setindex!(a::AbstractBlockSparseArray{<:Any,0}, value, ::Block{0}) - blocks(a)[] = value - return a +function Base.setindex!(a::AbstractBlockSparseArray{<:Any, 0}, value, ::Block{0}) + blocks(a)[] = value + return a end # Custom `_convert` works around the issue that @@ -72,83 +72,83 @@ _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)) + 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} - blocksize = ntuple(dim -> length(axes(a, dim)[I[dim]]), N) - if size(value) ≠ blocksize - throw( - DimensionMismatch( - "Trying to set block $(Block(Int.(I)...)), which has a size $blocksize, with data of size $(size(value)).", - ), - ) - end - if isstored(a, I...) - # This writes into existing blocks, or constructs blocks - # using the axes. - aI = @view! a[I...] - aI .= value - else - # 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 `_convert` once we drop support for Julia v1.10. - blocks(a)[Int.(I)...] = _convert(blocktype(a), value) - end - return a + a::AbstractBlockSparseArray{<:Any, N}, value, I::Vararg{Block{1}, N} + ) where {N} + blocksize = ntuple(dim -> length(axes(a, dim)[I[dim]]), N) + if size(value) ≠ blocksize + throw( + DimensionMismatch( + "Trying to set block $(Block(Int.(I)...)), which has a size $blocksize, with data of size $(size(value)).", + ), + ) + end + if isstored(a, I...) + # This writes into existing blocks, or constructs blocks + # using the axes. + aI = @view! a[I...] + aI .= value + else + # 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 `_convert` once we drop support for Julia v1.10. + blocks(a)[Int.(I)...] = _convert(blocktype(a), value) + end + return a end # Copy of `Base.dims2string` defined in `show.jl`. function dims_to_string(d) - isempty(d) && return "0-dimensional" - length(d) == 1 && return "$(d[1])-element" - return join(map(string, d), '×') + isempty(d) && return "0-dimensional" + length(d) == 1 && return "$(d[1])-element" + return join(map(string, d), '×') end # Copy of `BlockArrays.block2string` from `BlockArrays.jl`. block_to_string(b, s) = string(join(map(string, b), '×'), "-blocked ", dims_to_string(s)) using TypeParameterAccessors: type_parameters, unspecify_type_parameters -function concretetype_to_string_truncated(type::Type; param_truncation_length=typemax(Int)) - isconcretetype(type) || throw(ArgumentError("Type must be concrete.")) - alias = Base.make_typealias(type) - base_type, params = if isnothing(alias) - unspecify_type_parameters(type), type_parameters(type) - else - base_type_globalref, params_svec = alias - base_type_globalref.name, params_svec - end - str = string(base_type) - if isempty(params) - return str - end - str *= '{' - param_strings = map(params) do param - param_string = string(param) - if length(param_string) > param_truncation_length - return "…" +function concretetype_to_string_truncated(type::Type; param_truncation_length = typemax(Int)) + isconcretetype(type) || throw(ArgumentError("Type must be concrete.")) + alias = Base.make_typealias(type) + base_type, params = if isnothing(alias) + unspecify_type_parameters(type), type_parameters(type) + else + base_type_globalref, params_svec = alias + base_type_globalref.name, params_svec end - return param_string - end - str *= join(param_strings, ", ") - str *= '}' - return str + str = string(base_type) + if isempty(params) + return str + end + str *= '{' + param_strings = map(params) do param + param_string = string(param) + if length(param_string) > param_truncation_length + return "…" + end + return param_string + end + str *= join(param_strings, ", ") + str *= '}' + return str end function Base.summary(io::IO, a::AbstractBlockSparseArray) - print(io, block_to_string(blocksize(a), size(a))) - print(io, ' ') - print(io, concretetype_to_string_truncated(typeof(a); param_truncation_length=40)) - return nothing + print(io, block_to_string(blocksize(a), size(a))) + print(io, ' ') + print(io, concretetype_to_string_truncated(typeof(a); param_truncation_length = 40)) + return nothing end function Base.showarg(io::IO, a::AbstractBlockSparseArray, toplevel::Bool) - !toplevel && print(io, "::") - print(io, concretetype_to_string_truncated(typeof(a); param_truncation_length=40)) - return nothing + !toplevel && print(io, "::") + print(io, concretetype_to_string_truncated(typeof(a); param_truncation_length = 40)) + return nothing end diff --git a/src/abstractblocksparsearray/abstractblocksparsematrix.jl b/src/abstractblocksparsearray/abstractblocksparsematrix.jl index 376f408e..929a7e90 100644 --- a/src/abstractblocksparsearray/abstractblocksparsematrix.jl +++ b/src/abstractblocksparsearray/abstractblocksparsematrix.jl @@ -1,84 +1,84 @@ -const AbstractBlockSparseMatrix{T} = AbstractBlockSparseArray{T,2} +const AbstractBlockSparseMatrix{T} = AbstractBlockSparseArray{T, 2} # SVD is implemented by trying to # 1. Attempt to find a block-diagonal implementation by permuting # 2. Fallback to AbstractBlockArray implementation via BlockedArray function eigencopy_oftype(A::AbstractBlockSparseMatrix, T) - if is_block_permutation_matrix(A) - Acopy = similar(A, T) - for bI in eachblockstoredindex(A) - Acopy[bI] = eigencopy_oftype(@view!(A[bI]), T) + if is_block_permutation_matrix(A) + Acopy = similar(A, T) + for bI in eachblockstoredindex(A) + Acopy[bI] = eigencopy_oftype(@view!(A[bI]), T) + end + return Acopy + else + return BlockedMatrix{T}(A) end - return Acopy - else - return BlockedMatrix{T}(A) - end end function is_block_permutation_matrix(a::AbstractBlockSparseMatrix) - return allunique(first ∘ Tuple, eachblockstoredindex(a)) && - allunique(last ∘ Tuple, eachblockstoredindex(a)) + return allunique(first ∘ Tuple, eachblockstoredindex(a)) && + allunique(last ∘ Tuple, eachblockstoredindex(a)) end function _allocate_svd_output(A::AbstractBlockSparseMatrix, full::Bool, ::Algorithm) - @assert !full "TODO" - bm, bn = blocksize(A) - bmn = min(bm, bn) + @assert !full "TODO" + bm, bn = blocksize(A) + bmn = min(bm, bn) - brows = blocklengths(axes(A, 1)) - bcols = blocklengths(axes(A, 2)) - slengths = Vector{Int}(undef, bmn) + brows = blocklengths(axes(A, 1)) + bcols = blocklengths(axes(A, 2)) + slengths = Vector{Int}(undef, bmn) - # fill in values for blocks that are present - bIs = collect(eachblockstoredindex(A)) - browIs = Int.(first.(Tuple.(bIs))) - bcolIs = Int.(last.(Tuple.(bIs))) - for bI in eachblockstoredindex(A) - row, col = Int.(Tuple(bI)) - nrows = brows[row] - ncols = bcols[col] - slengths[col] = min(nrows, ncols) - end + # fill in values for blocks that are present + bIs = collect(eachblockstoredindex(A)) + browIs = Int.(first.(Tuple.(bIs))) + bcolIs = Int.(last.(Tuple.(bIs))) + for bI in eachblockstoredindex(A) + row, col = Int.(Tuple(bI)) + nrows = brows[row] + ncols = bcols[col] + slengths[col] = min(nrows, ncols) + end - # fill in values for blocks that aren't present, pairing them in order of occurence - # this is a convention, which at least gives the expected results for blockdiagonal - emptyrows = setdiff(1:bm, browIs) - emptycols = setdiff(1:bn, bcolIs) - for (row, col) in zip(emptyrows, emptycols) - slengths[col] = min(brows[row], bcols[col]) - end + # fill in values for blocks that aren't present, pairing them in order of occurence + # this is a convention, which at least gives the expected results for blockdiagonal + emptyrows = setdiff(1:bm, browIs) + emptycols = setdiff(1:bn, bcolIs) + for (row, col) in zip(emptyrows, emptycols) + slengths[col] = min(brows[row], bcols[col]) + end - s_axis = blockedrange(slengths) - U = similar(A, axes(A, 1), s_axis) - S = similar(A, real(eltype(A)), s_axis) - Vt = similar(A, s_axis, axes(A, 2)) + s_axis = blockedrange(slengths) + U = similar(A, axes(A, 1), s_axis) + S = similar(A, real(eltype(A)), s_axis) + Vt = similar(A, s_axis, axes(A, 2)) - # also fill in identities for blocks that aren't present - for (row, col) in zip(emptyrows, emptycols) - copyto!(@view!(U[Block(row, col)]), LinearAlgebra.I) - copyto!(@view!(Vt[Block(col, col)]), LinearAlgebra.I) - end + # also fill in identities for blocks that aren't present + for (row, col) in zip(emptyrows, emptycols) + copyto!(@view!(U[Block(row, col)]), LinearAlgebra.I) + copyto!(@view!(Vt[Block(col, col)]), LinearAlgebra.I) + end - return U, S, Vt + return U, S, Vt end function svd(A::AbstractBlockSparseMatrix; kwargs...) - return svd!(eigencopy_oftype(A, LinearAlgebra.eigtype(eltype(A))); kwargs...) + return svd!(eigencopy_oftype(A, LinearAlgebra.eigtype(eltype(A))); kwargs...) end function svd!( - A::AbstractBlockSparseMatrix; full::Bool=false, alg::Algorithm=default_svd_alg(A) -) - @assert is_block_permutation_matrix(A) "Cannot keep sparsity: use `svd` to convert to `BlockedMatrix" - U, S, Vt = _allocate_svd_output(A, full, alg) - for bI in eachblockstoredindex(A) - bUSV = svd!(@view!(A[bI]); full, alg) - brow, bcol = Tuple(bI) - U[brow, bcol] = bUSV.U - S[bcol] = bUSV.S - Vt[bcol, bcol] = bUSV.Vt - end + A::AbstractBlockSparseMatrix; full::Bool = false, alg::Algorithm = default_svd_alg(A) + ) + @assert is_block_permutation_matrix(A) "Cannot keep sparsity: use `svd` to convert to `BlockedMatrix" + U, S, Vt = _allocate_svd_output(A, full, alg) + for bI in eachblockstoredindex(A) + bUSV = svd!(@view!(A[bI]); full, alg) + brow, bcol = Tuple(bI) + U[brow, bcol] = bUSV.U + S[bcol] = bUSV.S + Vt[bcol, bcol] = bUSV.Vt + end - return SVD(U, S, Vt) + return SVD(U, S, Vt) end diff --git a/src/abstractblocksparsearray/abstractblocksparsevector.jl b/src/abstractblocksparsearray/abstractblocksparsevector.jl index ae1441c5..0e1b4de4 100644 --- a/src/abstractblocksparsearray/abstractblocksparsevector.jl +++ b/src/abstractblocksparsearray/abstractblocksparsevector.jl @@ -1 +1 @@ -const AbstractBlockSparseVector{T} = AbstractBlockSparseArray{T,1} +const AbstractBlockSparseVector{T} = AbstractBlockSparseArray{T, 1} diff --git a/src/abstractblocksparsearray/arraylayouts.jl b/src/abstractblocksparsearray/arraylayouts.jl index 44764385..fd3079d5 100644 --- a/src/abstractblocksparsearray/arraylayouts.jl +++ b/src/abstractblocksparsearray/arraylayouts.jl @@ -4,82 +4,82 @@ using SparseArraysBase: SparseLayout using TypeParameterAccessors: parenttype, similartype function ArrayLayouts.MemoryLayout(arraytype::Type{<:AnyAbstractBlockSparseArray}) - outer_layout = typeof(MemoryLayout(blockstype(arraytype))) - inner_layout = typeof(MemoryLayout(blocktype(arraytype))) - return BlockLayout{outer_layout,inner_layout}() + outer_layout = typeof(MemoryLayout(blockstype(arraytype))) + inner_layout = typeof(MemoryLayout(blocktype(arraytype))) + return BlockLayout{outer_layout, inner_layout}() end # TODO: Generalize to `BlockSparseVectorLike`/`AnyBlockSparseVector`. function ArrayLayouts.MemoryLayout( - arraytype::Type{<:Adjoint{<:Any,<:AbstractBlockSparseVector}} -) - return DualLayout{typeof(MemoryLayout(parenttype(arraytype)))}() + arraytype::Type{<:Adjoint{<:Any, <:AbstractBlockSparseVector}} + ) + return DualLayout{typeof(MemoryLayout(parenttype(arraytype)))}() end # TODO: Generalize to `BlockSparseVectorLike`/`AnyBlockSparseVector`. function ArrayLayouts.MemoryLayout( - arraytype::Type{<:Transpose{<:Any,<:AbstractBlockSparseVector}} -) - return DualLayout{typeof(MemoryLayout(parenttype(arraytype)))}() + arraytype::Type{<:Transpose{<:Any, <:AbstractBlockSparseVector}} + ) + return DualLayout{typeof(MemoryLayout(parenttype(arraytype)))}() end function Base.similar( - mul::MulAdd{ - <:BlockLayout{<:SparseLayout,BlockLayoutA}, - <:BlockLayout{<:SparseLayout,BlockLayoutB}, - LayoutC, - T, - A, - B, - C, - }, - elt::Type, - axes, -) where {BlockLayoutA,BlockLayoutB,LayoutC,T,A,B,C} + mul::MulAdd{ + <:BlockLayout{<:SparseLayout, BlockLayoutA}, + <:BlockLayout{<:SparseLayout, BlockLayoutB}, + LayoutC, + T, + A, + B, + C, + }, + elt::Type, + 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: 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.β) - # ) + # 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) + 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. function ArrayLayouts.sub_materialize(layout::BlockLayout{<:SparseLayout}, a, axes) - # TODO: Define `blocktype`/`blockstype` for `SubArray` wrapping `BlockSparseArray`. - # TODO: Use `similar`? - blocktype_a = blocktype(parent(a)) - a_dest = BlockSparseArray{eltype(a),length(axes),blocktype_a}(undef, axes) - a_dest .= a - return a_dest + # TODO: Define `blocktype`/`blockstype` for `SubArray` wrapping `BlockSparseArray`. + # TODO: Use `similar`? + blocktype_a = blocktype(parent(a)) + a_dest = BlockSparseArray{eltype(a), length(axes), blocktype_a}(undef, axes) + a_dest .= a + return a_dest end function _similar(arraytype::Type{<:AbstractArray}, size::Tuple) - return similar(arraytype, size) + return similar(arraytype, size) end function _similar( - ::Type{<:SubArray{<:Any,<:Any,<:ArrayType}}, size::Tuple -) where {ArrayType} - return similar(ArrayType, size) + ::Type{<:SubArray{<:Any, <:Any, <:ArrayType}}, size::Tuple + ) where {ArrayType} + return similar(ArrayType, size) end # Materialize a SubArray view. function ArrayLayouts.sub_materialize( - layout::BlockLayout{<:SparseLayout}, a, axes::Tuple{Vararg{Base.OneTo}} -) - a_dest = _similar(blocktype(a), length.(axes)) - a_dest .= a - return a_dest + layout::BlockLayout{<:SparseLayout}, a, axes::Tuple{Vararg{Base.OneTo}} + ) + a_dest = _similar(blocktype(a), length.(axes)) + a_dest .= a + return a_dest end diff --git a/src/abstractblocksparsearray/broadcast.jl b/src/abstractblocksparsearray/broadcast.jl index ecef502c..d2868ee8 100644 --- a/src/abstractblocksparsearray/broadcast.jl +++ b/src/abstractblocksparsearray/broadcast.jl @@ -2,67 +2,67 @@ using BlockArrays: AbstractBlockedUnitRange, BlockSlice using Base.Broadcast: Broadcast, BroadcastStyle function Broadcast.BroadcastStyle(arraytype::Type{<:AnyAbstractBlockSparseArray}) - return BlockSparseArrayStyle(BroadcastStyle(blocktype(arraytype))) + return BlockSparseArrayStyle(BroadcastStyle(blocktype(arraytype))) end # Fix ambiguity error with `BlockArrays`. function Broadcast.BroadcastStyle( - arraytype::Type{ - <:SubArray{ - <:Any, - <:Any, - <:AbstractBlockSparseArray, - <:Tuple{BlockSlice{<:Any,<:Any,<:AbstractBlockedUnitRange},Vararg{Any}}, - }, - }, -) - return BlockSparseArrayStyle{ndims(arraytype)}() + arraytype::Type{ + <:SubArray{ + <:Any, + <:Any, + <:AbstractBlockSparseArray, + <:Tuple{BlockSlice{<:Any, <:Any, <:AbstractBlockedUnitRange}, Vararg{Any}}, + }, + }, + ) + return BlockSparseArrayStyle{ndims(arraytype)}() end function Broadcast.BroadcastStyle( - arraytype::Type{ - <:SubArray{ - <:Any, - <:Any, - <:AbstractBlockSparseArray, - <:Tuple{ - BlockSlice{<:Any,<:Any,<:AbstractBlockedUnitRange}, - BlockSlice{<:Any,<:Any,<:AbstractBlockedUnitRange}, - Vararg{Any}, - }, - }, - }, -) - return BlockSparseArrayStyle{ndims(arraytype)}() + arraytype::Type{ + <:SubArray{ + <:Any, + <:Any, + <:AbstractBlockSparseArray, + <:Tuple{ + BlockSlice{<:Any, <:Any, <:AbstractBlockedUnitRange}, + BlockSlice{<:Any, <:Any, <:AbstractBlockedUnitRange}, + Vararg{Any}, + }, + }, + }, + ) + return BlockSparseArrayStyle{ndims(arraytype)}() end function Broadcast.BroadcastStyle( - arraytype::Type{ - <:SubArray{ - <:Any, - <:Any, - <:AbstractBlockSparseArray, - <:Tuple{Any,BlockSlice{<:Any,<:Any,<:AbstractBlockedUnitRange},Vararg{Any}}, - }, - }, -) - return BlockSparseArrayStyle{ndims(arraytype)}() + arraytype::Type{ + <:SubArray{ + <:Any, + <:Any, + <:AbstractBlockSparseArray, + <:Tuple{Any, BlockSlice{<:Any, <:Any, <:AbstractBlockedUnitRange}, Vararg{Any}}, + }, + }, + ) + return BlockSparseArrayStyle{ndims(arraytype)}() end # These catch cases that aren't caught by the standard # `BlockSparseArrayStyle` definition, and also fix # ambiguity issues. function Base.copyto!(dest::AnyAbstractBlockSparseArray, bc::Broadcasted) - copyto_blocksparse!(dest, bc) - return dest + copyto_blocksparse!(dest, bc) + return dest end function Base.copyto!( - dest::AnyAbstractBlockSparseArray, bc::Broadcasted{<:Base.Broadcast.AbstractArrayStyle{0}} -) - copyto_blocksparse!(dest, bc) - return dest + dest::AnyAbstractBlockSparseArray, bc::Broadcasted{<:Base.Broadcast.AbstractArrayStyle{0}} + ) + copyto_blocksparse!(dest, bc) + return dest end function Base.copyto!( - dest::AnyAbstractBlockSparseArray{<:Any,N}, bc::Broadcasted{BlockSparseArrayStyle{N}} -) where {N} - copyto_blocksparse!(dest, bc) - return dest + dest::AnyAbstractBlockSparseArray{<:Any, N}, bc::Broadcasted{BlockSparseArrayStyle{N}} + ) where {N} + copyto_blocksparse!(dest, bc) + return dest end diff --git a/src/abstractblocksparsearray/cat.jl b/src/abstractblocksparsearray/cat.jl index cdefdcf9..2e31cb3c 100644 --- a/src/abstractblocksparsearray/cat.jl +++ b/src/abstractblocksparsearray/cat.jl @@ -2,5 +2,5 @@ using DerivableInterfaces: @interface, interface using DerivableInterfaces.Concatenate: concatenate function Base._cat(dims, as::AnyAbstractBlockSparseArray...) - return concatenate(dims, as...) + return concatenate(dims, as...) end diff --git a/src/abstractblocksparsearray/linearalgebra.jl b/src/abstractblocksparsearray/linearalgebra.jl index 2b03de9e..d87a50e2 100644 --- a/src/abstractblocksparsearray/linearalgebra.jl +++ b/src/abstractblocksparsearray/linearalgebra.jl @@ -4,160 +4,162 @@ using LinearAlgebra: LinearAlgebra, Adjoint, Transpose, diag, norm, tr # but also takes the dual of the axes. # Fixes an issue raised in: # https://github.com/ITensor/ITensors.jl/issues/1336#issuecomment-2353434147 -function Base.copy(a::Adjoint{T,<:AbstractBlockSparseMatrix{T}}) where {T} - a_dest = similar(parent(a), axes(a)) - a_dest .= a - return a_dest +function Base.copy(a::Adjoint{T, <:AbstractBlockSparseMatrix{T}}) where {T} + a_dest = similar(parent(a), axes(a)) + a_dest .= a + return a_dest end # More efficient than the generic `LinearAlgebra` version. -function Base.copy(a::Transpose{T,<:AbstractBlockSparseMatrix{T}}) where {T} - a_dest = similar(parent(a), axes(a)) - a_dest .= a - return a_dest +function Base.copy(a::Transpose{T, <:AbstractBlockSparseMatrix{T}}) where {T} + a_dest = similar(parent(a), axes(a)) + a_dest .= a + return a_dest end -function LinearAlgebra.norm(a::AnyAbstractBlockSparseArray, p::Real=2) - nrmᵖ = float(norm(zero(eltype(a)))) - for I in eachblockstoredindex(a) - nrmᵖ += norm(@view(a[I]), p)^p - end - return nrmᵖ^(1/p) +function LinearAlgebra.norm(a::AnyAbstractBlockSparseArray, p::Real = 2) + nrmᵖ = float(norm(zero(eltype(a)))) + for I in eachblockstoredindex(a) + nrmᵖ += norm(@view(a[I]), p)^p + end + return nrmᵖ^(1 / p) end function LinearAlgebra.tr(a::AnyAbstractBlockSparseMatrix) - tr_a = zero(eltype(a)) - for I in eachstoredblockdiagindex(a) - tr_a += tr(@view(a[I])) - end - return tr_a + tr_a = zero(eltype(a)) + for I in eachstoredblockdiagindex(a) + tr_a += tr(@view(a[I])) + end + return tr_a end # TODO: Define in DiagonalArrays.jl. function diagaxis(a::AbstractArray) - LinearAlgebra.checksquare(a) - return axes(a, 1) + LinearAlgebra.checksquare(a) + return axes(a, 1) end function LinearAlgebra.diag(a::AnyAbstractBlockSparseMatrix) - # TODO: Add `checkblocksquare` to also check it is square blockwise. - LinearAlgebra.checksquare(a) - diagaxes = map(blockdiagindices(a)) do I - return diagaxis(@view(a[I])) - end - r = blockrange(diagaxes) - stored_blocks = Dict(( - Tuple(I)[1] => diag(@view!(a[I])) for I in eachstoredblockdiagindex(a) - )) - return blocksparse(stored_blocks, (r,)) + # TODO: Add `checkblocksquare` to also check it is square blockwise. + LinearAlgebra.checksquare(a) + diagaxes = map(blockdiagindices(a)) do I + return diagaxis(@view(a[I])) + end + r = blockrange(diagaxes) + stored_blocks = Dict( + ( + Tuple(I)[1] => diag(@view!(a[I])) for I in eachstoredblockdiagindex(a) + ) + ) + return blocksparse(stored_blocks, (r,)) end # TODO: Define `SparseArraysBase.isdiag`, define as # `isdiag(blocks(a))`. function blockisdiag(a::AbstractArray) - return all(eachblockstoredindex(a)) do I - return allequal(Tuple(I)) - end + return all(eachblockstoredindex(a)) do I + return allequal(Tuple(I)) + end end const MATRIX_FUNCTIONS = [ - :exp, - :cis, - :log, - :sqrt, - :cbrt, - :cos, - :sin, - :tan, - :csc, - :sec, - :cot, - :cosh, - :sinh, - :tanh, - :csch, - :sech, - :coth, - :acos, - :asin, - :atan, - :acsc, - :asec, - :acot, - :acosh, - :asinh, - :atanh, - :acsch, - :asech, - :acoth, + :exp, + :cis, + :log, + :sqrt, + :cbrt, + :cos, + :sin, + :tan, + :csc, + :sec, + :cot, + :cosh, + :sinh, + :tanh, + :csch, + :sech, + :coth, + :acos, + :asin, + :atan, + :acsc, + :asec, + :acot, + :acosh, + :asinh, + :atanh, + :acsch, + :asech, + :acoth, ] # Functions where the dense implementations in `LinearAlgebra` are # not type stable. const MATRIX_FUNCTIONS_UNSTABLE = [ - :log, - :sqrt, - :acos, - :asin, - :atan, - :acsc, - :asec, - :acot, - :acosh, - :asinh, - :atanh, - :acsch, - :asech, - :acoth, + :log, + :sqrt, + :acos, + :asin, + :atan, + :acsc, + :asec, + :acot, + :acosh, + :asinh, + :atanh, + :acsch, + :asech, + :acoth, ] function initialize_output_blocksparse(f::F, a::AbstractMatrix) where {F} - blockt = Base.promote_op(f, blocktype(a)) - elt′ = Base.promote_op(f, eltype(a)) - blockt′ = if !(blockt <: AbstractMatrix{elt′}) || blockt === Union{} - AbstractMatrix{elt′} - else - blockt - end - return similar(a, BlockType(blockt′)) + blockt = Base.promote_op(f, blocktype(a)) + elt′ = Base.promote_op(f, eltype(a)) + blockt′ = if !(blockt <: AbstractMatrix{elt′}) || blockt === Union{} + AbstractMatrix{elt′} + else + blockt + end + return similar(a, BlockType(blockt′)) end function matrix_function_blocksparse(f::F, a::AbstractMatrix; kwargs...) where {F} - blockisdiag(a) || throw(ArgumentError("`$f` only defined for block-diagonal matrices")) - fa = initialize_output_blocksparse(f, a) - for I in blockdiagindices(a) - fa[I] = f(a[I]; kwargs...) - end - return fa + blockisdiag(a) || throw(ArgumentError("`$f` only defined for block-diagonal matrices")) + fa = initialize_output_blocksparse(f, a) + for I in blockdiagindices(a) + fa[I] = f(a[I]; kwargs...) + end + return fa end for f in MATRIX_FUNCTIONS - @eval begin - function Base.$f(a::AnyAbstractBlockSparseMatrix) - return matrix_function_blocksparse($f, a) + @eval begin + function Base.$f(a::AnyAbstractBlockSparseMatrix) + return matrix_function_blocksparse($f, a) + end end - end end for f in MATRIX_FUNCTIONS_UNSTABLE - @eval begin - function initialize_output_blocksparse(::typeof($f), a::AbstractMatrix) - elt′ = complex(eltype(a)) - blockt = Base.promote_op(similar, blocktype(a), elt′) - blockt′ = if !(blockt <: AbstractMatrix{elt′}) || blockt === Union{} - AbstractMatrix{elt′} - else - blockt - end - return similar(a, BlockType(blockt′)) + @eval begin + function initialize_output_blocksparse(::typeof($f), a::AbstractMatrix) + elt′ = complex(eltype(a)) + blockt = Base.promote_op(similar, blocktype(a), elt′) + blockt′ = if !(blockt <: AbstractMatrix{elt′}) || blockt === Union{} + AbstractMatrix{elt′} + else + blockt + end + return similar(a, BlockType(blockt′)) + end end - end end function LinearAlgebra.inv(a::AnyAbstractBlockSparseMatrix) - return matrix_function_blocksparse(inv, a) + return matrix_function_blocksparse(inv, a) end using LinearAlgebra: LinearAlgebra, pinv function LinearAlgebra.pinv(a::AnyAbstractBlockSparseMatrix; kwargs...) - return matrix_function_blocksparse(pinv, a; kwargs...) + return matrix_function_blocksparse(pinv, a; kwargs...) end diff --git a/src/abstractblocksparsearray/map.jl b/src/abstractblocksparsearray/map.jl index 04513fc4..2d969f3f 100644 --- a/src/abstractblocksparsearray/map.jl +++ b/src/abstractblocksparsearray/map.jl @@ -6,77 +6,77 @@ using LinearAlgebra: Adjoint, Transpose # If the blocking of the slice doesn't match the blocking of the # parent array, reblock according to the blocking of the parent array. function reblock( - a::SubArray{<:Any,<:Any,<:AbstractBlockSparseArray,<:Tuple{Vararg{AbstractUnitRange}}} -) - # TODO: This relies on the behavior that slicing a block sparse - # array with a UnitRange inherits the blocking of the underlying - # block sparse array, we might change that default behavior - # so this might become something like `@blocked parent(a)[...]`. - return @view parent(a)[UnitRange{Int}.(parentindices(a))...] + a::SubArray{<:Any, <:Any, <:AbstractBlockSparseArray, <:Tuple{Vararg{AbstractUnitRange}}} + ) + # TODO: This relies on the behavior that slicing a block sparse + # array with a UnitRange inherits the blocking of the underlying + # block sparse array, we might change that default behavior + # so this might become something like `@blocked parent(a)[...]`. + return @view parent(a)[UnitRange{Int}.(parentindices(a))...] end # TODO: Make this more general, independent of `AbstractBlockSparseArray`. function reblock( - a::SubArray{<:Any,<:Any,<:AbstractBlockSparseArray,<:Tuple{Vararg{NonBlockedArray}}} -) - return @view parent(a)[map(I -> I.array, parentindices(a))...] + a::SubArray{<:Any, <:Any, <:AbstractBlockSparseArray, <:Tuple{Vararg{NonBlockedArray}}} + ) + return @view parent(a)[map(I -> I.array, parentindices(a))...] end # TODO: Make this more general, independent of `AbstractBlockSparseArray`. function reblock( - a::SubArray{ - <:Any, - <:Any, - <:AbstractBlockSparseArray, - <:Tuple{Vararg{BlockIndices{<:AbstractBlockVector{<:Block{1}}}}}, - }, -) - # Remove the blocking. - return @view parent(a)[map(I -> Vector(I.blocks), parentindices(a))...] + a::SubArray{ + <:Any, + <:Any, + <:AbstractBlockSparseArray, + <:Tuple{Vararg{BlockIndices{<:AbstractBlockVector{<:Block{1}}}}}, + }, + ) + # Remove the blocking. + return @view parent(a)[map(I -> Vector(I.blocks), parentindices(a))...] end function Base.map!(f, a_dest::AbstractArray, a_srcs::AnyAbstractBlockSparseArray...) - @interface interface(a_dest, a_srcs...) map!(f, a_dest, a_srcs...) - return a_dest + @interface interface(a_dest, a_srcs...) map!(f, a_dest, a_srcs...) + return a_dest end function Base.map!(f, a_dest::AnyAbstractBlockSparseArray, a_srcs::AbstractArray...) - @interface interface(a_dest, a_srcs...) map!(f, a_dest, a_srcs...) - return a_dest + @interface interface(a_dest, a_srcs...) map!(f, a_dest, a_srcs...) + return a_dest end function Base.map!( - f, a_dest::AnyAbstractBlockSparseArray, a_srcs::AnyAbstractBlockSparseArray... -) - @interface interface(a_dest, a_srcs...) map!(f, a_dest, a_srcs...) - return a_dest + f, a_dest::AnyAbstractBlockSparseArray, a_srcs::AnyAbstractBlockSparseArray... + ) + @interface interface(a_dest, a_srcs...) map!(f, a_dest, a_srcs...) + return a_dest end function Base.map(f, as::Vararg{AnyAbstractBlockSparseArray}) - return f.(as...) + return f.(as...) end function Base.copy!(a_dest::AbstractArray, a_src::AnyAbstractBlockSparseArray) - return @interface interface(a_src) copy!(a_dest, a_src) + return @interface interface(a_src) copy!(a_dest, a_src) end function Base.copyto!(a_dest::AbstractArray, a_src::AnyAbstractBlockSparseArray) - return @interface interface(a_src) copyto!(a_dest, a_src) + return @interface interface(a_src) copyto!(a_dest, a_src) end # Fix ambiguity error function Base.copyto!(a_dest::LayoutArray, a_src::AnyAbstractBlockSparseArray) - return @interface interface(a_src) copyto!(a_dest, a_src) + return @interface interface(a_src) copyto!(a_dest, a_src) end function Base.copyto!( - a_dest::AbstractMatrix, a_src::Transpose{T,<:AbstractBlockSparseMatrix{T}} -) where {T} - return @interface interface(a_src) copyto!(a_dest, a_src) + a_dest::AbstractMatrix, a_src::Transpose{T, <:AbstractBlockSparseMatrix{T}} + ) where {T} + return @interface interface(a_src) copyto!(a_dest, a_src) end function Base.copyto!( - a_dest::AbstractMatrix, a_src::Adjoint{T,<:AbstractBlockSparseMatrix{T}} -) where {T} - return @interface interface(a_src) copyto!(a_dest, a_src) + a_dest::AbstractMatrix, a_src::Adjoint{T, <:AbstractBlockSparseMatrix{T}} + ) where {T} + return @interface interface(a_src) copyto!(a_dest, a_src) end # This avoids going through the generic version that calls `Base.permutedims!`, @@ -85,7 +85,7 @@ end # `PermutedDimsArray`). # TODO: Handle slicing better in `map!` so that this can be removed. function Base.permutedims(a::AnyAbstractBlockSparseArray, perm) - @interface interface(a) permutedims(a, perm) + return @interface interface(a) permutedims(a, perm) end # The `::AbstractBlockSparseArrayInterface` version @@ -96,30 +96,30 @@ end # ``` # TODO: Handle slicing better in `map!` so that this can be removed. function Base.permutedims!(a_dest, a_src::AnyAbstractBlockSparseArray, perm) - return @interface interface(a_src) permutedims!(a_dest, a_src, perm) + return @interface interface(a_src) permutedims!(a_dest, a_src, perm) end function Base.mapreduce(f, op, as::AnyAbstractBlockSparseArray...; kwargs...) - return @interface interface(as...) mapreduce(f, op, as...; kwargs...) + return @interface interface(as...) mapreduce(f, op, as...; kwargs...) end function Base.iszero(a::AnyAbstractBlockSparseArray) - return @interface interface(a) iszero(a) + return @interface interface(a) iszero(a) end function Base.isreal(a::AnyAbstractBlockSparseArray) - return @interface interface(a) isreal(a) + return @interface interface(a) isreal(a) end # Helps with specialization of block operations by avoiding # having anonymous functions constructed inside the map/broadcast # code logic. function Base.:*(x::Number, a::AnyAbstractBlockSparseArray) - return map(Base.Fix1(*, x), a) + return map(Base.Fix1(*, x), a) end function Base.:*(a::AnyAbstractBlockSparseArray, x::Number) - return map(Base.Fix2(*, x), a) + return map(Base.Fix2(*, x), a) end function Base.:/(a::AnyAbstractBlockSparseArray, x::Number) - return map(Base.Fix2(/, x), a) + return map(Base.Fix2(/, x), a) end diff --git a/src/abstractblocksparsearray/sparsearrayinterface.jl b/src/abstractblocksparsearray/sparsearrayinterface.jl index bbb008ab..23dd9df7 100644 --- a/src/abstractblocksparsearray/sparsearrayinterface.jl +++ b/src/abstractblocksparsearray/sparsearrayinterface.jl @@ -4,31 +4,31 @@ using SparseArraysBase: SparseArraysBase, eachstoredindex, storedlength, storedv # Structure storing the block sparse storage # TODO: Delete this in favor of `storedvalues(blocks(a))`, # and rename `storedblocks(a)` and/or `eachstoredblock(a)`. -struct BlockSparseStorage{Arr<:AbstractBlockSparseArray} - array::Arr +struct BlockSparseStorage{Arr <: AbstractBlockSparseArray} + array::Arr end function blockindex_to_cartesianindex(a::AbstractArray, blockindex) - return CartesianIndex(getindex.(axes(a), getindex.(Block.(blockindex.I), blockindex.α))) + return CartesianIndex(getindex.(axes(a), getindex.(Block.(blockindex.I), blockindex.α))) end function Base.keys(s::BlockSparseStorage) - stored_blockindices = Iterators.map(eachstoredindex(blocks(s.array))) do I - block_axes = axes(blocks(s.array)[I]) - blockindices = Block(Tuple(I))[block_axes...] - return Iterators.map( - blockindex -> blockindex_to_cartesianindex(s.array, blockindex), blockindices - ) - end - return Iterators.flatten(stored_blockindices) + stored_blockindices = Iterators.map(eachstoredindex(blocks(s.array))) do I + block_axes = axes(blocks(s.array)[I]) + blockindices = Block(Tuple(I))[block_axes...] + return Iterators.map( + blockindex -> blockindex_to_cartesianindex(s.array, blockindex), blockindices + ) + end + return Iterators.flatten(stored_blockindices) end function Base.values(s::BlockSparseStorage) - return Iterators.map(I -> s.array[I], eachindex(s)) + return Iterators.map(I -> s.array[I], eachindex(s)) end function Base.iterate(s::BlockSparseStorage, args...) - return iterate(values(s), args...) + return iterate(values(s), args...) end ## TODO: Bring back this deifinition but check that it makes sense. @@ -38,5 +38,5 @@ end # TODO: Turn this into an `@interface ::AbstractBlockSparseArrayInterface` function. function SparseArraysBase.storedlength(a::AnyAbstractBlockSparseArray) - return sum(storedlength, storedvalues(blocks(a)); init=zero(Int)) + return sum(storedlength, storedvalues(blocks(a)); init = zero(Int)) end diff --git a/src/abstractblocksparsearray/unblockedsubarray.jl b/src/abstractblocksparsearray/unblockedsubarray.jl index a10d5419..4180745b 100644 --- a/src/abstractblocksparsearray/unblockedsubarray.jl +++ b/src/abstractblocksparsearray/unblockedsubarray.jl @@ -4,82 +4,82 @@ using BlockArrays: BlockArrays, Block, BlockIndexRange, BlockSlice using TypeParameterAccessors: TypeParameterAccessors, parenttype, similartype const UnblockedIndices = Union{ - Vector{<:Integer}, - BlockSlice{<:Block{1}}, - BlockSlice{<:BlockIndexRange{1}}, - BlockSlice{<:BlockIndexVector}, + Vector{<:Integer}, + BlockSlice{<:Block{1}}, + BlockSlice{<:BlockIndexRange{1}}, + BlockSlice{<:BlockIndexVector}, } -const UnblockedSubArray{T,N} = SubArray{ - T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{UnblockedIndices}} +const UnblockedSubArray{T, N} = SubArray{ + T, N, <:AbstractBlockSparseArray{T, N}, <:Tuple{Vararg{UnblockedIndices}}, } function BlockArrays.blocks(a::UnblockedSubArray) - return SingleBlockView(a) + return SingleBlockView(a) end function DerivableInterfaces.interface(arraytype::Type{<:UnblockedSubArray}) - return interface(blocktype(parenttype(arraytype))) + return interface(blocktype(parenttype(arraytype))) end function ArrayLayouts.MemoryLayout(arraytype::Type{<:UnblockedSubArray}) - return MemoryLayout(blocktype(parenttype(arraytype))) + return MemoryLayout(blocktype(parenttype(arraytype))) end function Broadcast.BroadcastStyle(arraytype::Type{<:UnblockedSubArray}) - return BroadcastStyle(blocktype(parenttype(arraytype))) + return BroadcastStyle(blocktype(parenttype(arraytype))) end function Base.similar( - a::UnblockedSubArray, elt::Type, axes::Tuple{Base.OneTo,Vararg{Base.OneTo}} -) - return similar(similartype(blocktype(parenttype(a)), elt), axes) + a::UnblockedSubArray, elt::Type, axes::Tuple{Base.OneTo, Vararg{Base.OneTo}} + ) + return similar(similartype(blocktype(parenttype(a)), elt), axes) end -function Base.similar(a::UnblockedSubArray, elt::Type, size::Tuple{Int,Vararg{Int}}) - return similar(a, elt, Base.OneTo.(size)) +function Base.similar(a::UnblockedSubArray, elt::Type, size::Tuple{Int, Vararg{Int}}) + return similar(a, elt, Base.OneTo.(size)) end function ArrayLayouts.sub_materialize(a::UnblockedSubArray) - a_cpu = adapt(Array, a) - a_cpu′ = similar(a_cpu) - a_cpu′ .= a_cpu - if typeof(a) === typeof(a_cpu) - return a_cpu′ - end - a′ = similar(a) - a′ .= a_cpu′ - return a′ + a_cpu = adapt(Array, a) + a_cpu′ = similar(a_cpu) + a_cpu′ .= a_cpu + if typeof(a) === typeof(a_cpu) + return a_cpu′ + end + a′ = similar(a) + a′ .= a_cpu′ + return a′ end function Base.map!( - f, a_dest::AbstractArray, a_src1::UnblockedSubArray, a_src_rest::UnblockedSubArray... -) - return invoke( - map!, - Tuple{Any,AbstractArray,AbstractArray,Vararg{AbstractArray}}, - f, - a_dest, - a_src1, - a_src_rest..., - ) + f, a_dest::AbstractArray, a_src1::UnblockedSubArray, a_src_rest::UnblockedSubArray... + ) + return invoke( + map!, + Tuple{Any, AbstractArray, AbstractArray, Vararg{AbstractArray}}, + f, + a_dest, + a_src1, + a_src_rest..., + ) end # Fix ambiguity and scalar indexing errors with GPUArrays. using Adapt: adapt using GPUArraysCore: GPUArraysCore function Base.map!( - f, - a_dest::GPUArraysCore.AnyGPUArray, - a_src1::UnblockedSubArray, - a_src_rest::UnblockedSubArray..., -) - a_dest_cpu = adapt(Array, a_dest) - a_srcs_cpu = map(adapt(Array), (a_src1, a_src_rest...)) - map!(f, a_dest_cpu, a_srcs_cpu...) - a_dest .= a_dest_cpu - return a_dest + f, + a_dest::GPUArraysCore.AnyGPUArray, + a_src1::UnblockedSubArray, + a_src_rest::UnblockedSubArray..., + ) + a_dest_cpu = adapt(Array, a_dest) + a_srcs_cpu = map(adapt(Array), (a_src1, a_src_rest...)) + map!(f, a_dest_cpu, a_srcs_cpu...) + a_dest .= a_dest_cpu + return a_dest end function Base.iszero(a::UnblockedSubArray) - return invoke(iszero, Tuple{AbstractArray}, adapt(Array, a)) + return invoke(iszero, Tuple{AbstractArray}, adapt(Array, a)) end diff --git a/src/abstractblocksparsearray/views.jl b/src/abstractblocksparsearray/views.jl index 19de8a31..edddc77b 100644 --- a/src/abstractblocksparsearray/views.jl +++ b/src/abstractblocksparsearray/views.jl @@ -1,92 +1,92 @@ using BlockArrays: - AbstractBlockedUnitRange, - BlockArrays, - Block, - BlockIndexRange, - BlockedVector, - blocklength, - blocksize, - viewblock + AbstractBlockedUnitRange, + BlockArrays, + Block, + BlockIndexRange, + BlockedVector, + blocklength, + blocksize, + viewblock # This splits `BlockIndexRange{N}` into # `NTuple{N,BlockIndexRange{1}}`. # TODO: Move to `BlockArraysExtensions`. to_tuple(x) = Tuple(x) function to_tuple(x::BlockIndexRange{N}) where {N} - blocks = Tuple(Block(x)) - n = length(blocks) - return ntuple(dim -> blocks[dim][x.indices[dim]], n) + blocks = Tuple(Block(x)) + n = length(blocks) + return ntuple(dim -> blocks[dim][x.indices[dim]], n) end # Override the default definition of `BlockArrays.blocksize`, # which is incorrect for certain slices. -function BlockArrays.blocksize(a::SubArray{<:Any,<:Any,<:AnyAbstractBlockSparseArray}) - return blocklength.(axes(a)) +function BlockArrays.blocksize(a::SubArray{<:Any, <:Any, <:AnyAbstractBlockSparseArray}) + return blocklength.(axes(a)) end function BlockArrays.blocksize( - a::SubArray{<:Any,<:Any,<:AnyAbstractBlockSparseArray}, i::Int -) - # TODO: Maybe use `blocklength(axes(a, i))` which would be a bit faster. - return blocksize(a)[i] + a::SubArray{<:Any, <:Any, <:AnyAbstractBlockSparseArray}, i::Int + ) + # TODO: Maybe use `blocklength(axes(a, i))` which would be a bit faster. + return blocksize(a)[i] end # These definitions circumvent some generic definitions in BlockArrays.jl: # https://github.com/JuliaArrays/BlockArrays.jl/blob/master/src/views.jl # which don't handle subslices of blocks properly. function Base.view( - a::SubArray{ - <:Any,N,<:AnyAbstractBlockSparseArray,<:Tuple{Vararg{BlockSlice{<:BlockRange{1}},N}} - }, - I::Block{N}, -) where {N} - return @interface interface(a) view(a, I) + a::SubArray{ + <:Any, N, <:AnyAbstractBlockSparseArray, <:Tuple{Vararg{BlockSlice{<:BlockRange{1}}, N}}, + }, + I::Block{N}, + ) where {N} + return @interface interface(a) view(a, I) end function Base.view( - a::SubArray{ - <:Any,N,<:AnyAbstractBlockSparseArray,<:Tuple{Vararg{BlockSlice{<:BlockRange{1}},N}} - }, - I::Vararg{Block{1},N}, -) where {N} - return @interface interface(a) view(a, I...) + a::SubArray{ + <:Any, N, <:AnyAbstractBlockSparseArray, <:Tuple{Vararg{BlockSlice{<:BlockRange{1}}, N}}, + }, + I::Vararg{Block{1}, N}, + ) where {N} + return @interface interface(a) view(a, I...) end function Base.view( - V::SubArray{<:Any,1,<:AnyAbstractBlockSparseArray,<:Tuple{BlockSlice{<:BlockRange{1}}}}, - I::Block{1}, -) - return @interface interface(a) view(a, I) + V::SubArray{<:Any, 1, <:AnyAbstractBlockSparseArray, <:Tuple{BlockSlice{<:BlockRange{1}}}}, + I::Block{1}, + ) + return @interface interface(a) view(a, I) end # Specialized code for getting the view of a block. function BlockArrays.viewblock( - a::AbstractBlockSparseArray{<:Any,N}, block::Block{N} -) where {N} - return viewblock(a, Tuple(block)...) + a::AbstractBlockSparseArray{<:Any, N}, block::Block{N} + ) where {N} + return viewblock(a, Tuple(block)...) end # TODO: Define `@interface interface(a) viewblock`. function BlockArrays.viewblock( - a::AbstractBlockSparseArray{<:Any,N}, block::Vararg{Block{1},N} -) where {N} - I = CartesianIndex(Int.(block)) - # TODO: isblockstored - return isstored(blocks(a), I) ? blocks(a)[I] : BlockView(a, block) + a::AbstractBlockSparseArray{<:Any, N}, block::Vararg{Block{1}, N} + ) where {N} + I = CartesianIndex(Int.(block)) + # TODO: isblockstored + return isstored(blocks(a), I) ? blocks(a)[I] : BlockView(a, block) end # Specialized code for getting the view of a subblock. function Base.view( - a::AbstractBlockSparseArray{<:Any,N}, block::BlockIndexRange{N} -) where {N} - return view(a, to_tuple(block)...) + a::AbstractBlockSparseArray{<:Any, N}, block::BlockIndexRange{N} + ) where {N} + return view(a, to_tuple(block)...) end # Specialized code for getting the view of a subblock. function Base.view( - a::SubArray{T,N,<:AbstractBlockSparseArray{T,N}}, I::BlockIndexRange{N} -) where {T,N} - return view(a, to_tuple(I)...) + a::SubArray{T, N, <:AbstractBlockSparseArray{T, N}}, I::BlockIndexRange{N} + ) where {T, N} + return view(a, to_tuple(I)...) end -function Base.view(a::AbstractBlockSparseArray{<:Any,N}, I::Vararg{Block{1},N}) where {N} - return viewblock(a, I...) +function Base.view(a::AbstractBlockSparseArray{<:Any, N}, I::Vararg{Block{1}, N}) where {N} + return viewblock(a, I...) end # TODO: Move to `GradedUnitRanges` or `BlockArraysExtensions`. @@ -98,116 +98,116 @@ to_block_indices(I::BlockIndexRange{1}) = only(I.indices) to_block_indices(I::BlockIndexVector{1}) = only(I.indices) function Base.view( - a::AbstractBlockSparseArray{<:Any,N}, - I::Vararg{Union{Block{1},BlockIndexRange{1},BlockIndexVector},N}, -) where {N} - return @views a[to_block.(I)...][to_block_indices.(I)...] + a::AbstractBlockSparseArray{<:Any, N}, + I::Vararg{Union{Block{1}, BlockIndexRange{1}, BlockIndexVector}, N}, + ) where {N} + return @views a[to_block.(I)...][to_block_indices.(I)...] end function Base.view( - a::SubArray{T,N,<:AbstractBlockSparseArray{T,N}}, I::Vararg{Block{1},N} -) where {T,N} - return viewblock(a, I...) + a::SubArray{T, N, <:AbstractBlockSparseArray{T, N}}, I::Vararg{Block{1}, N} + ) where {T, N} + return viewblock(a, I...) end function Base.view( - a::SubArray{T,N,<:AbstractBlockSparseArray{T,N}}, - I::Vararg{Union{Block{1},BlockIndexRange{1},BlockIndexVector},N}, -) where {T,N} - return @views a[to_block.(I)...][to_block_indices.(I)...] + a::SubArray{T, N, <:AbstractBlockSparseArray{T, N}}, + I::Vararg{Union{Block{1}, BlockIndexRange{1}, BlockIndexVector}, N}, + ) where {T, N} + return @views a[to_block.(I)...][to_block_indices.(I)...] end # Generic fallback. function BlockArrays.viewblock( - a::SubArray{T,N,<:AbstractBlockSparseArray{T,N}}, I::Vararg{Block{1},N} -) where {T,N} - return Base.invoke(view, Tuple{AbstractArray,Vararg{Any}}, a, I...) + a::SubArray{T, N, <:AbstractBlockSparseArray{T, N}}, I::Vararg{Block{1}, N} + ) where {T, N} + return Base.invoke(view, Tuple{AbstractArray, Vararg{Any}}, a, I...) end function Base.view( - a::SubArray{ - T, - N, - <:AbstractBlockSparseArray{T,N}, - <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}}, - }, - block::Union{Block{N},BlockIndexRange{N}}, -) where {T,N} - return viewblock(a, block) + a::SubArray{ + T, + N, + <:AbstractBlockSparseArray{T, N}, + <:Tuple{Vararg{Union{BlockSliceCollection, SubBlockSliceCollection}, N}}, + }, + block::Union{Block{N}, BlockIndexRange{N}}, + ) where {T, N} + return viewblock(a, block) end # Fix ambiguity error with BlockArrays.jl for slices like # `a = BlockSparseArray{Float64}(undef, [2, 2], [2, 2]); @view a[:, :]`. function Base.view( - a::SubArray{ - T, - N, - <:AbstractBlockSparseArray{T,N}, - <:Tuple{Vararg{Union{Base.Slice,BlockSlice{<:BlockRange{1}}},N}}, - }, - block::Block{N}, -) where {T,N} - return viewblock(a, block) + a::SubArray{ + T, + N, + <:AbstractBlockSparseArray{T, N}, + <:Tuple{Vararg{Union{Base.Slice, BlockSlice{<:BlockRange{1}}}, N}}, + }, + block::Block{N}, + ) where {T, N} + return viewblock(a, block) end function Base.view( - a::SubArray{ - T, - N, - <:AbstractBlockSparseArray{T,N}, - <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}}, - }, - block::Vararg{Union{Block{1},BlockIndexRange{1}},N}, -) where {T,N} - return viewblock(a, block...) + a::SubArray{ + T, + N, + <:AbstractBlockSparseArray{T, N}, + <:Tuple{Vararg{Union{BlockSliceCollection, SubBlockSliceCollection}, N}}, + }, + block::Vararg{Union{Block{1}, BlockIndexRange{1}}, N}, + ) where {T, N} + return viewblock(a, block...) end function BlockArrays.viewblock( - a::SubArray{ - T, - N, - <:AbstractBlockSparseArray{T,N}, - <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}}, - }, - block::Union{Block{N},BlockIndexRange{N},BlockIndexVector{N}}, -) where {T,N} - return viewblock(a, to_tuple(block)...) + a::SubArray{ + T, + N, + <:AbstractBlockSparseArray{T, N}, + <:Tuple{Vararg{Union{BlockSliceCollection, SubBlockSliceCollection}, N}}, + }, + block::Union{Block{N}, BlockIndexRange{N}, BlockIndexVector{N}}, + ) where {T, N} + return viewblock(a, to_tuple(block)...) end # Fixes ambiguity error with `AnyAbstractBlockSparseArray` definition. function Base.view( - a::SubArray{ - T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockSlice{<:BlockRange{1}},N}} - }, - block::Block{N}, -) where {T,N} - return viewblock(a, block) + a::SubArray{ + T, N, <:AbstractBlockSparseArray{T, N}, <:Tuple{Vararg{BlockSlice{<:BlockRange{1}}, N}}, + }, + block::Block{N}, + ) where {T, N} + return viewblock(a, block) end # Fixes ambiguity error with `AnyAbstractBlockSparseArray` definition. function Base.view( - a::SubArray{ - T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockSlice{<:BlockRange{1}},N}} - }, - block::Vararg{Block{1},N}, -) where {T,N} - return viewblock(a, block...) + a::SubArray{ + T, N, <:AbstractBlockSparseArray{T, N}, <:Tuple{Vararg{BlockSlice{<:BlockRange{1}}, N}}, + }, + block::Vararg{Block{1}, N}, + ) where {T, N} + return viewblock(a, block...) end # Disambiguate between block reindexing of blockwise views # (`BlockSliceCollection`) and subblockwise views (`SubBlockSliceCollection`), # which both include `Base.Slice`. function Base.view( - a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{Base.Slice,N}}}, - block::Block{N}, -) where {T,N} - return viewblock(a, block) + a::SubArray{T, N, <:AbstractBlockSparseArray{T, N}, <:Tuple{Vararg{Base.Slice, N}}}, + block::Block{N}, + ) where {T, N} + return viewblock(a, block) end # Block reindexing of blockwise views (`BlockSliceCollection`). -function viewblock_blockslice(a::SubArray{<:Any,N}, block::Vararg{Block{1},N}) where {N} - I = CartesianIndex(Int.(block)) - # TODO: Use `eachblockstoredindex`. - if I ∈ eachstoredindex(blocks(a)) - return blocks(a)[I] - end - return BlockView(parent(a), Block.(Base.reindex(parentindices(blocks(a)), Tuple(I)))) +function viewblock_blockslice(a::SubArray{<:Any, N}, block::Vararg{Block{1}, N}) where {N} + I = CartesianIndex(Int.(block)) + # TODO: Use `eachblockstoredindex`. + if I ∈ eachstoredindex(blocks(a)) + return blocks(a)[I] + end + return BlockView(parent(a), Block.(Base.reindex(parentindices(blocks(a)), Tuple(I)))) end # XXX: TODO: Distinguish if a sub-view of the block needs to be taken! @@ -216,172 +216,172 @@ end # in `blocksparsearrayinterface/blocksparsearrayinterface.jl`. # TODO: Define `@interface interface(a) viewblock`. function BlockArrays.viewblock( - a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockSliceCollection,N}}}, - block::Vararg{Block{1},N}, -) where {T,N} - return viewblock_blockslice(a, block...) + a::SubArray{T, N, <:AbstractBlockSparseArray{T, N}, <:Tuple{Vararg{BlockSliceCollection, N}}}, + block::Vararg{Block{1}, N}, + ) where {T, N} + return viewblock_blockslice(a, block...) end # Disambiguate between block reindexing of blockwise views # (`BlockSliceCollection`) and subblockwise views (`SubBlockSliceCollection`), # which both include `Base.Slice`. function BlockArrays.viewblock( - a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{Base.Slice,N}}}, - block::Vararg{Block{1},N}, -) where {T,N} - return viewblock_blockslice(a, block...) + a::SubArray{T, N, <:AbstractBlockSparseArray{T, N}, <:Tuple{Vararg{Base.Slice, N}}}, + block::Vararg{Block{1}, N}, + ) where {T, N} + return viewblock_blockslice(a, block...) end function to_blockindexrange( - a::BlockSlice{<:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}}, I::Block{1} -) - # TODO: Ideally we would just use `a.blocks[I]` but that doesn't - # work right now. - return blocks(a.block)[Int(I)] + a::BlockSlice{<:BlockVector{<:BlockIndex{1}, <:Vector{<:BlockIndexRange{1}}}}, I::Block{1} + ) + # TODO: Ideally we would just use `a.blocks[I]` but that doesn't + # work right now. + return blocks(a.block)[Int(I)] end function to_blockindexrange( - a::BlockIndices{<:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange}}}, I::Block{1} -) - # TODO: Ideally we would just use `a.blocks[I]` but that doesn't - # work right now. - return blocks(a.blocks)[Int(I)] + a::BlockIndices{<:BlockVector{<:BlockIndex{1}, <:Vector{<:BlockIndexRange}}}, I::Block{1} + ) + # TODO: Ideally we would just use `a.blocks[I]` but that doesn't + # work right now. + return blocks(a.blocks)[Int(I)] end function to_blockindexrange( - a::BlockIndices{<:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexVector}}}, I::Block{1} -) - # TODO: Ideally we would just use `a.blocks[I]` but that doesn't - # work right now. - return blocks(a.blocks)[Int(I)] + a::BlockIndices{<:BlockVector{<:BlockIndex{1}, <:Vector{<:BlockIndexVector}}}, I::Block{1} + ) + # TODO: Ideally we would just use `a.blocks[I]` but that doesn't + # work right now. + return blocks(a.blocks)[Int(I)] end function to_blockindexrange( - a::BlockIndices{<:BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector}}}, - I::Block{1}, -) - # TODO: Ideally we would just use `a.blocks[I]` but that doesn't - # work right now. - return blocks(a.blocks)[Int(I)] + a::BlockIndices{<:BlockVector{<:GenericBlockIndex{1}, <:Vector{<:BlockIndexVector}}}, + I::Block{1}, + ) + # TODO: Ideally we would just use `a.blocks[I]` but that doesn't + # work right now. + return blocks(a.blocks)[Int(I)] end function to_blockindexrange( - a::Base.Slice{<:AbstractBlockedUnitRange{<:Integer}}, I::Block{1} -) - @assert I in only(blockaxes(a.indices)) - return I + a::Base.Slice{<:AbstractBlockedUnitRange{<:Integer}}, I::Block{1} + ) + @assert I in only(blockaxes(a.indices)) + return I end function BlockArrays.viewblock( - a::SubArray{ - T, - N, - <:AbstractBlockSparseArray{T,N}, - <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}}, - }, - block::Vararg{Block{1},N}, -) where {T,N} - brs = ntuple(dim -> to_blockindexrange(parentindices(a)[dim], block[dim]), ndims(a)) - return @view parent(a)[brs...] + a::SubArray{ + T, + N, + <:AbstractBlockSparseArray{T, N}, + <:Tuple{Vararg{Union{BlockSliceCollection, SubBlockSliceCollection}, N}}, + }, + block::Vararg{Block{1}, N}, + ) where {T, N} + brs = ntuple(dim -> to_blockindexrange(parentindices(a)[dim], block[dim]), ndims(a)) + return @view parent(a)[brs...] end # TODO: Define `@interface interface() viewblock`. function BlockArrays.viewblock( - a::SubArray{ - T, - N, - <:AbstractBlockSparseArray{T,N}, - <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}}, - }, - block::Vararg{BlockIndexRange{1},N}, -) where {T,N} - return view(viewblock(a, Block.(block)...), map(b -> only(b.indices), block)...) + a::SubArray{ + T, + N, + <:AbstractBlockSparseArray{T, N}, + <:Tuple{Vararg{Union{BlockSliceCollection, SubBlockSliceCollection}, N}}, + }, + block::Vararg{BlockIndexRange{1}, N}, + ) where {T, N} + return view(viewblock(a, Block.(block)...), map(b -> only(b.indices), block)...) end function Base.view( - a::SubArray{ - T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{SubBlockSliceCollection,N}} - }, - block::Union{Block{N},BlockIndexRange{N}}, -) where {T,N} - return viewblock(a, block) + a::SubArray{ + T, N, <:AbstractBlockSparseArray{T, N}, <:Tuple{Vararg{SubBlockSliceCollection, N}}, + }, + block::Union{Block{N}, BlockIndexRange{N}}, + ) where {T, N} + return viewblock(a, block) end function Base.view( - a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockIndexRangeSlice,N}}}, - block::Union{Block{N},BlockIndexRange{N}}, -) where {T,N} - return viewblock(a, block) + a::SubArray{T, N, <:AbstractBlockSparseArray{T, N}, <:Tuple{Vararg{BlockIndexRangeSlice, N}}}, + block::Union{Block{N}, BlockIndexRange{N}}, + ) where {T, N} + return viewblock(a, block) end function Base.view( - a::SubArray{ - T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{SubBlockSliceCollection,N}} - }, - block::Vararg{Union{Block{1},BlockIndexRange{1},BlockIndexVector},N}, -) where {T,N} - return viewblock(a, block...) + a::SubArray{ + T, N, <:AbstractBlockSparseArray{T, N}, <:Tuple{Vararg{SubBlockSliceCollection, N}}, + }, + block::Vararg{Union{Block{1}, BlockIndexRange{1}, BlockIndexVector}, N}, + ) where {T, N} + return viewblock(a, block...) end # Fix ambiguity error. function Base.view( - a::SubArray{ - T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{SubBlockSliceCollection,N}} - }, - block::Vararg{Block{1},N}, -) where {T,N} - return viewblock(a, block...) + a::SubArray{ + T, N, <:AbstractBlockSparseArray{T, N}, <:Tuple{Vararg{SubBlockSliceCollection, N}}, + }, + block::Vararg{Block{1}, N}, + ) where {T, N} + return viewblock(a, block...) end function BlockArrays.viewblock( - a::SubArray{ - T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{SubBlockSliceCollection,N}} - }, - block::Union{Block{N},BlockIndexRange{N}}, -) where {T,N} - return viewblock(a, to_tuple(block)...) + a::SubArray{ + T, N, <:AbstractBlockSparseArray{T, N}, <:Tuple{Vararg{SubBlockSliceCollection, N}}, + }, + block::Union{Block{N}, BlockIndexRange{N}}, + ) where {T, N} + return viewblock(a, to_tuple(block)...) end blockedslice_blocks(x::BlockSlice) = x.block blockedslice_blocks(x::BlockIndices) = x.blocks # Reinterpret the slice blockwise. function blockedslice_blocks(x::Base.Slice) - return mortar( - map(BlockRange(x.indices)) do b - return BlockIndexRange(b, Base.Slice(Base.axes1(x.indices[b]))) - end, - ) + return mortar( + map(BlockRange(x.indices)) do b + return BlockIndexRange(b, Base.Slice(Base.axes1(x.indices[b]))) + end, + ) end # TODO: Define `@interface interface(a) viewblock`. function BlockArrays.viewblock( - a::SubArray{ - T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{SubBlockSliceCollection,N}} - }, - I::Vararg{Block{1},N}, -) where {T,N} - # TODO: Use `reindex`, `to_indices`, etc. - brs = ntuple(ndims(a)) do dim - # TODO: Ideally we would use this but it outputs a Vector, - # not a range: - # return parentindices(a)[dim].block[I[dim]] - return blocks(blockedslice_blocks(parentindices(a)[dim]))[Int(I[dim])] - end - return @view parent(a)[brs...] + a::SubArray{ + T, N, <:AbstractBlockSparseArray{T, N}, <:Tuple{Vararg{SubBlockSliceCollection, N}}, + }, + I::Vararg{Block{1}, N}, + ) where {T, N} + # TODO: Use `reindex`, `to_indices`, etc. + brs = ntuple(ndims(a)) do dim + # TODO: Ideally we would use this but it outputs a Vector, + # not a range: + # return parentindices(a)[dim].block[I[dim]] + return blocks(blockedslice_blocks(parentindices(a)[dim]))[Int(I[dim])] + end + return @view parent(a)[brs...] end # TODO: Define `@interface interface(a) viewblock`. function BlockArrays.viewblock( - a::SubArray{ - T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{SubBlockSliceCollection,N}} - }, - block::Vararg{BlockIndexRange{1},N}, -) where {T,N} - return view(viewblock(a, Block.(block)...), map(b -> only(b.indices), block)...) + a::SubArray{ + T, N, <:AbstractBlockSparseArray{T, N}, <:Tuple{Vararg{SubBlockSliceCollection, N}}, + }, + block::Vararg{BlockIndexRange{1}, N}, + ) where {T, N} + return view(viewblock(a, Block.(block)...), map(b -> only(b.indices), block)...) end # migrate wrapper layer for viewing `adjoint` and `transpose`. for (f, F) in ((:adjoint, :Adjoint), (:transpose, :Transpose)) - @eval begin - function Base.view(A::$F{<:Any,<:AbstractBlockSparseVector}, b::Block{1}) - return $f(view(parent(A), b)) - end + @eval begin + function Base.view(A::$F{<:Any, <:AbstractBlockSparseVector}, b::Block{1}) + return $f(view(parent(A), b)) + end - Base.view(A::$F{<:Any,<:AbstractBlockSparseMatrix}, b::Block{2}) = view(A, Tuple(b)...) - function Base.view(A::$F{<:Any,<:AbstractBlockSparseMatrix}, b1::Block{1}, b2::Block{1}) - return $f(view(parent(A), b2, b1)) + Base.view(A::$F{<:Any, <:AbstractBlockSparseMatrix}, b::Block{2}) = view(A, Tuple(b)...) + function Base.view(A::$F{<:Any, <:AbstractBlockSparseMatrix}, b1::Block{1}, b2::Block{1}) + return $f(view(parent(A), b2, b1)) + end end - end end diff --git a/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl index 207d2105..d85fbb47 100644 --- a/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl +++ b/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl @@ -1,129 +1,129 @@ using Adapt: Adapt, WrappedArray, adapt using ArrayLayouts: ArrayLayouts using BlockArrays: - BlockArrays, - AbstractBlockVector, - AbstractBlockedUnitRange, - BlockIndexRange, - BlockRange, - blockedrange, - mortar, - unblock + BlockArrays, + AbstractBlockVector, + AbstractBlockedUnitRange, + BlockIndexRange, + BlockRange, + blockedrange, + mortar, + unblock using DerivableInterfaces: DerivableInterfaces, @interface, DefaultArrayInterface, zero! using GPUArraysCore: @allowscalar using SplitApplyCombine: groupcount using TypeParameterAccessors: similartype -const WrappedAbstractBlockSparseArray{T,N} = WrappedArray{ - T,N,AbstractBlockSparseArray,AbstractBlockSparseArray{T,N} +const WrappedAbstractBlockSparseArray{T, N} = WrappedArray{ + T, N, AbstractBlockSparseArray, AbstractBlockSparseArray{T, N}, } -const AnyAbstractBlockSparseArray{T,N} = Union{ - <:AbstractBlockSparseArray{T,N},<:WrappedAbstractBlockSparseArray{T,N} +const AnyAbstractBlockSparseArray{T, N} = Union{ + <:AbstractBlockSparseArray{T, N}, <:WrappedAbstractBlockSparseArray{T, N}, } -const AnyAbstractBlockSparseVector{T} = AnyAbstractBlockSparseArray{T,1} -const AnyAbstractBlockSparseMatrix{T} = AnyAbstractBlockSparseArray{T,2} -const AnyAbstractBlockSparseVecOrMat{T,N} = Union{ - AnyAbstractBlockSparseVector{T},AnyAbstractBlockSparseMatrix{T} +const AnyAbstractBlockSparseVector{T} = AnyAbstractBlockSparseArray{T, 1} +const AnyAbstractBlockSparseMatrix{T} = AnyAbstractBlockSparseArray{T, 2} +const AnyAbstractBlockSparseVecOrMat{T, N} = Union{ + AnyAbstractBlockSparseVector{T}, AnyAbstractBlockSparseMatrix{T}, } function DerivableInterfaces.interface(arrayt::Type{<:AnyAbstractBlockSparseArray}) - return BlockSparseArrayInterface(interface(blocktype(arrayt))) + return BlockSparseArrayInterface(interface(blocktype(arrayt))) end # a[1:2, 1:2] function Base.to_indices( - a::AnyAbstractBlockSparseArray, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}} -) - return @interface interface(a) to_indices(a, inds, I) + a::AnyAbstractBlockSparseArray, inds, I::Tuple{UnitRange{<:Integer}, Vararg{Any}} + ) + return @interface interface(a) to_indices(a, inds, I) end function Base.to_indices( - a::AnyAbstractBlockSparseArray, inds, I::Tuple{AbstractArray{Bool},Vararg{Any}} -) - return @interface interface(a) to_indices(a, inds, I) + a::AnyAbstractBlockSparseArray, inds, I::Tuple{AbstractArray{Bool}, Vararg{Any}} + ) + return @interface interface(a) to_indices(a, inds, I) end # Fix ambiguity error with Base for logical indexing in Julia 1.10. # TODO: Delete this once we drop support for Julia 1.10. function Base.to_indices( - a::AnyAbstractBlockSparseArray, inds, I::Union{Tuple{BitArray{N}},Tuple{Array{Bool,N}}} -) where {N} - return @interface interface(a) to_indices(a, inds, I) + a::AnyAbstractBlockSparseArray, inds, I::Union{Tuple{BitArray{N}}, Tuple{Array{Bool, N}}} + ) where {N} + return @interface interface(a) to_indices(a, inds, I) end # a[[Block(2), Block(1)], [Block(2), Block(1)]] function Base.to_indices( - a::AnyAbstractBlockSparseArray, inds, I::Tuple{Vector{<:Block{1}},Vararg{Any}} -) - return @interface interface(a) to_indices(a, inds, I) + a::AnyAbstractBlockSparseArray, inds, I::Tuple{Vector{<:Block{1}}, Vararg{Any}} + ) + return @interface interface(a) to_indices(a, inds, I) end # a[BlockVector([Block(2), Block(1)], [2]), BlockVector([Block(2), Block(1)], [2])] # a[BlockedVector([Block(2), Block(1)], [2]), BlockedVector([Block(2), Block(1)], [2])] function Base.to_indices( - a::AnyAbstractBlockSparseArray, - inds, - I::Tuple{AbstractBlockVector{<:Block{1}},Vararg{Any}}, -) - return @interface interface(a) to_indices(a, inds, I) + a::AnyAbstractBlockSparseArray, + inds, + I::Tuple{AbstractBlockVector{<:Block{1}}, Vararg{Any}}, + ) + return @interface interface(a) to_indices(a, inds, I) end # a[mortar([Block(1)[1:2], Block(2)[1:3]])] function Base.to_indices( - a::AnyAbstractBlockSparseArray, - inds, - I::Tuple{BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}},Vararg{Any}}, -) - return @interface interface(a) to_indices(a, inds, I) + a::AnyAbstractBlockSparseArray, + inds, + I::Tuple{BlockVector{<:BlockIndex{1}, <:Vector{<:BlockIndexRange{1}}}, Vararg{Any}}, + ) + return @interface interface(a) to_indices(a, inds, I) end # a[mortar([Block(1)[[1, 2]], Block(2)[[1, 3]]])] function Base.to_indices( - a::AnyAbstractBlockSparseArray, - inds, - I::Tuple{BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexVector{1}}},Vararg{Any}}, -) - return @interface interface(a) to_indices(a, inds, I) + a::AnyAbstractBlockSparseArray, + inds, + I::Tuple{BlockVector{<:BlockIndex{1}, <:Vector{<:BlockIndexVector{1}}}, Vararg{Any}}, + ) + return @interface interface(a) to_indices(a, inds, I) end function Base.to_indices( - a::AnyAbstractBlockSparseArray, - inds, - I::Tuple{BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector{1}}},Vararg{Any}}, -) - return @interface interface(a) to_indices(a, inds, I) + a::AnyAbstractBlockSparseArray, + inds, + I::Tuple{BlockVector{<:GenericBlockIndex{1}, <:Vector{<:BlockIndexVector{1}}}, Vararg{Any}}, + ) + return @interface interface(a) to_indices(a, inds, I) end # a[[Block(1)[1:2], Block(2)[1:2]], [Block(1)[1:2], Block(2)[1:2]]] function Base.to_indices( - a::AnyAbstractBlockSparseArray, inds, I::Tuple{Vector{<:BlockIndexRange{1}},Vararg{Any}} -) - return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...)) + a::AnyAbstractBlockSparseArray, inds, I::Tuple{Vector{<:BlockIndexRange{1}}, Vararg{Any}} + ) + return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...)) end # a[[Block(1)[[1, 2]], Block(2)[[1, 2]]], [Block(1)[[1, 2]], Block(2)[[1, 2]]]] function Base.to_indices( - a::AnyAbstractBlockSparseArray, inds, I::Tuple{Vector{<:BlockIndexVector{1}},Vararg{Any}} -) - return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...)) + a::AnyAbstractBlockSparseArray, inds, I::Tuple{Vector{<:BlockIndexVector{1}}, Vararg{Any}} + ) + return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...)) end # BlockArrays `AbstractBlockArray` interface function BlockArrays.blocks(a::AnyAbstractBlockSparseArray) - @interface interface(a) blocks(a) + return @interface interface(a) blocks(a) end # Fix ambiguity error with `BlockArrays` using BlockArrays: BlockSlice function BlockArrays.blocks( - a::SubArray{<:Any,<:Any,<:AbstractBlockSparseArray,<:Tuple{Vararg{BlockSlice}}} -) - return @interface interface(a) blocks(a) + a::SubArray{<:Any, <:Any, <:AbstractBlockSparseArray, <:Tuple{Vararg{BlockSlice}}} + ) + return @interface interface(a) blocks(a) end using TypeParameterAccessors: parenttype function blockstype(arraytype::Type{<:WrappedAbstractBlockSparseArray}) - return blockstype(parenttype(arraytype)) + return blockstype(parenttype(arraytype)) end blocktype(a::AnyAbstractBlockSparseArray) = eltype(blocks(a)) @@ -131,86 +131,86 @@ blocktype(arraytype::Type{<:AnyAbstractBlockSparseArray}) = eltype(blockstype(ar using ArrayLayouts: ArrayLayouts function Base.getindex( - a::AnyAbstractBlockSparseArray{<:Any,N}, I::CartesianIndices{N} -) where {N} - return ArrayLayouts.layout_getindex(a, I) + a::AnyAbstractBlockSparseArray{<:Any, N}, I::CartesianIndices{N} + ) where {N} + return ArrayLayouts.layout_getindex(a, I) end function Base.getindex( - a::AnyAbstractBlockSparseArray{<:Any,N}, I::Vararg{AbstractUnitRange{<:Integer},N} -) where {N} - return ArrayLayouts.layout_getindex(a, I...) + a::AnyAbstractBlockSparseArray{<:Any, N}, I::Vararg{AbstractUnitRange{<:Integer}, N} + ) where {N} + return ArrayLayouts.layout_getindex(a, I...) end # TODO: Define `AnyBlockSparseMatrix`. function Base.getindex( - a::AnyAbstractBlockSparseArray{<:Any,2}, I::Vararg{AbstractUnitRange{<:Integer},2} -) - return ArrayLayouts.layout_getindex(a, I...) + a::AnyAbstractBlockSparseArray{<:Any, 2}, I::Vararg{AbstractUnitRange{<:Integer}, 2} + ) + return ArrayLayouts.layout_getindex(a, I...) end # Fixes ambiguity error. -function Base.getindex(a::AnyAbstractBlockSparseArray{<:Any,0}) - return ArrayLayouts.layout_getindex(a) +function Base.getindex(a::AnyAbstractBlockSparseArray{<:Any, 0}) + return ArrayLayouts.layout_getindex(a) end # TODO: Define `@interface interface(a) isassigned`. function Base.isassigned( - a::AnyAbstractBlockSparseArray{<:Any,N}, index::Vararg{Block{1},N} -) where {N} - return isassigned(blocks(a), Int.(index)...) + a::AnyAbstractBlockSparseArray{<:Any, N}, index::Vararg{Block{1}, N} + ) where {N} + return isassigned(blocks(a), Int.(index)...) end # Fix ambiguity error. -function Base.isassigned(a::AnyAbstractBlockSparseArray{<:Any,0}) - return isassigned(blocks(a)) +function Base.isassigned(a::AnyAbstractBlockSparseArray{<:Any, 0}) + return isassigned(blocks(a)) end -function Base.isassigned(a::AnyAbstractBlockSparseArray{<:Any,N}, index::Block{N}) where {N} - return isassigned(a, Tuple(index)...) +function Base.isassigned(a::AnyAbstractBlockSparseArray{<:Any, N}, index::Block{N}) where {N} + return isassigned(a, Tuple(index)...) end # TODO: Define `@interface interface(a) isassigned`. function Base.isassigned( - a::AnyAbstractBlockSparseArray{<:Any,N}, index::Vararg{BlockIndex{1},N} -) where {N} - b = block.(index) - return isassigned(a, b...) && isassigned(@view(a[b...]), blockindex.(index)...) + a::AnyAbstractBlockSparseArray{<:Any, N}, index::Vararg{BlockIndex{1}, N} + ) where {N} + b = block.(index) + return isassigned(a, b...) && isassigned(@view(a[b...]), blockindex.(index)...) end function Base.setindex!( - a::AnyAbstractBlockSparseArray{<:Any,N}, value, I::BlockIndex{N} -) where {N} - # TODO: Use `@interface interface(a) setindex!(...)`. - @interface interface(a) setindex!(a, value, I) - return a + a::AnyAbstractBlockSparseArray{<:Any, N}, value, I::BlockIndex{N} + ) where {N} + # TODO: Use `@interface interface(a) setindex!(...)`. + @interface interface(a) setindex!(a, value, I) + return a end # Fixes ambiguity error with BlockArrays.jl -function Base.setindex!(a::AnyAbstractBlockSparseArray{<:Any,1}, value, I::BlockIndex{1}) - # TODO: Use `@interface interface(a) setindex!(...)`. - @interface interface(a) setindex!(a, value, I) - return a +function Base.setindex!(a::AnyAbstractBlockSparseArray{<:Any, 1}, value, I::BlockIndex{1}) + # TODO: Use `@interface interface(a) setindex!(...)`. + @interface interface(a) setindex!(a, value, I) + return a end function ArrayLayouts.zero!(a::AnyAbstractBlockSparseArray) - return zero!(a) + return zero!(a) end # TODO: Use `@derive`. function Base.fill!(a::AnyAbstractBlockSparseArray, value) - return @interface interface(a) fill!(a, value) + return @interface interface(a) fill!(a, value) end # Needed by `BlockArrays` matrix multiplication interface function Base.similar( - arraytype::Type{<:AnyAbstractBlockSparseArray}, - axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}, -) - return similar(arraytype, eltype(arraytype), axes) + arraytype::Type{<:AnyAbstractBlockSparseArray}, + axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}, + ) + return similar(arraytype, eltype(arraytype), axes) end # Fixes ambiguity error. function Base.similar( - arraytype::Type{<:AnyAbstractBlockSparseArray}, axes::Tuple{Base.OneTo,Vararg{Base.OneTo}} -) - return similar(arraytype, eltype(arraytype), axes) + arraytype::Type{<:AnyAbstractBlockSparseArray}, axes::Tuple{Base.OneTo, Vararg{Base.OneTo}} + ) + return similar(arraytype, eltype(arraytype), axes) end # Needed by `BlockArrays` matrix multiplication interface @@ -218,171 +218,171 @@ end # is only appears to be needed in older versions of Julia like v1.6. # Delete once we drop support for older versions of Julia. function Base.similar( - arraytype::Type{<:AnyAbstractBlockSparseArray}, - axes::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}, -) - return similar(arraytype, eltype(arraytype), axes) + arraytype::Type{<:AnyAbstractBlockSparseArray}, + axes::Tuple{AbstractUnitRange{<:Integer}, Vararg{AbstractUnitRange{<:Integer}}}, + ) + return similar(arraytype, eltype(arraytype), axes) end # Fixes ambiguity error with `BlockArrays`. function Base.similar( - arraytype::Type{<:AnyAbstractBlockSparseArray}, - axes::Tuple{AbstractBlockedUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}, -) - return similar(arraytype, eltype(arraytype), axes) + arraytype::Type{<:AnyAbstractBlockSparseArray}, + axes::Tuple{AbstractBlockedUnitRange{<:Integer}, Vararg{AbstractUnitRange{<:Integer}}}, + ) + return similar(arraytype, eltype(arraytype), axes) end # Fixes ambiguity error with `BlockArrays`. function Base.similar( - arraytype::Type{<:AnyAbstractBlockSparseArray}, - axes::Tuple{ - AbstractUnitRange{<:Integer}, - AbstractBlockedUnitRange{<:Integer}, - Vararg{AbstractUnitRange{<:Integer}}, - }, -) - return similar(arraytype, eltype(arraytype), axes) + arraytype::Type{<:AnyAbstractBlockSparseArray}, + axes::Tuple{ + AbstractUnitRange{<:Integer}, + AbstractBlockedUnitRange{<:Integer}, + Vararg{AbstractUnitRange{<:Integer}}, + }, + ) + return similar(arraytype, eltype(arraytype), axes) end # Needed for disambiguation function Base.similar( - arraytype::Type{<:AnyAbstractBlockSparseArray}, - axes::Tuple{Vararg{AbstractBlockedUnitRange{<:Integer}}}, -) - return similar(arraytype, eltype(arraytype), axes) + arraytype::Type{<:AnyAbstractBlockSparseArray}, + axes::Tuple{Vararg{AbstractBlockedUnitRange{<:Integer}}}, + ) + return similar(arraytype, eltype(arraytype), axes) end function blocksparse_similar(a, elt::Type, axes::Tuple) - ndims = length(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) + ndims = length(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}} -) - return blocksparse_similar(a, elt, axes) + a::AbstractArray, elt::Type, axes::Tuple{Vararg{Int}} + ) + return blocksparse_similar(a, elt, axes) end @interface ::AbstractBlockSparseArrayInterface function Base.similar( - a::AbstractArray, elt::Type, axes::Tuple -) - return blocksparse_similar(a, elt, axes) + a::AbstractArray, elt::Type, axes::Tuple + ) + return blocksparse_similar(a, elt, axes) end # Fix ambiguity error when non-blocked ranges are passed. @interface ::AbstractBlockSparseArrayInterface function Base.similar( - a::AbstractArray, elt::Type, axes::Tuple{Base.OneTo,Vararg{Base.OneTo}} -) - return blocksparse_similar(a, elt, axes) + a::AbstractArray, elt::Type, axes::Tuple{Base.OneTo, Vararg{Base.OneTo}} + ) + return blocksparse_similar(a, elt, axes) end # Fix ambiguity error when empty axes are passed. @interface ::AbstractBlockSparseArrayInterface function Base.similar( - a::AbstractArray, elt::Type, axes::Tuple{} -) - return blocksparse_similar(a, elt, axes) + a::AbstractArray, elt::Type, axes::Tuple{} + ) + return blocksparse_similar(a, elt, axes) end @interface ::AbstractBlockSparseArrayInterface function Base.similar( - a::Type{<:AbstractArray}, elt::Type, axes::Tuple{Vararg{Int}} -) - return blocksparse_similar(a, elt, axes) + a::Type{<:AbstractArray}, elt::Type, axes::Tuple{Vararg{Int}} + ) + return blocksparse_similar(a, elt, axes) end @interface ::AbstractBlockSparseArrayInterface function Base.similar( - a::Type{<:AbstractArray}, elt::Type, axes::Tuple -) - return blocksparse_similar(a, elt, axes) + a::Type{<:AbstractArray}, elt::Type, axes::Tuple + ) + return blocksparse_similar(a, elt, axes) end # Needed by `BlockArrays` matrix multiplication interface # TODO: Define a `@interface interface(a) similar` function. function Base.similar( - arraytype::Type{<:AnyAbstractBlockSparseArray}, - elt::Type, - axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}, -) - return @interface interface(arraytype) similar(arraytype, elt, axes) + arraytype::Type{<:AnyAbstractBlockSparseArray}, + elt::Type, + axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}, + ) + return @interface interface(arraytype) similar(arraytype, elt, axes) end # TODO: Define a `@interface interface(a) similar` function. function Base.similar( - a::AnyAbstractBlockSparseArray, - elt::Type, - axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}, -) - return @interface interface(a) similar(a, elt, axes) + a::AnyAbstractBlockSparseArray, + elt::Type, + axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}, + ) + return @interface interface(a) similar(a, elt, axes) end # Fixes ambiguity error. function Base.similar(a::AnyAbstractBlockSparseArray, elt::Type, axes::Tuple{}) - return @interface interface(a) similar(a, elt, axes) + return @interface interface(a) similar(a, elt, axes) end # Fixes ambiguity error with `BlockArrays`. function Base.similar( - a::AnyAbstractBlockSparseArray, - elt::Type, - axes::Tuple{ - AbstractBlockedUnitRange{<:Integer},Vararg{AbstractBlockedUnitRange{<:Integer}} - }, -) - return @interface interface(a) similar(a, elt, axes) + a::AnyAbstractBlockSparseArray, + elt::Type, + axes::Tuple{ + AbstractBlockedUnitRange{<:Integer}, Vararg{AbstractBlockedUnitRange{<:Integer}}, + }, + ) + return @interface interface(a) similar(a, elt, axes) end # Fixes ambiguity error with `OffsetArrays`. function Base.similar( - a::AnyAbstractBlockSparseArray, - elt::Type, - axes::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}, -) - return @interface interface(a) similar(a, elt, axes) + a::AnyAbstractBlockSparseArray, + elt::Type, + axes::Tuple{AbstractUnitRange{<:Integer}, Vararg{AbstractUnitRange{<:Integer}}}, + ) + return @interface interface(a) similar(a, elt, axes) end # Fixes ambiguity error with `BlockArrays`. function Base.similar( - a::AnyAbstractBlockSparseArray, - elt::Type, - axes::Tuple{AbstractBlockedUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}, -) - return @interface interface(a) similar(a, elt, axes) + a::AnyAbstractBlockSparseArray, + elt::Type, + axes::Tuple{AbstractBlockedUnitRange{<:Integer}, Vararg{AbstractUnitRange{<:Integer}}}, + ) + 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)) + 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) + 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( - a::AnyAbstractBlockSparseArray, - elt::Type, - axes::Tuple{ - AbstractUnitRange{<:Integer}, - AbstractBlockedUnitRange{<:Integer}, - Vararg{AbstractUnitRange{<:Integer}}, - }, -) - # TODO: Use `@interface interface(a) similar(...)`. - return @interface interface(a) similar(a, elt, axes) + a::AnyAbstractBlockSparseArray, + elt::Type, + axes::Tuple{ + AbstractUnitRange{<:Integer}, + AbstractBlockedUnitRange{<:Integer}, + Vararg{AbstractUnitRange{<:Integer}}, + }, + ) + # TODO: Use `@interface interface(a) similar(...)`. + return @interface interface(a) similar(a, elt, axes) end # Fixes ambiguity error with `StaticArrays`. function Base.similar( - a::AnyAbstractBlockSparseArray, elt::Type, axes::Tuple{Base.OneTo,Vararg{Base.OneTo}} -) - return @interface interface(a) similar(a, elt, axes) + a::AnyAbstractBlockSparseArray, elt::Type, axes::Tuple{Base.OneTo, Vararg{Base.OneTo}} + ) + return @interface interface(a) similar(a, elt, axes) end struct BlockType{T} end BlockType(x) = BlockType{x}() function Base.similar(a::AbstractBlockSparseArray, ::BlockType{T}, ax) where {T} - return BlockSparseArray{eltype(T),ndims(T),T}(undef, ax) + return BlockSparseArray{eltype(T), ndims(T), T}(undef, ax) end function Base.similar(a::AbstractBlockSparseArray, T::BlockType) - return similar(a, T, axes(a)) + return similar(a, T, axes(a)) end # TODO: Implement this in a more generic way using a smarter `copyto!`, @@ -390,59 +390,59 @@ end # These are defined for now to avoid scalar indexing issues when there # are blocks on GPU, and also work with exotic block types like # KroneckerArrays. -function Base.Array{T,N}(a::AnyAbstractBlockSparseArray{<:Any,N}) where {T,N} - a_dest = zeros(T, size(a)) - for I in eachblockstoredindex(a) - # TODO: Use: `I′ = CartesianIndices(axes(a))[I]`, unfortunately this - # outputs `Matrix{CartesianIndex}` instead of `CartesianIndices`. - I′ = CartesianIndices(ntuple(dim -> axes(a, dim)[Tuple(I)[dim]], ndims(a))) - a_dest[I′] = Array{T,N}(@view(a[I])) - end - return a_dest +function Base.Array{T, N}(a::AnyAbstractBlockSparseArray{<:Any, N}) where {T, N} + a_dest = zeros(T, size(a)) + for I in eachblockstoredindex(a) + # TODO: Use: `I′ = CartesianIndices(axes(a))[I]`, unfortunately this + # outputs `Matrix{CartesianIndex}` instead of `CartesianIndices`. + I′ = CartesianIndices(ntuple(dim -> axes(a, dim)[Tuple(I)[dim]], ndims(a))) + a_dest[I′] = Array{T, N}(@view(a[I])) + end + return a_dest end function Base.Array{T}(a::AnyAbstractBlockSparseArray) where {T} - return Array{T,ndims(a)}(a) + return Array{T, ndims(a)}(a) end function Base.Array(a::AnyAbstractBlockSparseArray) - return Array{eltype(a)}(a) + return Array{eltype(a)}(a) end function SparseArraysBase.isstored( - a::AbstractBlockSparseArray{<:Any,N}, I::Vararg{Int,N} -) where {N} - return @interface interface(a) isstored(a, I...) + a::AbstractBlockSparseArray{<:Any, N}, I::Vararg{Int, N} + ) where {N} + return @interface interface(a) isstored(a, I...) end function Base.replace_in_print_matrix( - a::AnyAbstractBlockSparseVecOrMat, i::Integer, j::Integer, s::AbstractString -) - return isstored(a, i, j) ? s : Base.replace_with_centered_mark(s) + a::AnyAbstractBlockSparseVecOrMat, i::Integer, j::Integer, s::AbstractString + ) + return isstored(a, i, j) ? s : Base.replace_with_centered_mark(s) end # attempt to catch things that wrap GPU arrays function Base.print_array(io::IO, a::AnyAbstractBlockSparseArray) - a_cpu = adapt(Array, a) - if typeof(a_cpu) === typeof(a) # prevent infinite recursion - # need to specify ndims to allow specialized code for vector/matrix - @allowscalar @invoke Base.print_array( - io, a_cpu::AbstractArray{eltype(a_cpu),ndims(a_cpu)} - ) + a_cpu = adapt(Array, a) + if typeof(a_cpu) === typeof(a) # prevent infinite recursion + # need to specify ndims to allow specialized code for vector/matrix + @allowscalar @invoke Base.print_array( + io, a_cpu::AbstractArray{eltype(a_cpu), ndims(a_cpu)} + ) + return nothing + end + Base.print_array(io, a_cpu) return nothing - end - Base.print_array(io, a_cpu) - return nothing end using Adapt: Adapt, adapt -function Adapt.adapt_structure(to, a::SubArray{<:Any,<:Any,<:AbstractBlockSparseArray}) - # In the generic definition in Adapt.jl, `parentindices(a)` are also - # adapted, but is broken when the parent indices contained blocked unit - # ranges since `adapt` is broken on blocked unit ranges. - # TODO: Fix adapt for blocked unit ranges by making an AdaptExt for - # BlockArrays.jl. - return SubArray(adapt(to, parent(a)), parentindices(a)) +function Adapt.adapt_structure(to, a::SubArray{<:Any, <:Any, <:AbstractBlockSparseArray}) + # In the generic definition in Adapt.jl, `parentindices(a)` are also + # adapted, but is broken when the parent indices contained blocked unit + # ranges since `adapt` is broken on blocked unit ranges. + # TODO: Fix adapt for blocked unit ranges by making an AdaptExt for + # BlockArrays.jl. + return SubArray(adapt(to, parent(a)), parentindices(a)) end function Base.show(io::IO, a::AnyAbstractBlockSparseArray) - return show(io, convert(Array, a)) + return show(io, convert(Array, a)) end diff --git a/src/blocksparsearray/blockdiagonalarray.jl b/src/blocksparsearray/blockdiagonalarray.jl index 05345ac3..e3bcb96c 100644 --- a/src/blocksparsearray/blockdiagonalarray.jl +++ b/src/blocksparsearray/blockdiagonalarray.jl @@ -3,25 +3,25 @@ using DiagonalArrays: DiagonalArrays, diagonal using LinearAlgebra: Diagonal # type alias for block-diagonal -const BlockDiagonal{T,A,Axes,V<:AbstractVector{A}} = BlockSparseMatrix{ - T,A,Diagonal{A,V},Axes +const BlockDiagonal{T, A, Axes, V <: AbstractVector{A}} = BlockSparseMatrix{ + T, A, Diagonal{A, V}, Axes, } -const BlockSparseDiagonal{T,A<:AbstractBlockSparseVector{T}} = Diagonal{T,A} +const BlockSparseDiagonal{T, A <: AbstractBlockSparseVector{T}} = Diagonal{T, A} @interface interface::BlockSparseArrayInterface function blocks(a::BlockSparseDiagonal) - return Diagonal(Diagonal.(blocks(a.diag))) + return Diagonal(Diagonal.(blocks(a.diag))) end function BlockDiagonal(blocks::AbstractVector{<:AbstractMatrix}) - return _BlockSparseArray( - Diagonal(blocks), blockedrange.((size.(blocks, 1), size.(blocks, 2))) - ) + return _BlockSparseArray( + Diagonal(blocks), blockedrange.((size.(blocks, 1), size.(blocks, 2))) + ) end function DiagonalArrays.diagonal(S::BlockSparseVector) - D = similar(S, (axes(S, 1), axes(S, 1))) - for bI in eachblockstoredindex(S) - D[bI, bI] = diagonal(@view!(S[bI])) - end - return D + D = similar(S, (axes(S, 1), axes(S, 1))) + for bI in eachblockstoredindex(S) + D[bI, bI] = diagonal(@view!(S[bI])) + end + return D end diff --git a/src/blocksparsearray/blocksparsearray.jl b/src/blocksparsearray/blocksparsearray.jl index f708f345..7cf149fa 100644 --- a/src/blocksparsearray/blocksparsearray.jl +++ b/src/blocksparsearray/blocksparsearray.jl @@ -1,11 +1,11 @@ using BlockArrays: - BlockArrays, - Block, - BlockedUnitRange, - UndefBlocksInitializer, - blockedrange, - blocklength, - undef_blocks + BlockArrays, + Block, + BlockedUnitRange, + UndefBlocksInitializer, + blockedrange, + blocklength, + undef_blocks using DerivableInterfaces: @interface using Dictionaries: Dictionary using SparseArraysBase: SparseArrayDOK, Unstored @@ -22,83 +22,83 @@ Note that `undef_blocks` is defined in [BlockArrays.jl](https://juliaarrays.github.io/BlockArrays.jl/stable/lib/public/#BlockArrays.undef_blocks) and should be imported from that package to use it as an input to this constructor. """ -function SparseArraysBase.SparseArrayDOK{T,N}( - ::UndefBlocksInitializer, ax::Tuple{Vararg{AbstractUnitRange{<:Integer},N}} -) where {T,N} - return SparseArrayDOK{T,N}(undef, Unstored(ZeroBlocks{N,T}(ax))) -end -function SparseArraysBase.SparseArrayDOK{T,N}( - ::UndefBlocksInitializer, ax::Vararg{AbstractUnitRange{<:Integer},N} -) where {T,N} - return SparseArrayDOK{T,N}(undef_blocks, ax) -end -function SparseArraysBase.SparseArrayDOK{T,N}( - ::UndefBlocksInitializer, - dims::Tuple{AbstractVector{<:Integer},Vararg{AbstractVector{<:Integer}}}, -) where {T,N} - return SparseArrayDOK{T,N}(undef_blocks, blockedrange.(dims)) -end -function SparseArraysBase.SparseArrayDOK{T,N}( - ::UndefBlocksInitializer, - dim1::AbstractVector{<:Integer}, - dim_rest::AbstractVector{<:Integer}..., -) where {T,N} - return SparseArrayDOK{T,N}(undef_blocks, (dim1, dim_rest...)) +function SparseArraysBase.SparseArrayDOK{T, N}( + ::UndefBlocksInitializer, ax::Tuple{Vararg{AbstractUnitRange{<:Integer}, N}} + ) where {T, N} + return SparseArrayDOK{T, N}(undef, Unstored(ZeroBlocks{N, T}(ax))) +end +function SparseArraysBase.SparseArrayDOK{T, N}( + ::UndefBlocksInitializer, ax::Vararg{AbstractUnitRange{<:Integer}, N} + ) where {T, N} + return SparseArrayDOK{T, N}(undef_blocks, ax) +end +function SparseArraysBase.SparseArrayDOK{T, N}( + ::UndefBlocksInitializer, + dims::Tuple{AbstractVector{<:Integer}, Vararg{AbstractVector{<:Integer}}}, + ) where {T, N} + return SparseArrayDOK{T, N}(undef_blocks, blockedrange.(dims)) +end +function SparseArraysBase.SparseArrayDOK{T, N}( + ::UndefBlocksInitializer, + dim1::AbstractVector{<:Integer}, + dim_rest::AbstractVector{<:Integer}..., + ) where {T, N} + return SparseArrayDOK{T, N}(undef_blocks, (dim1, dim_rest...)) end function SparseArraysBase.SparseArrayDOK{T}( - ::UndefBlocksInitializer, ax::Tuple{Vararg{AbstractUnitRange{<:Integer},N}} -) where {T,N} - return SparseArrayDOK{T,N}(undef_blocks, ax) + ::UndefBlocksInitializer, ax::Tuple{Vararg{AbstractUnitRange{<:Integer}, N}} + ) where {T, N} + return SparseArrayDOK{T, N}(undef_blocks, ax) end function SparseArraysBase.SparseArrayDOK{T}( - ::UndefBlocksInitializer, ax::Vararg{AbstractUnitRange{<:Integer},N} -) where {T,N} - return SparseArrayDOK{T,N}(undef_blocks, ax) + ::UndefBlocksInitializer, ax::Vararg{AbstractUnitRange{<:Integer}, N} + ) where {T, N} + return SparseArrayDOK{T, N}(undef_blocks, ax) end function SparseArraysBase.SparseArrayDOK{T}( - ::UndefBlocksInitializer, - dims::Tuple{AbstractVector{<:Integer},Vararg{AbstractVector{<:Integer}}}, -) where {T} - return SparseArrayDOK{T}(undef_blocks, blockedrange.(dims)) + ::UndefBlocksInitializer, + dims::Tuple{AbstractVector{<:Integer}, Vararg{AbstractVector{<:Integer}}}, + ) where {T} + return SparseArrayDOK{T}(undef_blocks, blockedrange.(dims)) end function SparseArraysBase.SparseArrayDOK{T}( - ::UndefBlocksInitializer, - dim1::AbstractVector{<:Integer}, - dim_rest::AbstractVector{<:Integer}..., -) where {T} - return SparseArrayDOK{T}(undef_blocks, (dim1, dim_rest...)) + ::UndefBlocksInitializer, + dim1::AbstractVector{<:Integer}, + dim_rest::AbstractVector{<:Integer}..., + ) where {T} + return SparseArrayDOK{T}(undef_blocks, (dim1, dim_rest...)) end function _BlockSparseArray end struct BlockSparseArray{ - T, - N, - A<:AbstractArray{T,N}, - Blocks<:AbstractArray{A,N}, - Axes<:Tuple{Vararg{AbstractUnitRange{<:Integer},N}}, -} <: AbstractBlockSparseArray{T,N} - blocks::Blocks - axes::Axes - global @inline function _BlockSparseArray( - blocks::AbstractArray{<:AbstractArray{T,N},N}, - axes::Tuple{Vararg{AbstractUnitRange{<:Integer},N}}, - ) where {T,N} - Base.require_one_based_indexing(axes...) - Base.require_one_based_indexing(blocks) - return new{T,N,eltype(blocks),typeof(blocks),typeof(axes)}(blocks, axes) - end + T, + N, + A <: AbstractArray{T, N}, + Blocks <: AbstractArray{A, N}, + Axes <: Tuple{Vararg{AbstractUnitRange{<:Integer}, N}}, + } <: AbstractBlockSparseArray{T, N} + blocks::Blocks + axes::Axes + global @inline function _BlockSparseArray( + blocks::AbstractArray{<:AbstractArray{T, N}, N}, + axes::Tuple{Vararg{AbstractUnitRange{<:Integer}, N}}, + ) where {T, N} + Base.require_one_based_indexing(axes...) + Base.require_one_based_indexing(blocks) + return new{T, N, eltype(blocks), typeof(blocks), typeof(axes)}(blocks, axes) + end end # TODO: Can this definition be shortened? -const BlockSparseMatrix{T,A<:AbstractMatrix{T},Blocks<:AbstractMatrix{A},Axes<:Tuple{AbstractUnitRange{<:Integer},AbstractUnitRange{<:Integer}}} = BlockSparseArray{ - T,2,A,Blocks,Axes +const BlockSparseMatrix{T, A <: AbstractMatrix{T}, Blocks <: AbstractMatrix{A}, Axes <: Tuple{AbstractUnitRange{<:Integer}, AbstractUnitRange{<:Integer}}} = BlockSparseArray{ + T, 2, A, Blocks, Axes, } # TODO: Can this definition be shortened? -const BlockSparseVector{T,A<:AbstractVector{T},Blocks<:AbstractVector{A},Axes<:Tuple{AbstractUnitRange{<:Integer}}} = BlockSparseArray{ - T,1,A,Blocks,Axes +const BlockSparseVector{T, A <: AbstractVector{T}, Blocks <: AbstractVector{A}, Axes <: Tuple{AbstractUnitRange{<:Integer}}} = BlockSparseArray{ + T, 1, A, Blocks, Axes, } """ @@ -108,29 +108,29 @@ Construct a block sparse array from a sparse array of arrays and specified block The block sizes must be commensurate with the blocks of the axes. """ function sparsemortar( - blocks::AbstractArray{<:AbstractArray{T,N},N}, - axes::Tuple{Vararg{AbstractUnitRange{<:Integer},N}}, -) where {T,N} - return _BlockSparseArray(blocks, axes) + blocks::AbstractArray{<:AbstractArray{T, N}, N}, + axes::Tuple{Vararg{AbstractUnitRange{<:Integer}, N}}, + ) where {T, N} + return _BlockSparseArray(blocks, axes) end function sparsemortar( - blocks::AbstractArray{<:AbstractArray{T,N},N}, - axes::Vararg{AbstractUnitRange{<:Integer},N}, -) where {T,N} - return sparsemortar(blocks, axes) + blocks::AbstractArray{<:AbstractArray{T, N}, N}, + axes::Vararg{AbstractUnitRange{<:Integer}, N}, + ) where {T, N} + return sparsemortar(blocks, axes) end function sparsemortar( - blocks::AbstractArray{<:AbstractArray{T,N},N}, - dims::Tuple{AbstractVector{<:Integer},Vararg{AbstractVector{<:Integer}}}, -) where {T,N} - return sparsemortar(blocks, blockedrange.(dims)) + blocks::AbstractArray{<:AbstractArray{T, N}, N}, + dims::Tuple{AbstractVector{<:Integer}, Vararg{AbstractVector{<:Integer}}}, + ) where {T, N} + return sparsemortar(blocks, blockedrange.(dims)) end function sparsemortar( - blocks::AbstractArray{<:AbstractArray{T,N},N}, - dim1::AbstractVector{<:Integer}, - dim_rest::AbstractVector{<:Integer}..., -) where {T,N} - return sparsemortar(blocks, (dim1, dim_rest...)) + blocks::AbstractArray{<:AbstractArray{T, N}, N}, + dim1::AbstractVector{<:Integer}, + dim_rest::AbstractVector{<:Integer}..., + ) where {T, N} + return sparsemortar(blocks, (dim1, dim_rest...)) end @doc """ @@ -142,137 +142,137 @@ Construct an uninitialized N-dimensional BlockSparseArray containing elements of of block lengths in each dimension or a list of blocked ranges representing the axes. """ BlockSparseArray -function BlockSparseArray{T,N,A}( - ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange{<:Integer},N}} -) where {T,N,A<:AbstractArray{T,N}} - return _BlockSparseArray(SparseArrayDOK{A}(undef_blocks, axes), axes) +function BlockSparseArray{T, N, A}( + ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}, N}} + ) where {T, N, A <: AbstractArray{T, N}} + return _BlockSparseArray(SparseArrayDOK{A}(undef_blocks, axes), axes) end -function BlockSparseArray{T,N,A}( - ::UndefInitializer, axes::Vararg{AbstractUnitRange{<:Integer},N} -) where {T,N,A<:AbstractArray{T,N}} - return BlockSparseArray{T,N,A}(undef, axes) +function BlockSparseArray{T, N, A}( + ::UndefInitializer, axes::Vararg{AbstractUnitRange{<:Integer}, N} + ) where {T, N, A <: AbstractArray{T, N}} + return BlockSparseArray{T, N, A}(undef, axes) end -function BlockSparseArray{T,N,A}( - ::UndefInitializer, - dims::Tuple{AbstractVector{<:Integer},Vararg{AbstractVector{<:Integer}}}, -) where {T,N,A<:AbstractArray{T,N}} - length(dims) == N || - throw(ArgumentError("Length of dims doesn't match number of dimensions.")) - return BlockSparseArray{T,N,A}(undef, blockedrange.(dims)) +function BlockSparseArray{T, N, A}( + ::UndefInitializer, + dims::Tuple{AbstractVector{<:Integer}, Vararg{AbstractVector{<:Integer}}}, + ) where {T, N, A <: AbstractArray{T, N}} + length(dims) == N || + throw(ArgumentError("Length of dims doesn't match number of dimensions.")) + return BlockSparseArray{T, N, A}(undef, blockedrange.(dims)) end -function BlockSparseArray{T,N,A}( - ::UndefInitializer, - dim1::AbstractVector{<:Integer}, - dim_rest::AbstractVector{<:Integer}..., -) where {T,N,A<:AbstractArray{T,N}} - return BlockSparseArray{T,N,A}(undef, (dim1, dim_rest...)) +function BlockSparseArray{T, N, A}( + ::UndefInitializer, + dim1::AbstractVector{<:Integer}, + dim_rest::AbstractVector{<:Integer}..., + ) where {T, N, A <: AbstractArray{T, N}} + return BlockSparseArray{T, N, A}(undef, (dim1, dim_rest...)) end function similartype_unchecked( - A::Type{<:AbstractArray{T}}, axt::Type{<:Tuple{Vararg{Any,N}}} -) where {T,N} - A′ = Base.promote_op(similar, A, axt) - return !isconcretetype(A′) ? Array{T,N} : A′ + A::Type{<:AbstractArray{T}}, axt::Type{<:Tuple{Vararg{Any, N}}} + ) where {T, N} + A′ = Base.promote_op(similar, A, axt) + return !isconcretetype(A′) ? Array{T, N} : A′ end -function BlockSparseArray{T,N}( - ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange{<:Integer},N}} -) where {T,N} - axt = Tuple{blockaxistype.(axes)...} - # Ideally we would use: - # ```julia - # A = similartype(Array{T}, axt) - # ``` - # but that doesn't work when `similar` isn't defined or - # isn't type stable. - A = similartype_unchecked(Array{T}, axt) - return BlockSparseArray{T,N,A}(undef, axes) +function BlockSparseArray{T, N}( + ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}, N}} + ) where {T, N} + axt = Tuple{blockaxistype.(axes)...} + # Ideally we would use: + # ```julia + # A = similartype(Array{T}, axt) + # ``` + # but that doesn't work when `similar` isn't defined or + # isn't type stable. + A = similartype_unchecked(Array{T}, axt) + return BlockSparseArray{T, N, A}(undef, axes) end -function BlockSparseArray{T,N}( - ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}} -) where {T,N} - return throw(ArgumentError("Length of axes doesn't match number of dimensions.")) +function BlockSparseArray{T, N}( + ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}} + ) where {T, N} + return throw(ArgumentError("Length of axes doesn't match number of dimensions.")) end -function BlockSparseArray{T,N}( - ::UndefInitializer, axes::Vararg{AbstractUnitRange{<:Integer},N} -) where {T,N} - return BlockSparseArray{T,N}(undef, axes) +function BlockSparseArray{T, N}( + ::UndefInitializer, axes::Vararg{AbstractUnitRange{<:Integer}, N} + ) where {T, N} + return BlockSparseArray{T, N}(undef, axes) end -function BlockSparseArray{T,N}( - ::UndefInitializer, - dims::Tuple{AbstractVector{<:Integer},Vararg{AbstractVector{<:Integer}}}, -) where {T,N} - return BlockSparseArray{T,N}(undef, blockedrange.(dims)) +function BlockSparseArray{T, N}( + ::UndefInitializer, + dims::Tuple{AbstractVector{<:Integer}, Vararg{AbstractVector{<:Integer}}}, + ) where {T, N} + return BlockSparseArray{T, N}(undef, blockedrange.(dims)) end -function BlockSparseArray{T,N}( - ::UndefInitializer, - dim1::AbstractVector{<:Integer}, - dim_rest::AbstractVector{<:Integer}..., -) where {T,N} - return BlockSparseArray{T,N}(undef, (dim1, dim_rest...)) +function BlockSparseArray{T, N}( + ::UndefInitializer, + dim1::AbstractVector{<:Integer}, + dim_rest::AbstractVector{<:Integer}..., + ) where {T, N} + return BlockSparseArray{T, N}(undef, (dim1, dim_rest...)) end function BlockSparseArray{T}( - ::UndefInitializer, - dims::Tuple{AbstractVector{<:Integer},Vararg{AbstractVector{<:Integer}}}, -) where {T} - return BlockSparseArray{T,length(dims)}(undef, dims) + ::UndefInitializer, + dims::Tuple{AbstractVector{<:Integer}, Vararg{AbstractVector{<:Integer}}}, + ) where {T} + return BlockSparseArray{T, length(dims)}(undef, dims) end function BlockSparseArray{T}( - ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}} -) where {T} - return BlockSparseArray{T,length(axes)}(undef, axes) + ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}} + ) where {T} + return BlockSparseArray{T, length(axes)}(undef, axes) end function BlockSparseArray{T}( - ::UndefInitializer, - dim1::AbstractVector{<:Integer}, - dim_rest::AbstractVector{<:Integer}..., -) where {T} - return BlockSparseArray{T}(undef, (dim1, dim_rest...)) + ::UndefInitializer, + dim1::AbstractVector{<:Integer}, + dim_rest::AbstractVector{<:Integer}..., + ) where {T} + return BlockSparseArray{T}(undef, (dim1, dim_rest...)) end function BlockSparseArray{T}( - ::UndefInitializer, axes::Vararg{AbstractUnitRange{<:Integer}} -) where {T} - return BlockSparseArray{T}(undef, axes) + ::UndefInitializer, axes::Vararg{AbstractUnitRange{<:Integer}} + ) where {T} + return BlockSparseArray{T}(undef, axes) end # Convenient constructors. function blocksparsezeros(elt::Type, axes...) - return BlockSparseArray{elt}(undef, axes...) + return BlockSparseArray{elt}(undef, axes...) end function blocksparsezeros(axes...) - return blocksparsezeros(Float64, axes...) -end -function blocksparsezeros(::BlockType{A}, axes...) where {A<:AbstractArray} - # TODO: Use: - # ```julia - # B = similartype(A, Type{eltype(A)}, Tuple{blockaxistype.(axes)...}) - # BlockSparseArray{eltype(A),length(axes),B}(undef, axes...) - # ``` - # to make a bit more generic. - return BlockSparseArray{eltype(A),ndims(A),A}(undef, axes...) -end -function blocksparse(d::Dict{<:Block,<:AbstractArray}, ax::Tuple) - a = blocksparsezeros(BlockType(valtype(d)), ax...) - for I in eachindex(d) - a[I] = d[I] - end - return a + return blocksparsezeros(Float64, axes...) +end +function blocksparsezeros(::BlockType{A}, axes...) where {A <: AbstractArray} + # TODO: Use: + # ```julia + # B = similartype(A, Type{eltype(A)}, Tuple{blockaxistype.(axes)...}) + # BlockSparseArray{eltype(A),length(axes),B}(undef, axes...) + # ``` + # to make a bit more generic. + return BlockSparseArray{eltype(A), ndims(A), A}(undef, axes...) +end +function blocksparse(d::Dict{<:Block, <:AbstractArray}, ax::Tuple) + a = blocksparsezeros(BlockType(valtype(d)), ax...) + for I in eachindex(d) + a[I] = d[I] + end + return a end function blocksparse( - d::Dict{<:Block,<:AbstractArray}, blocklens::AbstractVector{<:Integer}... -) - return blocksparse(d, map(blockedrange, blocklens)) + d::Dict{<:Block, <:AbstractArray}, blocklens::AbstractVector{<:Integer}... + ) + return blocksparse(d, map(blockedrange, blocklens)) end # Base `AbstractArray` interface @@ -281,30 +281,30 @@ Base.axes(a::BlockSparseArray) = a.axes # BlockArrays `AbstractBlockArray` interface. # This is used by `blocks(::AnyAbstractBlockSparseArray)`. @interface ::AbstractBlockSparseArrayInterface BlockArrays.blocks(a::BlockSparseArray) = - a.blocks + a.blocks function blocktype( - arraytype::Type{<:BlockSparseArray{T,N,A}} -) where {T,N,A<:AbstractArray{T,N}} - return A + arraytype::Type{<:BlockSparseArray{T, N, A}} + ) where {T, N, A <: AbstractArray{T, N}} + return A end # TODO: Use `TypeParameterAccessors`. function blockstype( - arraytype::Type{<:BlockSparseArray{T,N,A,Blocks}} -) where {T,N,A<:AbstractArray{T,N},Blocks<:AbstractArray{A,N}} - return Blocks + arraytype::Type{<:BlockSparseArray{T, N, A, Blocks}} + ) where {T, N, A <: AbstractArray{T, N}, Blocks <: AbstractArray{A, N}} + return Blocks end function blockstype( - arraytype::Type{<:BlockSparseArray{T,N,A}} -) where {T,N,A<:AbstractArray{T,N}} - return SparseArrayDOK{A,N} + arraytype::Type{<:BlockSparseArray{T, N, A}} + ) where {T, N, A <: AbstractArray{T, N}} + return SparseArrayDOK{A, N} end -function blockstype(arraytype::Type{<:BlockSparseArray{T,N}}) where {T,N} - return SparseArrayDOK{AbstractArray{T,N},N} +function blockstype(arraytype::Type{<:BlockSparseArray{T, N}}) where {T, N} + return SparseArrayDOK{AbstractArray{T, N}, N} end function blockstype(arraytype::Type{<:BlockSparseArray{T}}) where {T} - return SparseArrayDOK{AbstractArray{T}} + return SparseArrayDOK{AbstractArray{T}} end blockstype(arraytype::Type{<:BlockSparseArray}) = SparseArrayDOK{AbstractArray} @@ -322,5 +322,5 @@ TypeParameterAccessors.position(::Type{BlockSparseArray}, ::typeof(eltype)) = Po TypeParameterAccessors.position(::Type{BlockSparseArray}, ::typeof(ndims)) = Position(2) TypeParameterAccessors.position(::Type{BlockSparseArray}, ::typeof(blocktype)) = Position(3) function TypeParameterAccessors.position(::Type{BlockSparseArray}, ::typeof(blockstype)) - return Position(4) + return Position(4) end diff --git a/src/blocksparsearrayinterface/arraylayouts.jl b/src/blocksparsearrayinterface/arraylayouts.jl index 14d3d563..f1511248 100644 --- a/src/blocksparsearrayinterface/arraylayouts.jl +++ b/src/blocksparsearrayinterface/arraylayouts.jl @@ -5,45 +5,45 @@ using SparseArraysBase: SparseLayout using LinearAlgebra: LinearAlgebra, dot, mul! @interface ::AbstractBlockSparseArrayInterface function BlockArrays.muladd!( - α::Number, a1::AbstractArray, a2::AbstractArray, β::Number, a_dest::AbstractArray -) - mul!(blocks(a_dest), blocks(a1), blocks(a2), α, β) - return a_dest + α::Number, a1::AbstractArray, a2::AbstractArray, β::Number, a_dest::AbstractArray + ) + mul!(blocks(a_dest), blocks(a1), blocks(a2), α, β) + return a_dest end function DerivableInterfaces.interface(m::MulAdd) - return interface(m.A, m.B, m.C) + return interface(m.A, m.B, m.C) end function ArrayLayouts.materialize!( - m::MatMulMatAdd{ - <:BlockLayout{<:SparseLayout}, - <:BlockLayout{<:SparseLayout}, - <:BlockLayout{<:SparseLayout}, - }, -) - @interface interface(m) muladd!(m.α, m.A, m.B, m.β, m.C) - return m.C + m::MatMulMatAdd{ + <:BlockLayout{<:SparseLayout}, + <:BlockLayout{<:SparseLayout}, + <:BlockLayout{<:SparseLayout}, + }, + ) + @interface interface(m) muladd!(m.α, m.A, m.B, m.β, m.C) + return m.C end function ArrayLayouts.materialize!( - m::MatMulVecAdd{ - <:BlockLayout{<:SparseLayout}, - <:BlockLayout{<:SparseLayout}, - <:BlockLayout{<:SparseLayout}, - }, -) - @interface interface(m) matmul!(m) - return m.C + m::MatMulVecAdd{ + <:BlockLayout{<:SparseLayout}, + <:BlockLayout{<:SparseLayout}, + <:BlockLayout{<:SparseLayout}, + }, + ) + @interface interface(m) matmul!(m) + return m.C end @interface ::AbstractBlockSparseArrayInterface function LinearAlgebra.dot( - a1::AbstractArray, a2::AbstractArray -) - # TODO: Add a check that the blocking of `a1` and `a2` are - # the same, or the same up to a reshape. - return dot(blocks(a1), blocks(a2)) + a1::AbstractArray, a2::AbstractArray + ) + # TODO: Add a check that the blocking of `a1` and `a2` are + # the same, or the same up to a reshape. + return dot(blocks(a1), blocks(a2)) end -function Base.copy(d::Dot{<:BlockLayout{<:SparseLayout},<:BlockLayout{<:SparseLayout}}) - return @interface interface(d.A, d.B) dot(d.A, d.B) +function Base.copy(d::Dot{<:BlockLayout{<:SparseLayout}, <:BlockLayout{<:SparseLayout}}) + return @interface interface(d.A, d.B) dot(d.A, d.B) end diff --git a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index bdda71e0..02e4fec4 100644 --- a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -1,77 +1,77 @@ using ArrayLayouts: ArrayLayouts using BlockArrays: - BlockArrays, - AbstractBlockVector, - Block, - BlockIndex, - BlockRange, - BlockSlice, - BlockVector, - BlockedUnitRange, - BlockedVector, - block, - blockcheckbounds, - blockisequal, - blocklengths, - blocklength, - blocks, - findblockindex + BlockArrays, + AbstractBlockVector, + Block, + BlockIndex, + BlockRange, + BlockSlice, + BlockVector, + BlockedUnitRange, + BlockedVector, + block, + blockcheckbounds, + blockisequal, + blocklengths, + blocklength, + blocks, + findblockindex using DerivableInterfaces: - DerivableInterfaces, - @interface, - AbstractArrayInterface, - DefaultArrayInterface, - interface, - permuteddims, - zero! + DerivableInterfaces, + @interface, + AbstractArrayInterface, + DefaultArrayInterface, + interface, + permuteddims, + zero! using LinearAlgebra: Adjoint, Transpose using SparseArraysBase: - AbstractSparseArrayInterface, - getstoredindex, - getunstoredindex, - eachstoredindex, - perm, - iperm, - storedlength, - storedvalues + AbstractSparseArrayInterface, + getstoredindex, + getunstoredindex, + eachstoredindex, + perm, + iperm, + storedlength, + storedvalues # Like `SparseArraysBase.eachstoredindex` but # at the block level, i.e. iterates over the # stored block locations. function eachblockstoredindex(a::AbstractArray) - # TODO: Use `Iterators.map`. - return Block.(Tuple.(eachstoredindex(blocks(a)))) + # TODO: Use `Iterators.map`. + 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)...) + 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)...) +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) - return Block.(Tuple.(diagindices(blocks(a)))) + return Block.(Tuple.(diagindices(blocks(a)))) end function eachstoredblockdiagindex(a::AbstractArray) - return eachblockstoredindex(a) ∩ blockdiagindices(a) + return eachblockstoredindex(a) ∩ blockdiagindices(a) end function eachunstoredblockdiagindex(a::AbstractArray) - return setdiff(blockdiagindices(a), eachblockstoredindex(a)) + return setdiff(blockdiagindices(a), eachblockstoredindex(a)) end # Like `BlockArrays.eachblock` but only iterating # over stored blocks. function eachstoredblock(a::AbstractArray) - return storedvalues(blocks(a)) + return storedvalues(blocks(a)) end function blockstype(a::AbstractArray) - return typeof(blocks(a)) + return typeof(blocks(a)) end #= @@ -104,82 +104,82 @@ julia> length(blocks(BlockedVector{Float64}(randn(0)))) ``` =# function blocktype(a::AbstractArray) - if isempty(blocks(a)) - error("`blocktype` can't be determined if `isempty(blocks(a))`.") - end - return mapreduce(typeof, promote_type, blocks(a)) + if isempty(blocks(a)) + error("`blocktype` can't be determined if `isempty(blocks(a))`.") + end + return mapreduce(typeof, promote_type, blocks(a)) end using BlockArrays: BlockArray -blockstype(::Type{<:BlockArray{<:Any,<:Any,B}}) where {B} = B +blockstype(::Type{<:BlockArray{<:Any, <:Any, B}}) where {B} = B blockstype(a::BlockArray) = blockstype(typeof(a)) blocktype(arraytype::Type{<:BlockArray}) = eltype(blockstype(arraytype)) blocktype(a::BlockArray) = eltype(blocks(a)) -abstract type AbstractBlockSparseArrayInterface{N,B<:AbstractArrayInterface{N}} <: - AbstractSparseArrayInterface{N} end +abstract type AbstractBlockSparseArrayInterface{N, B <: AbstractArrayInterface{N}} <: +AbstractSparseArrayInterface{N} end -function blockinterface(interface::AbstractBlockSparseArrayInterface{<:Any,B}) where {B} - return B() +function blockinterface(interface::AbstractBlockSparseArrayInterface{<:Any, B}) where {B} + return B() end # TODO: Also support specifying the `blocktype` along with the `eltype`. function Base.similar(interface::AbstractBlockSparseArrayInterface, T::Type, ax::Tuple) - # TODO: Generalize by storing the block interface in the block sparse array interface. - N = length(ax) - B = similartype(typeof(blockinterface(interface)), Type{T}, Tuple{blockaxistype.(ax)...}) - return similar(BlockSparseArray{T,N,B}, ax) + # TODO: Generalize by storing the block interface in the block sparse array interface. + N = length(ax) + B = similartype(typeof(blockinterface(interface)), Type{T}, Tuple{blockaxistype.(ax)...}) + return similar(BlockSparseArray{T, N, B}, ax) end -struct BlockSparseArrayInterface{N,B<:AbstractArrayInterface{N}} <: - AbstractBlockSparseArrayInterface{N,B} - blockinterface::B +struct BlockSparseArrayInterface{N, B <: AbstractArrayInterface{N}} <: + AbstractBlockSparseArrayInterface{N, B} + blockinterface::B end function BlockSparseArrayInterface{N}(blockinterface::AbstractArrayInterface{N}) where {N} - return BlockSparseArrayInterface{N,typeof(blockinterface)}(blockinterface) + return BlockSparseArrayInterface{N, typeof(blockinterface)}(blockinterface) end -function BlockSparseArrayInterface{M,B}(::Val{N}) where {M,B<:AbstractArrayInterface{M},N} - B′ = B(Val(N)) - return BlockSparseArrayInterface(B′) +function BlockSparseArrayInterface{M, B}(::Val{N}) where {M, B <: AbstractArrayInterface{M}, N} + B′ = B(Val(N)) + return BlockSparseArrayInterface(B′) end function BlockSparseArrayInterface{N}() where {N} - return BlockSparseArrayInterface{N}(DefaultArrayInterface{N}()) + return BlockSparseArrayInterface{N}(DefaultArrayInterface{N}()) end BlockSparseArrayInterface(::Val{N}) where {N} = BlockSparseArrayInterface{N}() -BlockSparseArrayInterface{M}(::Val{N}) where {M,N} = BlockSparseArrayInterface{N}() +BlockSparseArrayInterface{M}(::Val{N}) where {M, N} = BlockSparseArrayInterface{N}() BlockSparseArrayInterface() = BlockSparseArrayInterface{Any}() function DerivableInterfaces.combine_interface_rule( - interface1::AbstractBlockSparseArrayInterface, - interface2::AbstractBlockSparseArrayInterface, -) - B = interface(blockinterface(interface1), blockinterface(interface2)) - return BlockSparseArrayInterface(B) + interface1::AbstractBlockSparseArrayInterface, + interface2::AbstractBlockSparseArrayInterface, + ) + B = interface(blockinterface(interface1), blockinterface(interface2)) + return BlockSparseArrayInterface(B) end @interface ::AbstractBlockSparseArrayInterface function BlockArrays.blocks(a::AbstractArray) - return error("Not implemented") + return error("Not implemented") end @interface ::AbstractBlockSparseArrayInterface function SparseArraysBase.isstored( - a::AbstractArray{<:Any,N}, I::Vararg{Int,N} -) where {N} - bI = BlockIndex(findblockindex.(axes(a), I)) - return isstored(blocks(a), bI.I...) && isstored(blocks(a)[bI.I...], bI.α...) + a::AbstractArray{<:Any, N}, I::Vararg{Int, N} + ) where {N} + bI = BlockIndex(findblockindex.(axes(a), I)) + return isstored(blocks(a), bI.I...) && isstored(blocks(a)[bI.I...], bI.α...) end @interface ::AbstractBlockSparseArrayInterface function Base.getindex( - a::AbstractArray{<:Any,N}, I::Vararg{Int,N} -) where {N} - @boundscheck checkbounds(a, I...) - return a[findblockindex.(axes(a), I)...] + a::AbstractArray{<:Any, N}, I::Vararg{Int, N} + ) where {N} + @boundscheck checkbounds(a, I...) + return a[findblockindex.(axes(a), I)...] end # Fix ambiguity error. @interface ::AbstractBlockSparseArrayInterface function Base.getindex( - a::AbstractArray{<:Any,0} -) - return a[Block()[]] + a::AbstractArray{<:Any, 0} + ) + return a[Block()[]] end # a[1:2, 1:2] @@ -189,122 +189,122 @@ end # https://github.com/JuliaArrays/BlockArrays.jl/issues/347 and also # https://github.com/ITensor/ITensors.jl/issues/1336. @interface ::AbstractBlockSparseArrayInterface function Base.to_indices( - a, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}} -) - bs1 = to_blockindices(inds[1], I[1]) - I1 = BlockSlice(bs1, blockedunitrange_getindices(inds[1], I[1])) - return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) + a, inds, I::Tuple{UnitRange{<:Integer}, Vararg{Any}} + ) + bs1 = to_blockindices(inds[1], I[1]) + I1 = BlockSlice(bs1, blockedunitrange_getindices(inds[1], I[1])) + return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) end @interface ::AbstractBlockSparseArrayInterface function Base.to_indices( - a, inds, I::Tuple{AbstractArray{Bool},Vararg{Any}} -) - bs1 = to_blockindices(inds[1], I[1]) - I1 = BlockIndices(bs1, blockedunitrange_getindices(inds[1], I[1])) - return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) + a, inds, I::Tuple{AbstractArray{Bool}, Vararg{Any}} + ) + bs1 = to_blockindices(inds[1], I[1]) + I1 = BlockIndices(bs1, blockedunitrange_getindices(inds[1], I[1])) + return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) end # Special case when there is no blocking. @interface ::AbstractBlockSparseArrayInterface function Base.to_indices( - a, - inds::Tuple{Base.OneTo{<:Integer},Vararg{Any}}, - I::Tuple{UnitRange{<:Integer},Vararg{Any}}, -) - return (inds[1][I[1]], to_indices(a, Base.tail(inds), Base.tail(I))...) + a, + inds::Tuple{Base.OneTo{<:Integer}, Vararg{Any}}, + I::Tuple{UnitRange{<:Integer}, Vararg{Any}}, + ) + return (inds[1][I[1]], to_indices(a, Base.tail(inds), Base.tail(I))...) end # a[[Block(2), Block(1)], [Block(2), Block(1)]] @interface ::AbstractBlockSparseArrayInterface function Base.to_indices( - a, inds, I::Tuple{Vector{<:Block{1}},Vararg{Any}} -) - I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1])) - return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) + a, inds, I::Tuple{Vector{<:Block{1}}, Vararg{Any}} + ) + I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1])) + return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) end # a[mortar([Block(1)[1:2], Block(2)[1:3]]), mortar([Block(1)[1:2], Block(2)[1:3]])] # a[[Block(1)[1:2], Block(2)[1:3]], [Block(1)[1:2], Block(2)[1:3]]] @interface ::AbstractBlockSparseArrayInterface function Base.to_indices( - a, - inds, - I::Tuple{ - BlockVector{<:BlockIndex{1},<:Vector{<:Union{BlockIndexRange{1},BlockIndexVector{1}}}}, - Vararg{Any}, - }, -) - # Index the `inds` by the `BlockIndexRange`/`BlockIndexVector`s on each block - # in order to canonicalize the indices and preserve metadata, - # such as sector data for symmetric tensors. - bs = mortar( - map(blocks(I[1])) do bi - b = Block(bi) - binds = only(bi.indices) - return BlockIndexVector(b, Base.axes1(inds[1][b])[binds]) - end, - ) - I1 = BlockIndices(bs, blockedunitrange_getindices(inds[1], I[1])) - return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) + a, + inds, + I::Tuple{ + BlockVector{<:BlockIndex{1}, <:Vector{<:Union{BlockIndexRange{1}, BlockIndexVector{1}}}}, + Vararg{Any}, + }, + ) + # Index the `inds` by the `BlockIndexRange`/`BlockIndexVector`s on each block + # in order to canonicalize the indices and preserve metadata, + # such as sector data for symmetric tensors. + bs = mortar( + map(blocks(I[1])) do bi + b = Block(bi) + binds = only(bi.indices) + return BlockIndexVector(b, Base.axes1(inds[1][b])[binds]) + end, + ) + I1 = BlockIndices(bs, blockedunitrange_getindices(inds[1], I[1])) + return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) end @interface ::AbstractBlockSparseArrayInterface function Base.to_indices( - a, - inds, - I::Tuple{BlockVector{<:GenericBlockIndex{1},<:Vector{<:BlockIndexVector{1}}},Vararg{Any}}, -) - I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1])) - return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) + a, + inds, + I::Tuple{BlockVector{<:GenericBlockIndex{1}, <:Vector{<:BlockIndexVector{1}}}, Vararg{Any}}, + ) + I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1])) + return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) end # a[BlockVector([Block(2), Block(1)], [2]), BlockVector([Block(2), Block(1)], [2])] # Permute and merge blocks. # TODO: This isn't merging blocks yet, that needs to be implemented that. @interface ::AbstractBlockSparseArrayInterface function Base.to_indices( - a, inds, I::Tuple{AbstractBlockVector{<:Block{1}},Vararg{Any}} -) - I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1])) - return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) + a, inds, I::Tuple{AbstractBlockVector{<:Block{1}}, Vararg{Any}} + ) + I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1])) + return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) end # TODO: Need to implement this! function block_merge end @interface ::AbstractBlockSparseArrayInterface function Base.setindex!( - a::AbstractArray{<:Any,N}, value, I::Vararg{Int,N} -) where {N} - @boundscheck checkbounds(a, I...) - a[findblockindex.(axes(a), I)...] = value - return a + a::AbstractArray{<:Any, N}, value, I::Vararg{Int, N} + ) where {N} + @boundscheck checkbounds(a, I...) + a[findblockindex.(axes(a), I)...] = value + return a end # Fix ambiguity error. @interface ::AbstractBlockSparseArrayInterface function Base.setindex!( - a::AbstractArray{<:Any,0}, value -) - # TODO: Use `Block()[]` once https://github.com/JuliaArrays/BlockArrays.jl/issues/430 - # is fixed. - a[BlockIndex()] = value - return a + a::AbstractArray{<:Any, 0}, value + ) + # TODO: Use `Block()[]` once https://github.com/JuliaArrays/BlockArrays.jl/issues/430 + # is fixed. + a[BlockIndex()] = value + return a end @interface ::AbstractBlockSparseArrayInterface function Base.setindex!( - a::AbstractArray{<:Any,N}, value, I::BlockIndex{N} -) where {N} - i = Int.(Tuple(block(I))) - a_b = blocks(a)[i...] - a_b[I.α...] = value - # Set the block, required if it is structurally zero. - blocks(a)[i...] = a_b - return a + a::AbstractArray{<:Any, N}, value, I::BlockIndex{N} + ) where {N} + i = Int.(Tuple(block(I))) + a_b = blocks(a)[i...] + a_b[I.α...] = value + # Set the block, required if it is structurally zero. + blocks(a)[i...] = a_b + return a end # Fix ambiguity error. @interface ::AbstractBlockSparseArrayInterface function Base.setindex!( - a::AbstractArray{<:Any,0}, value, I::BlockIndex{0} -) - a_b = blocks(a)[] - # `value[]` handles scalars and 0-dimensional arrays. - a_b[] = value[] - # Set the block, required if it is structurally zero. - blocks(a)[] = a_b - return a + a::AbstractArray{<:Any, 0}, value, I::BlockIndex{0} + ) + a_b = blocks(a)[] + # `value[]` handles scalars and 0-dimensional arrays. + a_b[] = value[] + # Set the block, required if it is structurally zero. + blocks(a)[] = a_b + return a end # Version of `permutedims!` that assumes the destination and source @@ -313,70 +313,70 @@ end # `blockisequal_map[!]`. # TODO: Maybe define a `BlockIsEqualInterface` for these kinds of functions. function blockisequal_permutedims!(a_dest::AbstractArray, a_src::AbstractArray, perm) - blocks(a_dest) .= blocks(permuteddims(a_src, perm)) - return a_dest + blocks(a_dest) .= blocks(permuteddims(a_src, perm)) + return a_dest end # We overload `permutedims` here so that we can assume the destination and source # have the same blocking and avoid non-GPU friendly slicing operations in block sparse `map!`. # TODO: Delete this and handle this logic in block sparse `map!`. @interface ::AbstractBlockSparseArrayInterface function Base.permutedims( - a::AbstractArray, perm -) - a_dest = similar(permuteddims(a, perm)) - # TODO: Maybe define this as `@interface BlockIsEqualInterface() permutedims!(...)`. - blockisequal_permutedims!(a_dest, a, perm) - return a_dest + a::AbstractArray, perm + ) + a_dest = similar(permuteddims(a, perm)) + # TODO: Maybe define this as `@interface BlockIsEqualInterface() permutedims!(...)`. + blockisequal_permutedims!(a_dest, a, perm) + return a_dest end # We overload `permutedims!` here so that we can special case when the destination and source # have the same blocking and avoid non-GPU friendly slicing operations in block sparse `map!`. # TODO: Delete this and handle this logic in block sparse `map!`. @interface ::AbstractBlockSparseArrayInterface function Base.permutedims!( - a_dest::AbstractArray, a_src::AbstractArray, perm -) - if all(blockisequal.(axes(a_dest), axes(permuteddims(a_src, perm)))) - # TODO: Maybe define this as `@interface BlockIsEqualInterface() permutedims!(...)`. - blockisequal_permutedims!(a_dest, a_src, perm) + a_dest::AbstractArray, a_src::AbstractArray, perm + ) + if all(blockisequal.(axes(a_dest), axes(permuteddims(a_src, perm)))) + # TODO: Maybe define this as `@interface BlockIsEqualInterface() permutedims!(...)`. + blockisequal_permutedims!(a_dest, a_src, perm) + return a_dest + end + @interface DefaultArrayInterface() permutedims!(a_dest, a_src, perm) return a_dest - end - @interface DefaultArrayInterface() permutedims!(a_dest, a_src, perm) - return a_dest end @interface ::AbstractBlockSparseArrayInterface function Base.fill!(a::AbstractArray, value) - # TODO: Only do this check if `value isa Number`? - if iszero(value) - zero!(a) + # TODO: Only do this check if `value isa Number`? + if iszero(value) + zero!(a) + return a + end + # TODO: Maybe use `map` over `blocks(a)` or something + # like that. + for b in BlockRange(a) + fill!(@view!(a[b]), value) + end return a - end - # TODO: Maybe use `map` over `blocks(a)` or something - # like that. - for b in BlockRange(a) - fill!(@view!(a[b]), value) - end - return a end @interface ::AbstractBlockSparseArrayInterface function DerivableInterfaces.zero!( - a::AbstractArray -) - # This will try to empty the storage if possible. - zero!(blocks(a)) - return a + a::AbstractArray + ) + # This will try to empty the storage if possible. + zero!(blocks(a)) + return a end # TODO: Rename to `blockstoredlength`. function blockstoredlength(a::AbstractArray) - return storedlength(blocks(a)) + return storedlength(blocks(a)) end # BlockArrays using SparseArraysBase: SparseArraysBase, AbstractSparseArray, AbstractSparseMatrix -_perm(::PermutedDimsArray{<:Any,<:Any,perm}) where {perm} = perm -_invperm(::PermutedDimsArray{<:Any,<:Any,<:Any,invperm}) where {invperm} = invperm +_perm(::PermutedDimsArray{<:Any, <:Any, perm}) where {perm} = perm +_invperm(::PermutedDimsArray{<:Any, <:Any, <:Any, invperm}) where {invperm} = invperm _getindices(t::Tuple, indices) = map(i -> t[i], indices) _getindices(i::CartesianIndex, indices) = CartesianIndex(_getindices(Tuple(i), indices)) @@ -384,43 +384,43 @@ _getindices(i::CartesianIndex, indices) = CartesianIndex(_getindices(Tuple(i), i # wrapping a block spare array, i.e. `blocks(array)` where `a` is a `PermutedDimsArray`. # TODO: Delete this in favor of `NestedPermutedDimsArrays.NestedPermutedDimsArray`. struct SparsePermutedDimsArrayBlocks{ - T,N,BlockType<:AbstractArray{T,N},Array<:PermutedDimsArray{T,N} -} <: AbstractSparseArray{BlockType,N} - array::Array + T, N, BlockType <: AbstractArray{T, N}, Array <: PermutedDimsArray{T, N}, + } <: AbstractSparseArray{BlockType, N} + array::Array end @interface ::AbstractBlockSparseArrayInterface function BlockArrays.blocks( - a::PermutedDimsArray -) - return SparsePermutedDimsArrayBlocks{eltype(a),ndims(a),blocktype(parent(a)),typeof(a)}(a) + a::PermutedDimsArray + ) + return SparsePermutedDimsArrayBlocks{eltype(a), ndims(a), blocktype(parent(a)), typeof(a)}(a) end function Base.size(a::SparsePermutedDimsArrayBlocks) - return _getindices(size(blocks(parent(a.array))), _perm(a.array)) + return _getindices(size(blocks(parent(a.array))), _perm(a.array)) end function SparseArraysBase.isstored( - a::SparsePermutedDimsArrayBlocks{<:Any,N}, index::Vararg{Int,N} -) where {N} - return isstored(blocks(parent(a.array)), _getindices(index, _invperm(a.array))...) + a::SparsePermutedDimsArrayBlocks{<:Any, N}, index::Vararg{Int, N} + ) where {N} + return isstored(blocks(parent(a.array)), _getindices(index, _invperm(a.array))...) end function SparseArraysBase.getstoredindex( - a::SparsePermutedDimsArrayBlocks{<:Any,N}, index::Vararg{Int,N} -) where {N} - return permuteddims( - getstoredindex(blocks(parent(a.array)), _getindices(index, _invperm(a.array))...), - _perm(a.array), - ) + a::SparsePermutedDimsArrayBlocks{<:Any, N}, index::Vararg{Int, N} + ) where {N} + return permuteddims( + getstoredindex(blocks(parent(a.array)), _getindices(index, _invperm(a.array))...), + _perm(a.array), + ) end function SparseArraysBase.getunstoredindex( - a::SparsePermutedDimsArrayBlocks{<:Any,N}, index::Vararg{Int,N} -) where {N} - return permuteddims( - getunstoredindex(blocks(parent(a.array)), _getindices(index, _invperm(a.array))...), - _perm(a.array), - ) + a::SparsePermutedDimsArrayBlocks{<:Any, N}, index::Vararg{Int, N} + ) where {N} + return permuteddims( + getunstoredindex(blocks(parent(a.array)), _getindices(index, _invperm(a.array))...), + _perm(a.array), + ) end function SparseArraysBase.eachstoredindex( - ::IndexCartesian, a::SparsePermutedDimsArrayBlocks -) - return map(I -> _getindices(I, _perm(a.array)), eachstoredindex(blocks(parent(a.array)))) + ::IndexCartesian, a::SparsePermutedDimsArrayBlocks + ) + return map(I -> _getindices(I, _perm(a.array)), eachstoredindex(blocks(parent(a.array)))) end ## TODO: Define `storedvalues` instead. ## function SparseArraysBase.sparse_storage(a::SparsePermutedDimsArrayBlocks) @@ -431,91 +431,91 @@ reverse_index(index) = reverse(index) reverse_index(index::CartesianIndex) = CartesianIndex(reverse(Tuple(index))) @interface ::AbstractBlockSparseArrayInterface BlockArrays.blocks(a::Transpose) = transpose( - blocks(parent(a)) + blocks(parent(a)) ) @interface ::AbstractBlockSparseArrayInterface BlockArrays.blocks(a::Adjoint) = adjoint( - blocks(parent(a)) + blocks(parent(a)) ) # Represents the array of arrays of a `SubArray` # wrapping a block spare array, i.e. `blocks(array)` where `a` is a `SubArray`. -struct SparseSubArrayBlocks{T,N,BlockType<:AbstractArray{T,N},Array<:SubArray{T,N}} <: - AbstractSparseArray{BlockType,N} - array::Array +struct SparseSubArrayBlocks{T, N, BlockType <: AbstractArray{T, N}, Array <: SubArray{T, N}} <: + AbstractSparseArray{BlockType, N} + array::Array end @interface ::AbstractBlockSparseArrayInterface function BlockArrays.blocks(a::SubArray) - return SparseSubArrayBlocks{eltype(a),ndims(a),blocktype(parent(a)),typeof(a)}(a) + return SparseSubArrayBlocks{eltype(a), ndims(a), blocktype(parent(a)), typeof(a)}(a) end # TODO: Define this as `blockrange(a::AbstractArray, indices::Tuple{Vararg{AbstractUnitRange}})`. function blockrange(a::SparseSubArrayBlocks) - blockranges = blockrange.(axes(parent(a.array)), a.array.indices) - return map(blockrange -> Int.(blockrange), blockranges) + blockranges = blockrange.(axes(parent(a.array)), a.array.indices) + return map(blockrange -> Int.(blockrange), blockranges) end function Base.axes(a::SparseSubArrayBlocks) - return Base.OneTo.(length.(blockrange(a))) + return Base.OneTo.(length.(blockrange(a))) end function Base.size(a::SparseSubArrayBlocks) - return length.(axes(a)) + return length.(axes(a)) end # TODO: Make a faster version for when the slice is blockwise. function SparseArraysBase.isstored( - a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N} -) where {N} - J = Base.reindex(parentindices(a.array), to_indices(a.array, Block.(I))) - # TODO: Try doing this blockwise when possible rather - # than elementwise. - return any(Iterators.product(J...)) do K - return isstored(parent(a.array), K...) - end + a::SparseSubArrayBlocks{<:Any, N}, I::Vararg{Int, N} + ) where {N} + J = Base.reindex(parentindices(a.array), to_indices(a.array, Block.(I))) + # TODO: Try doing this blockwise when possible rather + # than elementwise. + return any(Iterators.product(J...)) do K + return isstored(parent(a.array), K...) + end end # TODO: Define `getstoredindex`, `getunstoredindex` instead. -function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) where {N} - # TODO: Should this be defined as `@view a.array[Block(I)]` instead? - return @view a.array[Block(I)] +function Base.getindex(a::SparseSubArrayBlocks{<:Any, N}, I::Vararg{Int, N}) where {N} + # TODO: Should this be defined as `@view a.array[Block(I)]` instead? + return @view a.array[Block(I)] - ## parent_blocks = @view blocks(parent(a.array))[blockrange(a)...] - ## parent_block = parent_blocks[I...] - ## # TODO: Define this using `blockrange(a::AbstractArray, indices::Tuple{Vararg{AbstractUnitRange}})`. - ## block = Block(ntuple(i -> blockrange(a)[i][I[i]], ndims(a))) - ## return @view parent_block[blockindices(parent(a.array), block, a.array.indices)...] + ## parent_blocks = @view blocks(parent(a.array))[blockrange(a)...] + ## parent_block = parent_blocks[I...] + ## # TODO: Define this using `blockrange(a::AbstractArray, indices::Tuple{Vararg{AbstractUnitRange}})`. + ## block = Block(ntuple(i -> blockrange(a)[i][I[i]], ndims(a))) + ## return @view parent_block[blockindices(parent(a.array), block, a.array.indices)...] end # TODO: This should be handled by generic `AbstractSparseArray` code. # TODO: Define `getstoredindex`, `getunstoredindex` instead. -function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::CartesianIndex{N}) where {N} - return a[Tuple(I)...] +function Base.getindex(a::SparseSubArrayBlocks{<:Any, N}, I::CartesianIndex{N}) where {N} + return a[Tuple(I)...] end # TODO: Define `setstoredindex!`, `setunstoredindex!` instead. -function Base.setindex!(a::SparseSubArrayBlocks{<:Any,N}, value, I::Vararg{Int,N}) where {N} - parent_blocks = @view blocks(parent(a.array))[blockrange(a)...] - # TODO: The following line is required to instantiate - # uninstantiated blocks, maybe use `@view!` instead, - # or some other code pattern. - parent_blocks[I...] = parent_blocks[I...] - # TODO: Define this using `blockrange(a::AbstractArray, indices::Tuple{Vararg{AbstractUnitRange}})`. - block = Block(ntuple(i -> blockrange(a)[i][I[i]], ndims(a))) - return parent_blocks[I...][blockindices(parent(a.array), block, a.array.indices)...] = - value -end -function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) where {N} - if CartesianIndex(I) ∉ CartesianIndices(a) - return false - end - # TODO: Implement this properly. - return true +function Base.setindex!(a::SparseSubArrayBlocks{<:Any, N}, value, I::Vararg{Int, N}) where {N} + parent_blocks = @view blocks(parent(a.array))[blockrange(a)...] + # TODO: The following line is required to instantiate + # uninstantiated blocks, maybe use `@view!` instead, + # or some other code pattern. + parent_blocks[I...] = parent_blocks[I...] + # TODO: Define this using `blockrange(a::AbstractArray, indices::Tuple{Vararg{AbstractUnitRange}})`. + block = Block(ntuple(i -> blockrange(a)[i][I[i]], ndims(a))) + return parent_blocks[I...][blockindices(parent(a.array), block, a.array.indices)...] = + value +end +function Base.isassigned(a::SparseSubArrayBlocks{<:Any, N}, I::Vararg{Int, N}) where {N} + if CartesianIndex(I) ∉ CartesianIndices(a) + return false + end + # TODO: Implement this properly. + return true end function SparseArraysBase.eachstoredindex(::IndexCartesian, a::SparseSubArrayBlocks) - isempty(a) && return CartesianIndex{ndims(a)}[] - inds = filter(eachindex(a)) do I - return isstored(a, I) - end - return inds + isempty(a) && return CartesianIndex{ndims(a)}[] + inds = filter(eachindex(a)) do I + return isstored(a, I) + end + return inds - ## # TODO: This only works for blockwise slices, i.e. slices using - ## # `BlockSliceCollection`. - ## return eachstoredindex(view(blocks(parent(a.array)), blockrange(a)...)) + ## # TODO: This only works for blockwise slices, i.e. slices using + ## # `BlockSliceCollection`. + ## return eachstoredindex(view(blocks(parent(a.array)), blockrange(a)...)) end # TODO: Either make this the generic interface or define @@ -533,9 +533,9 @@ SparseArraysBase.storedlength(a::SparseSubArrayBlocks) = length(eachstoredindex( ## end function SparseArraysBase.getunstoredindex( - a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N} -) where {N} - return error("Not implemented.") + a::SparseSubArrayBlocks{<:Any, N}, I::Vararg{Int, N} + ) where {N} + return error("Not implemented.") end # Convert a blockwise slice on a block sparse array @@ -547,9 +547,9 @@ to_blocks_indices(I::BlockIndices{<:Vector{<:Block{1}}}) = Int.(I.blocks) to_blocks_indices(I::Base.Slice) = Base.OneTo(blocklength(I.indices)) @interface ::AbstractBlockSparseArrayInterface function BlockArrays.blocks( - a::SubArray{<:Any,<:Any,<:Any,<:Tuple{Vararg{BlockSliceCollection}}} -) - return @view blocks(parent(a))[map(to_blocks_indices, parentindices(a))...] + a::SubArray{<:Any, <:Any, <:Any, <:Tuple{Vararg{BlockSliceCollection}}} + ) + return @view blocks(parent(a))[map(to_blocks_indices, parentindices(a))...] end using BlockArrays: BlocksView diff --git a/src/blocksparsearrayinterface/broadcast.jl b/src/blocksparsearrayinterface/broadcast.jl index a4ee8ae4..dd24d884 100644 --- a/src/blocksparsearrayinterface/broadcast.jl +++ b/src/blocksparsearrayinterface/broadcast.jl @@ -1,86 +1,86 @@ using Base.Broadcast: - Broadcast, BroadcastStyle, AbstractArrayStyle, DefaultArrayStyle, Broadcasted + Broadcast, BroadcastStyle, AbstractArrayStyle, DefaultArrayStyle, Broadcasted using GPUArraysCore: @allowscalar using MapBroadcast: Mapped using DerivableInterfaces: DerivableInterfaces, @interface -abstract type AbstractBlockSparseArrayStyle{N,B<:AbstractArrayStyle{N}} <: - AbstractArrayStyle{N} end +abstract type AbstractBlockSparseArrayStyle{N, B <: AbstractArrayStyle{N}} <: +AbstractArrayStyle{N} end -blockstyle(::AbstractBlockSparseArrayStyle{N,B}) where {N,B<:AbstractArrayStyle{N}} = B() +blockstyle(::AbstractBlockSparseArrayStyle{N, B}) where {N, B <: AbstractArrayStyle{N}} = B() function Broadcast.BroadcastStyle( - style1::AbstractBlockSparseArrayStyle, style2::AbstractBlockSparseArrayStyle -) - style = Broadcast.result_style(blockstyle(style1), blockstyle(style2)) - return BlockSparseArrayStyle(style) + style1::AbstractBlockSparseArrayStyle, style2::AbstractBlockSparseArrayStyle + ) + style = Broadcast.result_style(blockstyle(style1), blockstyle(style2)) + return BlockSparseArrayStyle(style) end function DerivableInterfaces.interface( - ::Type{<:AbstractBlockSparseArrayStyle{N,B}} -) where {N,B<:AbstractArrayStyle{N}} - return BlockSparseArrayInterface(interface(B)) + ::Type{<:AbstractBlockSparseArrayStyle{N, B}} + ) where {N, B <: AbstractArrayStyle{N}} + return BlockSparseArrayInterface(interface(B)) end -struct BlockSparseArrayStyle{N,B<:AbstractArrayStyle{N}} <: - AbstractBlockSparseArrayStyle{N,B} - blockstyle::B +struct BlockSparseArrayStyle{N, B <: AbstractArrayStyle{N}} <: + AbstractBlockSparseArrayStyle{N, B} + blockstyle::B end function BlockSparseArrayStyle{N}(blockstyle::AbstractArrayStyle{N}) where {N} - return BlockSparseArrayStyle{N,typeof(blockstyle)}(blockstyle) + return BlockSparseArrayStyle{N, typeof(blockstyle)}(blockstyle) end -function BlockSparseArrayStyle{N,B}() where {N,B<:AbstractArrayStyle{N}} - return BlockSparseArrayStyle{N,B}(B()) +function BlockSparseArrayStyle{N, B}() where {N, B <: AbstractArrayStyle{N}} + return BlockSparseArrayStyle{N, B}(B()) end BlockSparseArrayStyle{N}() where {N} = BlockSparseArrayStyle{N}(DefaultArrayStyle{N}()) BlockSparseArrayStyle(::Val{N}) where {N} = BlockSparseArrayStyle{N}() -BlockSparseArrayStyle{M}(::Val{N}) where {M,N} = BlockSparseArrayStyle{N}() -function BlockSparseArrayStyle{M,B}(::Val{N}) where {M,B<:AbstractArrayStyle{M},N} - return BlockSparseArrayStyle{N}(B(Val(N))) +BlockSparseArrayStyle{M}(::Val{N}) where {M, N} = BlockSparseArrayStyle{N}() +function BlockSparseArrayStyle{M, B}(::Val{N}) where {M, B <: AbstractArrayStyle{M}, N} + return BlockSparseArrayStyle{N}(B(Val(N))) end Broadcast.BroadcastStyle(a::BlockSparseArrayStyle, ::DefaultArrayStyle{0}) = a function Broadcast.BroadcastStyle( - ::BlockSparseArrayStyle{N}, a::DefaultArrayStyle -) where {N} - return BroadcastStyle(DefaultArrayStyle{N}(), a) + ::BlockSparseArrayStyle{N}, a::DefaultArrayStyle + ) where {N} + return BroadcastStyle(DefaultArrayStyle{N}(), a) end function Broadcast.BroadcastStyle( - ::BlockSparseArrayStyle{N}, ::Broadcast.Style{Tuple} -) where {N} - return DefaultArrayStyle{N}() + ::BlockSparseArrayStyle{N}, ::Broadcast.Style{Tuple} + ) where {N} + return DefaultArrayStyle{N}() end function Base.similar(bc::Broadcasted{<:BlockSparseArrayStyle}, elt::Type, ax) - # TODO: Make this more generic, base it off sure this handles GPU arrays properly. - m = Mapped(bc) - return similar(first(m.args), elt, ax) + # TODO: Make this more generic, base it off sure this handles GPU arrays properly. + m = Mapped(bc) + return similar(first(m.args), elt, ax) end # Catches cases like `dest .= value` or `dest .= value1 .+ value2`. # If the RHS is zero, this makes sure that the storage is emptied, # which is logic that is handled by `fill!`. function copyto_blocksparse!(dest::AbstractArray, bc::Broadcasted{<:AbstractArrayStyle{0}}) - # `[]` is used to unwrap zero-dimensional arrays. - bcf = Broadcast.flatten(bc) - value = @allowscalar bcf.f(map(arg -> arg[], bcf.args)...) - return @interface BlockSparseArrayInterface() fill!(dest, value) + # `[]` is used to unwrap zero-dimensional arrays. + bcf = Broadcast.flatten(bc) + value = @allowscalar bcf.f(map(arg -> arg[], bcf.args)...) + return @interface BlockSparseArrayInterface() fill!(dest, value) end # Broadcasting implementation # TODO: Delete this in favor of `DerivableInterfaces` version. function copyto_blocksparse!(dest::AbstractArray, bc::Broadcasted) - # convert to map - # flatten and only keep the AbstractArray arguments - m = Mapped(bc) - @interface interface(dest, bc) map!(m.f, dest, m.args...) - return dest + # convert to map + # flatten and only keep the AbstractArray arguments + m = Mapped(bc) + @interface interface(dest, bc) map!(m.f, dest, m.args...) + return dest end function Base.copyto!( - dest::AbstractArray{<:Any,N}, bc::Broadcasted{BlockSparseArrayStyle{N}} -) where {N} - copyto_blocksparse!(dest, bc) - return dest + dest::AbstractArray{<:Any, N}, bc::Broadcasted{BlockSparseArrayStyle{N}} + ) where {N} + copyto_blocksparse!(dest, bc) + return dest end diff --git a/src/blocksparsearrayinterface/cat.jl b/src/blocksparsearrayinterface/cat.jl index 2c0178a6..851d161d 100644 --- a/src/blocksparsearrayinterface/cat.jl +++ b/src/blocksparsearrayinterface/cat.jl @@ -2,13 +2,13 @@ using BlockArrays: blocks using DerivableInterfaces.Concatenate: Concatenated, cat! function Base.copyto!( - dest::AbstractArray, concat::Concatenated{<:BlockSparseArrayInterface} -) - # TODO: This assumes the destination blocking is commensurate with - # the blocking of the sources, for example because it was constructed - # based on the input arguments. Maybe check that explicitly. - # This should mostly just get called from `cat` anyway and not get - # called explicitly. - cat!(blocks(dest), blocks.(concat.args)...; dims=concat.dims) - return dest + dest::AbstractArray, concat::Concatenated{<:BlockSparseArrayInterface} + ) + # TODO: This assumes the destination blocking is commensurate with + # the blocking of the sources, for example because it was constructed + # based on the input arguments. Maybe check that explicitly. + # This should mostly just get called from `cat` anyway and not get + # called explicitly. + cat!(blocks(dest), blocks.(concat.args)...; dims = concat.dims) + return dest end diff --git a/src/blocksparsearrayinterface/getunstoredblock.jl b/src/blocksparsearrayinterface/getunstoredblock.jl index c5ae1dca..60cf6a88 100644 --- a/src/blocksparsearrayinterface/getunstoredblock.jl +++ b/src/blocksparsearrayinterface/getunstoredblock.jl @@ -2,44 +2,44 @@ using BlockArrays: Block using DerivableInterfaces: zero! struct ZeroBlocks{ - N,A<:AbstractArray{<:Any,N},ParentAxes<:Tuple{Vararg{AbstractUnitRange{<:Integer},N}} -} <: AbstractArray{A,N} - parentaxes::ParentAxes + N, A <: AbstractArray{<:Any, N}, ParentAxes <: Tuple{Vararg{AbstractUnitRange{<:Integer}, N}}, + } <: AbstractArray{A, N} + parentaxes::ParentAxes end -function ZeroBlocks{N,A}( - ax::Ax -) where {N,A<:AbstractArray{<:Any,N},Ax<:Tuple{Vararg{AbstractUnitRange{<:Integer},N}}} - return ZeroBlocks{N,A,Ax}(ax) +function ZeroBlocks{N, A}( + ax::Ax + ) where {N, A <: AbstractArray{<:Any, N}, Ax <: Tuple{Vararg{AbstractUnitRange{<:Integer}, N}}} + return ZeroBlocks{N, A, Ax}(ax) end Base.size(a::ZeroBlocks) = map(blocklength, a.parentaxes) -function Base.AbstractArray{A}(a::ZeroBlocks{N}) where {N,A} - return ZeroBlocks{N,A}(a.parentaxes) +function Base.AbstractArray{A}(a::ZeroBlocks{N}) where {N, A} + return ZeroBlocks{N, A}(a.parentaxes) end -@inline function Base.getindex(a::ZeroBlocks{N,A}, I::Vararg{Int,N}) where {N,A} - # TODO: Use `BlockArrays.eachblockaxes`. - ax = ntuple(N) do d - return eachblockaxis(a.parentaxes[d])[I[d]] - end - !isconcretetype(A) && return zero!(similar(Array{eltype(A),N}, ax)) - return zero!(similar(A, ax)) +@inline function Base.getindex(a::ZeroBlocks{N, A}, I::Vararg{Int, N}) where {N, A} + # TODO: Use `BlockArrays.eachblockaxes`. + ax = ntuple(N) do d + return eachblockaxis(a.parentaxes[d])[I[d]] + end + !isconcretetype(A) && return zero!(similar(Array{eltype(A), N}, ax)) + return zero!(similar(A, ax)) end # TODO: Use `Base.to_indices`. -@inline function Base.getindex(a::ZeroBlocks{N,A}, I::CartesianIndex{N}) where {N,A} - return a[Tuple(I)...] +@inline function Base.getindex(a::ZeroBlocks{N, A}, I::CartesianIndex{N}) where {N, A} + return a[Tuple(I)...] end # TODO: this is a hack and is also type-unstable using LinearAlgebra: Diagonal using TypeParameterAccessors: similartype -function Base.getindex(a::ZeroBlocks{2,A}, I::Vararg{Int,2}) where {V,A<:Diagonal{<:Any,V}} - ax = ntuple(2) do d - return only(axes(a.parentaxes[d][Block(I[d])])) - end - if allequal(I) - return Diagonal(zero!(similar(V, first(ax)))) - else - return zero!(similar(similartype(V, typeof(ax)), ax)) - end +function Base.getindex(a::ZeroBlocks{2, A}, I::Vararg{Int, 2}) where {V, A <: Diagonal{<:Any, V}} + ax = ntuple(2) do d + return only(axes(a.parentaxes[d][Block(I[d])])) + end + if allequal(I) + return Diagonal(zero!(similar(V, first(ax)))) + else + return zero!(similar(similartype(V, typeof(ax)), ax)) + end end diff --git a/src/blocksparsearrayinterface/linearalgebra.jl b/src/blocksparsearrayinterface/linearalgebra.jl index 702e7613..837b18b8 100644 --- a/src/blocksparsearrayinterface/linearalgebra.jl +++ b/src/blocksparsearrayinterface/linearalgebra.jl @@ -1,12 +1,12 @@ using LinearAlgebra: LinearAlgebra, mul! @interface ::AbstractBlockSparseArrayInterface function LinearAlgebra.mul!( - a_dest::AbstractMatrix, - a1::AbstractMatrix, - a2::AbstractMatrix, - α::Number=true, - β::Number=false, -) - mul!(blocks(a_dest), blocks(a1), blocks(a2), α, β) - return a_dest + a_dest::AbstractMatrix, + a1::AbstractMatrix, + a2::AbstractMatrix, + α::Number = true, + β::Number = false, + ) + mul!(blocks(a_dest), blocks(a1), blocks(a2), α, β) + return a_dest end diff --git a/src/blocksparsearrayinterface/map.jl b/src/blocksparsearrayinterface/map.jl index 03b5ec08..c243443d 100644 --- a/src/blocksparsearrayinterface/map.jl +++ b/src/blocksparsearrayinterface/map.jl @@ -4,64 +4,64 @@ using GPUArraysCore: @allowscalar # Check if the block structures are the same. function same_block_structure(as::AbstractArray...) - isempty(as) && return true - return all( - ntuple(ndims(first(as))) do dim - ax = map(Base.Fix2(axes, dim), as) - return blockisequal(ax...) - end, - ) + isempty(as) && return true + return all( + ntuple(ndims(first(as))) do dim + ax = map(Base.Fix2(axes, dim), as) + return blockisequal(ax...) + end, + ) end # Find the common stored blocks, assuming the block structures are the same. function union_eachblockstoredindex(as::AbstractArray...) - return ∪(map(eachblockstoredindex, as)...) + 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)...] +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)...) +function viewblock_stored(a::AbstractArray{<:Any, N}, I::Block{N}) where {N} + return viewblock_stored(a, Tuple(I)...) end using FillArrays: Zeros, fillsimilar # 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 fillsimilar(Zeros{eltype(a)}(block_ax), block_ax) - end +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 fillsimilar(Zeros{eltype(a)}(block_ax), block_ax) + end end -function viewblock_or_zeros(a::AbstractArray{<:Any,N}, I::Block{N}) where {N} - return viewblock_or_zeros(a, Tuple(I)...) +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 + 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)) - Is = if f_preserves_zeros - union_eachblockstoredindex(a_dest, a_srcs...) - else - BlockRange(a_dest) - end - for I in Is - map_block!(f, a_dest, I, a_srcs...) - end - return a_dest + # TODO: This assumes element types are numbers, generalize this logic. + f_preserves_zeros = f(zero.(eltype.(a_srcs))...) == zero(eltype(a_dest)) + Is = if f_preserves_zeros + union_eachblockstoredindex(a_dest, a_srcs...) + else + BlockRange(a_dest) + end + for I in Is + map_block!(f, a_dest, I, a_srcs...) + end + return a_dest end # TODO: Rewrite this so that it takes the blocking structure @@ -70,79 +70,79 @@ end # `reblock` is a partial solution to that, but a bit ad-hoc. ## TODO: Make this an `@interface AbstractBlockSparseArrayInterface` function. @interface interface::AbstractBlockSparseArrayInterface function Base.map!( - f, a_dest::AbstractArray, a_srcs::AbstractArray... -) - if isempty(a_srcs) - error("Can't call `map!` with zero source terms.") - end - if iszero(ndims(a_dest)) - @interface interface map_zero_dim!(f, a_dest, a_srcs...) - return a_dest - end - if same_block_structure(a_dest, a_srcs...) - map_blockwise!(f, a_dest, a_srcs...) - return a_dest - end - # TODO: This assumes element types are numbers, generalize this logic. - f_preserves_zeros = f(zero.(eltype.(a_srcs))...) == zero(eltype(a_dest)) - a_dest, a_srcs = reblock(a_dest), reblock.(a_srcs) - for I in union_stored_blocked_cartesianindices(a_dest, a_srcs...) - BI_dest = blockindexrange(a_dest, I) - BI_srcs = map(a_src -> blockindexrange(a_src, I), a_srcs) - # TODO: Investigate why this doesn't work: - # block_dest = @view a_dest[_block(BI_dest)] - block_dest = blocks_maybe_single(a_dest)[Int.(Tuple(_block(BI_dest)))...] - # TODO: Investigate why this doesn't work: - # block_srcs = ntuple(i -> @view(a_srcs[i][_block(BI_srcs[i])]), length(a_srcs)) - block_srcs = ntuple(length(a_srcs)) do i - return blocks_maybe_single(a_srcs[i])[Int.(Tuple(_block(BI_srcs[i])))...] + f, a_dest::AbstractArray, a_srcs::AbstractArray... + ) + if isempty(a_srcs) + error("Can't call `map!` with zero source terms.") end - subblock_dest = @view block_dest[BI_dest.indices...] - subblock_srcs = ntuple(i -> @view(block_srcs[i][BI_srcs[i].indices...]), length(a_srcs)) - I_dest = CartesianIndex(Int.(Tuple(_block(BI_dest)))) - # If the function preserves zero values and all of the source blocks are zero, - # the output block will be zero. In that case, if the block isn't stored yet, - # don't do anything. - if f_preserves_zeros && all(iszero, subblock_srcs) && !isstored(blocks(a_dest), I_dest) - continue + if iszero(ndims(a_dest)) + @interface interface map_zero_dim!(f, a_dest, a_srcs...) + return a_dest end - # TODO: Use `map!!` to handle immutable blocks. - map!(f, subblock_dest, subblock_srcs...) - # Replace the entire block, handles initializing new blocks - # or if blocks are immutable. - blocks(a_dest)[I_dest] = block_dest - end - return a_dest + if same_block_structure(a_dest, a_srcs...) + map_blockwise!(f, a_dest, a_srcs...) + return a_dest + end + # TODO: This assumes element types are numbers, generalize this logic. + f_preserves_zeros = f(zero.(eltype.(a_srcs))...) == zero(eltype(a_dest)) + a_dest, a_srcs = reblock(a_dest), reblock.(a_srcs) + for I in union_stored_blocked_cartesianindices(a_dest, a_srcs...) + BI_dest = blockindexrange(a_dest, I) + BI_srcs = map(a_src -> blockindexrange(a_src, I), a_srcs) + # TODO: Investigate why this doesn't work: + # block_dest = @view a_dest[_block(BI_dest)] + block_dest = blocks_maybe_single(a_dest)[Int.(Tuple(_block(BI_dest)))...] + # TODO: Investigate why this doesn't work: + # block_srcs = ntuple(i -> @view(a_srcs[i][_block(BI_srcs[i])]), length(a_srcs)) + block_srcs = ntuple(length(a_srcs)) do i + return blocks_maybe_single(a_srcs[i])[Int.(Tuple(_block(BI_srcs[i])))...] + end + subblock_dest = @view block_dest[BI_dest.indices...] + subblock_srcs = ntuple(i -> @view(block_srcs[i][BI_srcs[i].indices...]), length(a_srcs)) + I_dest = CartesianIndex(Int.(Tuple(_block(BI_dest)))) + # If the function preserves zero values and all of the source blocks are zero, + # the output block will be zero. In that case, if the block isn't stored yet, + # don't do anything. + if f_preserves_zeros && all(iszero, subblock_srcs) && !isstored(blocks(a_dest), I_dest) + continue + end + # TODO: Use `map!!` to handle immutable blocks. + map!(f, subblock_dest, subblock_srcs...) + # Replace the entire block, handles initializing new blocks + # or if blocks are immutable. + blocks(a_dest)[I_dest] = block_dest + end + return a_dest end @interface ::AbstractBlockSparseArrayInterface function Base.mapreduce( - f, op, as::AbstractArray...; kwargs... -) - # TODO: Define an `init` value based on the element type. - return @interface interface(blocks.(as)...) mapreduce( - block -> mapreduce(f, op, block), op, blocks.(as)...; kwargs... - ) + f, op, as::AbstractArray...; kwargs... + ) + # TODO: Define an `init` value based on the element type. + return @interface interface(blocks.(as)...) mapreduce( + block -> mapreduce(f, op, block), op, blocks.(as)...; kwargs... + ) end @interface ::AbstractBlockSparseArrayInterface function Base.iszero(a::AbstractArray) - # TODO: Just call `iszero(blocks(a))`? - return @interface interface(blocks(a)) iszero(blocks(a)) + # TODO: Just call `iszero(blocks(a))`? + return @interface interface(blocks(a)) iszero(blocks(a)) end @interface ::AbstractBlockSparseArrayInterface function Base.isreal(a::AbstractArray) - # TODO: Just call `isreal(blocks(a))`? - return @interface interface(blocks(a)) isreal(blocks(a)) + # TODO: Just call `isreal(blocks(a))`? + return @interface interface(blocks(a)) isreal(blocks(a)) end # Helper functions for block sparse map. # Returns `Vector{<:CartesianIndices}` function union_stored_blocked_cartesianindices(as::Vararg{AbstractArray}) - combined_axes = combine_axes(axes.(as)...) - stored_blocked_cartesianindices_as = map(as) do a - return blocked_cartesianindices(axes(a), combined_axes, eachblockstoredindex(a)) - end - return ∪(stored_blocked_cartesianindices_as...) + combined_axes = combine_axes(axes.(as)...) + stored_blocked_cartesianindices_as = map(as) do a + return blocked_cartesianindices(axes(a), combined_axes, eachblockstoredindex(a)) + end + return ∪(stored_blocked_cartesianindices_as...) end # This is used by `map` to get the output axes. @@ -155,10 +155,10 @@ reblock(a) = a function map_zero_dim! end @interface ::AbstractArrayInterface function map_zero_dim!( - f, a_dest::AbstractArray, a_srcs::AbstractArray... -) - @allowscalar a_dest[] = f.(map(a_src -> a_src[], a_srcs)...) - return a_dest + f, a_dest::AbstractArray, a_srcs::AbstractArray... + ) + @allowscalar a_dest[] = f.(map(a_src -> a_src[], a_srcs)...) + return a_dest end # TODO: Do we need this function or can we just use `map`? @@ -166,20 +166,20 @@ end # specify the function preserves zeros, i.e. # `map(f, a; preserves_zeros=true)` or `@preserves_zeros map(f, a)`. 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)) - 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) - a′ = BlockSparseArray{eltype(blocktype_a′),ndims(a),blocktype_a′}(undef, axes(a)) - for (B, b) in zip(block_stored_indices, stored_blocks) - a′[B] = b - end - return a′ + 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)) + 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) + a′ = BlockSparseArray{eltype(blocktype_a′), ndims(a), blocktype_a′}(undef, axes(a)) + for (B, b) in zip(block_stored_indices, stored_blocks) + a′[B] = b + end + return a′ end diff --git a/src/blocksparsearrayinterface/views.jl b/src/blocksparsearrayinterface/views.jl index 56d17b02..7445684e 100644 --- a/src/blocksparsearrayinterface/views.jl +++ b/src/blocksparsearrayinterface/views.jl @@ -1,3 +1,3 @@ @interface ::AbstractBlockSparseArrayInterface function Base.view(a, I...) - return Base.invoke(view, Tuple{AbstractArray,Vararg{Any}}, a, I...) + return Base.invoke(view, Tuple{AbstractArray, Vararg{Any}}, a, I...) end diff --git a/src/factorizations/eig.jl b/src/factorizations/eig.jl index afc8be2f..1f51ced4 100644 --- a/src/factorizations/eig.jl +++ b/src/factorizations/eig.jl @@ -6,122 +6,122 @@ using MatrixAlgebraKit: default_eig_algorithm, eig_full!, eig_vals! using MatrixAlgebraKit: default_eigh_algorithm, eigh_full!, eigh_vals! for f in [:default_eig_algorithm, :default_eigh_algorithm] - @eval begin - function MatrixAlgebraKit.$f(::Type{<:AbstractBlockSparseMatrix}; kwargs...) - return BlockDiagonalAlgorithm() do block - return $f(block; kwargs...) - end + @eval begin + function MatrixAlgebraKit.$f(::Type{<:AbstractBlockSparseMatrix}; kwargs...) + return BlockDiagonalAlgorithm() do block + return $f(block; kwargs...) + end + end end - end end function output_type(::typeof(eig_full!), A::Type{<:AbstractMatrix{T}}) where {T} - DV = Base.promote_op(eig_full!, A) - return if isconcretetype(DV) - DV - else - Tuple{AbstractMatrix{complex(T)},AbstractMatrix{complex(T)}} - end + DV = Base.promote_op(eig_full!, A) + return if isconcretetype(DV) + DV + else + Tuple{AbstractMatrix{complex(T)}, AbstractMatrix{complex(T)}} + end end function output_type(::typeof(eigh_full!), A::Type{<:AbstractMatrix{T}}) where {T} - DV = Base.promote_op(eigh_full!, A) - return isconcretetype(DV) ? DV : Tuple{AbstractMatrix{real(T)},AbstractMatrix{T}} + DV = Base.promote_op(eigh_full!, A) + return isconcretetype(DV) ? DV : Tuple{AbstractMatrix{real(T)}, AbstractMatrix{T}} end function MatrixAlgebraKit.check_input( - ::typeof(eig_full!), A::AbstractBlockSparseMatrix, (D, V), ::BlockDiagonalAlgorithm -) - @assert isa(D, AbstractBlockSparseMatrix) && isa(V, AbstractBlockSparseMatrix) - @assert eltype(V) === eltype(D) === complex(eltype(A)) - @assert axes(A, 1) == axes(A, 2) - @assert axes(A) == axes(D) == axes(V) - @assert isblockdiagonal(A) - return nothing + ::typeof(eig_full!), A::AbstractBlockSparseMatrix, (D, V), ::BlockDiagonalAlgorithm + ) + @assert isa(D, AbstractBlockSparseMatrix) && isa(V, AbstractBlockSparseMatrix) + @assert eltype(V) === eltype(D) === complex(eltype(A)) + @assert axes(A, 1) == axes(A, 2) + @assert axes(A) == axes(D) == axes(V) + @assert isblockdiagonal(A) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(eigh_full!), A::AbstractBlockSparseMatrix, (D, V), ::BlockDiagonalAlgorithm -) - @assert isa(D, AbstractBlockSparseMatrix) && isa(V, AbstractBlockSparseMatrix) - @assert eltype(V) === eltype(A) - @assert eltype(D) === real(eltype(A)) - @assert axes(A, 1) == axes(A, 2) - @assert axes(A) == axes(D) == axes(V) - @assert isblockdiagonal(A) - return nothing + ::typeof(eigh_full!), A::AbstractBlockSparseMatrix, (D, V), ::BlockDiagonalAlgorithm + ) + @assert isa(D, AbstractBlockSparseMatrix) && isa(V, AbstractBlockSparseMatrix) + @assert eltype(V) === eltype(A) + @assert eltype(D) === real(eltype(A)) + @assert axes(A, 1) == axes(A, 2) + @assert axes(A) == axes(D) == axes(V) + @assert isblockdiagonal(A) + return nothing end for f in [:eig_full!, :eigh_full!] - @eval begin - function MatrixAlgebraKit.initialize_output( - ::typeof($f), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm - ) - Td, Tv = fieldtypes(output_type($f, blocktype(A))) - D = similar(A, BlockType(Td)) - V = similar(A, BlockType(Tv)) - return (D, V) - end - function MatrixAlgebraKit.$f( - A::AbstractBlockSparseMatrix, (D, V), alg::BlockDiagonalAlgorithm - ) - MatrixAlgebraKit.check_input($f, A, (D, V), alg) + @eval begin + function MatrixAlgebraKit.initialize_output( + ::typeof($f), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm + ) + Td, Tv = fieldtypes(output_type($f, blocktype(A))) + D = similar(A, BlockType(Td)) + V = similar(A, BlockType(Tv)) + return (D, V) + end + function MatrixAlgebraKit.$f( + A::AbstractBlockSparseMatrix, (D, V), alg::BlockDiagonalAlgorithm + ) + MatrixAlgebraKit.check_input($f, A, (D, V), alg) - # do decomposition on each block - for bI in blockdiagindices(A) - if isstored(A, bI) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - bD, bV = $f(block, block_alg) - D[bI] = bD - V[bI] = bV - else - # TODO: this should be `V[bI] = LinearAlgebra.I` - copyto!(@view!(V[bI]), LinearAlgebra.I) + # do decomposition on each block + for bI in blockdiagindices(A) + if isstored(A, bI) + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bD, bV = $f(block, block_alg) + D[bI] = bD + V[bI] = bV + else + # TODO: this should be `V[bI] = LinearAlgebra.I` + copyto!(@view!(V[bI]), LinearAlgebra.I) + end + end + return (D, V) end - end - return (D, V) end - end end function output_type(f::typeof(eig_vals!), A::Type{<:AbstractMatrix{T}}) where {T} - D = Base.promote_op(f, A) - !isconcretetype(D) && return AbstractVector{complex(T)} - return D + D = Base.promote_op(f, A) + !isconcretetype(D) && return AbstractVector{complex(T)} + return D end function output_type(f::typeof(eigh_vals!), A::Type{<:AbstractMatrix{T}}) where {T} - D = Base.promote_op(f, A) - !isconcretetype(D) && return AbstractVector{real(T)} - return D + D = Base.promote_op(f, A) + !isconcretetype(D) && return AbstractVector{real(T)} + return D end for f in [:eig_vals!, :eigh_vals!] - @eval begin - function MatrixAlgebraKit.initialize_output( - ::typeof($f), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm - ) - T = output_type($f, blocktype(A)) - return similar(A, BlockType(T), axes(A, 1)) - end - function MatrixAlgebraKit.check_input( - ::typeof($f), A::AbstractBlockSparseMatrix, D, ::BlockDiagonalAlgorithm - ) - @assert isa(D, AbstractBlockSparseVector) - @assert eltype(D) === $(f == :eig_vals! ? complex : real)(eltype(A)) - @assert axes(A, 1) == axes(A, 2) - @assert (axes(A, 1),) == axes(D) - @assert isblockdiagonal(A) - return nothing - end + @eval begin + function MatrixAlgebraKit.initialize_output( + ::typeof($f), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm + ) + T = output_type($f, blocktype(A)) + return similar(A, BlockType(T), axes(A, 1)) + end + function MatrixAlgebraKit.check_input( + ::typeof($f), A::AbstractBlockSparseMatrix, D, ::BlockDiagonalAlgorithm + ) + @assert isa(D, AbstractBlockSparseVector) + @assert eltype(D) === $(f == :eig_vals! ? complex : real)(eltype(A)) + @assert axes(A, 1) == axes(A, 2) + @assert (axes(A, 1),) == axes(D) + @assert isblockdiagonal(A) + return nothing + end - function MatrixAlgebraKit.$f( - A::AbstractBlockSparseMatrix, D, alg::BlockDiagonalAlgorithm - ) - MatrixAlgebraKit.check_input($f, A, D, alg) - for I in eachblockstoredindex(A) - block = @view!(A[I]) - D[Tuple(I)[1]] = $f(block, block_algorithm(alg, block)) - end - return D + function MatrixAlgebraKit.$f( + A::AbstractBlockSparseMatrix, D, alg::BlockDiagonalAlgorithm + ) + MatrixAlgebraKit.check_input($f, A, D, alg) + for I in eachblockstoredindex(A) + block = @view!(A[I]) + D[Tuple(I)[1]] = $f(block, block_algorithm(alg, block)) + end + return D + end end - end end diff --git a/src/factorizations/lq.jl b/src/factorizations/lq.jl index 412fbe16..8ae937f6 100644 --- a/src/factorizations/lq.jl +++ b/src/factorizations/lq.jl @@ -1,147 +1,147 @@ using MatrixAlgebraKit: MatrixAlgebraKit, default_lq_algorithm, lq_compact!, lq_full! function MatrixAlgebraKit.default_lq_algorithm( - A::Type{<:AbstractBlockSparseMatrix}; kwargs... -) - return BlockPermutedDiagonalAlgorithm() do block - return default_lq_algorithm(block; kwargs...) - end + A::Type{<:AbstractBlockSparseMatrix}; kwargs... + ) + return BlockPermutedDiagonalAlgorithm() do block + return default_lq_algorithm(block; kwargs...) + end end function output_type( - f::Union{typeof(lq_compact!),typeof(lq_full!)}, A::Type{<:AbstractMatrix{T}} -) where {T} - LQ = Base.promote_op(f, A) - return isconcretetype(LQ) ? LQ : Tuple{AbstractMatrix{T},AbstractMatrix{T}} + f::Union{typeof(lq_compact!), typeof(lq_full!)}, A::Type{<:AbstractMatrix{T}} + ) where {T} + LQ = Base.promote_op(f, A) + return isconcretetype(LQ) ? LQ : Tuple{AbstractMatrix{T}, AbstractMatrix{T}} end function MatrixAlgebraKit.initialize_output( - ::typeof(lq_compact!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm -) - return nothing + ::typeof(lq_compact!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm + ) + return nothing end function MatrixAlgebraKit.initialize_output( - ::typeof(lq_compact!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm -) - brows = eachblockaxis(axes(A, 1)) - bcols = eachblockaxis(axes(A, 2)) - # using the property that zip stops as soon as one of the iterators is exhausted - l_axes = map(splat(infimum), zip(brows, bcols)) - l_axis = mortar_axis(l_axes) + ::typeof(lq_compact!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm + ) + brows = eachblockaxis(axes(A, 1)) + bcols = eachblockaxis(axes(A, 2)) + # using the property that zip stops as soon as one of the iterators is exhausted + l_axes = map(splat(infimum), zip(brows, bcols)) + l_axis = mortar_axis(l_axes) - BL, BQ = fieldtypes(output_type(lq_compact!, blocktype(A))) - L = similar(A, BlockType(BL), (axes(A, 1), l_axis)) - Q = similar(A, BlockType(BQ), (l_axis, axes(A, 2))) + BL, BQ = fieldtypes(output_type(lq_compact!, blocktype(A))) + L = similar(A, BlockType(BL), (axes(A, 1), l_axis)) + Q = similar(A, BlockType(BQ), (l_axis, axes(A, 2))) - return L, Q + return L, Q end function MatrixAlgebraKit.initialize_output( - ::typeof(lq_full!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm -) - return nothing + ::typeof(lq_full!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm + ) + return nothing end function MatrixAlgebraKit.initialize_output( - ::typeof(lq_full!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm -) - BL, BQ = fieldtypes(output_type(lq_full!, blocktype(A))) - L = similar(A, BlockType(BL), (axes(A, 1), axes(A, 2))) - Q = similar(A, BlockType(BQ), (axes(A, 2), axes(A, 2))) - return L, Q + ::typeof(lq_full!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm + ) + BL, BQ = fieldtypes(output_type(lq_full!, blocktype(A))) + L = similar(A, BlockType(BL), (axes(A, 1), axes(A, 2))) + Q = similar(A, BlockType(BQ), (axes(A, 2), axes(A, 2))) + return L, Q end function MatrixAlgebraKit.check_input( - ::typeof(lq_compact!), A::AbstractBlockSparseMatrix, LQ, ::BlockPermutedDiagonalAlgorithm -) - @assert isblockpermuteddiagonal(A) - return nothing + ::typeof(lq_compact!), A::AbstractBlockSparseMatrix, LQ, ::BlockPermutedDiagonalAlgorithm + ) + @assert isblockpermuteddiagonal(A) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(lq_compact!), A::AbstractBlockSparseMatrix, (L, Q), ::BlockDiagonalAlgorithm -) - @assert isa(L, AbstractBlockSparseMatrix) && isa(Q, AbstractBlockSparseMatrix) - @assert eltype(A) == eltype(L) == eltype(Q) - @assert axes(A, 1) == axes(L, 1) && axes(A, 2) == axes(Q, 2) - @assert axes(L, 2) == axes(Q, 1) - @assert isblockdiagonal(A) - return nothing + ::typeof(lq_compact!), A::AbstractBlockSparseMatrix, (L, Q), ::BlockDiagonalAlgorithm + ) + @assert isa(L, AbstractBlockSparseMatrix) && isa(Q, AbstractBlockSparseMatrix) + @assert eltype(A) == eltype(L) == eltype(Q) + @assert axes(A, 1) == axes(L, 1) && axes(A, 2) == axes(Q, 2) + @assert axes(L, 2) == axes(Q, 1) + @assert isblockdiagonal(A) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(lq_full!), A::AbstractBlockSparseMatrix, LQ, ::BlockPermutedDiagonalAlgorithm -) - @assert isblockpermuteddiagonal(A) - return nothing + ::typeof(lq_full!), A::AbstractBlockSparseMatrix, LQ, ::BlockPermutedDiagonalAlgorithm + ) + @assert isblockpermuteddiagonal(A) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(lq_full!), A::AbstractBlockSparseMatrix, (L, Q), ::BlockDiagonalAlgorithm -) - @assert isa(L, AbstractBlockSparseMatrix) && isa(Q, AbstractBlockSparseMatrix) - @assert eltype(A) == eltype(L) == eltype(Q) - @assert axes(A, 1) == axes(L, 1) && axes(A, 2) == axes(Q, 2) - @assert axes(L, 2) == axes(Q, 1) - @assert isblockdiagonal(A) - return nothing + ::typeof(lq_full!), A::AbstractBlockSparseMatrix, (L, Q), ::BlockDiagonalAlgorithm + ) + @assert isa(L, AbstractBlockSparseMatrix) && isa(Q, AbstractBlockSparseMatrix) + @assert eltype(A) == eltype(L) == eltype(Q) + @assert axes(A, 1) == axes(L, 1) && axes(A, 2) == axes(Q, 2) + @assert axes(L, 2) == axes(Q, 1) + @assert isblockdiagonal(A) + return nothing end function MatrixAlgebraKit.lq_compact!( - A::AbstractBlockSparseMatrix, LQ, alg::BlockPermutedDiagonalAlgorithm -) - MatrixAlgebraKit.check_input(lq_compact!, A, LQ, alg) - Ad, (invrowperm, invcolperm) = blockdiagonalize(A) - Ld, Qd = lq_compact!(Ad, BlockDiagonalAlgorithm(alg)) - L = transform_rows(Ld, invrowperm) - Q = transform_cols(Qd, invcolperm) - return L, Q + A::AbstractBlockSparseMatrix, LQ, alg::BlockPermutedDiagonalAlgorithm + ) + MatrixAlgebraKit.check_input(lq_compact!, A, LQ, alg) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + Ld, Qd = lq_compact!(Ad, BlockDiagonalAlgorithm(alg)) + L = transform_rows(Ld, invrowperm) + Q = transform_cols(Qd, invcolperm) + return L, Q end function MatrixAlgebraKit.lq_compact!( - A::AbstractBlockSparseMatrix, (L, Q), alg::BlockDiagonalAlgorithm -) - MatrixAlgebraKit.check_input(lq_compact!, A, (L, Q), alg) + A::AbstractBlockSparseMatrix, (L, Q), alg::BlockDiagonalAlgorithm + ) + MatrixAlgebraKit.check_input(lq_compact!, A, (L, Q), alg) - # do decomposition on each block - for bI in blockdiagindices(A) - if isstored(A, bI) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - bL, bQ = lq_compact!(block, block_alg) - L[bI] = bL - Q[bI] = bQ - else - # TODO: this should be `Q[bI] = LinearAlgebra.I` - copyto!(@view!(Q[bI]), LinearAlgebra.I) + # do decomposition on each block + for bI in blockdiagindices(A) + if isstored(A, bI) + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bL, bQ = lq_compact!(block, block_alg) + L[bI] = bL + Q[bI] = bQ + else + # TODO: this should be `Q[bI] = LinearAlgebra.I` + copyto!(@view!(Q[bI]), LinearAlgebra.I) + end end - end - return L, Q + return L, Q end function MatrixAlgebraKit.lq_full!( - A::AbstractBlockSparseMatrix, LQ, alg::BlockPermutedDiagonalAlgorithm -) - MatrixAlgebraKit.check_input(lq_full!, A, LQ, alg) - Ad, (invrowperm, invcolperm) = blockdiagonalize(A) - Ld, Qd = lq_full!(Ad, BlockDiagonalAlgorithm(alg)) - L = transform_rows(Ld, invrowperm) - Q = transform_cols(Qd, invcolperm) - return L, Q + A::AbstractBlockSparseMatrix, LQ, alg::BlockPermutedDiagonalAlgorithm + ) + MatrixAlgebraKit.check_input(lq_full!, A, LQ, alg) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + Ld, Qd = lq_full!(Ad, BlockDiagonalAlgorithm(alg)) + L = transform_rows(Ld, invrowperm) + Q = transform_cols(Qd, invcolperm) + return L, Q end function MatrixAlgebraKit.lq_full!( - A::AbstractBlockSparseMatrix, (L, Q), alg::BlockDiagonalAlgorithm -) - MatrixAlgebraKit.check_input(lq_full!, A, (L, Q), alg) + A::AbstractBlockSparseMatrix, (L, Q), alg::BlockDiagonalAlgorithm + ) + MatrixAlgebraKit.check_input(lq_full!, A, (L, Q), alg) - for bI in blockdiagindices(A) - if isstored(A, bI) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - bL, bQ = lq_full!(block, block_alg) - L[bI] = bL - Q[bI] = bQ - else - # TODO: this should be `Q[bI] = LinearAlgebra.I` - copyto!(@view!(Q[bI]), LinearAlgebra.I) + for bI in blockdiagindices(A) + if isstored(A, bI) + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bL, bQ = lq_full!(block, block_alg) + L[bI] = bL + Q[bI] = bQ + else + # TODO: this should be `Q[bI] = LinearAlgebra.I` + copyto!(@view!(Q[bI]), LinearAlgebra.I) + end end - end - return L, Q + return L, Q end diff --git a/src/factorizations/orthnull.jl b/src/factorizations/orthnull.jl index dc8b1d86..f35921df 100644 --- a/src/factorizations/orthnull.jl +++ b/src/factorizations/orthnull.jl @@ -1,177 +1,177 @@ using MatrixAlgebraKit: - MatrixAlgebraKit, - default_svd_algorithm, - left_null!, - left_null_svd!, - left_orth!, - left_polar!, - lq_compact!, - null_truncation_strategy, - qr_compact!, - right_null!, - right_null_svd!, - right_orth!, - right_polar!, - select_algorithm, - svd_compact! + MatrixAlgebraKit, + default_svd_algorithm, + left_null!, + left_null_svd!, + left_orth!, + left_polar!, + lq_compact!, + null_truncation_strategy, + qr_compact!, + right_null!, + right_null_svd!, + right_orth!, + right_polar!, + select_algorithm, + svd_compact! function MatrixAlgebraKit.initialize_output( - ::typeof(left_orth!), A::AbstractBlockSparseMatrix -) - return nothing + ::typeof(left_orth!), A::AbstractBlockSparseMatrix + ) + return nothing end function MatrixAlgebraKit.check_input(::typeof(left_orth!), A::AbstractBlockSparseMatrix, F) - !isnothing(F) && throw( - ArgumentError( - "`left_orth!` on block sparse matrices does not support specifying the output" - ), - ) - return nothing + !isnothing(F) && throw( + ArgumentError( + "`left_orth!` on block sparse matrices does not support specifying the output" + ), + ) + return nothing end function MatrixAlgebraKit.left_orth_qr!(A::AbstractBlockSparseMatrix, F, alg) - !isnothing(F) && throw( - ArgumentError( - "`left_orth!` on block sparse matrices does not support specifying the output" - ), - ) - alg′ = select_algorithm(qr_compact!, A, alg) - return qr_compact!(A, alg′) + !isnothing(F) && throw( + ArgumentError( + "`left_orth!` on block sparse matrices does not support specifying the output" + ), + ) + alg′ = select_algorithm(qr_compact!, A, alg) + return qr_compact!(A, alg′) end function MatrixAlgebraKit.left_orth_polar!(A::AbstractBlockSparseMatrix, F, alg) - !isnothing(F) && throw( - ArgumentError( - "`left_orth!` on block sparse matrices does not support specifying the output" - ), - ) - alg′ = select_algorithm(left_polar!, A, alg) - return left_polar!(A, alg′) + !isnothing(F) && throw( + ArgumentError( + "`left_orth!` on block sparse matrices does not support specifying the output" + ), + ) + alg′ = select_algorithm(left_polar!, A, alg) + return left_polar!(A, alg′) end function MatrixAlgebraKit.left_orth_svd!( - A::AbstractBlockSparseMatrix, F, alg, trunc::Nothing=nothing -) - !isnothing(F) && throw( - ArgumentError( - "`left_orth!` on block sparse matrices does not support specifying the output" - ), - ) - alg′ = select_algorithm(svd_compact!, A, alg) - U, S, Vᴴ = svd_compact!(A, alg′) - return U, S * Vᴴ + A::AbstractBlockSparseMatrix, F, alg, trunc::Nothing = nothing + ) + !isnothing(F) && throw( + ArgumentError( + "`left_orth!` on block sparse matrices does not support specifying the output" + ), + ) + alg′ = select_algorithm(svd_compact!, A, alg) + U, S, Vᴴ = svd_compact!(A, alg′) + return U, S * Vᴴ end function MatrixAlgebraKit.left_orth_svd!(A::AbstractBlockSparseMatrix, F, alg, trunc) - !isnothing(F) && throw( - ArgumentError( - "`left_orth!` on block sparse matrices does not support specifying the output" - ), - ) - alg′ = select_algorithm(svd_compact!, A, alg) - alg_trunc = select_algorithm(svd_trunc!, A, alg′; trunc) - U, S, Vᴴ = svd_trunc!(A, alg_trunc) - return U, S * Vᴴ + !isnothing(F) && throw( + ArgumentError( + "`left_orth!` on block sparse matrices does not support specifying the output" + ), + ) + alg′ = select_algorithm(svd_compact!, A, alg) + alg_trunc = select_algorithm(svd_trunc!, A, alg′; trunc) + U, S, Vᴴ = svd_trunc!(A, alg_trunc) + return U, S * Vᴴ end function MatrixAlgebraKit.initialize_output( - ::typeof(right_orth!), A::AbstractBlockSparseMatrix -) - return nothing + ::typeof(right_orth!), A::AbstractBlockSparseMatrix + ) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(right_orth!), A::AbstractBlockSparseMatrix, F::Nothing -) - !isnothing(F) && throw( - ArgumentError( - "`right_orth!` on block sparse matrices does not support specifying the output" - ), - ) - return nothing + ::typeof(right_orth!), A::AbstractBlockSparseMatrix, F::Nothing + ) + !isnothing(F) && throw( + ArgumentError( + "`right_orth!` on block sparse matrices does not support specifying the output" + ), + ) + return nothing end function MatrixAlgebraKit.right_orth_lq!(A::AbstractBlockSparseMatrix, F, alg) - !isnothing(F) && throw( - ArgumentError( - "`right_orth!` on block sparse matrices does not support specifying the output" - ), - ) - alg′ = select_algorithm(lq_compact!, A, alg) - return lq_compact!(A, alg′) + !isnothing(F) && throw( + ArgumentError( + "`right_orth!` on block sparse matrices does not support specifying the output" + ), + ) + alg′ = select_algorithm(lq_compact!, A, alg) + return lq_compact!(A, alg′) end function MatrixAlgebraKit.right_orth_polar!(A::AbstractBlockSparseMatrix, F, alg) - !isnothing(F) && throw( - ArgumentError( - "`right_orth!` on block sparse matrices does not support specifying the output" - ), - ) - alg′ = select_algorithm(right_polar!, A, alg) - return right_polar!(A, alg′) + !isnothing(F) && throw( + ArgumentError( + "`right_orth!` on block sparse matrices does not support specifying the output" + ), + ) + alg′ = select_algorithm(right_polar!, A, alg) + return right_polar!(A, alg′) end function MatrixAlgebraKit.right_orth_svd!( - A::AbstractBlockSparseMatrix, F, alg, trunc::Nothing=nothing -) - !isnothing(F) && throw( - ArgumentError( - "`right_orth!` on block sparse matrices does not support specifying the output" - ), - ) - alg′ = select_algorithm(svd_compact!, A, alg) - U, S, Vᴴ = svd_compact!(A, alg′) - return U * S, Vᴴ + A::AbstractBlockSparseMatrix, F, alg, trunc::Nothing = nothing + ) + !isnothing(F) && throw( + ArgumentError( + "`right_orth!` on block sparse matrices does not support specifying the output" + ), + ) + alg′ = select_algorithm(svd_compact!, A, alg) + U, S, Vᴴ = svd_compact!(A, alg′) + return U * S, Vᴴ end function MatrixAlgebraKit.right_orth_svd!(A::AbstractBlockSparseMatrix, F, alg, trunc) - !isnothing(F) && throw( - ArgumentError( - "`right_orth!` on block sparse matrices does not support specifying the output" - ), - ) - alg′ = select_algorithm(svd_compact!, A, alg) - alg_trunc = select_algorithm(svd_trunc!, A, alg′; trunc) - U, S, Vᴴ = svd_trunc!(A, alg_trunc) - return U * S, Vᴴ + !isnothing(F) && throw( + ArgumentError( + "`right_orth!` on block sparse matrices does not support specifying the output" + ), + ) + alg′ = select_algorithm(svd_compact!, A, alg) + alg_trunc = select_algorithm(svd_trunc!, A, alg′; trunc) + U, S, Vᴴ = svd_trunc!(A, alg_trunc) + return U * S, Vᴴ end function MatrixAlgebraKit.initialize_output( - ::typeof(left_null!), A::AbstractBlockSparseMatrix -) - return nothing + ::typeof(left_null!), A::AbstractBlockSparseMatrix + ) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(left_null!), A::AbstractBlockSparseMatrix, N::Nothing -) - return nothing + ::typeof(left_null!), A::AbstractBlockSparseMatrix, N::Nothing + ) + return nothing end function MatrixAlgebraKit.left_null_qr!(A::AbstractBlockSparseMatrix, N, alg) - return left_null_svd!(A, N, default_svd_algorithm(A)) + return left_null_svd!(A, N, default_svd_algorithm(A)) end function MatrixAlgebraKit.left_null_svd!( - A::AbstractBlockSparseMatrix, N, alg, trunc::Nothing -) - return left_null_svd!(A, N, alg, null_truncation_strategy(; atol=0, rtol=0)) + A::AbstractBlockSparseMatrix, N, alg, trunc::Nothing + ) + return left_null_svd!(A, N, alg, null_truncation_strategy(; atol = 0, rtol = 0)) end function MatrixAlgebraKit.truncate!( - ::typeof(left_null!), - (U, S)::Tuple{AbstractBlockSparseMatrix,AbstractBlockSparseMatrix}, - strategy::TruncationStrategy, -) - return error("Not implemented.") + ::typeof(left_null!), + (U, S)::Tuple{AbstractBlockSparseMatrix, AbstractBlockSparseMatrix}, + strategy::TruncationStrategy, + ) + return error("Not implemented.") end function MatrixAlgebraKit.initialize_output( - ::typeof(right_null!), A::AbstractBlockSparseMatrix -) - return nothing + ::typeof(right_null!), A::AbstractBlockSparseMatrix + ) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(right_null!), A::AbstractBlockSparseMatrix, N::Nothing -) - return nothing + ::typeof(right_null!), A::AbstractBlockSparseMatrix, N::Nothing + ) + return nothing end function MatrixAlgebraKit.right_null_lq!(A::AbstractBlockSparseMatrix, N, alg) - return error("Not implement.") + return error("Not implement.") end function MatrixAlgebraKit.truncate!( - ::typeof(right_null!), - (S, Vᴴ)::Tuple{AbstractBlockSparseMatrix,AbstractBlockSparseMatrix}, - strategy::TruncationStrategy, -) - return error("Not implemented.") + ::typeof(right_null!), + (S, Vᴴ)::Tuple{AbstractBlockSparseMatrix, AbstractBlockSparseMatrix}, + strategy::TruncationStrategy, + ) + return error("Not implemented.") end diff --git a/src/factorizations/polar.jl b/src/factorizations/polar.jl index f123662f..c217600b 100644 --- a/src/factorizations/polar.jl +++ b/src/factorizations/polar.jl @@ -1,48 +1,48 @@ using MatrixAlgebraKit: - MatrixAlgebraKit, - PolarViaSVD, - check_input, - default_algorithm, - left_polar!, - right_polar!, - svd_compact! + MatrixAlgebraKit, + PolarViaSVD, + check_input, + default_algorithm, + left_polar!, + right_polar!, + svd_compact! function MatrixAlgebraKit.check_input(::typeof(left_polar!), A::AbstractBlockSparseMatrix) - @views for I in eachblockstoredindex(A) - m, n = size(A[I]) - m >= n || - throw(ArgumentError("each input matrix block needs at least as many rows as columns")) - end - return nothing + @views for I in eachblockstoredindex(A) + m, n = size(A[I]) + m >= n || + throw(ArgumentError("each input matrix block needs at least as many rows as columns")) + end + return nothing end function MatrixAlgebraKit.check_input(::typeof(right_polar!), A::AbstractBlockSparseMatrix) - @views for I in eachblockstoredindex(A) - m, n = size(A[I]) - m <= n || - throw(ArgumentError("each input matrix block needs at least as many columns as rows")) - end - return nothing + @views for I in eachblockstoredindex(A) + m, n = size(A[I]) + m <= n || + throw(ArgumentError("each input matrix block needs at least as many columns as rows")) + end + return nothing end function MatrixAlgebraKit.left_polar!(A::AbstractBlockSparseMatrix, alg::PolarViaSVD) - check_input(left_polar!, A) - # TODO: Use more in-place operations here, avoid `copy`. - U, S, Vᴴ = svd_compact!(A, alg.svdalg) - W = U * Vᴴ - # TODO: `copy` is required for now because of: - # https://github.com/ITensor/BlockSparseArrays.jl/issues/24 - # Remove when that is fixed. - P = copy(Vᴴ') * S * Vᴴ - return (W, P) + check_input(left_polar!, A) + # TODO: Use more in-place operations here, avoid `copy`. + U, S, Vᴴ = svd_compact!(A, alg.svdalg) + W = U * Vᴴ + # TODO: `copy` is required for now because of: + # https://github.com/ITensor/BlockSparseArrays.jl/issues/24 + # Remove when that is fixed. + P = copy(Vᴴ') * S * Vᴴ + return (W, P) end function MatrixAlgebraKit.right_polar!(A::AbstractBlockSparseMatrix, alg::PolarViaSVD) - check_input(right_polar!, A) - # TODO: Use more in-place operations here, avoid `copy`. - U, S, Vᴴ = svd_compact!(A, alg.svdalg) - Wᴴ = U * Vᴴ - # TODO: `copy` is required for now because of: - # https://github.com/ITensor/BlockSparseArrays.jl/issues/24 - # Remove when that is fixed. - P = U * S * copy(U') - return (P, Wᴴ) + check_input(right_polar!, A) + # TODO: Use more in-place operations here, avoid `copy`. + U, S, Vᴴ = svd_compact!(A, alg.svdalg) + Wᴴ = U * Vᴴ + # TODO: `copy` is required for now because of: + # https://github.com/ITensor/BlockSparseArrays.jl/issues/24 + # Remove when that is fixed. + P = U * S * copy(U') + return (P, Wᴴ) end diff --git a/src/factorizations/qr.jl b/src/factorizations/qr.jl index 81ab4d5d..474c345b 100644 --- a/src/factorizations/qr.jl +++ b/src/factorizations/qr.jl @@ -1,150 +1,150 @@ using MatrixAlgebraKit: MatrixAlgebraKit, default_qr_algorithm, qr_compact!, qr_full! function MatrixAlgebraKit.default_qr_algorithm( - ::Type{<:AbstractBlockSparseMatrix}; kwargs... -) - return BlockPermutedDiagonalAlgorithm() do block - return default_qr_algorithm(block; kwargs...) - end + ::Type{<:AbstractBlockSparseMatrix}; kwargs... + ) + return BlockPermutedDiagonalAlgorithm() do block + return default_qr_algorithm(block; kwargs...) + end end function output_type( - f::Union{typeof(qr_compact!),typeof(qr_full!)}, A::Type{<:AbstractMatrix{T}} -) where {T} - QR = Base.promote_op(f, A) - return isconcretetype(QR) ? QR : Tuple{AbstractMatrix{T},AbstractMatrix{T}} + f::Union{typeof(qr_compact!), typeof(qr_full!)}, A::Type{<:AbstractMatrix{T}} + ) where {T} + QR = Base.promote_op(f, A) + return isconcretetype(QR) ? QR : Tuple{AbstractMatrix{T}, AbstractMatrix{T}} end function MatrixAlgebraKit.initialize_output( - ::typeof(qr_compact!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm -) - return nothing + ::typeof(qr_compact!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm + ) + return nothing end function MatrixAlgebraKit.initialize_output( - ::typeof(qr_compact!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm -) - brows = eachblockaxis(axes(A, 1)) - bcols = eachblockaxis(axes(A, 2)) - # using the property that zip stops as soon as one of the iterators is exhausted - r_axes = map(splat(infimum), zip(brows, bcols)) - r_axis = mortar_axis(r_axes) - - BQ, BR = fieldtypes(output_type(qr_compact!, blocktype(A))) - Q = similar(A, BlockType(BQ), (axes(A, 1), r_axis)) - R = similar(A, BlockType(BR), (r_axis, axes(A, 2))) - - return Q, R + ::typeof(qr_compact!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm + ) + brows = eachblockaxis(axes(A, 1)) + bcols = eachblockaxis(axes(A, 2)) + # using the property that zip stops as soon as one of the iterators is exhausted + r_axes = map(splat(infimum), zip(brows, bcols)) + r_axis = mortar_axis(r_axes) + + BQ, BR = fieldtypes(output_type(qr_compact!, blocktype(A))) + Q = similar(A, BlockType(BQ), (axes(A, 1), r_axis)) + R = similar(A, BlockType(BR), (r_axis, axes(A, 2))) + + return Q, R end function MatrixAlgebraKit.initialize_output( - ::typeof(qr_full!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm -) - return nothing + ::typeof(qr_full!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm + ) + return nothing end function MatrixAlgebraKit.initialize_output( - ::typeof(qr_full!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm -) - BQ, BR = fieldtypes(output_type(qr_full!, blocktype(A))) - Q = similar(A, BlockType(BQ), (axes(A, 1), axes(A, 1))) - R = similar(A, BlockType(BR), (axes(A, 1), axes(A, 2))) - return Q, R + ::typeof(qr_full!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm + ) + BQ, BR = fieldtypes(output_type(qr_full!, blocktype(A))) + Q = similar(A, BlockType(BQ), (axes(A, 1), axes(A, 1))) + R = similar(A, BlockType(BR), (axes(A, 1), axes(A, 2))) + return Q, R end function MatrixAlgebraKit.check_input( - ::typeof(qr_compact!), A::AbstractBlockSparseMatrix, QR, ::BlockPermutedDiagonalAlgorithm -) - @assert isblockpermuteddiagonal(A) - return nothing + ::typeof(qr_compact!), A::AbstractBlockSparseMatrix, QR, ::BlockPermutedDiagonalAlgorithm + ) + @assert isblockpermuteddiagonal(A) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(qr_compact!), A::AbstractBlockSparseMatrix, (Q, R), ::BlockDiagonalAlgorithm -) - @assert isa(Q, AbstractBlockSparseMatrix) && isa(R, AbstractBlockSparseMatrix) - @assert eltype(A) == eltype(Q) == eltype(R) - @assert axes(A, 1) == axes(Q, 1) && axes(A, 2) == axes(R, 2) - @assert axes(Q, 2) == axes(R, 1) - @assert isblockdiagonal(A) - return nothing + ::typeof(qr_compact!), A::AbstractBlockSparseMatrix, (Q, R), ::BlockDiagonalAlgorithm + ) + @assert isa(Q, AbstractBlockSparseMatrix) && isa(R, AbstractBlockSparseMatrix) + @assert eltype(A) == eltype(Q) == eltype(R) + @assert axes(A, 1) == axes(Q, 1) && axes(A, 2) == axes(R, 2) + @assert axes(Q, 2) == axes(R, 1) + @assert isblockdiagonal(A) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(qr_full!), A::AbstractBlockSparseMatrix, QR, ::BlockPermutedDiagonalAlgorithm -) - @assert isblockpermuteddiagonal(A) - return nothing + ::typeof(qr_full!), A::AbstractBlockSparseMatrix, QR, ::BlockPermutedDiagonalAlgorithm + ) + @assert isblockpermuteddiagonal(A) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(qr_full!), A::AbstractBlockSparseMatrix, (Q, R), ::BlockDiagonalAlgorithm -) - @assert isa(Q, AbstractBlockSparseMatrix) && isa(R, AbstractBlockSparseMatrix) - @assert eltype(A) == eltype(Q) == eltype(R) - @assert axes(A, 1) == axes(Q, 1) && axes(A, 2) == axes(R, 2) - @assert axes(Q, 2) == axes(R, 1) - @assert isblockdiagonal(A) - return nothing + ::typeof(qr_full!), A::AbstractBlockSparseMatrix, (Q, R), ::BlockDiagonalAlgorithm + ) + @assert isa(Q, AbstractBlockSparseMatrix) && isa(R, AbstractBlockSparseMatrix) + @assert eltype(A) == eltype(Q) == eltype(R) + @assert axes(A, 1) == axes(Q, 1) && axes(A, 2) == axes(R, 2) + @assert axes(Q, 2) == axes(R, 1) + @assert isblockdiagonal(A) + return nothing end function MatrixAlgebraKit.qr_compact!( - A::AbstractBlockSparseMatrix, QR, alg::BlockPermutedDiagonalAlgorithm -) - check_input(qr_compact!, A, QR, alg) - Ad, (invrowperm, invcolperm) = blockdiagonalize(A) - Qd, Rd = qr_compact!(Ad, BlockDiagonalAlgorithm(alg)) - Q = transform_rows(Qd, invrowperm) - R = transform_cols(Rd, invcolperm) - return Q, R + A::AbstractBlockSparseMatrix, QR, alg::BlockPermutedDiagonalAlgorithm + ) + check_input(qr_compact!, A, QR, alg) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + Qd, Rd = qr_compact!(Ad, BlockDiagonalAlgorithm(alg)) + Q = transform_rows(Qd, invrowperm) + R = transform_cols(Rd, invcolperm) + return Q, R end function MatrixAlgebraKit.qr_compact!( - A::AbstractBlockSparseMatrix, (Q, R), alg::BlockDiagonalAlgorithm -) - MatrixAlgebraKit.check_input(qr_compact!, A, (Q, R), alg) - - # do decomposition on each block - for bI in blockdiagindices(A) - if isstored(A, bI) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - bQ, bR = qr_compact!(block, block_alg) - Q[bI] = bQ - R[bI] = bR - else - # TODO: this should be `Q[bI] = LinearAlgebra.I` - copyto!(@view!(Q[bI]), LinearAlgebra.I) + A::AbstractBlockSparseMatrix, (Q, R), alg::BlockDiagonalAlgorithm + ) + MatrixAlgebraKit.check_input(qr_compact!, A, (Q, R), alg) + + # do decomposition on each block + for bI in blockdiagindices(A) + if isstored(A, bI) + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bQ, bR = qr_compact!(block, block_alg) + Q[bI] = bQ + R[bI] = bR + else + # TODO: this should be `Q[bI] = LinearAlgebra.I` + copyto!(@view!(Q[bI]), LinearAlgebra.I) + end end - end - return Q, R + return Q, R end function MatrixAlgebraKit.qr_full!( - A::AbstractBlockSparseMatrix, QR, alg::BlockPermutedDiagonalAlgorithm -) - check_input(qr_full!, A, QR, alg) - Ad, (invrowperm, invcolperm) = blockdiagonalize(A) - Qd, Rd = qr_full!(Ad, BlockDiagonalAlgorithm(alg)) - Q = transform_rows(Qd, invrowperm) - R = transform_cols(Rd, invcolperm) - return Q, R + A::AbstractBlockSparseMatrix, QR, alg::BlockPermutedDiagonalAlgorithm + ) + check_input(qr_full!, A, QR, alg) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + Qd, Rd = qr_full!(Ad, BlockDiagonalAlgorithm(alg)) + Q = transform_rows(Qd, invrowperm) + R = transform_cols(Rd, invcolperm) + return Q, R end function MatrixAlgebraKit.qr_full!( - A::AbstractBlockSparseMatrix, (Q, R), alg::BlockDiagonalAlgorithm -) - MatrixAlgebraKit.check_input(qr_full!, A, (Q, R), alg) - - for bI in blockdiagindices(A) - if isstored(A, bI) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - bQ, bR = qr_full!(block, block_alg) - Q[bI] = bQ - R[bI] = bR - else - # TODO: this should be `Q[bI] = LinearAlgebra.I` - copyto!(@view!(Q[bI]), LinearAlgebra.I) + A::AbstractBlockSparseMatrix, (Q, R), alg::BlockDiagonalAlgorithm + ) + MatrixAlgebraKit.check_input(qr_full!, A, (Q, R), alg) + + for bI in blockdiagindices(A) + if isstored(A, bI) + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bQ, bR = qr_full!(block, block_alg) + Q[bI] = bQ + R[bI] = bR + else + # TODO: this should be `Q[bI] = LinearAlgebra.I` + copyto!(@view!(Q[bI]), LinearAlgebra.I) + end end - end - return Q, R + return Q, R end diff --git a/src/factorizations/svd.jl b/src/factorizations/svd.jl index 13f5bc0e..d3c2fa28 100644 --- a/src/factorizations/svd.jl +++ b/src/factorizations/svd.jl @@ -1,241 +1,241 @@ using DiagonalArrays: diagonaltype using MatrixAlgebraKit: - MatrixAlgebraKit, check_input, default_svd_algorithm, svd_compact!, svd_full!, svd_vals! + MatrixAlgebraKit, check_input, default_svd_algorithm, svd_compact!, svd_full!, svd_vals! using TypeParameterAccessors: realtype function MatrixAlgebraKit.default_svd_algorithm( - ::Type{<:AbstractBlockSparseMatrix}; kwargs... -) - return BlockPermutedDiagonalAlgorithm() do block - return default_svd_algorithm(block; kwargs...) - end + ::Type{<:AbstractBlockSparseMatrix}; kwargs... + ) + return BlockPermutedDiagonalAlgorithm() do block + return default_svd_algorithm(block; kwargs...) + end end function output_type( - f::Union{typeof(svd_compact!),typeof(svd_full!)}, A::Type{<:AbstractMatrix{T}} -) where {T} - USVᴴ = Base.promote_op(f, A) - return if isconcretetype(USVᴴ) - USVᴴ - else - Tuple{AbstractMatrix{T},AbstractMatrix{realtype(T)},AbstractMatrix{T}} - end + f::Union{typeof(svd_compact!), typeof(svd_full!)}, A::Type{<:AbstractMatrix{T}} + ) where {T} + USVᴴ = Base.promote_op(f, A) + return if isconcretetype(USVᴴ) + USVᴴ + else + Tuple{AbstractMatrix{T}, AbstractMatrix{realtype(T)}, AbstractMatrix{T}} + end end function output_type(::typeof(svd_vals!), A::Type{<:AbstractMatrix{T}}) where {T} - S = Base.promote_op(svd_vals!, A) - return isconcretetype(S) ? S : AbstractVector{real(T)} + S = Base.promote_op(svd_vals!, A) + return isconcretetype(S) ? S : AbstractVector{real(T)} end function MatrixAlgebraKit.initialize_output( - ::typeof(svd_compact!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm -) - return nothing + ::typeof(svd_compact!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm + ) + return nothing end function MatrixAlgebraKit.initialize_output( - ::typeof(svd_compact!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm -) - brows = eachblockaxis(axes(A, 1)) - bcols = eachblockaxis(axes(A, 2)) - # using the property that zip stops as soon as one of the iterators is exhausted - s_axes = map(splat(infimum), zip(brows, bcols)) - s_axis = mortar_axis(s_axes) - S_axes = (s_axis, s_axis) + ::typeof(svd_compact!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm + ) + brows = eachblockaxis(axes(A, 1)) + bcols = eachblockaxis(axes(A, 2)) + # using the property that zip stops as soon as one of the iterators is exhausted + s_axes = map(splat(infimum), zip(brows, bcols)) + s_axis = mortar_axis(s_axes) + S_axes = (s_axis, s_axis) - BU, BS, BVᴴ = fieldtypes(output_type(svd_compact!, blocktype(A))) - 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))) + BU, BS, BVᴴ = fieldtypes(output_type(svd_compact!, blocktype(A))) + 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))) - return U, S, Vᴴ + return U, S, Vᴴ end function MatrixAlgebraKit.initialize_output( - ::typeof(svd_full!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm -) - return nothing + ::typeof(svd_full!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm + ) + return nothing end function MatrixAlgebraKit.initialize_output( - ::typeof(svd_full!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm -) - BU, BS, BVᴴ = fieldtypes(output_type(svd_full!, blocktype(A))) - U = similar(A, BlockType(BU), (axes(A, 1), axes(A, 1))) - S = similar(A, BlockType(BS), axes(A)) - Vᴴ = similar(A, BlockType(BVᴴ), (axes(A, 2), axes(A, 2))) + ::typeof(svd_full!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm + ) + BU, BS, BVᴴ = fieldtypes(output_type(svd_full!, blocktype(A))) + U = similar(A, BlockType(BU), (axes(A, 1), axes(A, 1))) + S = similar(A, BlockType(BS), axes(A)) + Vᴴ = similar(A, BlockType(BVᴴ), (axes(A, 2), axes(A, 2))) - return U, S, Vᴴ + return U, S, Vᴴ end function MatrixAlgebraKit.initialize_output( - ::typeof(svd_vals!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm -) - return nothing + ::typeof(svd_vals!), ::AbstractBlockSparseMatrix, ::BlockPermutedDiagonalAlgorithm + ) + return nothing end function MatrixAlgebraKit.initialize_output( - ::typeof(svd_vals!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm -) - brows = eachblockaxis(axes(A, 1)) - bcols = eachblockaxis(axes(A, 2)) - # using the property that zip stops as soon as one of the iterators is exhausted - s_axes = map(splat(infimum), zip(brows, bcols)) - s_axis = mortar_axis(s_axes) + ::typeof(svd_vals!), A::AbstractBlockSparseMatrix, alg::BlockDiagonalAlgorithm + ) + brows = eachblockaxis(axes(A, 1)) + bcols = eachblockaxis(axes(A, 2)) + # using the property that zip stops as soon as one of the iterators is exhausted + s_axes = map(splat(infimum), zip(brows, bcols)) + s_axis = mortar_axis(s_axes) - BS = output_type(svd_vals!, blocktype(A)) - return similar(A, BlockType(BS), S_axes) + BS = output_type(svd_vals!, blocktype(A)) + return similar(A, BlockType(BS), S_axes) end function MatrixAlgebraKit.check_input( - ::typeof(svd_compact!), - A::AbstractBlockSparseMatrix, - USVᴴ, - ::BlockPermutedDiagonalAlgorithm, -) - @assert isblockpermuteddiagonal(A) + ::typeof(svd_compact!), + A::AbstractBlockSparseMatrix, + USVᴴ, + ::BlockPermutedDiagonalAlgorithm, + ) + return @assert isblockpermuteddiagonal(A) end function MatrixAlgebraKit.check_input( - ::typeof(svd_compact!), A::AbstractBlockSparseMatrix, (U, S, Vᴴ), ::BlockDiagonalAlgorithm -) - @assert isa(U, AbstractBlockSparseMatrix) && - isa(S, AbstractBlockSparseMatrix) && - isa(Vᴴ, AbstractBlockSparseMatrix) - @assert eltype(A) == eltype(U) == eltype(Vᴴ) - @assert real(eltype(A)) == eltype(S) - @assert axes(A, 1) == axes(U, 1) && axes(A, 2) == axes(Vᴴ, 2) - @assert axes(S, 1) == axes(S, 2) - @assert isblockdiagonal(A) - return nothing + ::typeof(svd_compact!), A::AbstractBlockSparseMatrix, (U, S, Vᴴ), ::BlockDiagonalAlgorithm + ) + @assert isa(U, AbstractBlockSparseMatrix) && + isa(S, AbstractBlockSparseMatrix) && + isa(Vᴴ, AbstractBlockSparseMatrix) + @assert eltype(A) == eltype(U) == eltype(Vᴴ) + @assert real(eltype(A)) == eltype(S) + @assert axes(A, 1) == axes(U, 1) && axes(A, 2) == axes(Vᴴ, 2) + @assert axes(S, 1) == axes(S, 2) + @assert isblockdiagonal(A) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(svd_full!), A::AbstractBlockSparseMatrix, USVᴴ, ::BlockPermutedDiagonalAlgorithm -) - @assert isblockpermuteddiagonal(A) - return nothing + ::typeof(svd_full!), A::AbstractBlockSparseMatrix, USVᴴ, ::BlockPermutedDiagonalAlgorithm + ) + @assert isblockpermuteddiagonal(A) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(svd_full!), A::AbstractBlockSparseMatrix, (U, S, Vᴴ), ::BlockDiagonalAlgorithm -) - @assert isa(U, AbstractBlockSparseMatrix) && - isa(S, AbstractBlockSparseMatrix) && - isa(Vᴴ, AbstractBlockSparseMatrix) - @assert eltype(A) == eltype(U) == eltype(Vᴴ) - @assert real(eltype(A)) == eltype(S) - @assert axes(A, 1) == axes(U, 1) && axes(A, 2) == axes(Vᴴ, 1) == axes(Vᴴ, 2) - @assert axes(S, 2) == axes(A, 2) - @assert isblockdiagonal(A) - return nothing + ::typeof(svd_full!), A::AbstractBlockSparseMatrix, (U, S, Vᴴ), ::BlockDiagonalAlgorithm + ) + @assert isa(U, AbstractBlockSparseMatrix) && + isa(S, AbstractBlockSparseMatrix) && + isa(Vᴴ, AbstractBlockSparseMatrix) + @assert eltype(A) == eltype(U) == eltype(Vᴴ) + @assert real(eltype(A)) == eltype(S) + @assert axes(A, 1) == axes(U, 1) && axes(A, 2) == axes(Vᴴ, 1) == axes(Vᴴ, 2) + @assert axes(S, 2) == axes(A, 2) + @assert isblockdiagonal(A) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(svd_vals!), A::AbstractBlockSparseMatrix, S, ::BlockPermutedDiagonalAlgorithm -) - @assert isblockpermuteddiagonal(A) - return nothing + ::typeof(svd_vals!), A::AbstractBlockSparseMatrix, S, ::BlockPermutedDiagonalAlgorithm + ) + @assert isblockpermuteddiagonal(A) + return nothing end function MatrixAlgebraKit.check_input( - ::typeof(svd_vals!), A::AbstractBlockSparseMatrix, S, ::BlockDiagonalAlgorithm -) - @assert isa(S, AbstractBlockSparseVector) - @assert real(eltype(A)) == eltype(S) - @assert isblockdiagonal(A) - return nothing + ::typeof(svd_vals!), A::AbstractBlockSparseMatrix, S, ::BlockDiagonalAlgorithm + ) + @assert isa(S, AbstractBlockSparseVector) + @assert real(eltype(A)) == eltype(S) + @assert isblockdiagonal(A) + return nothing end function MatrixAlgebraKit.svd_compact!( - A::AbstractBlockSparseMatrix, USVᴴ, alg::BlockPermutedDiagonalAlgorithm -) - check_input(svd_compact!, A, USVᴴ, alg) + A::AbstractBlockSparseMatrix, USVᴴ, alg::BlockPermutedDiagonalAlgorithm + ) + check_input(svd_compact!, A, USVᴴ, alg) - Ad, (invrowperm, invcolperm) = blockdiagonalize(A) - Ud, S, Vᴴd = svd_compact!(Ad, BlockDiagonalAlgorithm(alg)) - U = transform_rows(Ud, invrowperm) - Vᴴ = transform_cols(Vᴴd, invcolperm) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + Ud, S, Vᴴd = svd_compact!(Ad, BlockDiagonalAlgorithm(alg)) + U = transform_rows(Ud, invrowperm) + Vᴴ = transform_cols(Vᴴd, invcolperm) - return U, S, Vᴴ + return U, S, Vᴴ end function MatrixAlgebraKit.svd_compact!( - A::AbstractBlockSparseMatrix, (U, S, Vᴴ), alg::BlockDiagonalAlgorithm -) - check_input(svd_compact!, A, (U, S, Vᴴ), alg) - - for bI in blockdiagindices(A) - if isstored(A, bI) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - bU, bS, bVᴴ = svd_compact!(block, block_alg) - U[bI] = bU - S[bI] = bS - Vᴴ[bI] = bVᴴ - else - # TODO: this should be `U[bI] = LinearAlgebra.I` and `Vᴴ[bI] = LinearAlgebra.I` - copyto!(@view!(U[bI]), LinearAlgebra.I) - copyto!(@view!(Vᴴ[bI]), LinearAlgebra.I) + A::AbstractBlockSparseMatrix, (U, S, Vᴴ), alg::BlockDiagonalAlgorithm + ) + check_input(svd_compact!, A, (U, S, Vᴴ), alg) + + for bI in blockdiagindices(A) + if isstored(A, bI) + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bU, bS, bVᴴ = svd_compact!(block, block_alg) + U[bI] = bU + S[bI] = bS + Vᴴ[bI] = bVᴴ + else + # TODO: this should be `U[bI] = LinearAlgebra.I` and `Vᴴ[bI] = LinearAlgebra.I` + copyto!(@view!(U[bI]), LinearAlgebra.I) + copyto!(@view!(Vᴴ[bI]), LinearAlgebra.I) + end end - end - return U, S, Vᴴ + return U, S, Vᴴ end function MatrixAlgebraKit.svd_full!( - A::AbstractBlockSparseMatrix, USVᴴ, alg::BlockPermutedDiagonalAlgorithm -) - check_input(svd_full!, A, USVᴴ, alg) + A::AbstractBlockSparseMatrix, USVᴴ, alg::BlockPermutedDiagonalAlgorithm + ) + check_input(svd_full!, A, USVᴴ, alg) - Ad, (invrowperm, invcolperm) = blockdiagonalize(A) - Ud, S, Vᴴd = svd_full!(Ad, BlockDiagonalAlgorithm(alg)) - U = transform_rows(Ud, invrowperm) - Vᴴ = transform_cols(Vᴴd, invcolperm) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + Ud, S, Vᴴd = svd_full!(Ad, BlockDiagonalAlgorithm(alg)) + U = transform_rows(Ud, invrowperm) + Vᴴ = transform_cols(Vᴴd, invcolperm) - return U, S, Vᴴ + return U, S, Vᴴ end function MatrixAlgebraKit.svd_full!( - A::AbstractBlockSparseMatrix, (U, S, Vᴴ), alg::BlockDiagonalAlgorithm -) - check_input(svd_full!, A, (U, S, Vᴴ), alg) - - for bI in blockdiagindices(A) - if isstored(A, bI) - block = @view!(A[bI]) - block_alg = block_algorithm(alg, block) - bU, bS, bVᴴ = svd_full!(block, block_alg) - U[bI] = bU - S[bI] = bS - Vᴴ[bI] = bVᴴ - else - # TODO: this should be `U[bI] = LinearAlgebra.I` and `Vᴴ[bI] = LinearAlgebra.I` - copyto!(@view!(U[bI]), LinearAlgebra.I) - copyto!(@view!(Vᴴ[bI]), LinearAlgebra.I) + A::AbstractBlockSparseMatrix, (U, S, Vᴴ), alg::BlockDiagonalAlgorithm + ) + check_input(svd_full!, A, (U, S, Vᴴ), alg) + + for bI in blockdiagindices(A) + if isstored(A, bI) + block = @view!(A[bI]) + block_alg = block_algorithm(alg, block) + bU, bS, bVᴴ = svd_full!(block, block_alg) + U[bI] = bU + S[bI] = bS + Vᴴ[bI] = bVᴴ + else + # TODO: this should be `U[bI] = LinearAlgebra.I` and `Vᴴ[bI] = LinearAlgebra.I` + copyto!(@view!(U[bI]), LinearAlgebra.I) + copyto!(@view!(Vᴴ[bI]), LinearAlgebra.I) + end end - end - # Complete the unitaries for rectangular inputs - # TODO: this should be `U[Block(I, I)] = LinearAlgebra.I` - for I in (blocksize(A, 2) + 1):blocksize(A, 1) - copyto!(@view!(U[Block(I, I)]), LinearAlgebra.I) - end - # TODO: this should be `Vᴴ[Block(I, I)] = LinearAlgebra.I` - for I in (blocksize(A, 1) + 1):blocksize(A, 2) - copyto!(@view!(Vᴴ[Block(I, I)]), LinearAlgebra.I) - end + # Complete the unitaries for rectangular inputs + # TODO: this should be `U[Block(I, I)] = LinearAlgebra.I` + for I in (blocksize(A, 2) + 1):blocksize(A, 1) + copyto!(@view!(U[Block(I, I)]), LinearAlgebra.I) + end + # TODO: this should be `Vᴴ[Block(I, I)] = LinearAlgebra.I` + for I in (blocksize(A, 1) + 1):blocksize(A, 2) + copyto!(@view!(Vᴴ[Block(I, I)]), LinearAlgebra.I) + end - return U, S, Vᴴ + return U, S, Vᴴ end function MatrixAlgebraKit.svd_vals!( - A::AbstractBlockSparseMatrix, S, alg::BlockPermutedDiagonalAlgorithm -) - MatrixAlgebraKit.check_input(svd_vals!, A, S, alg) - Ad, _ = blockdiagonalize(A) - return svd_vals!(Ad, BlockDiagonalAlgorithm(alg)) + A::AbstractBlockSparseMatrix, S, alg::BlockPermutedDiagonalAlgorithm + ) + MatrixAlgebraKit.check_input(svd_vals!, A, S, alg) + Ad, _ = blockdiagonalize(A) + return svd_vals!(Ad, BlockDiagonalAlgorithm(alg)) end function MatrixAlgebraKit.svd_vals!( - A::AbstractBlockSparseMatrix, S, alg::BlockDiagonalAlgorithm -) - MatrixAlgebraKit.check_input(svd_vals!, A, S, alg) - for I in eachblockstoredindex(A) - block = @view!(A[I]) - S[Tuple(I)[1]] = svd_vals!(block, block_algorithm(alg, block)) - end - return S + A::AbstractBlockSparseMatrix, S, alg::BlockDiagonalAlgorithm + ) + MatrixAlgebraKit.check_input(svd_vals!, A, S, alg) + for I in eachblockstoredindex(A) + block = @view!(A[I]) + S[Tuple(I)[1]] = svd_vals!(block, block_algorithm(alg, block)) + end + return S end diff --git a/src/factorizations/truncation.jl b/src/factorizations/truncation.jl index eb04195e..ee822e03 100644 --- a/src/factorizations/truncation.jl +++ b/src/factorizations/truncation.jl @@ -1,13 +1,13 @@ using MatrixAlgebraKit: - MatrixAlgebraKit, - TruncatedAlgorithm, - TruncationStrategy, - diagview, - eig_trunc!, - eigh_trunc!, - findtruncated, - svd_trunc!, - truncate! + MatrixAlgebraKit, + TruncatedAlgorithm, + TruncationStrategy, + diagview, + eig_trunc!, + eigh_trunc!, + findtruncated, + svd_trunc!, + truncate! """ BlockDiagonalTruncationStrategy(strategy::TruncationStrategy) @@ -15,51 +15,51 @@ using MatrixAlgebraKit: A wrapper for `TruncationStrategy` that implements the wrapped strategy on a block-by-block basis, which is possible if the input matrix is a block-diagonal matrix. """ -struct BlockDiagonalTruncationStrategy{T<:TruncationStrategy} <: TruncationStrategy - strategy::T +struct BlockDiagonalTruncationStrategy{T <: TruncationStrategy} <: TruncationStrategy + strategy::T end function BlockDiagonalTruncationStrategy(alg::BlockPermutedDiagonalAlgorithm) - return BlockDiagonalTruncationStrategy(alg.strategy) + return BlockDiagonalTruncationStrategy(alg.strategy) end function MatrixAlgebraKit.svd_trunc!( - A::AbstractBlockSparseMatrix, - out, - alg::TruncatedAlgorithm{<:BlockPermutedDiagonalAlgorithm}, -) - Ad, (invrowperm, invcolperm) = blockdiagonalize(A) - blockalg = BlockDiagonalAlgorithm(alg.alg) - blockstrategy = BlockDiagonalTruncationStrategy(alg.trunc) - Ud, S, Vᴴd = svd_trunc!(Ad, TruncatedAlgorithm(blockalg, blockstrategy)) + A::AbstractBlockSparseMatrix, + out, + alg::TruncatedAlgorithm{<:BlockPermutedDiagonalAlgorithm}, + ) + Ad, (invrowperm, invcolperm) = blockdiagonalize(A) + blockalg = BlockDiagonalAlgorithm(alg.alg) + blockstrategy = BlockDiagonalTruncationStrategy(alg.trunc) + Ud, S, Vᴴd = svd_trunc!(Ad, TruncatedAlgorithm(blockalg, blockstrategy)) - U = transform_rows(Ud, invrowperm) - Vᴴ = transform_cols(Vᴴd, invcolperm) + U = transform_rows(Ud, invrowperm) + Vᴴ = transform_cols(Vᴴd, invcolperm) - return U, S, Vᴴ + return U, S, Vᴴ end for f in [:eig_trunc!, :eigh_trunc!] - @eval begin - function MatrixAlgebraKit.truncate!( - ::typeof($f), - (D, V)::NTuple{2,AbstractBlockSparseMatrix}, - strategy::TruncationStrategy, - ) - return truncate!($f, (D, V), BlockDiagonalTruncationStrategy(strategy)) + @eval begin + function MatrixAlgebraKit.truncate!( + ::typeof($f), + (D, V)::NTuple{2, AbstractBlockSparseMatrix}, + strategy::TruncationStrategy, + ) + return truncate!($f, (D, V), BlockDiagonalTruncationStrategy(strategy)) + end end - end end # cannot use regular slicing here: I want to slice without altering blockstructure # solution: use boolean indexing and slice the mask, effectively cheaply inverting the map function MatrixAlgebraKit.findtruncated( - values::AbstractVector, strategy::BlockDiagonalTruncationStrategy -) - ind = findtruncated(Vector(values), strategy.strategy) - indexmask = falses(length(values)) - indexmask[ind] .= true - return to_truncated_indices(values, indexmask) + values::AbstractVector, strategy::BlockDiagonalTruncationStrategy + ) + ind = findtruncated(Vector(values), strategy.strategy) + indexmask = falses(length(values)) + indexmask[ind] .= true + return to_truncated_indices(values, indexmask) end # Allow customizing the indices output by `findtruncated` @@ -67,30 +67,30 @@ end # a block or Kronecker structure. to_truncated_indices(values, I) = I function to_truncated_indices(values::AbstractBlockVector, I::AbstractVector{Bool}) - I′ = BlockedVector(I, blocklengths(axis(values))) - blocks = map(BlockRange(values)) do b - return _getindex(b, to_truncated_indices(values[b], I′[b])) - end - return blocks + I′ = BlockedVector(I, blocklengths(axis(values))) + blocks = map(BlockRange(values)) do b + return _getindex(b, to_truncated_indices(values[b], I′[b])) + end + return blocks end function MatrixAlgebraKit.truncate!( - ::typeof(svd_trunc!), - (U, S, Vᴴ)::NTuple{3,AbstractBlockSparseMatrix}, - strategy::BlockDiagonalTruncationStrategy, -) - I = findtruncated(diag(S), strategy) - return (U[:, I], S[I, I], Vᴴ[I, :]) + ::typeof(svd_trunc!), + (U, S, Vᴴ)::NTuple{3, AbstractBlockSparseMatrix}, + strategy::BlockDiagonalTruncationStrategy, + ) + I = findtruncated(diag(S), strategy) + return (U[:, I], S[I, I], Vᴴ[I, :]) end for f in [:eig_trunc!, :eigh_trunc!] - @eval begin - function MatrixAlgebraKit.truncate!( - ::typeof($f), - (D, V)::NTuple{2,AbstractBlockSparseMatrix}, - strategy::BlockDiagonalTruncationStrategy, - ) - I = findtruncated(diag(D), strategy) - return (D[I, I], V[:, I]) + @eval begin + function MatrixAlgebraKit.truncate!( + ::typeof($f), + (D, V)::NTuple{2, AbstractBlockSparseMatrix}, + strategy::BlockDiagonalTruncationStrategy, + ) + I = findtruncated(diag(D), strategy) + return (D[I, I], V[:, I]) + end end - end end diff --git a/src/factorizations/utility.jl b/src/factorizations/utility.jl index b269c0fc..f0e73233 100644 --- a/src/factorizations/utility.jl +++ b/src/factorizations/utility.jl @@ -1,23 +1,23 @@ using MatrixAlgebraKit: MatrixAlgebraKit function infimum(r1::AbstractUnitRange, r2::AbstractUnitRange) - (isone(first(r1)) && isone(first(r2))) || - throw(ArgumentError("infimum only defined for ranges starting at 1")) - if length(r1) ≤ length(r2) - return r1 - else - return r1[r2] - end + (isone(first(r1)) && isone(first(r2))) || + throw(ArgumentError("infimum only defined for ranges starting at 1")) + if length(r1) ≤ length(r2) + return r1 + else + return r1[r2] + end end function supremum(r1::AbstractUnitRange, r2::AbstractUnitRange) - (isone(first(r1)) && isone(first(r2))) || - throw(ArgumentError("supremum only defined for ranges starting at 1")) - if length(r1) ≥ length(r2) - return r1 - else - return r2 - end + (isone(first(r1)) && isone(first(r2))) || + throw(ArgumentError("supremum only defined for ranges starting at 1")) + if length(r1) ≥ length(r2) + return r1 + else + return r2 + end end transform_rows(A::AbstractMatrix, invrowperm) = A[invrowperm, :] @@ -25,44 +25,44 @@ transform_rows(A::AbstractVector, invrowperm) = A[invrowperm] transform_cols(A::AbstractMatrix, invcolperm) = A[:, invcolperm] function blockdiagonalize(A::AbstractBlockSparseMatrix) - # sort in order to avoid depending on internal details such as dictionary order - bIs = sort!(collect(eachblockstoredindex(A)); by=Int ∘ last ∘ Tuple) + # sort in order to avoid depending on internal details such as dictionary order + bIs = sort!(collect(eachblockstoredindex(A)); by = Int ∘ last ∘ Tuple) - # obtain permutation for the blocks that are present - rowperm = map(first ∘ Tuple, bIs) - colperm = map(last ∘ Tuple, bIs) + # obtain permutation for the blocks that are present + rowperm = map(first ∘ Tuple, bIs) + colperm = map(last ∘ Tuple, bIs) - # These checks might be expensive but are convenient for now - allunique(rowperm) || throw(ArgumentError("input contains more than one block per row")) - allunique(colperm) || - throw(ArgumentError("input contains more than one block per column")) + # These checks might be expensive but are convenient for now + allunique(rowperm) || throw(ArgumentError("input contains more than one block per row")) + allunique(colperm) || + throw(ArgumentError("input contains more than one block per column")) - # post-process empty rows and columns: this pairs them up in order of occurance, - # putting empty rows and columns due to rectangularity at the end - emptyrows = setdiff(Block.(1:blocksize(A, 1)), rowperm) - append!(rowperm, emptyrows) - emptycols = setdiff(Block.(1:blocksize(A, 2)), colperm) - append!(colperm, emptycols) + # post-process empty rows and columns: this pairs them up in order of occurance, + # putting empty rows and columns due to rectangularity at the end + emptyrows = setdiff(Block.(1:blocksize(A, 1)), rowperm) + append!(rowperm, emptyrows) + emptycols = setdiff(Block.(1:blocksize(A, 2)), colperm) + append!(colperm, emptycols) - invrowperm = Block.(invperm(Int.(rowperm))) - invcolperm = Block.(invperm(Int.(colperm))) + invrowperm = Block.(invperm(Int.(rowperm))) + invcolperm = Block.(invperm(Int.(colperm))) - return A[rowperm, colperm], (invrowperm, invcolperm) + return A[rowperm, colperm], (invrowperm, invcolperm) end function isblockdiagonal(A::AbstractBlockSparseMatrix) - for bI in eachblockstoredindex(A) - row, col = Tuple(bI) - row == col || return false - end - # don't need to check for rows and cols appearing only once - # this is guaranteed as long as eachblockstoredindex is unique, which we assume - return true + for bI in eachblockstoredindex(A) + row, col = Tuple(bI) + row == col || return false + end + # don't need to check for rows and cols appearing only once + # this is guaranteed as long as eachblockstoredindex is unique, which we assume + return true end function isblockpermuteddiagonal(A::AbstractBlockSparseMatrix) - return allunique(first.(eachblockstoredindex(A))) && - allunique(last.(eachblockstoredindex(A))) + return allunique(first.(eachblockstoredindex(A))) && + allunique(last.(eachblockstoredindex(A))) end """ @@ -75,14 +75,14 @@ the algorithm that will be used. This can be leveraged to allow for different al each block. """ struct BlockDiagonalAlgorithm{F} <: MatrixAlgebraKit.AbstractAlgorithm - falg::F + falg::F end function block_algorithm(alg::BlockDiagonalAlgorithm, a::AbstractMatrix) - return block_algorithm(alg, typeof(a)) + return block_algorithm(alg, typeof(a)) end function block_algorithm(alg::BlockDiagonalAlgorithm, A::Type{<:AbstractMatrix}) - return alg.falg(A) + return alg.falg(A) end """ @@ -96,15 +96,15 @@ the algorithm that will be used. This can be leveraged to allow for different al each block. """ struct BlockPermutedDiagonalAlgorithm{F} <: MatrixAlgebraKit.AbstractAlgorithm - falg::F + falg::F end function block_algorithm(alg::BlockPermutedDiagonalAlgorithm, a::AbstractMatrix) - return block_algorithm(alg, typeof(a)) + return block_algorithm(alg, typeof(a)) end function block_algorithm(alg::BlockPermutedDiagonalAlgorithm, A::Type{<:AbstractMatrix}) - return alg.falg(A) + return alg.falg(A) end function BlockDiagonalAlgorithm(alg::BlockPermutedDiagonalAlgorithm) - return BlockDiagonalAlgorithm(alg.falg) + return BlockDiagonalAlgorithm(alg.falg) end diff --git a/test/TestBlockSparseArraysUtils.jl b/test/TestBlockSparseArraysUtils.jl index 36f45398..de20f7ae 100644 --- a/test/TestBlockSparseArraysUtils.jl +++ b/test/TestBlockSparseArraysUtils.jl @@ -2,14 +2,14 @@ module TestBlockSparseArraysUtils using BlockArrays: BlockRange function set_blocks!(a::AbstractArray, f::Function, blocks::Function) - set_blocks!(a, f, filter(blocks, BlockRange(a))) - return a + set_blocks!(a, f, filter(blocks, BlockRange(a))) + return a end function set_blocks!(a::AbstractArray, f::Function, blocks::Vector) - for b in blocks - a[b] = f(eltype(a), size(@view(a[b]))) - end - return a + for b in blocks + a[b] = f(eltype(a), size(@view(a[b]))) + end + return a end end diff --git a/test/runtests.jl b/test/runtests.jl index 98b2d2b8..00080503 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -6,60 +6,62 @@ using Suppressor: Suppressor const pat = r"(?:--group=)(\w+)" arg_id = findfirst(contains(pat), ARGS) const GROUP = uppercase( - if isnothing(arg_id) - get(ENV, "GROUP", "ALL") - else - only(match(pat, ARGS[arg_id]).captures) - end, + if isnothing(arg_id) + get(ENV, "GROUP", "ALL") + else + only(match(pat, ARGS[arg_id]).captures) + end, ) "match files of the form `test_*.jl`, but exclude `*setup*.jl`" function istestfile(fn) - return endswith(fn, ".jl") && startswith(basename(fn), "test_") && !contains(fn, "setup") + return endswith(fn, ".jl") && startswith(basename(fn), "test_") && !contains(fn, "setup") end "match files of the form `*.jl`, but exclude `*_notest.jl` and `*setup*.jl`" function isexamplefile(fn) - return endswith(fn, ".jl") && !endswith(fn, "_notest.jl") && !contains(fn, "setup") + return endswith(fn, ".jl") && !endswith(fn, "_notest.jl") && !contains(fn, "setup") end @time begin - # tests in groups based on folder structure - for testgroup in filter(isdir, readdir(@__DIR__)) - if GROUP == "ALL" || GROUP == uppercase(testgroup) - groupdir = joinpath(@__DIR__, testgroup) - for file in filter(istestfile, readdir(groupdir)) - filename = joinpath(groupdir, file) - @eval @safetestset $file begin - include($filename) + # tests in groups based on folder structure + for testgroup in filter(isdir, readdir(@__DIR__)) + if GROUP == "ALL" || GROUP == uppercase(testgroup) + groupdir = joinpath(@__DIR__, testgroup) + for file in filter(istestfile, readdir(groupdir)) + filename = joinpath(groupdir, file) + @eval @safetestset $file begin + include($filename) + end + end end - end end - end - # single files in top folder - for file in filter(istestfile, readdir(@__DIR__)) - (file == basename(@__FILE__)) && continue # exclude this file to avoid infinite recursion - @eval @safetestset $file begin - include($file) + # single files in top folder + for file in filter(istestfile, readdir(@__DIR__)) + (file == basename(@__FILE__)) && continue # exclude this file to avoid infinite recursion + @eval @safetestset $file begin + include($file) + end end - end - # test examples - examplepath = joinpath(@__DIR__, "..", "examples") - for (root, _, files) in walkdir(examplepath) - contains(chopprefix(root, @__DIR__), "setup") && continue - for file in filter(isexamplefile, files) - filename = joinpath(root, file) - @eval begin - @safetestset $file begin - $(Expr( - :macrocall, - GlobalRef(Suppressor, Symbol("@suppress")), - LineNumberNode(@__LINE__, @__FILE__), - :(include($filename)), - )) + # test examples + examplepath = joinpath(@__DIR__, "..", "examples") + for (root, _, files) in walkdir(examplepath) + contains(chopprefix(root, @__DIR__), "setup") && continue + for file in filter(isexamplefile, files) + filename = joinpath(root, file) + @eval begin + @safetestset $file begin + $( + Expr( + :macrocall, + GlobalRef(Suppressor, Symbol("@suppress")), + LineNumberNode(@__LINE__, @__FILE__), + :(include($filename)), + ) + ) + end + end end - end end - end end diff --git a/test/test_abstract_blocktype.jl b/test/test_abstract_blocktype.jl index d66d4b22..297b54f1 100644 --- a/test/test_abstract_blocktype.jl +++ b/test/test_abstract_blocktype.jl @@ -4,141 +4,141 @@ using BlockSparseArrays: BlockSparseMatrix, blockstoredlength using JLArrays: JLArray using LinearAlgebra: hermitianpart, norm using MatrixAlgebraKit: - eig_full, - eig_trunc, - eig_vals, - eigh_full, - eigh_trunc, - eigh_vals, - isisometry, - left_orth, - left_polar, - lq_compact, - lq_full, - qr_compact, - qr_full, - right_orth, - right_polar, - svd_compact, - svd_full, - svd_trunc + eig_full, + eig_trunc, + eig_vals, + eigh_full, + eigh_trunc, + eigh_vals, + isisometry, + left_orth, + left_polar, + lq_compact, + lq_full, + qr_compact, + qr_full, + right_orth, + right_polar, + svd_compact, + svd_full, + svd_trunc 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 + elt in elts - dev = adapt(arrayt) + 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 = 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 = BlockSparseMatrix{elt,AbstractMatrix{elt}}(undef, [2, 3], [2, 3]) - a[Block(1, 1)] = dev(randn(elt, 2, 2)) - @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 Matrix{elt} - @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} + a = BlockSparseMatrix{elt, AbstractMatrix{elt}}(undef, [2, 3], [2, 3]) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + @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 Matrix{elt} + @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} - a = BlockSparseMatrix{elt,AbstractMatrix{elt}}(undef, [2, 3], [2, 3]) - a[Block(1, 1)] = dev(randn(elt, 2, 2)) - a′ = BlockSparseMatrix{elt,AbstractMatrix{elt}}(undef, [2, 3], [2, 3]) - a′[Block(2, 2)] = dev(randn(elt, 3, 3)) + a = BlockSparseMatrix{elt, AbstractMatrix{elt}}(undef, [2, 3], [2, 3]) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + a′ = BlockSparseMatrix{elt, AbstractMatrix{elt}}(undef, [2, 3], [2, 3]) + a′[Block(2, 2)] = dev(randn(elt, 3, 3)) - b = copy(a) - @test Array(b) ≈ Array(a) + b = copy(a) + @test Array(b) ≈ Array(a) - b = a + a′ - @test Array(b) ≈ Array(a) + Array(a′) + b = a + a′ + @test Array(b) ≈ Array(a) + Array(a′) - b = 3a - @test Array(b) ≈ 3Array(a) + b = 3a + @test Array(b) ≈ 3Array(a) - b = a * a - @test Array(b) ≈ Array(a) * Array(a) + b = a * a + @test Array(b) ≈ Array(a) * Array(a) - b = a * a′ - @test Array(b) ≈ Array(a) * Array(a′) - @test norm(b) ≈ 0 + b = a * a′ + @test Array(b) ≈ Array(a) * Array(a′) + @test norm(b) ≈ 0 - a = BlockSparseMatrix{elt,AbstractMatrix{elt}}(undef, [2, 3], [2, 3]) - a[Block(1, 1)] = dev(randn(elt, 2, 2)) - for f in (eig_full, eig_trunc) + a = BlockSparseMatrix{elt, AbstractMatrix{elt}}(undef, [2, 3], [2, 3]) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + for f in (eig_full, eig_trunc) + if arrayt === Array + d, v = f(a) + @test a * v ≈ v * d + else + @test_broken f(a) + end + end if arrayt === Array - d, v = f(a) - @test a * v ≈ v * d + d = eig_vals(a) + @test sort(Vector(d); by = abs) ≈ sort(eig_vals(Matrix(a)); by = abs) else - @test_broken f(a) + @test_broken eig_vals(a) end - end - if arrayt === Array - d = eig_vals(a) - @test sort(Vector(d); by=abs) ≈ sort(eig_vals(Matrix(a)); by=abs) - else - @test_broken eig_vals(a) - end - a = BlockSparseMatrix{elt,AbstractMatrix{elt}}(undef, [2, 3], [2, 3]) - a[Block(1, 1)] = dev(parent(hermitianpart(randn(elt, 2, 2)))) - for f in (eigh_full, eigh_trunc) + a = BlockSparseMatrix{elt, AbstractMatrix{elt}}(undef, [2, 3], [2, 3]) + a[Block(1, 1)] = dev(parent(hermitianpart(randn(elt, 2, 2)))) + for f in (eigh_full, eigh_trunc) + if arrayt === Array + d, v = f(a) + @test a * v ≈ v * d + else + @test_broken f(a) + end + end if arrayt === Array - d, v = f(a) - @test a * v ≈ v * d + d = eigh_vals(a) + @test sort(Vector(d); by = abs) ≈ sort(eig_vals(Matrix(a)); by = abs) else - @test_broken f(a) + @test_broken eigh_vals(a) end - end - if arrayt === Array - d = eigh_vals(a) - @test sort(Vector(d); by=abs) ≈ sort(eig_vals(Matrix(a)); by=abs) - else - @test_broken eigh_vals(a) - end - a = BlockSparseMatrix{elt,AbstractMatrix{elt}}(undef, [2, 3], [2, 3]) - a[Block(1, 1)] = dev(randn(elt, 2, 2)) - for f in (left_orth, left_polar, qr_compact, qr_full) - if arrayt ≢ Array && f ≡ left_orth - @test_broken f(a) - else - u, c = f(a) - @test u * c ≈ a - if arrayt ≡ Array - @test isisometry(u; side=:left) - else - # TODO: Fix comparison with UniformScaling on GPU. - @test_broken isisometry(u; side=:left) - end + a = BlockSparseMatrix{elt, AbstractMatrix{elt}}(undef, [2, 3], [2, 3]) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + for f in (left_orth, left_polar, qr_compact, qr_full) + if arrayt ≢ Array && f ≡ left_orth + @test_broken f(a) + else + u, c = f(a) + @test u * c ≈ a + if arrayt ≡ Array + @test isisometry(u; side = :left) + else + # TODO: Fix comparison with UniformScaling on GPU. + @test_broken isisometry(u; side = :left) + end + end end - end - for f in (right_orth, right_polar, lq_compact, lq_full) - if arrayt ≢ Array && f ≡ right_orth - @test_broken f(a) - else - c, u = f(a) - @test c * u ≈ a - if arrayt ≡ Array - @test isisometry(u; side=:right) - else - # TODO: Fix comparison with UniformScaling on GPU. - @test_broken isisometry(u; side=:right) - end + for f in (right_orth, right_polar, lq_compact, lq_full) + if arrayt ≢ Array && f ≡ right_orth + @test_broken f(a) + else + c, u = f(a) + @test c * u ≈ a + if arrayt ≡ Array + @test isisometry(u; side = :right) + else + # TODO: Fix comparison with UniformScaling on GPU. + @test_broken isisometry(u; side = :right) + end + end end - end - for f in (svd_compact, svd_full, svd_trunc) - if arrayt ≢ Array && (f ≡ svd_full || f ≡ svd_trunc) - @test_broken f(a) - else - u, s, v = f(a) - @test u * s * v ≈ a + for f in (svd_compact, svd_full, svd_trunc) + if arrayt ≢ Array && (f ≡ svd_full || f ≡ svd_trunc) + @test_broken f(a) + else + u, s, v = f(a) + @test u * s * v ≈ a + end end - end end diff --git a/test/test_aqua.jl b/test/test_aqua.jl index 0463d2db..cede4561 100644 --- a/test/test_aqua.jl +++ b/test/test_aqua.jl @@ -3,6 +3,6 @@ using Aqua: Aqua using Test: @testset @testset "Code quality (Aqua.jl)" begin - # TODO: Reenable ambiguity and piracy checks. - Aqua.test_all(BlockSparseArrays; ambiguities=false, piracies=false) + # TODO: Reenable ambiguity and piracy checks. + Aqua.test_all(BlockSparseArrays; ambiguities = false, piracies = false) end diff --git a/test/test_basics.jl b/test/test_basics.jl index a464fa33..7627386f 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -1,46 +1,46 @@ using Adapt: adapt using ArrayLayouts: zero! using BlockArrays: - BlockArrays, - Block, - BlockArray, - BlockRange, - BlockVector, - BlockedOneTo, - BlockedArray, - BlockedVector, - blockedrange, - blocklength, - blocklengths, - blocksize, - blocksizes, - mortar, - undef_blocks + BlockArrays, + Block, + BlockArray, + BlockRange, + BlockVector, + BlockedOneTo, + BlockedArray, + BlockedVector, + blockedrange, + blocklength, + blocklengths, + blocksize, + blocksizes, + mortar, + undef_blocks using BlockSparseArrays: - @view!, - BlockSparseArray, - BlockSparseMatrix, - BlockSparseVector, - BlockType, - BlockView, - blockdiagindices, - blockreshape, - blocksparse, - blocksparsezeros, - blockstoredlength, - blockstype, - blocktype, - eachblockstoredindex, - eachstoredblock, - eachstoredblockdiagindex, - similartype_unchecked, - sparsemortar, - view! + @view!, + BlockSparseArray, + BlockSparseMatrix, + BlockSparseVector, + BlockType, + BlockView, + blockdiagindices, + blockreshape, + blocksparse, + blocksparsezeros, + blockstoredlength, + blockstype, + blocktype, + eachblockstoredindex, + eachstoredblock, + eachstoredblockdiagindex, + similartype_unchecked, + sparsemortar, + view! using GPUArraysCore: @allowscalar using JLArrays: JLArray, JLMatrix using LinearAlgebra: Adjoint, Transpose, dot, norm, tr using SparseArraysBase: - SparseArrayDOK, SparseMatrixDOK, SparseVectorDOK, isstored, storedlength + SparseArrayDOK, SparseMatrixDOK, SparseVectorDOK, isstored, storedlength using Test: @test, @test_broken, @test_throws, @testset, @inferred using TestExtras: @constinferred using TypeParameterAccessors: TypeParameterAccessors, Position @@ -49,490 +49,490 @@ include("TestBlockSparseArraysUtils.jl") elts = (Float32, Float64, ComplexF32) arrayts = (Array, JLArray) @testset "BlockSparseArrays basics (arraytype=$arrayt, eltype=$elt)" for arrayt in arrayts, - elt in elts - - dev = adapt(arrayt) - - @testset "Broken" begin - # TODO: Fix this and turn it into a proper test. - 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)) - @test_broken a[:, 4] - - # TODO: Fix this and turn it into a proper test. - 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)) - @allowscalar @test a[2:4, 4] == Array(a)[2:4, 4] - @test_broken a[4, 2:4] - - @test a[Block(1), :] isa BlockSparseArray{elt} - @test adjoint(a) isa Adjoint{elt,<:BlockSparseArray} - @test_broken adjoint(a)[Block(1), :] isa Adjoint{elt,<:BlockSparseArray} - # could also be directly a BlockSparseArray - end - @testset "Constructors" begin - # BlockSparseMatrix - bs = ([2, 3], [3, 4]) - for T in ( - BlockSparseArray{elt}, - BlockSparseArray{elt,2}, - BlockSparseMatrix{elt}, - BlockSparseArray{elt,2,Matrix{elt}}, - BlockSparseMatrix{elt,Matrix{elt}}, - ## BlockSparseArray{elt,2,Matrix{elt},SparseMatrixDOK{Matrix{elt}}}, # TODO - ## BlockSparseMatrix{elt,Matrix{elt},SparseMatrixDOK{Matrix{elt}}}, # TODO - ) - for args in ( - (undef, bs), - (undef, bs...), - (undef, blockedrange.(bs)), - (undef, blockedrange.(bs)...), - ) - a = T(args...) - @test eltype(a) == elt - @test blocktype(a) == Matrix{elt} - @test blockstype(a) <: SparseMatrixDOK{Matrix{elt}} - @test blocklengths.(axes(a)) == ([2, 3], [3, 4]) + elt in elts + + dev = adapt(arrayt) + + @testset "Broken" begin + # TODO: Fix this and turn it into a proper test. + 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)) + @test_broken a[:, 4] + + # TODO: Fix this and turn it into a proper test. + 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)) + @allowscalar @test a[2:4, 4] == Array(a)[2:4, 4] + @test_broken a[4, 2:4] + + @test a[Block(1), :] isa BlockSparseArray{elt} + @test adjoint(a) isa Adjoint{elt, <:BlockSparseArray} + @test_broken adjoint(a)[Block(1), :] isa Adjoint{elt, <:BlockSparseArray} + # could also be directly a BlockSparseArray + end + @testset "Constructors" begin + # BlockSparseMatrix + bs = ([2, 3], [3, 4]) + for T in ( + BlockSparseArray{elt}, + BlockSparseArray{elt, 2}, + BlockSparseMatrix{elt}, + BlockSparseArray{elt, 2, Matrix{elt}}, + BlockSparseMatrix{elt, Matrix{elt}}, + ## BlockSparseArray{elt,2,Matrix{elt},SparseMatrixDOK{Matrix{elt}}}, # TODO + ## BlockSparseMatrix{elt,Matrix{elt},SparseMatrixDOK{Matrix{elt}}}, # TODO + ) + for args in ( + (undef, bs), + (undef, bs...), + (undef, blockedrange.(bs)), + (undef, blockedrange.(bs)...), + ) + a = T(args...) + @test eltype(a) == elt + @test blocktype(a) == Matrix{elt} + @test blockstype(a) <: SparseMatrixDOK{Matrix{elt}} + @test blocklengths.(axes(a)) == ([2, 3], [3, 4]) + @test iszero(a) + @test iszero(blockstoredlength(a)) + @test iszero(storedlength(a)) + end + T != BlockSparseArray{elt} && @test_throws ArgumentError T(undef, bs[1:1]) + end + + # BlockSparseVector + bs = ([2, 3],) + for T in ( + BlockSparseArray{elt}, + BlockSparseArray{elt, 1}, + BlockSparseVector{elt}, + BlockSparseArray{elt, 1, Vector{elt}}, + BlockSparseVector{elt, Vector{elt}}, + ## BlockSparseArray{elt,1,Vector{elt},SparseVectorDOK{Vector{elt}}}, # TODO + ## BlockSparseVector{elt,Vector{elt},SparseVectorDOK{Vector{elt}}}, # TODO + ) + for args in ( + (undef, bs), + (undef, bs...), + (undef, blockedrange.(bs)), + (undef, blockedrange.(bs)...), + ) + a = T(args...) + @test eltype(a) == elt + @test blocktype(a) == Vector{elt} + @test blockstype(a) <: SparseVectorDOK{Vector{elt}} + @test blocklengths.(axes(a)) == ([2, 3],) + @test iszero(a) + @test iszero(blockstoredlength(a)) + @test iszero(storedlength(a)) + end + end + + for dims in ( + ([2, 2], [2, 2]), + (([2, 2], [2, 2]),), + blockedrange.(([2, 2], [2, 2])), + (blockedrange.(([2, 2], [2, 2])),), + ) + @test_throws ArgumentError BlockSparseVector{elt}(undef, dims...) + end + + # Convenient constructors. + a = blocksparsezeros([2, 3], [2, 3]) @test iszero(a) @test iszero(blockstoredlength(a)) - @test iszero(storedlength(a)) - end - T != BlockSparseArray{elt} && @test_throws ArgumentError T(undef, bs[1:1]) - end + @test a isa BlockSparseMatrix{Float64, Matrix{Float64}} + @test blocktype(a) == Matrix{Float64} + @test blockstype(a) <: SparseMatrixDOK{Matrix{Float64}} + @test blocksize(a) == (2, 2) + @test blocksizes(a) == [(2, 2) (2, 3); (3, 2) (3, 3)] - # BlockSparseVector - bs = ([2, 3],) - for T in ( - BlockSparseArray{elt}, - BlockSparseArray{elt,1}, - BlockSparseVector{elt}, - BlockSparseArray{elt,1,Vector{elt}}, - BlockSparseVector{elt,Vector{elt}}, - ## BlockSparseArray{elt,1,Vector{elt},SparseVectorDOK{Vector{elt}}}, # TODO - ## BlockSparseVector{elt,Vector{elt},SparseVectorDOK{Vector{elt}}}, # TODO - ) - for args in ( - (undef, bs), - (undef, bs...), - (undef, blockedrange.(bs)), - (undef, blockedrange.(bs)...), - ) - a = T(args...) - @test eltype(a) == elt - @test blocktype(a) == Vector{elt} - @test blockstype(a) <: SparseVectorDOK{Vector{elt}} - @test blocklengths.(axes(a)) == ([2, 3],) + a = blocksparsezeros(elt, [2, 3], [2, 3]) @test iszero(a) @test iszero(blockstoredlength(a)) - @test iszero(storedlength(a)) - end - end - - for dims in ( - ([2, 2], [2, 2]), - (([2, 2], [2, 2]),), - blockedrange.(([2, 2], [2, 2])), - (blockedrange.(([2, 2], [2, 2])),), - ) - @test_throws ArgumentError BlockSparseVector{elt}(undef, dims...) - end + @test a isa BlockSparseMatrix{elt, Matrix{elt}} + @test blocktype(a) == Matrix{elt} + @test blockstype(a) <: SparseMatrixDOK{Matrix{elt}} + @test blocksize(a) == (2, 2) + @test blocksizes(a) == [(2, 2) (2, 3); (3, 2) (3, 3)] - # Convenient constructors. - a = blocksparsezeros([2, 3], [2, 3]) - @test iszero(a) - @test iszero(blockstoredlength(a)) - @test a isa BlockSparseMatrix{Float64,Matrix{Float64}} - @test blocktype(a) == Matrix{Float64} - @test blockstype(a) <: SparseMatrixDOK{Matrix{Float64}} - @test blocksize(a) == (2, 2) - @test blocksizes(a) == [(2, 2) (2, 3); (3, 2) (3, 3)] - - a = blocksparsezeros(elt, [2, 3], [2, 3]) - @test iszero(a) - @test iszero(blockstoredlength(a)) - @test a isa BlockSparseMatrix{elt,Matrix{elt}} - @test blocktype(a) == Matrix{elt} - @test blockstype(a) <: SparseMatrixDOK{Matrix{elt}} - @test blocksize(a) == (2, 2) - @test blocksizes(a) == [(2, 2) (2, 3); (3, 2) (3, 3)] - - a = blocksparsezeros(BlockType(arrayt{elt,2}), [2, 3], [2, 3]) - @test iszero(a) - @test iszero(blockstoredlength(a)) - @test a isa BlockSparseMatrix{elt,arrayt{elt,2}} - @test blocktype(a) == arrayt{elt,2} - @test blockstype(a) <: SparseMatrixDOK{arrayt{elt,2}} - @test blocksize(a) == (2, 2) - @test blocksizes(a) == [(2, 2) (2, 3); (3, 2) (3, 3)] - - d = Dict(Block(1, 1) => dev(randn(elt, 2, 2)), Block(2, 2) => dev(randn(elt, 3, 3))) - a = blocksparse(d, [2, 3], [2, 3]) - @test !iszero(a) - @test a[Block(1, 1)] == d[Block(1, 1)] - @test a[Block(2, 2)] == d[Block(2, 2)] - @test blockstoredlength(a) == 2 - @test a isa BlockSparseMatrix{elt,arrayt{elt,2}} - @test blocktype(a) == arrayt{elt,2} - @test blockstype(a) <: SparseMatrixDOK{arrayt{elt,2}} - @test blocksize(a) == (2, 2) - @test blocksizes(a) == [(2, 2) (2, 3); (3, 2) (3, 3)] - end - @testset "blockstype, blocktype" begin - a = arrayt(randn(elt, 2, 2)) - @test (@constinferred blockstype(a)) <: BlockArrays.BlocksView{elt,2} - # TODO: This is difficult to determine just from type information. - @test_broken blockstype(typeof(a)) <: BlockArrays.BlocksView{elt,2} - @test (@constinferred blocktype(a)) <: SubArray{elt,2,arrayt{elt,2}} - # TODO: This is difficult to determine just from type information. - @test_broken blocktype(typeof(a)) <: SubArray{elt,2,arrayt{elt,2}} - - a = BlockSparseMatrix{elt,arrayt{elt,2}}(undef, [1, 1], [1, 1]) - @test (@constinferred blockstype(a)) <: SparseMatrixDOK{arrayt{elt,2}} - @test (@constinferred blockstype(typeof(a))) <: SparseMatrixDOK{arrayt{elt,2}} - @test (@constinferred blocktype(a)) <: arrayt{elt,2} - @test (@constinferred blocktype(typeof(a))) <: arrayt{elt,2} - - a = BlockArray(arrayt(randn(elt, (2, 2))), [1, 1], [1, 1]) - @test (@constinferred blockstype(a)) === Matrix{arrayt{elt,2}} - @test (@constinferred blockstype(typeof(a))) === Matrix{arrayt{elt,2}} - @test (@constinferred blocktype(a)) <: arrayt{elt,2} - @test (@constinferred blocktype(typeof(a))) <: arrayt{elt,2} - - a = BlockedArray(arrayt(randn(elt, 2, 2)), [1, 1], [1, 1]) - @test (@constinferred blockstype(a)) <: BlockArrays.BlocksView{elt,2} - # TODO: This is difficult to determine just from type information. - @test_broken blockstype(typeof(a)) <: BlockArrays.BlocksView{elt,2} - @test (@constinferred blocktype(a)) <: SubArray{elt,2,arrayt{elt,2}} - # TODO: This is difficult to determine just from type information. - @test_broken blocktype(typeof(a)) <: SubArray{elt,2,arrayt{elt,2}} - - # sparsemortar - for ax in ( - ([2, 3], [2, 3]), - (([2, 3], [2, 3]),), - blockedrange.(([2, 3], [2, 3])), - (blockedrange.(([2, 3], [2, 3])),), - ) - blocks = SparseArrayDOK{arrayt{elt,2}}(undef_blocks, ax...) - blocks[2, 1] = arrayt(randn(elt, 3, 2)) - blocks[1, 2] = arrayt(randn(elt, 2, 3)) - a = sparsemortar(blocks, ax...) - @test a isa BlockSparseArray{elt,2,arrayt{elt,2}} - @test iszero(a[Block(1, 1)]) - @test a[Block(2, 1)] == blocks[2, 1] - @test a[Block(1, 2)] == blocks[1, 2] - @test iszero(a[Block(2, 2)]) - end - end - @testset "Basics" begin - a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - @allowscalar @test a == dev( - BlockSparseArray{elt}(undef, blockedrange([2, 3]), blockedrange([2, 3])) - ) - @test eltype(a) === elt - @test axes(a) == (1:5, 1:5) - @test all(aᵢ -> aᵢ isa BlockedOneTo, axes(a)) - @test blocklength.(axes(a)) == (2, 2) - @test blocksize(a) == (2, 2) - @test size(a) == (5, 5) - @test blockstoredlength(a) == 0 - @test iszero(a) - @allowscalar @test all(I -> iszero(a[I]), eachindex(a)) - @test_throws DimensionMismatch a[Block(1, 1)] = randn(elt, 2, 3) - - a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) - a[3, 3] = 33 - @test eltype(a) === elt - @test axes(a) == (1:5, 1:5) - @test all(aᵢ -> aᵢ isa BlockedOneTo, axes(a)) - @test blocklength.(axes(a)) == (2, 2) - @test blocksize(a) == (2, 2) - @test size(a) == (5, 5) - @test blockstoredlength(a) == 1 - @test !iszero(a) - @test a[3, 3] == 33 - @test all(eachindex(a)) do I - if I == CartesianIndex(3, 3) - a[I] == 33 - else - iszero(a[I]) - end + a = blocksparsezeros(BlockType(arrayt{elt, 2}), [2, 3], [2, 3]) + @test iszero(a) + @test iszero(blockstoredlength(a)) + @test a isa BlockSparseMatrix{elt, arrayt{elt, 2}} + @test blocktype(a) == arrayt{elt, 2} + @test blockstype(a) <: SparseMatrixDOK{arrayt{elt, 2}} + @test blocksize(a) == (2, 2) + @test blocksizes(a) == [(2, 2) (2, 3); (3, 2) (3, 3)] + + d = Dict(Block(1, 1) => dev(randn(elt, 2, 2)), Block(2, 2) => dev(randn(elt, 3, 3))) + a = blocksparse(d, [2, 3], [2, 3]) + @test !iszero(a) + @test a[Block(1, 1)] == d[Block(1, 1)] + @test a[Block(2, 2)] == d[Block(2, 2)] + @test blockstoredlength(a) == 2 + @test a isa BlockSparseMatrix{elt, arrayt{elt, 2}} + @test blocktype(a) == arrayt{elt, 2} + @test blockstype(a) <: SparseMatrixDOK{arrayt{elt, 2}} + @test blocksize(a) == (2, 2) + @test blocksizes(a) == [(2, 2) (2, 3); (3, 2) (3, 3)] end - - a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) - a[Block(2, 1)] = randn(elt, 3, 2) - a[Block(1, 2)] = randn(elt, 2, 3) - @test issetequal(eachstoredblock(a), [a[Block(2, 1)], a[Block(1, 2)]]) - @test issetequal(eachblockstoredindex(a), [Block(2, 1), Block(1, 2)]) - @test issetequal(blockdiagindices(a), [Block(1, 1), Block(2, 2)]) - @test isempty(eachstoredblockdiagindex(a)) - @test norm(a) ≈ norm(Array(a)) - for p in 1:3 - @test norm(a, p) ≈ norm(Array(a), p) + @testset "blockstype, blocktype" begin + a = arrayt(randn(elt, 2, 2)) + @test (@constinferred blockstype(a)) <: BlockArrays.BlocksView{elt, 2} + # TODO: This is difficult to determine just from type information. + @test_broken blockstype(typeof(a)) <: BlockArrays.BlocksView{elt, 2} + @test (@constinferred blocktype(a)) <: SubArray{elt, 2, arrayt{elt, 2}} + # TODO: This is difficult to determine just from type information. + @test_broken blocktype(typeof(a)) <: SubArray{elt, 2, arrayt{elt, 2}} + + a = BlockSparseMatrix{elt, arrayt{elt, 2}}(undef, [1, 1], [1, 1]) + @test (@constinferred blockstype(a)) <: SparseMatrixDOK{arrayt{elt, 2}} + @test (@constinferred blockstype(typeof(a))) <: SparseMatrixDOK{arrayt{elt, 2}} + @test (@constinferred blocktype(a)) <: arrayt{elt, 2} + @test (@constinferred blocktype(typeof(a))) <: arrayt{elt, 2} + + a = BlockArray(arrayt(randn(elt, (2, 2))), [1, 1], [1, 1]) + @test (@constinferred blockstype(a)) === Matrix{arrayt{elt, 2}} + @test (@constinferred blockstype(typeof(a))) === Matrix{arrayt{elt, 2}} + @test (@constinferred blocktype(a)) <: arrayt{elt, 2} + @test (@constinferred blocktype(typeof(a))) <: arrayt{elt, 2} + + a = BlockedArray(arrayt(randn(elt, 2, 2)), [1, 1], [1, 1]) + @test (@constinferred blockstype(a)) <: BlockArrays.BlocksView{elt, 2} + # TODO: This is difficult to determine just from type information. + @test_broken blockstype(typeof(a)) <: BlockArrays.BlocksView{elt, 2} + @test (@constinferred blocktype(a)) <: SubArray{elt, 2, arrayt{elt, 2}} + # TODO: This is difficult to determine just from type information. + @test_broken blocktype(typeof(a)) <: SubArray{elt, 2, arrayt{elt, 2}} + + # sparsemortar + for ax in ( + ([2, 3], [2, 3]), + (([2, 3], [2, 3]),), + blockedrange.(([2, 3], [2, 3])), + (blockedrange.(([2, 3], [2, 3])),), + ) + blocks = SparseArrayDOK{arrayt{elt, 2}}(undef_blocks, ax...) + blocks[2, 1] = arrayt(randn(elt, 3, 2)) + blocks[1, 2] = arrayt(randn(elt, 2, 3)) + a = sparsemortar(blocks, ax...) + @test a isa BlockSparseArray{elt, 2, arrayt{elt, 2}} + @test iszero(a[Block(1, 1)]) + @test a[Block(2, 1)] == blocks[2, 1] + @test a[Block(1, 2)] == blocks[1, 2] + @test iszero(a[Block(2, 2)]) + end end - @test tr(a) ≈ tr(Array(a)) - - a[3, 3] = NaN - @test isnan(norm(a)) - - a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - a[Block(1, 1)] = dev(randn(elt, 2, 2)) - @test issetequal(eachstoredblockdiagindex(a), [Block(1, 1)]) - @test norm(a) ≈ norm(Array(a)) - for p in 1:3 - @test norm(a, p) ≈ norm(Array(a), p) + @testset "Basics" begin + a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + @allowscalar @test a == dev( + BlockSparseArray{elt}(undef, blockedrange([2, 3]), blockedrange([2, 3])) + ) + @test eltype(a) === elt + @test axes(a) == (1:5, 1:5) + @test all(aᵢ -> aᵢ isa BlockedOneTo, axes(a)) + @test blocklength.(axes(a)) == (2, 2) + @test blocksize(a) == (2, 2) + @test size(a) == (5, 5) + @test blockstoredlength(a) == 0 + @test iszero(a) + @allowscalar @test all(I -> iszero(a[I]), eachindex(a)) + @test_throws DimensionMismatch a[Block(1, 1)] = randn(elt, 2, 3) + + a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) + a[3, 3] = 33 + @test eltype(a) === elt + @test axes(a) == (1:5, 1:5) + @test all(aᵢ -> aᵢ isa BlockedOneTo, axes(a)) + @test blocklength.(axes(a)) == (2, 2) + @test blocksize(a) == (2, 2) + @test size(a) == (5, 5) + @test blockstoredlength(a) == 1 + @test !iszero(a) + @test a[3, 3] == 33 + @test all(eachindex(a)) do I + if I == CartesianIndex(3, 3) + a[I] == 33 + else + iszero(a[I]) + end + end + + a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) + a[Block(2, 1)] = randn(elt, 3, 2) + a[Block(1, 2)] = randn(elt, 2, 3) + @test issetequal(eachstoredblock(a), [a[Block(2, 1)], a[Block(1, 2)]]) + @test issetequal(eachblockstoredindex(a), [Block(2, 1), Block(1, 2)]) + @test issetequal(blockdiagindices(a), [Block(1, 1), Block(2, 2)]) + @test isempty(eachstoredblockdiagindex(a)) + @test norm(a) ≈ norm(Array(a)) + for p in 1:3 + @test norm(a, p) ≈ norm(Array(a), p) + end + @test tr(a) ≈ tr(Array(a)) + + a[3, 3] = NaN + @test isnan(norm(a)) + + a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + @test issetequal(eachstoredblockdiagindex(a), [Block(1, 1)]) + @test norm(a) ≈ norm(Array(a)) + for p in 1:3 + @test norm(a, p) ≈ norm(Array(a), p) + end + @test tr(a) ≈ tr(Array(a)) + + # Empty constructor + for a in (dev(BlockSparseArray{elt}(undef)),) + @test size(a) == () + @test isone(length(a)) + @test blocksize(a) == () + @test blocksizes(a) == fill(()) + @test iszero(blockstoredlength(a)) + @test iszero(@allowscalar(a[])) + @test iszero(@allowscalar(a[CartesianIndex()])) + @test a[Block()] == dev(fill(0)) + @test iszero(@allowscalar(a[Block()][])) + @test iszero(@allowscalar(a[Block()[]])) + @test Array(a) isa Array{elt, 0} + @test Array(a) == fill(0) + for b in ( + (b = copy(a); @allowscalar(b[] = 2); b), + (b = copy(a); @allowscalar(b[CartesianIndex()] = 2); b), + (b = copy(a); @allowscalar(b[Block()[]] = 2); b), + # Regression test for https://github.com/ITensor/BlockSparseArrays.jl/issues/27. + (b = copy(a); b[Block()] = dev(fill(2)); b), + ) + @test size(b) == () + @test isone(length(b)) + @test blocksize(b) == () + @test blocksizes(b) == fill(()) + @test isone(blockstoredlength(b)) + @test @allowscalar(b[]) == 2 + @test @allowscalar(b[CartesianIndex()]) == 2 + @test b[Block()] == dev(fill(2)) + @test @allowscalar(b[Block()][]) == 2 + @test @allowscalar(b[Block()[]]) == 2 + @test Array(b) isa Array{elt, 0} + @test Array(b) == fill(2) + end + end + + @testset "Transpose" begin + a = dev(BlockSparseArray{elt}(undef, [2, 2], [3, 3, 1])) + a[Block(1, 1)] = dev(randn(elt, 2, 3)) + a[Block(2, 3)] = dev(randn(elt, 2, 1)) + + at = @inferred transpose(a) + @test at isa Transpose + @test size(at) == reverse(size(a)) + @test blocksize(at) == reverse(blocksize(a)) + @test storedlength(at) == storedlength(a) + @test blockstoredlength(at) == blockstoredlength(a) + for bind in eachblockstoredindex(a) + bindt = Block(reverse(Int.(Tuple(bind)))) + @test bindt in eachblockstoredindex(at) + end + + @test @views(at[Block(1, 1)]) == transpose(a[Block(1, 1)]) + @test @views(at[Block(1, 1)]) isa Transpose + @test @views(at[Block(3, 2)]) == transpose(a[Block(2, 3)]) + # TODO: BlockView == AbstractArray calls scalar code + @test @allowscalar @views(at[Block(1, 2)]) == transpose(a[Block(2, 1)]) + @test @views(at[Block(1, 2)]) isa Transpose + end + + @testset "Adjoint" begin + a = dev(BlockSparseArray{elt}(undef, [2, 2], [3, 3, 1])) + a[Block(1, 1)] = dev(randn(elt, 2, 3)) + a[Block(2, 3)] = dev(randn(elt, 2, 1)) + + at = @inferred adjoint(a) + @test at isa Adjoint + @test size(at) == reverse(size(a)) + @test blocksize(at) == reverse(blocksize(a)) + @test storedlength(at) == storedlength(a) + @test blockstoredlength(at) == blockstoredlength(a) + for bind in eachblockstoredindex(a) + bindt = Block(reverse(Int.(Tuple(bind)))) + @test bindt in eachblockstoredindex(at) + end + + @test @views(at[Block(1, 1)]) == adjoint(a[Block(1, 1)]) + @test @views(at[Block(1, 1)]) isa Adjoint + @test @views(at[Block(3, 2)]) == adjoint(a[Block(2, 3)]) + # TODO: BlockView == AbstractArray calls scalar code + @test @allowscalar @views(at[Block(1, 2)]) == adjoint(a[Block(2, 1)]) + @test @views(at[Block(1, 2)]) isa Adjoint + end end - @test tr(a) ≈ tr(Array(a)) - - # Empty constructor - for a in (dev(BlockSparseArray{elt}(undef)),) - @test size(a) == () - @test isone(length(a)) - @test blocksize(a) == () - @test blocksizes(a) == fill(()) - @test iszero(blockstoredlength(a)) - @test iszero(@allowscalar(a[])) - @test iszero(@allowscalar(a[CartesianIndex()])) - @test a[Block()] == dev(fill(0)) - @test iszero(@allowscalar(a[Block()][])) - @test iszero(@allowscalar(a[Block()[]])) - @test Array(a) isa Array{elt,0} - @test Array(a) == fill(0) - for b in ( - (b=copy(a); @allowscalar(b[] = 2); b), - (b=copy(a); @allowscalar(b[CartesianIndex()] = 2); b), - (b=copy(a); @allowscalar(b[Block()[]] = 2); b), - # Regression test for https://github.com/ITensor/BlockSparseArrays.jl/issues/27. - (b=copy(a); b[Block()]=dev(fill(2)); b), - ) - @test size(b) == () - @test isone(length(b)) - @test blocksize(b) == () - @test blocksizes(b) == fill(()) - @test isone(blockstoredlength(b)) - @test @allowscalar(b[]) == 2 - @test @allowscalar(b[CartesianIndex()]) == 2 - @test b[Block()] == dev(fill(2)) - @test @allowscalar(b[Block()][]) == 2 - @test @allowscalar(b[Block()[]]) == 2 - @test Array(b) isa Array{elt,0} - @test Array(b) == fill(2) - end + @testset "adapt" begin + a = BlockSparseArray{elt}(undef, [2, 2], [2, 2]) + a_12 = randn(elt, 2, 2) + a[Block(1, 2)] = a_12 + a_jl = adapt(JLArray, a) + @test a_jl isa BlockSparseMatrix{elt, JLMatrix{elt}} + @test blocktype(a_jl) == JLMatrix{elt} + @test blockstoredlength(a_jl) == 1 + @test a_jl[Block(1, 2)] isa JLMatrix{elt} + @test adapt(Array, a_jl[Block(1, 2)]) == a_12 + + a = BlockSparseArray{elt}(undef, [2, 2], [2, 2]) + a_12 = randn(elt, 2, 2) + a[Block(1, 2)] = a_12 + a_jl = adapt(JLArray, @view(a[:, :])) + @test a_jl isa SubArray{elt, 2, <:BlockSparseMatrix{elt, JLMatrix{elt}}} + @test blocktype(a_jl) == JLMatrix{elt} + @test blockstoredlength(a_jl) == 1 + @test a_jl[Block(1, 2)] isa JLMatrix{elt} + @test adapt(Array, a_jl[Block(1, 2)]) == a_12 end - - @testset "Transpose" begin - a = dev(BlockSparseArray{elt}(undef, [2, 2], [3, 3, 1])) - a[Block(1, 1)] = dev(randn(elt, 2, 3)) - a[Block(2, 3)] = dev(randn(elt, 2, 1)) - - at = @inferred transpose(a) - @test at isa Transpose - @test size(at) == reverse(size(a)) - @test blocksize(at) == reverse(blocksize(a)) - @test storedlength(at) == storedlength(a) - @test blockstoredlength(at) == blockstoredlength(a) - for bind in eachblockstoredindex(a) - bindt = Block(reverse(Int.(Tuple(bind)))) - @test bindt in eachblockstoredindex(at) - end - - @test @views(at[Block(1, 1)]) == transpose(a[Block(1, 1)]) - @test @views(at[Block(1, 1)]) isa Transpose - @test @views(at[Block(3, 2)]) == transpose(a[Block(2, 3)]) - # TODO: BlockView == AbstractArray calls scalar code - @test @allowscalar @views(at[Block(1, 2)]) == transpose(a[Block(2, 1)]) - @test @views(at[Block(1, 2)]) isa Transpose + @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)])))) + a2 = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + a2[Block(1, 1)] = dev(randn(elt, size(@view(a1[Block(1, 1)])))) + a_dest = a1 * a2 + @allowscalar @test Array(a_dest) ≈ Array(a1) * Array(a2) + @test a_dest isa BlockSparseArray{elt} + @test blockstoredlength(a_dest) == 1 end - - @testset "Adjoint" begin - a = dev(BlockSparseArray{elt}(undef, [2, 2], [3, 3, 1])) - a[Block(1, 1)] = dev(randn(elt, 2, 3)) - a[Block(2, 3)] = dev(randn(elt, 2, 1)) - - at = @inferred adjoint(a) - @test at isa Adjoint - @test size(at) == reverse(size(a)) - @test blocksize(at) == reverse(blocksize(a)) - @test storedlength(at) == storedlength(a) - @test blockstoredlength(at) == blockstoredlength(a) - for bind in eachblockstoredindex(a) - bindt = Block(reverse(Int.(Tuple(bind)))) - @test bindt in eachblockstoredindex(at) - end - - @test @views(at[Block(1, 1)]) == adjoint(a[Block(1, 1)]) - @test @views(at[Block(1, 1)]) isa Adjoint - @test @views(at[Block(3, 2)]) == adjoint(a[Block(2, 3)]) - # TODO: BlockView == AbstractArray calls scalar code - @test @allowscalar @views(at[Block(1, 2)]) == adjoint(a[Block(2, 1)]) - @test @views(at[Block(1, 2)]) isa Adjoint + @testset "Matrix multiplication" begin + a1 = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + a1[Block(1, 2)] = dev(randn(elt, size(@view(a1[Block(1, 2)])))) + a1[Block(2, 1)] = dev(randn(elt, size(@view(a1[Block(2, 1)])))) + a2 = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + a2[Block(1, 2)] = dev(randn(elt, size(@view(a2[Block(1, 2)])))) + a2[Block(2, 1)] = dev(randn(elt, size(@view(a2[Block(2, 1)])))) + for (a1′, a2′) in ((a1, a2), (a1', a2), (a1, a2'), (a1', a2')) + a_dest = a1′ * a2′ + @allowscalar @test Array(a_dest) ≈ Array(a1′) * Array(a2′) + end end - end - @testset "adapt" begin - a = BlockSparseArray{elt}(undef, [2, 2], [2, 2]) - a_12 = randn(elt, 2, 2) - a[Block(1, 2)] = a_12 - a_jl = adapt(JLArray, a) - @test a_jl isa BlockSparseMatrix{elt,JLMatrix{elt}} - @test blocktype(a_jl) == JLMatrix{elt} - @test blockstoredlength(a_jl) == 1 - @test a_jl[Block(1, 2)] isa JLMatrix{elt} - @test adapt(Array, a_jl[Block(1, 2)]) == a_12 - - a = BlockSparseArray{elt}(undef, [2, 2], [2, 2]) - a_12 = randn(elt, 2, 2) - a[Block(1, 2)] = a_12 - a_jl = adapt(JLArray, @view(a[:, :])) - @test a_jl isa SubArray{elt,2,<:BlockSparseMatrix{elt,JLMatrix{elt}}} - @test blocktype(a_jl) == JLMatrix{elt} - @test blockstoredlength(a_jl) == 1 - @test a_jl[Block(1, 2)] isa JLMatrix{elt} - @test adapt(Array, a_jl[Block(1, 2)]) == a_12 - 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)])))) - a2 = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - a2[Block(1, 1)] = dev(randn(elt, size(@view(a1[Block(1, 1)])))) - a_dest = a1 * a2 - @allowscalar @test Array(a_dest) ≈ Array(a1) * Array(a2) - @test a_dest isa BlockSparseArray{elt} - @test blockstoredlength(a_dest) == 1 - end - @testset "Matrix multiplication" begin - a1 = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - a1[Block(1, 2)] = dev(randn(elt, size(@view(a1[Block(1, 2)])))) - a1[Block(2, 1)] = dev(randn(elt, size(@view(a1[Block(2, 1)])))) - a2 = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - a2[Block(1, 2)] = dev(randn(elt, size(@view(a2[Block(1, 2)])))) - a2[Block(2, 1)] = dev(randn(elt, size(@view(a2[Block(2, 1)])))) - for (a1′, a2′) in ((a1, a2), (a1', a2), (a1, a2'), (a1', a2')) - a_dest = a1′ * a2′ - @allowscalar @test Array(a_dest) ≈ Array(a1′) * Array(a2′) + @testset "Dot product" begin + a1 = dev(BlockSparseArray{elt}(undef, [2, 3, 4])) + a1[Block(1)] = dev(randn(elt, size(@view(a1[Block(1)])))) + a1[Block(3)] = dev(randn(elt, size(@view(a1[Block(3)])))) + a2 = dev(BlockSparseArray{elt}(undef, [2, 3, 4])) + a2[Block(2)] = dev(randn(elt, size(@view(a1[Block(2)])))) + a2[Block(3)] = dev(randn(elt, size(@view(a1[Block(3)])))) + @test a1' * a2 ≈ Array(a1)' * Array(a2) + @test dot(a1, a2) ≈ a1' * a2 end - end - @testset "Dot product" begin - a1 = dev(BlockSparseArray{elt}(undef, [2, 3, 4])) - a1[Block(1)] = dev(randn(elt, size(@view(a1[Block(1)])))) - a1[Block(3)] = dev(randn(elt, size(@view(a1[Block(3)])))) - a2 = dev(BlockSparseArray{elt}(undef, [2, 3, 4])) - a2[Block(2)] = dev(randn(elt, size(@view(a1[Block(2)])))) - a2[Block(3)] = dev(randn(elt, size(@view(a1[Block(3)])))) - @test a1' * a2 ≈ Array(a1)' * Array(a2) - @test dot(a1, a2) ≈ a1' * a2 - end - @testset "cat" begin - a1 = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - a1[Block(2, 1)] = dev(randn(elt, size(@view(a1[Block(2, 1)])))) - a2 = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - a2[Block(1, 2)] = dev(randn(elt, size(@view(a2[Block(1, 2)])))) - - a_dest = cat(a1, a2; dims=1) - @test blockstoredlength(a_dest) == 2 - @test blocklengths.(axes(a_dest)) == ([2, 3, 2, 3], [2, 3]) - @test issetequal(eachblockstoredindex(a_dest), [Block(2, 1), Block(3, 2)]) - @test a_dest[Block(2, 1)] == a1[Block(2, 1)] - @test a_dest[Block(3, 2)] == a2[Block(1, 2)] - - a_dest = cat(a1, a2; dims=2) - @test blockstoredlength(a_dest) == 2 - @test blocklengths.(axes(a_dest)) == ([2, 3], [2, 3, 2, 3]) - @test issetequal(eachblockstoredindex(a_dest), [Block(2, 1), Block(1, 4)]) - @test a_dest[Block(2, 1)] == a1[Block(2, 1)] - @test a_dest[Block(1, 4)] == a2[Block(1, 2)] - - a_dest = cat(a1, a2; dims=(1, 2)) - @test blockstoredlength(a_dest) == 2 - @test blocklengths.(axes(a_dest)) == ([2, 3, 2, 3], [2, 3, 2, 3]) - @test issetequal(eachblockstoredindex(a_dest), [Block(2, 1), Block(3, 4)]) - @test a_dest[Block(2, 1)] == a1[Block(2, 1)] - @test a_dest[Block(3, 4)] == a2[Block(1, 2)] - end - - @testset "blockreshape" begin - a = dev(BlockSparseArray{elt}(undef, ([3, 4], [2, 3]))) - a[Block(1, 2)] = dev(randn(elt, size(@view(a[Block(1, 2)])))) - a[Block(2, 1)] = dev(randn(elt, size(@view(a[Block(2, 1)])))) - b = blockreshape(a, [6, 8, 9, 12]) - @test reshape(a[Block(1, 2)], 9) == b[Block(3)] - @test reshape(a[Block(2, 1)], 8) == b[Block(2)] - @test blockstoredlength(b) == 2 - @test storedlength(b) == 17 - - # Zero-dimensional limit (check for ambiguity errors). - # Regression test for https://github.com/ITensor/BlockSparseArrays.jl/issues/98. - a = dev(BlockSparseArray{elt}(undef, ())) - a[Block()] = dev(randn(elt, ())) - b = blockreshape(a) - @test a[Block()] == b[Block()] - @test blockstoredlength(b) == 1 - @test storedlength(b) == 1 - end - @testset "show" begin - vectort_elt = arrayt{elt,1} - matrixt_elt = arrayt{elt,2} - arrayt_elt = arrayt{elt,3} - - a = BlockSparseVector{elt,arrayt{elt,1}}(undef, [2, 2]) - res = sprint(summary, a) - function ref_vec(elt, arrayt, prefix="") - return "2-blocked 4-element $(prefix)BlockSparseVector{$(elt), $(arrayt), …, …}" + @testset "cat" begin + a1 = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + a1[Block(2, 1)] = dev(randn(elt, size(@view(a1[Block(2, 1)])))) + a2 = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + a2[Block(1, 2)] = dev(randn(elt, size(@view(a2[Block(1, 2)])))) + + a_dest = cat(a1, a2; dims = 1) + @test blockstoredlength(a_dest) == 2 + @test blocklengths.(axes(a_dest)) == ([2, 3, 2, 3], [2, 3]) + @test issetequal(eachblockstoredindex(a_dest), [Block(2, 1), Block(3, 2)]) + @test a_dest[Block(2, 1)] == a1[Block(2, 1)] + @test a_dest[Block(3, 2)] == a2[Block(1, 2)] + + a_dest = cat(a1, a2; dims = 2) + @test blockstoredlength(a_dest) == 2 + @test blocklengths.(axes(a_dest)) == ([2, 3], [2, 3, 2, 3]) + @test issetequal(eachblockstoredindex(a_dest), [Block(2, 1), Block(1, 4)]) + @test a_dest[Block(2, 1)] == a1[Block(2, 1)] + @test a_dest[Block(1, 4)] == a2[Block(1, 2)] + + a_dest = cat(a1, a2; dims = (1, 2)) + @test blockstoredlength(a_dest) == 2 + @test blocklengths.(axes(a_dest)) == ([2, 3, 2, 3], [2, 3, 2, 3]) + @test issetequal(eachblockstoredindex(a_dest), [Block(2, 1), Block(3, 4)]) + @test a_dest[Block(2, 1)] == a1[Block(2, 1)] + @test a_dest[Block(3, 4)] == a2[Block(1, 2)] end - # Either option is possible depending on namespacing. - @test (res == ref_vec(elt, vectort_elt)) || - (res == ref_vec(elt, vectort_elt, "BlockSparseArrays.")) - - a = BlockSparseMatrix{elt,arrayt{elt,2}}(undef, [2, 2], [2, 2]) - res = sprint(summary, a) - function ref_mat(elt, arrayt, prefix="") - return "2×2-blocked 4×4 $(prefix)BlockSparseMatrix{$(elt), $(arrayt), …, …}" + + @testset "blockreshape" begin + a = dev(BlockSparseArray{elt}(undef, ([3, 4], [2, 3]))) + a[Block(1, 2)] = dev(randn(elt, size(@view(a[Block(1, 2)])))) + a[Block(2, 1)] = dev(randn(elt, size(@view(a[Block(2, 1)])))) + b = blockreshape(a, [6, 8, 9, 12]) + @test reshape(a[Block(1, 2)], 9) == b[Block(3)] + @test reshape(a[Block(2, 1)], 8) == b[Block(2)] + @test blockstoredlength(b) == 2 + @test storedlength(b) == 17 + + # Zero-dimensional limit (check for ambiguity errors). + # Regression test for https://github.com/ITensor/BlockSparseArrays.jl/issues/98. + a = dev(BlockSparseArray{elt}(undef, ())) + a[Block()] = dev(randn(elt, ())) + b = blockreshape(a) + @test a[Block()] == b[Block()] + @test blockstoredlength(b) == 1 + @test storedlength(b) == 1 end - # Either option is possible depending on namespacing. - @test (res == ref_mat(elt, matrixt_elt)) || - (res == ref_mat(elt, matrixt_elt, "BlockSparseArrays.")) - - a = BlockSparseArray{elt,3,arrayt{elt,3}}(undef, [2, 2], [2, 2], [2, 2]) - res = sprint(summary, a) - function ref_arr(elt, arrayt, prefix="") - return "2×2×2-blocked 4×4×4 $(prefix)BlockSparseArray{$(elt), 3, $(arrayt), …, …}" + @testset "show" begin + vectort_elt = arrayt{elt, 1} + matrixt_elt = arrayt{elt, 2} + arrayt_elt = arrayt{elt, 3} + + a = BlockSparseVector{elt, arrayt{elt, 1}}(undef, [2, 2]) + res = sprint(summary, a) + function ref_vec(elt, arrayt, prefix = "") + return "2-blocked 4-element $(prefix)BlockSparseVector{$(elt), $(arrayt), …, …}" + end + # Either option is possible depending on namespacing. + @test (res == ref_vec(elt, vectort_elt)) || + (res == ref_vec(elt, vectort_elt, "BlockSparseArrays.")) + + a = BlockSparseMatrix{elt, arrayt{elt, 2}}(undef, [2, 2], [2, 2]) + res = sprint(summary, a) + function ref_mat(elt, arrayt, prefix = "") + return "2×2-blocked 4×4 $(prefix)BlockSparseMatrix{$(elt), $(arrayt), …, …}" + end + # Either option is possible depending on namespacing. + @test (res == ref_mat(elt, matrixt_elt)) || + (res == ref_mat(elt, matrixt_elt, "BlockSparseArrays.")) + + a = BlockSparseArray{elt, 3, arrayt{elt, 3}}(undef, [2, 2], [2, 2], [2, 2]) + res = sprint(summary, a) + function ref_arr(elt, arrayt, prefix = "") + return "2×2×2-blocked 4×4×4 $(prefix)BlockSparseArray{$(elt), 3, $(arrayt), …, …}" + end + @test (res == ref_arr(elt, arrayt_elt)) || + (res == ref_arr(elt, arrayt_elt, "BlockSparseArrays.")) + + if elt === Float64 + # Not testing other element types since they change the + # spacing so it isn't easy to make the test general. + + a′ = BlockSparseVector{elt, arrayt{elt, 1}}(undef, [2, 2]) + @allowscalar a′[1] = 1 + a = a′ + @test sprint(show, "text/plain", a) == + "$(summary(a)):\n $(eltype(a)(1))\n $(zero(eltype(a)))\n ───\n ⋅ \n ⋅ " + a = @view a′[:] + @test sprint(show, "text/plain", a) == + "$(summary(a)):\n $(eltype(a)(1))\n $(zero(eltype(a)))\n ⋅ \n ⋅ " + + a′ = BlockSparseMatrix{elt, arrayt{elt, 2}}(undef, [2, 2], [2, 2]) + @allowscalar a′[1, 2] = 12 + for a in (a′, @view(a′[:, :])) + @test sprint(show, "text/plain", a) == + "$(summary(a)):\n $(zero(eltype(a))) $(eltype(a)(12)) │ ⋅ ⋅ \n $(zero(eltype(a))) $(zero(eltype(a))) │ ⋅ ⋅ \n ───────────┼──────────\n ⋅ ⋅ │ ⋅ ⋅ \n ⋅ ⋅ │ ⋅ ⋅ " + end + + a′ = BlockSparseArray{elt, 3, arrayt{elt, 3}}(undef, [2, 2], [2, 2], [2, 2]) + @allowscalar a′[1, 2, 1] = 121 + for a in (a′, @view(a′[:, :, :])) + @test sprint(show, "text/plain", a) == + "$(summary(a)):\n[:, :, 1] =\n $(zero(eltype(a))) $(eltype(a)(121)) ⋅ ⋅ \n $(zero(eltype(a))) $(zero(eltype(a))) ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n\n[:, :, 2] =\n $(zero(eltype(a))) $(zero(eltype(a))) ⋅ ⋅ \n $(zero(eltype(a))) $(zero(eltype(a))) ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n\n[:, :, 3] =\n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n\n[:, :, 4] =\n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ " + end + end end - @test (res == ref_arr(elt, arrayt_elt)) || - (res == ref_arr(elt, arrayt_elt, "BlockSparseArrays.")) - - if elt === Float64 - # Not testing other element types since they change the - # spacing so it isn't easy to make the test general. - - a′ = BlockSparseVector{elt,arrayt{elt,1}}(undef, [2, 2]) - @allowscalar a′[1] = 1 - a = a′ - @test sprint(show, "text/plain", a) == - "$(summary(a)):\n $(eltype(a)(1))\n $(zero(eltype(a)))\n ───\n ⋅ \n ⋅ " - a = @view a′[:] - @test sprint(show, "text/plain", a) == - "$(summary(a)):\n $(eltype(a)(1))\n $(zero(eltype(a)))\n ⋅ \n ⋅ " - - a′ = BlockSparseMatrix{elt,arrayt{elt,2}}(undef, [2, 2], [2, 2]) - @allowscalar a′[1, 2] = 12 - for a in (a′, @view(a′[:, :])) - @test sprint(show, "text/plain", a) == - "$(summary(a)):\n $(zero(eltype(a))) $(eltype(a)(12)) │ ⋅ ⋅ \n $(zero(eltype(a))) $(zero(eltype(a))) │ ⋅ ⋅ \n ───────────┼──────────\n ⋅ ⋅ │ ⋅ ⋅ \n ⋅ ⋅ │ ⋅ ⋅ " - end - - a′ = BlockSparseArray{elt,3,arrayt{elt,3}}(undef, [2, 2], [2, 2], [2, 2]) - @allowscalar a′[1, 2, 1] = 121 - for a in (a′, @view(a′[:, :, :])) - @test sprint(show, "text/plain", a) == - "$(summary(a)):\n[:, :, 1] =\n $(zero(eltype(a))) $(eltype(a)(121)) ⋅ ⋅ \n $(zero(eltype(a))) $(zero(eltype(a))) ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n\n[:, :, 2] =\n $(zero(eltype(a))) $(zero(eltype(a))) ⋅ ⋅ \n $(zero(eltype(a))) $(zero(eltype(a))) ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n\n[:, :, 3] =\n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n\n[:, :, 4] =\n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ \n ⋅ ⋅ ⋅ ⋅ " - end + @testset "TypeParameterAccessors.position" begin + @test TypeParameterAccessors.position(BlockSparseArray, eltype) == Position(1) + @test TypeParameterAccessors.position(BlockSparseArray, ndims) == Position(2) + @test TypeParameterAccessors.position(BlockSparseArray, blocktype) == Position(3) + @test TypeParameterAccessors.position(BlockSparseArray, blockstype) == Position(4) end - end - @testset "TypeParameterAccessors.position" begin - @test TypeParameterAccessors.position(BlockSparseArray, eltype) == Position(1) - @test TypeParameterAccessors.position(BlockSparseArray, ndims) == Position(2) - @test TypeParameterAccessors.position(BlockSparseArray, blocktype) == Position(3) - @test TypeParameterAccessors.position(BlockSparseArray, blockstype) == Position(4) - end end diff --git a/test/test_blockrange.jl b/test/test_blockrange.jl index 35675655..e3ff72ca 100644 --- a/test/test_blockrange.jl +++ b/test/test_blockrange.jl @@ -3,14 +3,14 @@ using BlockSparseArrays: blockrange, eachblockaxis using Test: @test, @testset @testset "blockrange" begin - r = blockrange(AbstractUnitRange{Int}[Base.OneTo(3), 1:4]) - @test eachblockaxis(r) == [Base.OneTo(3), 1:4] - @test eachblockaxis(r)[1] === Base.OneTo(3) - @test eachblockaxis(r)[2] === 1:4 - @test r[Block(1)] == 1:3 - @test r[Block(2)] == 4:7 - @test first(r) == 1 - @test last(r) == 7 - @test blocklength(r) == 2 - @test r == 1:7 + r = blockrange(AbstractUnitRange{Int}[Base.OneTo(3), 1:4]) + @test eachblockaxis(r) == [Base.OneTo(3), 1:4] + @test eachblockaxis(r)[1] === Base.OneTo(3) + @test eachblockaxis(r)[2] === 1:4 + @test r[Block(1)] == 1:3 + @test r[Block(2)] == 4:7 + @test first(r) == 1 + @test last(r) == 7 + @test blocklength(r) == 2 + @test r == 1:7 end diff --git a/test/test_exports.jl b/test/test_exports.jl index 7121c797..bafaff2a 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -1,15 +1,15 @@ using BlockSparseArrays: BlockSparseArrays using Test: @test, @testset @testset "Test exports" begin - exports = [ - :BlockSparseArrays, - :BlockSparseArray, - :BlockSparseMatrix, - :BlockSparseVector, - :blockstoredlength, - :eachblockstoredindex, - :eachstoredblock, - :sparsemortar, - ] - @test issetequal(names(BlockSparseArrays), exports) + exports = [ + :BlockSparseArrays, + :BlockSparseArray, + :BlockSparseMatrix, + :BlockSparseVector, + :blockstoredlength, + :eachblockstoredindex, + :eachstoredblock, + :sparsemortar, + ] + @test issetequal(names(BlockSparseArrays), exports) end diff --git a/test/test_factorizations.jl b/test/test_factorizations.jl index e9d9dc38..9446bf57 100644 --- a/test/test_factorizations.jl +++ b/test/test_factorizations.jl @@ -1,134 +1,134 @@ using Adapt: adapt using BlockArrays: Block, BlockedMatrix, BlockedVector, blocks, mortar using BlockSparseArrays: - BlockSparseArrays, - BlockDiagonal, - BlockSparseArray, - BlockSparseMatrix, - blockstoredlength, - eachblockstoredindex + BlockSparseArrays, + BlockDiagonal, + BlockSparseArray, + BlockSparseMatrix, + blockstoredlength, + eachblockstoredindex using LinearAlgebra: LinearAlgebra, Diagonal, hermitianpart, pinv using MatrixAlgebraKit: - diagview, - eig_full, - eig_trunc, - eig_vals, - eigh_full, - eigh_trunc, - eigh_vals, - left_orth, - left_polar, - lq_compact, - lq_full, - qr_compact, - qr_full, - right_orth, - right_polar, - svd_compact, - svd_full, - svd_trunc, - truncrank, - trunctol + diagview, + eig_full, + eig_trunc, + eig_vals, + eigh_full, + eigh_trunc, + eigh_vals, + left_orth, + left_polar, + lq_compact, + lq_full, + qr_compact, + qr_full, + right_orth, + right_polar, + svd_compact, + svd_full, + svd_trunc, + truncrank, + trunctol using Random: Random using StableRNGs: StableRNG using Test: @inferred, @test, @test_throws, @testset @testset "Matrix functions (T=$elt)" for elt in (Float32, Float64, ComplexF64) - for matrixt in (Matrix, AbstractMatrix) - a = BlockSparseMatrix{elt,matrixt{elt}}(undef, [2, 3], [2, 3]) - rng = StableRNG(123) - a[Block(1, 1)] = randn(rng, elt, 2, 2) - a[Block(2, 2)] = randn(rng, elt, 3, 3) - MATRIX_FUNCTIONS = BlockSparseArrays.MATRIX_FUNCTIONS - MATRIX_FUNCTIONS = [MATRIX_FUNCTIONS; [:inv, :pinv]] - # Only works when real, also isn't defined in Julia 1.10. - MATRIX_FUNCTIONS = setdiff(MATRIX_FUNCTIONS, [:cbrt]) - MATRIX_FUNCTIONS_LOW_ACCURACY = [:acoth] - for f in setdiff(MATRIX_FUNCTIONS, MATRIX_FUNCTIONS_LOW_ACCURACY) - @eval begin - fa = $f($a) - @test Matrix(fa) ≈ $f(Matrix($a)) rtol = √(eps(real($elt))) - @test fa isa BlockSparseMatrix - @test issetequal(eachblockstoredindex(fa), [Block(1, 1), Block(2, 2)]) - end - end - for f in MATRIX_FUNCTIONS_LOW_ACCURACY - @eval begin - fa = $f($a) - # `acoth` appears to be broken on this matrix on Windows and Ubuntu - # for real matrices. - skip = !Sys.isapple() && ($elt <: Real) - @test Matrix(fa) ≈ $f(Matrix($a)) rtol = √eps(real($elt)) skip = skip - @test fa isa BlockSparseMatrix - @test issetequal(eachblockstoredindex(fa), [Block(1, 1), Block(2, 2)]) - end + for matrixt in (Matrix, AbstractMatrix) + a = BlockSparseMatrix{elt, matrixt{elt}}(undef, [2, 3], [2, 3]) + rng = StableRNG(123) + a[Block(1, 1)] = randn(rng, elt, 2, 2) + a[Block(2, 2)] = randn(rng, elt, 3, 3) + MATRIX_FUNCTIONS = BlockSparseArrays.MATRIX_FUNCTIONS + MATRIX_FUNCTIONS = [MATRIX_FUNCTIONS; [:inv, :pinv]] + # Only works when real, also isn't defined in Julia 1.10. + MATRIX_FUNCTIONS = setdiff(MATRIX_FUNCTIONS, [:cbrt]) + MATRIX_FUNCTIONS_LOW_ACCURACY = [:acoth] + for f in setdiff(MATRIX_FUNCTIONS, MATRIX_FUNCTIONS_LOW_ACCURACY) + @eval begin + fa = $f($a) + @test Matrix(fa) ≈ $f(Matrix($a)) rtol = √(eps(real($elt))) + @test fa isa BlockSparseMatrix + @test issetequal(eachblockstoredindex(fa), [Block(1, 1), Block(2, 2)]) + end + end + for f in MATRIX_FUNCTIONS_LOW_ACCURACY + @eval begin + fa = $f($a) + # `acoth` appears to be broken on this matrix on Windows and Ubuntu + # for real matrices. + skip = !Sys.isapple() && ($elt <: Real) + @test Matrix(fa) ≈ $f(Matrix($a)) rtol = √eps(real($elt)) skip = skip + @test fa isa BlockSparseMatrix + @test issetequal(eachblockstoredindex(fa), [Block(1, 1), Block(2, 2)]) + end + end end - end - # Catch case of off-diagonal blocks. - for matrixt in (Matrix, AbstractMatrix) - a = BlockSparseMatrix{elt,matrixt{elt}}(undef, [2, 3], [2, 3]) - rng = StableRNG(123) - a[Block(1, 1)] = randn(rng, elt, 2, 2) - a[Block(1, 2)] = randn(rng, elt, 2, 3) - for f in BlockSparseArrays.MATRIX_FUNCTIONS - @eval begin - @test_throws ArgumentError $f($a) - end + # Catch case of off-diagonal blocks. + for matrixt in (Matrix, AbstractMatrix) + a = BlockSparseMatrix{elt, matrixt{elt}}(undef, [2, 3], [2, 3]) + rng = StableRNG(123) + a[Block(1, 1)] = randn(rng, elt, 2, 2) + a[Block(1, 2)] = randn(rng, elt, 2, 3) + for f in BlockSparseArrays.MATRIX_FUNCTIONS + @eval begin + @test_throws ArgumentError $f($a) + end + end end - end - # Missing diagonal blocks. - for matrixt in (Matrix, AbstractMatrix) - a = BlockSparseMatrix{elt,matrixt{elt}}(undef, [2, 3], [2, 3]) - rng = StableRNG(123) - a[Block(2, 2)] = randn(rng, elt, 3, 3) - MATRIX_FUNCTIONS = BlockSparseArrays.MATRIX_FUNCTIONS - # These functions involve inverses so they break when there are zeros on the diagonal. - MATRIX_FUNCTIONS_SINGULAR = [ - :log, :acsc, :asec, :acot, :acsch, :asech, :acoth, :csc, :cot, :csch, :coth - ] - MATRIX_FUNCTIONS = setdiff(MATRIX_FUNCTIONS, MATRIX_FUNCTIONS_SINGULAR) - # Dense version is broken for some reason, investigate. - MATRIX_FUNCTIONS = setdiff(MATRIX_FUNCTIONS, [:cbrt]) - for f in MATRIX_FUNCTIONS - @eval begin - fa = $f($a) - @test Matrix(fa) ≈ $f(Matrix($a)) rtol = √(eps(real($elt))) - @test fa isa BlockSparseMatrix - @test issetequal(eachblockstoredindex(fa), [Block(1, 1), Block(2, 2)]) - end + # Missing diagonal blocks. + for matrixt in (Matrix, AbstractMatrix) + a = BlockSparseMatrix{elt, matrixt{elt}}(undef, [2, 3], [2, 3]) + rng = StableRNG(123) + a[Block(2, 2)] = randn(rng, elt, 3, 3) + MATRIX_FUNCTIONS = BlockSparseArrays.MATRIX_FUNCTIONS + # These functions involve inverses so they break when there are zeros on the diagonal. + MATRIX_FUNCTIONS_SINGULAR = [ + :log, :acsc, :asec, :acot, :acsch, :asech, :acoth, :csc, :cot, :csch, :coth, + ] + MATRIX_FUNCTIONS = setdiff(MATRIX_FUNCTIONS, MATRIX_FUNCTIONS_SINGULAR) + # Dense version is broken for some reason, investigate. + MATRIX_FUNCTIONS = setdiff(MATRIX_FUNCTIONS, [:cbrt]) + for f in MATRIX_FUNCTIONS + @eval begin + fa = $f($a) + @test Matrix(fa) ≈ $f(Matrix($a)) rtol = √(eps(real($elt))) + @test fa isa BlockSparseMatrix + @test issetequal(eachblockstoredindex(fa), [Block(1, 1), Block(2, 2)]) + end + end + + SINGULAR_EXCEPTION = if VERSION < v"1.11-" + # A different exception is thrown in older versions of Julia. + LinearAlgebra.LAPACKException + else + LinearAlgebra.SingularException + end + for f in setdiff(MATRIX_FUNCTIONS_SINGULAR, [:log]) + @eval begin + @test_throws $SINGULAR_EXCEPTION $f($a) + end + end end - - SINGULAR_EXCEPTION = if VERSION < v"1.11-" - # A different exception is thrown in older versions of Julia. - LinearAlgebra.LAPACKException - else - LinearAlgebra.SingularException - end - for f in setdiff(MATRIX_FUNCTIONS_SINGULAR, [:log]) - @eval begin - @test_throws $SINGULAR_EXCEPTION $f($a) - end - end - end end -function test_svd(a, (U, S, Vᴴ); full=false) - # Check that the SVD is correct - @test (U * S * Vᴴ ≈ a) - @test (U' * U ≈ LinearAlgebra.I) - @test (Vᴴ * Vᴴ' ≈ LinearAlgebra.I) - full || return nothing - - # Check factors are unitary - @test (U * U' ≈ LinearAlgebra.I) - @test (Vᴴ' * Vᴴ ≈ LinearAlgebra.I) - return nothing +function test_svd(a, (U, S, Vᴴ); full = false) + # Check that the SVD is correct + @test (U * S * Vᴴ ≈ a) + @test (U' * U ≈ LinearAlgebra.I) + @test (Vᴴ * Vᴴ' ≈ LinearAlgebra.I) + full || return nothing + + # Check factors are unitary + @test (U * U' ≈ LinearAlgebra.I) + @test (Vᴴ' * Vᴴ ≈ LinearAlgebra.I) + return nothing end blockszs = ( - ([2, 2], [2, 2]), ([2, 2], [2, 3]), ([2, 2, 1], [2, 3]), ([2, 3], [2]), ([2], [2, 3]) + ([2, 2], [2, 2]), ([2, 2], [2, 3]), ([2, 2, 1], [2, 3]), ([2, 3], [2]), ([2], [2, 3]), ) eltypes = (Float32, Float64, ComplexF64) test_params = Iterators.product(blockszs, eltypes) @@ -136,116 +136,88 @@ test_params = Iterators.product(blockszs, eltypes) # svd_compact! # ------------ @testset "svd_compact ($m, $n) BlockSparseMatrix{$T}" for ((m, n), T) in test_params - a = BlockSparseArray{T}(undef, m, n) - - # test empty matrix - usv_empty = svd_compact(a) - test_svd(a, usv_empty) - - # test blockdiagonal - rng = StableRNG(123) - for i in LinearAlgebra.diagind(blocks(a)) - I = CartesianIndices(blocks(a))[i] - a[Block(I.I...)] = rand(rng, T, size(blocks(a)[i])) - end - usv = svd_compact(a) - test_svd(a, usv) - - rng = StableRNG(123) - perm = Random.randperm(rng, length(m)) - b = a[Block.(perm), Block.(1:length(n))] - usv = svd_compact(b) - test_svd(b, usv) - - # test permuted blockdiagonal with missing row/col - rng = StableRNG(123) - I_removed = rand(rng, eachblockstoredindex(b)) - c = copy(b) - delete!(blocks(c).storage, CartesianIndex(Int.(Tuple(I_removed)))) - usv = svd_compact(c) - test_svd(c, usv) + a = BlockSparseArray{T}(undef, m, n) + + # test empty matrix + usv_empty = svd_compact(a) + test_svd(a, usv_empty) + + # test blockdiagonal + rng = StableRNG(123) + for i in LinearAlgebra.diagind(blocks(a)) + I = CartesianIndices(blocks(a))[i] + a[Block(I.I...)] = rand(rng, T, size(blocks(a)[i])) + end + usv = svd_compact(a) + test_svd(a, usv) + + rng = StableRNG(123) + perm = Random.randperm(rng, length(m)) + b = a[Block.(perm), Block.(1:length(n))] + usv = svd_compact(b) + test_svd(b, usv) + + # test permuted blockdiagonal with missing row/col + rng = StableRNG(123) + I_removed = rand(rng, eachblockstoredindex(b)) + c = copy(b) + delete!(blocks(c).storage, CartesianIndex(Int.(Tuple(I_removed)))) + usv = svd_compact(c) + test_svd(c, usv) end # svd_full! # --------- @testset "svd_full ($m, $n) BlockSparseMatrix{$T}" for ((m, n), T) in test_params - a = BlockSparseArray{T}(undef, m, n) - - # test empty matrix - usv_empty = svd_full(a) - test_svd(a, usv_empty; full=true) - - # test blockdiagonal - rng = StableRNG(123) - for i in LinearAlgebra.diagind(blocks(a)) - I = CartesianIndices(blocks(a))[i] - a[Block(I.I...)] = rand(rng, T, size(blocks(a)[i])) - end - usv = svd_full(a) - test_svd(a, usv; full=true) - - rng = StableRNG(123) - perm = Random.randperm(rng, length(m)) - b = a[Block.(perm), Block.(1:length(n))] - usv = svd_full(b) - test_svd(b, usv; full=true) - - # test permuted blockdiagonal with missing row/col - rng = StableRNG(123) - I_removed = rand(rng, eachblockstoredindex(b)) - c = copy(b) - delete!(blocks(c).storage, CartesianIndex(Int.(Tuple(I_removed)))) - usv = svd_full(c) - test_svd(c, usv; full=true) + a = BlockSparseArray{T}(undef, m, n) + + # test empty matrix + usv_empty = svd_full(a) + test_svd(a, usv_empty; full = true) + + # test blockdiagonal + rng = StableRNG(123) + for i in LinearAlgebra.diagind(blocks(a)) + I = CartesianIndices(blocks(a))[i] + a[Block(I.I...)] = rand(rng, T, size(blocks(a)[i])) + end + usv = svd_full(a) + test_svd(a, usv; full = true) + + rng = StableRNG(123) + perm = Random.randperm(rng, length(m)) + b = a[Block.(perm), Block.(1:length(n))] + usv = svd_full(b) + test_svd(b, usv; full = true) + + # test permuted blockdiagonal with missing row/col + rng = StableRNG(123) + I_removed = rand(rng, eachblockstoredindex(b)) + c = copy(b) + delete!(blocks(c).storage, CartesianIndex(Int.(Tuple(I_removed)))) + usv = svd_full(c) + test_svd(c, usv; full = true) end # svd_trunc! # ---------- @testset "svd_trunc ($m, $n) BlockSparseMatrix{$T}" for ((m, n), T) in test_params - a = BlockSparseArray{T}(undef, m, n) - - # test blockdiagonal - rng = StableRNG(123) - for i in LinearAlgebra.diagind(blocks(a)) - I = CartesianIndices(blocks(a))[i] - a[Block(I.I...)] = rand(rng, T, size(blocks(a)[i])) - end - - minmn = min(size(a)...) - r = max(1, minmn - 2) - trunc = truncrank(r) - - U1, S1, V1ᴴ = svd_trunc(a; trunc) - U2, S2, V2ᴴ = svd_trunc(Matrix(a); trunc) - @test size(U1) == size(U2) - @test size(S1) == size(S2) - @test size(V1ᴴ) == size(V2ᴴ) - @test Matrix(U1 * S1 * V1ᴴ) ≈ U2 * S2 * V2ᴴ - - @test (U1' * U1 ≈ LinearAlgebra.I) - @test (V1ᴴ * V1ᴴ' ≈ LinearAlgebra.I) - - atol = minimum(LinearAlgebra.diag(S1)) + 10 * eps(real(T)) - trunc = trunctol(; atol) - - U1, S1, V1ᴴ = svd_trunc(a; trunc) - U2, S2, V2ᴴ = svd_trunc(Matrix(a); trunc) - @test size(U1) == size(U2) - @test size(S1) == size(S2) - @test size(V1ᴴ) == size(V2ᴴ) - @test Matrix(U1 * S1 * V1ᴴ) ≈ U2 * S2 * V2ᴴ - - @test (U1' * U1 ≈ LinearAlgebra.I) - @test (V1ᴴ * V1ᴴ' ≈ LinearAlgebra.I) - - # test permuted blockdiagonal - rng = StableRNG(123) - perm = Random.randperm(rng, length(m)) - b = a[Block.(perm), Block.(1:length(n))] - for trunc in (truncrank(r), trunctol(; atol)) - U1, S1, V1ᴴ = svd_trunc(b; trunc) - U2, S2, V2ᴴ = svd_trunc(Matrix(b); trunc) + a = BlockSparseArray{T}(undef, m, n) + + # test blockdiagonal + rng = StableRNG(123) + for i in LinearAlgebra.diagind(blocks(a)) + I = CartesianIndices(blocks(a))[i] + a[Block(I.I...)] = rand(rng, T, size(blocks(a)[i])) + end + + minmn = min(size(a)...) + r = max(1, minmn - 2) + trunc = truncrank(r) + + U1, S1, V1ᴴ = svd_trunc(a; trunc) + U2, S2, V2ᴴ = svd_trunc(Matrix(a); trunc) @test size(U1) == size(U2) @test size(S1) == size(S2) @test size(V1ᴴ) == size(V2ᴴ) @@ -253,15 +225,12 @@ end @test (U1' * U1 ≈ LinearAlgebra.I) @test (V1ᴴ * V1ᴴ' ≈ LinearAlgebra.I) - end - - # test permuted blockdiagonal with missing row/col - I_removed = rand(eachblockstoredindex(b)) - c = copy(b) - delete!(blocks(c).storage, CartesianIndex(Int.(Tuple(I_removed)))) - for trunc in (truncrank(r), trunctol(; atol)) - U1, S1, V1ᴴ = svd_trunc(c; trunc) - U2, S2, V2ᴴ = svd_trunc(Matrix(c); trunc) + + atol = minimum(LinearAlgebra.diag(S1)) + 10 * eps(real(T)) + trunc = trunctol(; atol) + + U1, S1, V1ᴴ = svd_trunc(a; trunc) + U2, S2, V2ᴴ = svd_trunc(Matrix(a); trunc) @test size(U1) == size(U2) @test size(S1) == size(S2) @test size(V1ᴴ) == size(V2ᴴ) @@ -269,227 +238,258 @@ end @test (U1' * U1 ≈ LinearAlgebra.I) @test (V1ᴴ * V1ᴴ' ≈ LinearAlgebra.I) - end + + # test permuted blockdiagonal + rng = StableRNG(123) + perm = Random.randperm(rng, length(m)) + b = a[Block.(perm), Block.(1:length(n))] + for trunc in (truncrank(r), trunctol(; atol)) + U1, S1, V1ᴴ = svd_trunc(b; trunc) + U2, S2, V2ᴴ = svd_trunc(Matrix(b); trunc) + @test size(U1) == size(U2) + @test size(S1) == size(S2) + @test size(V1ᴴ) == size(V2ᴴ) + @test Matrix(U1 * S1 * V1ᴴ) ≈ U2 * S2 * V2ᴴ + + @test (U1' * U1 ≈ LinearAlgebra.I) + @test (V1ᴴ * V1ᴴ' ≈ LinearAlgebra.I) + end + + # test permuted blockdiagonal with missing row/col + I_removed = rand(eachblockstoredindex(b)) + c = copy(b) + delete!(blocks(c).storage, CartesianIndex(Int.(Tuple(I_removed)))) + for trunc in (truncrank(r), trunctol(; atol)) + U1, S1, V1ᴴ = svd_trunc(c; trunc) + U2, S2, V2ᴴ = svd_trunc(Matrix(c); trunc) + @test size(U1) == size(U2) + @test size(S1) == size(S2) + @test size(V1ᴴ) == size(V2ᴴ) + @test Matrix(U1 * S1 * V1ᴴ) ≈ U2 * S2 * V2ᴴ + + @test (U1' * U1 ≈ LinearAlgebra.I) + @test (V1ᴴ * V1ᴴ' ≈ LinearAlgebra.I) + end end @testset "qr_compact (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64) - for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3] - A = BlockSparseArray{T}(undef, ([i, j], [k, l])) - rng = StableRNG(123) - A[Block(1, 1)] = randn(rng, T, i, k) - A[Block(2, 2)] = randn(rng, T, j, l) - Q, R = qr_compact(A) - @test Matrix(Q'Q) ≈ LinearAlgebra.I - @test A ≈ Q * R - end + for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3] + A = BlockSparseArray{T}(undef, ([i, j], [k, l])) + rng = StableRNG(123) + A[Block(1, 1)] = randn(rng, T, i, k) + A[Block(2, 2)] = randn(rng, T, j, l) + Q, R = qr_compact(A) + @test Matrix(Q'Q) ≈ LinearAlgebra.I + @test A ≈ Q * R + end end @testset "qr_full (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64) - for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3] - A = BlockSparseArray{T}(undef, ([i, j], [k, l])) - rng = StableRNG(123) - A[Block(1, 1)] = randn(rng, T, i, k) - A[Block(2, 2)] = randn(rng, T, j, l) - Q, R = qr_full(A) - Q′, R′ = qr_full(Matrix(A)) - @test size(Q) == size(Q′) - @test size(R) == size(R′) - @test Matrix(Q'Q) ≈ LinearAlgebra.I - @test Matrix(Q * Q') ≈ LinearAlgebra.I - @test A ≈ Q * R - end + for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3] + A = BlockSparseArray{T}(undef, ([i, j], [k, l])) + rng = StableRNG(123) + A[Block(1, 1)] = randn(rng, T, i, k) + A[Block(2, 2)] = randn(rng, T, j, l) + Q, R = qr_full(A) + Q′, R′ = qr_full(Matrix(A)) + @test size(Q) == size(Q′) + @test size(R) == size(R′) + @test Matrix(Q'Q) ≈ LinearAlgebra.I + @test Matrix(Q * Q') ≈ LinearAlgebra.I + @test A ≈ Q * R + end end @testset "lq_compact" for T in (Float32, Float64, ComplexF32, ComplexF64) - for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3] - A = BlockSparseArray{T}(undef, ([i, j], [k, l])) - rng = StableRNG(123) - A[Block(1, 1)] = randn(rng, T, i, k) - A[Block(2, 2)] = randn(rng, T, j, l) - L, Q = lq_compact(A) - @test Matrix(Q * Q') ≈ LinearAlgebra.I - @test A ≈ L * Q - end + for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3] + A = BlockSparseArray{T}(undef, ([i, j], [k, l])) + rng = StableRNG(123) + A[Block(1, 1)] = randn(rng, T, i, k) + A[Block(2, 2)] = randn(rng, T, j, l) + L, Q = lq_compact(A) + @test Matrix(Q * Q') ≈ LinearAlgebra.I + @test A ≈ L * Q + end end @testset "lq_full" for T in (Float32, Float64, ComplexF32, ComplexF64) - for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3] - A = BlockSparseArray{T}(undef, ([i, j], [k, l])) - rng = StableRNG(123) - A[Block(1, 1)] = randn(rng, T, i, k) - A[Block(2, 2)] = randn(rng, T, j, l) - L, Q = lq_full(A) - L′, Q′ = lq_full(Matrix(A)) - @test size(L) == size(L′) - @test size(Q) == size(Q′) - @test Matrix(Q * Q') ≈ LinearAlgebra.I - @test Matrix(Q'Q) ≈ LinearAlgebra.I - @test A ≈ L * Q - end + for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3] + A = BlockSparseArray{T}(undef, ([i, j], [k, l])) + rng = StableRNG(123) + A[Block(1, 1)] = randn(rng, T, i, k) + A[Block(2, 2)] = randn(rng, T, j, l) + L, Q = lq_full(A) + L′, Q′ = lq_full(Matrix(A)) + @test size(L) == size(L′) + @test size(Q) == size(Q′) + @test Matrix(Q * Q') ≈ LinearAlgebra.I + @test Matrix(Q'Q) ≈ LinearAlgebra.I + @test A ≈ L * Q + end end @testset "left_polar (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64) - A = BlockSparseArray{T}(undef, ([3, 4], [2, 3])) - rng = StableRNG(123) - A[Block(1, 1)] = randn(rng, T, 3, 2) - A[Block(2, 2)] = randn(rng, T, 4, 3) - - U, C = left_polar(A) - @test U * C ≈ A - @test Matrix(U'U) ≈ LinearAlgebra.I + A = BlockSparseArray{T}(undef, ([3, 4], [2, 3])) + rng = StableRNG(123) + A[Block(1, 1)] = randn(rng, T, 3, 2) + A[Block(2, 2)] = randn(rng, T, 4, 3) + + U, C = left_polar(A) + @test U * C ≈ A + @test Matrix(U'U) ≈ LinearAlgebra.I end @testset "right_polar (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64) - A = BlockSparseArray{T}(undef, ([2, 3], [3, 4])) - rng = StableRNG(123) - A[Block(1, 1)] = randn(rng, T, 2, 3) - A[Block(2, 2)] = randn(rng, T, 3, 4) - - C, U = right_polar(A) - @test C * U ≈ A - @test Matrix(U * U') ≈ LinearAlgebra.I + A = BlockSparseArray{T}(undef, ([2, 3], [3, 4])) + rng = StableRNG(123) + A[Block(1, 1)] = randn(rng, T, 2, 3) + A[Block(2, 2)] = randn(rng, T, 3, 4) + + C, U = right_polar(A) + @test C * U ≈ A + @test Matrix(U * U') ≈ LinearAlgebra.I end @testset "left_orth (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64) - A = BlockSparseArray{T}(undef, ([3, 4], [2, 3])) - rng = StableRNG(123) - A[Block(1, 1)] = randn(rng, T, 3, 2) - A[Block(2, 2)] = randn(rng, T, 4, 3) + A = BlockSparseArray{T}(undef, ([3, 4], [2, 3])) + rng = StableRNG(123) + A[Block(1, 1)] = randn(rng, T, 3, 2) + A[Block(2, 2)] = randn(rng, T, 4, 3) - for kind in (:polar, :qr, :svd) - U, C = left_orth(A; kind) - @test U * C ≈ A - @test Matrix(U'U) ≈ LinearAlgebra.I - end + for kind in (:polar, :qr, :svd) + U, C = left_orth(A; kind) + @test U * C ≈ A + @test Matrix(U'U) ≈ LinearAlgebra.I + end - U, C = left_orth(A; trunc=(; maxrank=2)) - @test size(U, 2) ≤ 2 - @test size(C, 1) ≤ 2 - @test Matrix(U'U) ≈ LinearAlgebra.I + U, C = left_orth(A; trunc = (; maxrank = 2)) + @test size(U, 2) ≤ 2 + @test size(C, 1) ≤ 2 + @test Matrix(U'U) ≈ LinearAlgebra.I end @testset "right_orth (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64) - A = BlockSparseArray{T}(undef, ([2, 3], [3, 4])) - rng = StableRNG(123) - A[Block(1, 1)] = randn(rng, T, 2, 3) - A[Block(2, 2)] = randn(rng, T, 3, 4) + A = BlockSparseArray{T}(undef, ([2, 3], [3, 4])) + rng = StableRNG(123) + A[Block(1, 1)] = randn(rng, T, 2, 3) + A[Block(2, 2)] = randn(rng, T, 3, 4) - for kind in (:lq, :polar, :svd) - C, U = right_orth(A; kind) - @test C * U ≈ A - @test Matrix(U * U') ≈ LinearAlgebra.I - end + for kind in (:lq, :polar, :svd) + C, U = right_orth(A; kind) + @test C * U ≈ A + @test Matrix(U * U') ≈ LinearAlgebra.I + end - C, U = right_orth(A; trunc=(; maxrank=2)) - @test size(C, 2) ≤ 2 - @test size(U, 1) ≤ 2 - @test Matrix(U * U') ≈ LinearAlgebra.I + C, U = right_orth(A; trunc = (; maxrank = 2)) + @test size(C, 2) ≤ 2 + @test size(U, 1) ≤ 2 + @test Matrix(U * U') ≈ LinearAlgebra.I end @testset "eig_full (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64) - A = BlockSparseArray{T}(undef, ([2, 3], [2, 3])) - rng = StableRNG(123) - A[Block(1, 1)] = randn(rng, T, 2, 2) - A[Block(2, 2)] = randn(rng, T, 3, 3) - - D, V = eig_full(A) - @test size(D) == size(A) - @test size(D) == size(A) - @test blockstoredlength(D) == 2 - @test blockstoredlength(V) == 2 - @test issetequal(eachblockstoredindex(D), [Block(1, 1), Block(2, 2)]) - @test issetequal(eachblockstoredindex(V), [Block(1, 1), Block(2, 2)]) - @test A * V ≈ V * D + A = BlockSparseArray{T}(undef, ([2, 3], [2, 3])) + rng = StableRNG(123) + A[Block(1, 1)] = randn(rng, T, 2, 2) + A[Block(2, 2)] = randn(rng, T, 3, 3) + + D, V = eig_full(A) + @test size(D) == size(A) + @test size(D) == size(A) + @test blockstoredlength(D) == 2 + @test blockstoredlength(V) == 2 + @test issetequal(eachblockstoredindex(D), [Block(1, 1), Block(2, 2)]) + @test issetequal(eachblockstoredindex(V), [Block(1, 1), Block(2, 2)]) + @test A * V ≈ V * D end @testset "eig_vals (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64) - A = BlockSparseArray{T}(undef, ([2, 3], [2, 3])) - rng = StableRNG(123) - A[Block(1, 1)] = randn(rng, T, 2, 2) - A[Block(2, 2)] = randn(rng, T, 3, 3) - - D = eig_vals(A) - @test size(D) == (size(A, 1),) - @test blockstoredlength(D) == 2 - D′ = eig_vals(Matrix(A)) - @test sort(D; by=abs) ≈ sort(D′; by=abs) + A = BlockSparseArray{T}(undef, ([2, 3], [2, 3])) + rng = StableRNG(123) + A[Block(1, 1)] = randn(rng, T, 2, 2) + A[Block(2, 2)] = randn(rng, T, 3, 3) + + D = eig_vals(A) + @test size(D) == (size(A, 1),) + @test blockstoredlength(D) == 2 + D′ = eig_vals(Matrix(A)) + @test sort(D; by = abs) ≈ sort(D′; by = abs) end @testset "eig_trunc (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64) - A = BlockSparseArray{T}(undef, ([2, 3], [2, 3])) - rng = StableRNG(123) - D1 = [1.0, 0.1] - V1 = randn(rng, T, 2, 2) - A1 = V1 * Diagonal(D1) * inv(V1) - D2 = [1.0, 0.5, 0.1] - V2 = randn(rng, T, 3, 3) - A2 = V2 * Diagonal(D2) * inv(V2) - A[Block(1, 1)] = A1 - A[Block(2, 2)] = A2 - - D, V = eig_trunc(A; trunc=(; maxrank=3)) - @test size(D) == (3, 3) - @test size(D) == (3, 3) - @test blockstoredlength(D) == 2 - @test blockstoredlength(V) == 2 - @test issetequal(eachblockstoredindex(D), [Block(1, 1), Block(2, 2)]) - @test issetequal(eachblockstoredindex(V), [Block(1, 1), Block(2, 2)]) - @test A * V ≈ V * D - @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] + A = BlockSparseArray{T}(undef, ([2, 3], [2, 3])) + rng = StableRNG(123) + D1 = [1.0, 0.1] + V1 = randn(rng, T, 2, 2) + A1 = V1 * Diagonal(D1) * inv(V1) + D2 = [1.0, 0.5, 0.1] + V2 = randn(rng, T, 3, 3) + A2 = V2 * Diagonal(D2) * inv(V2) + A[Block(1, 1)] = A1 + A[Block(2, 2)] = A2 + + D, V = eig_trunc(A; trunc = (; maxrank = 3)) + @test size(D) == (3, 3) + @test size(D) == (3, 3) + @test blockstoredlength(D) == 2 + @test blockstoredlength(V) == 2 + @test issetequal(eachblockstoredindex(D), [Block(1, 1), Block(2, 2)]) + @test issetequal(eachblockstoredindex(V), [Block(1, 1), Block(2, 2)]) + @test A * V ≈ V * D + @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 herm(x) = parent(hermitianpart(x)) @testset "eigh_full (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64) - A = BlockSparseArray{T}(undef, ([2, 3], [2, 3])) - rng = StableRNG(123) - A[Block(1, 1)] = herm(randn(rng, T, 2, 2)) - A[Block(2, 2)] = herm(randn(rng, T, 3, 3)) - - D, V = eigh_full(A) - @test size(D) == size(A) - @test size(D) == size(A) - @test blockstoredlength(D) == 2 - @test blockstoredlength(V) == 2 - @test issetequal(eachblockstoredindex(D), [Block(1, 1), Block(2, 2)]) - @test issetequal(eachblockstoredindex(V), [Block(1, 1), Block(2, 2)]) - @test A * V ≈ V * D + A = BlockSparseArray{T}(undef, ([2, 3], [2, 3])) + rng = StableRNG(123) + A[Block(1, 1)] = herm(randn(rng, T, 2, 2)) + A[Block(2, 2)] = herm(randn(rng, T, 3, 3)) + + D, V = eigh_full(A) + @test size(D) == size(A) + @test size(D) == size(A) + @test blockstoredlength(D) == 2 + @test blockstoredlength(V) == 2 + @test issetequal(eachblockstoredindex(D), [Block(1, 1), Block(2, 2)]) + @test issetequal(eachblockstoredindex(V), [Block(1, 1), Block(2, 2)]) + @test A * V ≈ V * D end @testset "eigh_vals (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64) - A = BlockSparseArray{T}(undef, ([2, 3], [2, 3])) - rng = StableRNG(123) - A[Block(1, 1)] = herm(randn(rng, T, 2, 2)) - A[Block(2, 2)] = herm(randn(rng, T, 3, 3)) - - D = eigh_vals(A) - @test size(D) == (size(A, 1),) - @test blockstoredlength(D) == 2 - D′ = eigh_vals(Matrix(A)) - @test sort(D; by=abs) ≈ sort(D′; by=abs) + A = BlockSparseArray{T}(undef, ([2, 3], [2, 3])) + rng = StableRNG(123) + A[Block(1, 1)] = herm(randn(rng, T, 2, 2)) + A[Block(2, 2)] = herm(randn(rng, T, 3, 3)) + + D = eigh_vals(A) + @test size(D) == (size(A, 1),) + @test blockstoredlength(D) == 2 + D′ = eigh_vals(Matrix(A)) + @test sort(D; by = abs) ≈ sort(D′; by = abs) end @testset "eigh_trunc (T=$T)" for T in (Float32, Float64, ComplexF32, ComplexF64) - A = BlockSparseArray{T}(undef, ([2, 3], [2, 3])) - rng = StableRNG(123) - D1 = [1.0, 0.1] - V1, _ = qr_compact(randn(rng, T, 2, 2)) - A1 = V1 * Diagonal(D1) * V1' - D2 = [1.0, 0.5, 0.1] - V2, _ = qr_compact(randn(rng, T, 3, 3)) - A2 = V2 * Diagonal(D2) * V2' - A[Block(1, 1)] = herm(A1) - A[Block(2, 2)] = herm(A2) - - D, V = eigh_trunc(A; trunc=(; maxrank=3)) - @test size(D) == (3, 3) - @test size(D) == (3, 3) - @test blockstoredlength(D) == 2 - @test blockstoredlength(V) == 2 - @test issetequal(eachblockstoredindex(D), [Block(1, 1), Block(2, 2)]) - @test issetequal(eachblockstoredindex(V), [Block(1, 1), Block(2, 2)]) - @test A * V ≈ V * D - @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] + A = BlockSparseArray{T}(undef, ([2, 3], [2, 3])) + rng = StableRNG(123) + D1 = [1.0, 0.1] + V1, _ = qr_compact(randn(rng, T, 2, 2)) + A1 = V1 * Diagonal(D1) * V1' + D2 = [1.0, 0.5, 0.1] + V2, _ = qr_compact(randn(rng, T, 3, 3)) + A2 = V2 * Diagonal(D2) * V2' + A[Block(1, 1)] = herm(A1) + A[Block(2, 2)] = herm(A2) + + D, V = eigh_trunc(A; trunc = (; maxrank = 3)) + @test size(D) == (3, 3) + @test size(D) == (3, 3) + @test blockstoredlength(D) == 2 + @test blockstoredlength(V) == 2 + @test issetequal(eachblockstoredindex(D), [Block(1, 1), Block(2, 2)]) + @test issetequal(eachblockstoredindex(V), [Block(1, 1), Block(2, 2)]) + @test A * V ≈ V * D + @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 diff --git a/test/test_genericblockindex.jl b/test/test_genericblockindex.jl index cc159d92..4308d099 100644 --- a/test/test_genericblockindex.jl +++ b/test/test_genericblockindex.jl @@ -1,23 +1,23 @@ using BlockArrays: - Block, - BlockIndex, - BlockSlice, - BlockedArray, - BlockedVector, - block, - blockedrange, - blockindex, - mortar + Block, + BlockIndex, + BlockSlice, + BlockedArray, + BlockedVector, + block, + blockedrange, + blockindex, + mortar using BlockSparseArrays: - BlockSparseArrays, - BlockIndexVector, - BlockIndices, - GenericBlockIndex, - blocksparsezeros, - blockedunitrange_getindices, - to_block, - to_block_indices, - to_blockindexrange + BlockSparseArrays, + BlockIndexVector, + BlockIndices, + GenericBlockIndex, + blocksparsezeros, + blockedunitrange_getindices, + to_block, + to_block_indices, + to_blockindexrange using Test: @test, @test_broken, @testset # blockrange @@ -29,256 +29,260 @@ using Test: @test, @test_broken, @testset # to_blockindexrange @testset "GenericBlockIndex" begin - i1 = GenericBlockIndex(Block(1), ("x",)) - i2 = GenericBlockIndex(Block(2), ("y",)) - i = GenericBlockIndex(Block(1, 2), ("x", "y")) - @test sprint(show, i) == "Block(1, 2)[x, y]" - @test i isa GenericBlockIndex{2,Tuple{Int64,Int64},Tuple{String,String}} - @test GenericBlockIndex(Block(1), "x") === i1 - @test GenericBlockIndex(1, "x") === i1 - @test GenericBlockIndex(1, ("x",)) === i1 - @test GenericBlockIndex((1,), "x") === i1 - @test GenericBlockIndex((1, 2), ("x", "y")) === i - @test GenericBlockIndex((Block(1), Block(2)), ("x", "y")) === i - @test GenericBlockIndex((i1, i2)) === i - @test block(i1) == Block(1) - @test block(i) == Block(1, 2) - @test blockindex(i1) == "x" - @test GenericBlockIndex((), ()) == GenericBlockIndex(Block(), ()) - @test GenericBlockIndex(Block(1, 2), ("x",)) == GenericBlockIndex(Block(1, 2), ("x", 1)) + i1 = GenericBlockIndex(Block(1), ("x",)) + i2 = GenericBlockIndex(Block(2), ("y",)) + i = GenericBlockIndex(Block(1, 2), ("x", "y")) + @test sprint(show, i) == "Block(1, 2)[x, y]" + @test i isa GenericBlockIndex{2, Tuple{Int64, Int64}, Tuple{String, String}} + @test GenericBlockIndex(Block(1), "x") === i1 + @test GenericBlockIndex(1, "x") === i1 + @test GenericBlockIndex(1, ("x",)) === i1 + @test GenericBlockIndex((1,), "x") === i1 + @test GenericBlockIndex((1, 2), ("x", "y")) === i + @test GenericBlockIndex((Block(1), Block(2)), ("x", "y")) === i + @test GenericBlockIndex((i1, i2)) === i + @test block(i1) == Block(1) + @test block(i) == Block(1, 2) + @test blockindex(i1) == "x" + @test GenericBlockIndex((), ()) == GenericBlockIndex(Block(), ()) + @test GenericBlockIndex(Block(1, 2), ("x",)) == GenericBlockIndex(Block(1, 2), ("x", 1)) - i1 = GenericBlockIndex(Block(1), (1,)) - i2 = GenericBlockIndex(Block(2), (2,)) - i = GenericBlockIndex(Block(1, 2), (1, 2)) - v = BlockedVector(["a", "b", "c", "d"], [2, 2]) - @test v[i1] == "a" - @test v[i2] == "d" + i1 = GenericBlockIndex(Block(1), (1,)) + i2 = GenericBlockIndex(Block(2), (2,)) + i = GenericBlockIndex(Block(1, 2), (1, 2)) + v = BlockedVector(["a", "b", "c", "d"], [2, 2]) + @test v[i1] == "a" + @test v[i2] == "d" - a = collect(Iterators.product(v, v)) - @test a[i1, i1] == ("a", "a") - @test a[i2, i1] == ("d", "a") - @test a[i1, i2] == ("a", "d") - @test a[i] == ("a", "d") - @test a[i2, i2] == ("d", "d") + a = collect(Iterators.product(v, v)) + @test a[i1, i1] == ("a", "a") + @test a[i2, i1] == ("d", "a") + @test a[i1, i2] == ("a", "d") + @test a[i] == ("a", "d") + @test a[i2, i2] == ("d", "d") - I = BlockIndexVector(Block(1), [1, 2]) - @test eltype(I) === BlockIndex{1,Tuple{Int},Tuple{Int}} - @test ndims(I) === 1 - @test length(I) === 2 - @test size(I) === (2,) - @test I[1] === Block(1)[1] - @test I[2] === Block(1)[2] - @test block(I) === Block(1) - @test Block(I) === Block(1) - @test copy(I) == BlockIndexVector(Block(1), [1, 2]) + I = BlockIndexVector(Block(1), [1, 2]) + @test eltype(I) === BlockIndex{1, Tuple{Int}, Tuple{Int}} + @test ndims(I) === 1 + @test length(I) === 2 + @test size(I) === (2,) + @test I[1] === Block(1)[1] + @test I[2] === Block(1)[2] + @test block(I) === Block(1) + @test Block(I) === Block(1) + @test copy(I) == BlockIndexVector(Block(1), [1, 2]) - I = BlockIndexVector(Block(1, 2), ([1, 2], [3, 4])) - @test eltype(I) === BlockIndex{2,Tuple{Int,Int},Tuple{Int,Int}} - @test ndims(I) === 2 - @test length(I) === 4 - @test size(I) === (2, 2) - @test I[1, 1] === Block(1, 2)[1, 3] - @test I[2, 1] === Block(1, 2)[2, 3] - @test I[1, 2] === Block(1, 2)[1, 4] - @test I[2, 2] === Block(1, 2)[2, 4] - @test block(I) === Block(1, 2) - @test Block(I) === Block(1, 2) - @test copy(I) == BlockIndexVector(Block(1, 2), ([1, 2], [3, 4])) + I = BlockIndexVector(Block(1, 2), ([1, 2], [3, 4])) + @test eltype(I) === BlockIndex{2, Tuple{Int, Int}, Tuple{Int, Int}} + @test ndims(I) === 2 + @test length(I) === 4 + @test size(I) === (2, 2) + @test I[1, 1] === Block(1, 2)[1, 3] + @test I[2, 1] === Block(1, 2)[2, 3] + @test I[1, 2] === Block(1, 2)[1, 4] + @test I[2, 2] === Block(1, 2)[2, 4] + @test block(I) === Block(1, 2) + @test Block(I) === Block(1, 2) + @test copy(I) == BlockIndexVector(Block(1, 2), ([1, 2], [3, 4])) - I = BlockIndexVector(Block(1), ["x", "y"]) - @test eltype(I) === GenericBlockIndex{1,Tuple{Int},Tuple{String}} - @test ndims(I) === 1 - @test length(I) === 2 - @test size(I) === (2,) - @test I[1] === GenericBlockIndex(Block(1), "x") - @test I[2] === GenericBlockIndex(Block(1), "y") - @test block(I) === Block(1) - @test Block(I) === Block(1) - @test copy(I) == BlockIndexVector(Block(1), ["x", "y"]) + I = BlockIndexVector(Block(1), ["x", "y"]) + @test eltype(I) === GenericBlockIndex{1, Tuple{Int}, Tuple{String}} + @test ndims(I) === 1 + @test length(I) === 2 + @test size(I) === (2,) + @test I[1] === GenericBlockIndex(Block(1), "x") + @test I[2] === GenericBlockIndex(Block(1), "y") + @test block(I) === Block(1) + @test Block(I) === Block(1) + @test copy(I) == BlockIndexVector(Block(1), ["x", "y"]) - I = BlockIndexVector(Block(1, 2), (["x", "y"], ["z", "w"])) - @test eltype(I) === GenericBlockIndex{2,Tuple{Int,Int},Tuple{String,String}} - @test ndims(I) === 2 - @test length(I) === 4 - @test size(I) === (2, 2) - @test I[1, 1] === GenericBlockIndex(Block(1, 2), ("x", "z")) - @test I[2, 1] === GenericBlockIndex(Block(1, 2), ("y", "z")) - @test I[1, 2] === GenericBlockIndex(Block(1, 2), ("x", "w")) - @test I[2, 2] === GenericBlockIndex(Block(1, 2), ("y", "w")) - @test block(I) === Block(1, 2) - @test Block(I) === Block(1, 2) - @test copy(I) == BlockIndexVector(Block(1, 2), (["x", "y"], ["z", "w"])) + I = BlockIndexVector(Block(1, 2), (["x", "y"], ["z", "w"])) + @test eltype(I) === GenericBlockIndex{2, Tuple{Int, Int}, Tuple{String, String}} + @test ndims(I) === 2 + @test length(I) === 4 + @test size(I) === (2, 2) + @test I[1, 1] === GenericBlockIndex(Block(1, 2), ("x", "z")) + @test I[2, 1] === GenericBlockIndex(Block(1, 2), ("y", "z")) + @test I[1, 2] === GenericBlockIndex(Block(1, 2), ("x", "w")) + @test I[2, 2] === GenericBlockIndex(Block(1, 2), ("y", "w")) + @test block(I) === Block(1, 2) + @test Block(I) === Block(1, 2) + @test copy(I) == BlockIndexVector(Block(1, 2), (["x", "y"], ["z", "w"])) - v = BlockedVector(["a", "b", "c", "d"], [2, 2]) - i = BlockIndexVector(Block(1), [2, 1]) - @test v[i] == ["b", "a"] - i = BlockIndexVector(Block(2), [2, 1]) - @test v[i] == ["d", "c"] + v = BlockedVector(["a", "b", "c", "d"], [2, 2]) + i = BlockIndexVector(Block(1), [2, 1]) + @test v[i] == ["b", "a"] + i = BlockIndexVector(Block(2), [2, 1]) + @test v[i] == ["d", "c"] - v = BlockedVector(["a", "b", "c", "d"], [2, 2]) - i = BlockIndexVector{1,GenericBlockIndex{1,Tuple{Int},Tuple{String}}}(Block(1), [2, 1]) - @test v[i] == ["b", "a"] - i = BlockIndexVector(Block(2), [2, 1]) - @test v[i] == ["d", "c"] + v = BlockedVector(["a", "b", "c", "d"], [2, 2]) + i = BlockIndexVector{1, GenericBlockIndex{1, Tuple{Int}, Tuple{String}}}(Block(1), [2, 1]) + @test v[i] == ["b", "a"] + i = BlockIndexVector(Block(2), [2, 1]) + @test v[i] == ["d", "c"] - a = collect(Iterators.product(v, v)) - i1 = BlockIndexVector(Block(1), [2, 1]) - i2 = BlockIndexVector(Block(2), [1, 2]) - i = BlockIndexVector(Block(1, 2), ([2, 1], [1, 2])) - @test a[i1, i1] == [("b", "b") ("b", "a"); ("a", "b") ("a", "a")] - @test a[i2, i1] == [("c", "b") ("c", "a"); ("d", "b") ("d", "a")] - @test a[i1, i2] == [("b", "c") ("b", "d"); ("a", "c") ("a", "d")] - @test a[i] == [("b", "c") ("b", "d"); ("a", "c") ("a", "d")] - @test a[i2, i2] == [("c", "c") ("c", "d"); ("d", "c") ("d", "d")] + a = collect(Iterators.product(v, v)) + i1 = BlockIndexVector(Block(1), [2, 1]) + i2 = BlockIndexVector(Block(2), [1, 2]) + i = BlockIndexVector(Block(1, 2), ([2, 1], [1, 2])) + @test a[i1, i1] == [("b", "b") ("b", "a"); ("a", "b") ("a", "a")] + @test a[i2, i1] == [("c", "b") ("c", "a"); ("d", "b") ("d", "a")] + @test a[i1, i2] == [("b", "c") ("b", "d"); ("a", "c") ("a", "d")] + @test a[i] == [("b", "c") ("b", "d"); ("a", "c") ("a", "d")] + @test a[i2, i2] == [("c", "c") ("c", "d"); ("d", "c") ("d", "d")] - a = collect(Iterators.product(v, v)) - i1 = BlockIndexVector{1,GenericBlockIndex{1,Tuple{Int},Tuple{String}}}(Block(1), [2, 1]) - i2 = BlockIndexVector{1,GenericBlockIndex{1,Tuple{Int},Tuple{String}}}(Block(2), [1, 2]) - i = BlockIndexVector{2,GenericBlockIndex{2,Tuple{Int,Int},Tuple{String,String}}}( - Block(1, 2), ([2, 1], [1, 2]) - ) - @test a[i1, i1] == [("b", "b") ("b", "a"); ("a", "b") ("a", "a")] - @test a[i2, i1] == [("c", "b") ("c", "a"); ("d", "b") ("d", "a")] - @test a[i1, i2] == [("b", "c") ("b", "d"); ("a", "c") ("a", "d")] - @test a[i] == [("b", "c") ("b", "d"); ("a", "c") ("a", "d")] - @test a[i2, i2] == [("c", "c") ("c", "d"); ("d", "c") ("d", "d")] + a = collect(Iterators.product(v, v)) + i1 = BlockIndexVector{1, GenericBlockIndex{1, Tuple{Int}, Tuple{String}}}(Block(1), [2, 1]) + i2 = BlockIndexVector{1, GenericBlockIndex{1, Tuple{Int}, Tuple{String}}}(Block(2), [1, 2]) + i = BlockIndexVector{2, GenericBlockIndex{2, Tuple{Int, Int}, Tuple{String, String}}}( + Block(1, 2), ([2, 1], [1, 2]) + ) + @test a[i1, i1] == [("b", "b") ("b", "a"); ("a", "b") ("a", "a")] + @test a[i2, i1] == [("c", "b") ("c", "a"); ("d", "b") ("d", "a")] + @test a[i1, i2] == [("b", "c") ("b", "d"); ("a", "c") ("a", "d")] + @test a[i] == [("b", "c") ("b", "d"); ("a", "c") ("a", "d")] + @test a[i2, i2] == [("c", "c") ("c", "d"); ("d", "c") ("d", "d")] - r = blockedrange([2, 3]) - i = mortar([BlockIndexVector(Block(2), [1]), BlockIndexVector(Block(1), [1, 2])]) - # TODO: Check the indices make sense and are in bounds. - @test BlockSparseArrays.blockrange(r, i) == [Block(2), Block(1)] + r = blockedrange([2, 3]) + i = mortar([BlockIndexVector(Block(2), [1]), BlockIndexVector(Block(1), [1, 2])]) + # TODO: Check the indices make sense and are in bounds. + @test BlockSparseArrays.blockrange(r, i) == [Block(2), Block(1)] - r = blockedrange([2, 3]) - i = mortar([BlockIndexVector(Block(2), ["x"]), BlockIndexVector(Block(1), ["y", "z"])]) - # TODO: Check the indices make sense and are in bounds. - @test BlockSparseArrays.blockrange(r, i) == [Block(2), Block(1)] + r = blockedrange([2, 3]) + i = mortar([BlockIndexVector(Block(2), ["x"]), BlockIndexVector(Block(1), ["y", "z"])]) + # TODO: Check the indices make sense and are in bounds. + @test BlockSparseArrays.blockrange(r, i) == [Block(2), Block(1)] - i = GenericBlockIndex(Block(2), 1) - # TODO: Is this a good definition? - @test Base.to_index(i) === i + i = GenericBlockIndex(Block(2), 1) + # TODO: Is this a good definition? + @test Base.to_index(i) === i - r = blockedrange([2, 3]) - @test checkindex(Bool, r, GenericBlockIndex(Block(1), 1)) - @test checkindex(Bool, r, GenericBlockIndex(Block(1), 2)) - @test !checkindex(Bool, r, GenericBlockIndex(Block(1), 3)) - @test checkindex(Bool, r, GenericBlockIndex(Block(2), 1)) - @test checkindex(Bool, r, GenericBlockIndex(Block(2), 2)) - @test checkindex(Bool, r, GenericBlockIndex(Block(2), 3)) - @test !checkindex(Bool, r, GenericBlockIndex(Block(2), 4)) - @test !checkindex(Bool, r, GenericBlockIndex(Block(3), 1)) + r = blockedrange([2, 3]) + @test checkindex(Bool, r, GenericBlockIndex(Block(1), 1)) + @test checkindex(Bool, r, GenericBlockIndex(Block(1), 2)) + @test !checkindex(Bool, r, GenericBlockIndex(Block(1), 3)) + @test checkindex(Bool, r, GenericBlockIndex(Block(2), 1)) + @test checkindex(Bool, r, GenericBlockIndex(Block(2), 2)) + @test checkindex(Bool, r, GenericBlockIndex(Block(2), 3)) + @test !checkindex(Bool, r, GenericBlockIndex(Block(2), 4)) + @test !checkindex(Bool, r, GenericBlockIndex(Block(3), 1)) - a = BlockedArray(randn(5, 5), [2, 3], [2, 3]) - i = GenericBlockIndex(Block(1), 1) - @test to_indices(a, (i, i)) == (1, 1) - @test to_indices(a, axes(a), (i, i)) == (1, 1) - i = GenericBlockIndex(Block(2), 2) - @test to_indices(a, (i, i)) == (4, 4) - @test to_indices(a, axes(a), (i, i)) == (4, 4) + a = BlockedArray(randn(5, 5), [2, 3], [2, 3]) + i = GenericBlockIndex(Block(1), 1) + @test to_indices(a, (i, i)) == (1, 1) + @test to_indices(a, axes(a), (i, i)) == (1, 1) + i = GenericBlockIndex(Block(2), 2) + @test to_indices(a, (i, i)) == (4, 4) + @test to_indices(a, axes(a), (i, i)) == (4, 4) - r = blockedrange([2, 3]) - i = mortar([BlockIndexVector(Block(2), [1, 3]), BlockIndexVector(Block(1), [2])]) - @test blockedunitrange_getindices(r, i) == mortar([[3, 5], [2]]) + r = blockedrange([2, 3]) + i = mortar([BlockIndexVector(Block(2), [1, 3]), BlockIndexVector(Block(1), [2])]) + @test blockedunitrange_getindices(r, i) == mortar([[3, 5], [2]]) - r = blockedrange([2, 3]) - T = GenericBlockIndex{1,Tuple{Int},Tuple{Int}} - i = mortar([ - BlockIndexVector{1,T}(Block(2), [1, 3]), BlockIndexVector{1,T}(Block(1), [2]) - ]) - @test blockedunitrange_getindices(r, i) == mortar([[3, 5], [2]]) + r = blockedrange([2, 3]) + T = GenericBlockIndex{1, Tuple{Int}, Tuple{Int}} + i = mortar( + [ + BlockIndexVector{1, T}(Block(2), [1, 3]), BlockIndexVector{1, T}(Block(1), [2]), + ] + ) + @test blockedunitrange_getindices(r, i) == mortar([[3, 5], [2]]) - # Internal functions. - @test to_block(Block(2)) === Block(2) - @test to_block(Block(2)[2:3]) === Block(2) - @test to_block(BlockIndexVector(Block(2), [1, 3])) === Block(2) - @test to_block_indices(Block(2)) === (:) - @test to_block_indices(Block(2)[2:3]) === 2:3 - @test to_block_indices(BlockIndexVector(Block(2), [1, 3])) == [1, 3] + # Internal functions. + @test to_block(Block(2)) === Block(2) + @test to_block(Block(2)[2:3]) === Block(2) + @test to_block(BlockIndexVector(Block(2), [1, 3])) === Block(2) + @test to_block_indices(Block(2)) === (:) + @test to_block_indices(Block(2)[2:3]) === 2:3 + @test to_block_indices(BlockIndexVector(Block(2), [1, 3])) == [1, 3] - a = blocksparsezeros([2, 3], [2, 3]) - i = Block(2)[2:3] - @test to_indices(a, (i, i)) === - to_indices(a, axes(a), (i, i)) === - (BlockSlice(i, 4:5), BlockSlice(i, 4:5)) + a = blocksparsezeros([2, 3], [2, 3]) + i = Block(2)[2:3] + @test to_indices(a, (i, i)) === + to_indices(a, axes(a), (i, i)) === + (BlockSlice(i, 4:5), BlockSlice(i, 4:5)) - a = blocksparsezeros([2, 3], [2, 3]) - i = mortar([Block(2)[2:3], Block(1)[2:2]]) - for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) - for Iⱼ in I - @test Iⱼ == BlockIndices(i, mortar([4:5, 2:2])) - @test Iⱼ isa BlockIndices - @test Iⱼ.blocks == i - @test Iⱼ.indices == mortar([4:5, 2:2]) + a = blocksparsezeros([2, 3], [2, 3]) + i = mortar([Block(2)[2:3], Block(1)[2:2]]) + for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) + for Iⱼ in I + @test Iⱼ == BlockIndices(i, mortar([4:5, 2:2])) + @test Iⱼ isa BlockIndices + @test Iⱼ.blocks == i + @test Iⱼ.indices == mortar([4:5, 2:2]) + end end - end - a = blocksparsezeros([2, 3], [2, 3]) - i = BlockIndexVector(Block(2), [1, 3]) - for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) - for Iⱼ in I - @test Iⱼ == BlockIndices(i, [3, 5]) - @test Iⱼ isa BlockIndices - @test Iⱼ.blocks == i - @test Iⱼ.indices == [3, 5] + a = blocksparsezeros([2, 3], [2, 3]) + i = BlockIndexVector(Block(2), [1, 3]) + for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) + for Iⱼ in I + @test Iⱼ == BlockIndices(i, [3, 5]) + @test Iⱼ isa BlockIndices + @test Iⱼ.blocks == i + @test Iⱼ.indices == [3, 5] + end end - end - a = blocksparsezeros([2, 3], [2, 3]) - T = GenericBlockIndex{1,Tuple{Int},Tuple{Int}} - i = BlockIndexVector{1,T}(Block(2), [1, 3]) - for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) - for Iⱼ in I - @test Iⱼ == BlockIndices(i, [3, 5]) - @test Iⱼ isa BlockIndices - @test Iⱼ.blocks == i - @test Iⱼ.indices == [3, 5] + a = blocksparsezeros([2, 3], [2, 3]) + T = GenericBlockIndex{1, Tuple{Int}, Tuple{Int}} + i = BlockIndexVector{1, T}(Block(2), [1, 3]) + for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) + for Iⱼ in I + @test Iⱼ == BlockIndices(i, [3, 5]) + @test Iⱼ isa BlockIndices + @test Iⱼ.blocks == i + @test Iⱼ.indices == [3, 5] + end end - end - a = blocksparsezeros([2, 3], [2, 3]) - i = mortar([BlockIndexVector(Block(2), [1, 3]), BlockIndexVector(Block(1), [2])]) - for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) - for Iⱼ in I - @test Iⱼ == BlockIndices(i, mortar([[3, 5], [2]])) - @test Iⱼ isa BlockIndices - @test Iⱼ.blocks == i - @test Iⱼ.indices == mortar([[3, 5], [2]]) + a = blocksparsezeros([2, 3], [2, 3]) + i = mortar([BlockIndexVector(Block(2), [1, 3]), BlockIndexVector(Block(1), [2])]) + for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) + for Iⱼ in I + @test Iⱼ == BlockIndices(i, mortar([[3, 5], [2]])) + @test Iⱼ isa BlockIndices + @test Iⱼ.blocks == i + @test Iⱼ.indices == mortar([[3, 5], [2]]) + end end - end - a = blocksparsezeros([2, 3], [2, 3]) - i = [BlockIndexVector(Block(2), [1, 3]), BlockIndexVector(Block(1), [2])] - for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) - for Iⱼ in I - @test Iⱼ == BlockIndices(mortar(i), mortar([[3, 5], [2]])) - @test Iⱼ isa BlockIndices - @test Iⱼ.blocks == mortar(i) - @test Iⱼ.indices == mortar([[3, 5], [2]]) + a = blocksparsezeros([2, 3], [2, 3]) + i = [BlockIndexVector(Block(2), [1, 3]), BlockIndexVector(Block(1), [2])] + for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) + for Iⱼ in I + @test Iⱼ == BlockIndices(mortar(i), mortar([[3, 5], [2]])) + @test Iⱼ isa BlockIndices + @test Iⱼ.blocks == mortar(i) + @test Iⱼ.indices == mortar([[3, 5], [2]]) + end end - end - a = blocksparsezeros([2, 3], [2, 3]) - T = GenericBlockIndex{1,Tuple{Int},Tuple{Int}} - i = mortar([ - BlockIndexVector{1,T}(Block(2), [1, 3]), BlockIndexVector{1,T}(Block(1), [2]) - ]) - for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) - for Iⱼ in I - @test Iⱼ == BlockIndices(i, mortar([[3, 5], [2]])) - @test Iⱼ isa BlockIndices - @test Iⱼ.blocks == i - @test Iⱼ.indices == mortar([[3, 5], [2]]) + a = blocksparsezeros([2, 3], [2, 3]) + T = GenericBlockIndex{1, Tuple{Int}, Tuple{Int}} + i = mortar( + [ + BlockIndexVector{1, T}(Block(2), [1, 3]), BlockIndexVector{1, T}(Block(1), [2]), + ] + ) + for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) + for Iⱼ in I + @test Iⱼ == BlockIndices(i, mortar([[3, 5], [2]])) + @test Iⱼ isa BlockIndices + @test Iⱼ.blocks == i + @test Iⱼ.indices == mortar([[3, 5], [2]]) + end end - end - a = blocksparsezeros([2, 3], [2, 3]) - T = GenericBlockIndex{1,Tuple{Int},Tuple{Int}} - i = [BlockIndexVector{1,T}(Block(2), [1, 3]), BlockIndexVector{1,T}(Block(1), [2])] - for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) - for Iⱼ in I - @test Iⱼ == BlockIndices(mortar(i), mortar([[3, 5], [2]])) - @test Iⱼ isa BlockIndices - @test Iⱼ.blocks == mortar(i) - @test Iⱼ.indices == mortar([[3, 5], [2]]) + a = blocksparsezeros([2, 3], [2, 3]) + T = GenericBlockIndex{1, Tuple{Int}, Tuple{Int}} + i = [BlockIndexVector{1, T}(Block(2), [1, 3]), BlockIndexVector{1, T}(Block(1), [2])] + for I in (to_indices(a, (i, i)), to_indices(a, axes(a), (i, i))) + for Iⱼ in I + @test Iⱼ == BlockIndices(mortar(i), mortar([[3, 5], [2]])) + @test Iⱼ isa BlockIndices + @test Iⱼ.blocks == mortar(i) + @test Iⱼ.indices == mortar([[3, 5], [2]]) + end end - end end diff --git a/test/test_issues.jl b/test/test_issues.jl index 0ce92ca0..cb849486 100644 --- a/test/test_issues.jl +++ b/test/test_issues.jl @@ -6,20 +6,20 @@ using LinearAlgebra: LinearAlgebra using Test: @test, @testset @testset "Issue 162" begin - ax = (blockedrange([2, 2]), blockedrange([2, 2, 2])) - bs = Dict(Block(1, 1) => randn(2, 2), Block(2, 3) => randn(2, 2)) - a = blocksparse(bs, ax) - U, S, Vᴴ = svd_compact(a) + ax = (blockedrange([2, 2]), blockedrange([2, 2, 2])) + bs = Dict(Block(1, 1) => randn(2, 2), Block(2, 3) => randn(2, 2)) + a = blocksparse(bs, ax) + U, S, Vᴴ = svd_compact(a) - @test U * S * Vᴴ ≈ a - @test U' * U ≈ LinearAlgebra.I - @test Vᴴ * Vᴴ' ≈ LinearAlgebra.I + @test U * S * Vᴴ ≈ a + @test U' * U ≈ LinearAlgebra.I + @test Vᴴ * Vᴴ' ≈ LinearAlgebra.I - U, S, Vᴴ = svd_full(a); + U, S, Vᴴ = svd_full(a) - @test U * S * Vᴴ ≈ a - @test U' * U ≈ LinearAlgebra.I - @test U * U' ≈ LinearAlgebra.I - @test Vᴴ * Vᴴ' ≈ LinearAlgebra.I - @test Vᴴ' * Vᴴ ≈ LinearAlgebra.I + @test U * S * Vᴴ ≈ a + @test U' * U ≈ LinearAlgebra.I + @test U * U' ≈ LinearAlgebra.I + @test Vᴴ * Vᴴ' ≈ LinearAlgebra.I + @test Vᴴ' * Vᴴ ≈ LinearAlgebra.I end diff --git a/test/test_map.jl b/test/test_map.jl index 1f41b7af..34c5cfdd 100644 --- a/test/test_map.jl +++ b/test/test_map.jl @@ -1,15 +1,15 @@ using Adapt: adapt using BlockArrays: Block, BlockRange, blockedrange, blocksize, blocksizes, mortar using BlockSparseArrays: - BlockSparseArray, - BlockSparseVector, - BlockVector, - BlockView, - BlockedVector, - blocklengths, - blockstoredlength, - blocktype, - eachblockstoredindex + BlockSparseArray, + BlockSparseVector, + BlockVector, + BlockView, + BlockedVector, + blocklengths, + blockstoredlength, + blocktype, + eachblockstoredindex using DerivableInterfaces: zero! using GPUArraysCore: @allowscalar using JLArrays: JLArray @@ -20,711 +20,711 @@ 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 - - a = dev(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 - - a = dev(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])) - - # Regression test for 0-dimensional in-place broadcasting. - rng = StableRNG(123) - a = dev(BlockSparseArray{elt}(undef)) - @allowscalar a[] = randn(rng, elt) - b = dev(BlockSparseArray{elt}(undef)) - @allowscalar b[] = randn(rng, elt) - c = similar(a) - c .= 2 .* a .+ 3 .* b - @allowscalar @test c[] == 2 * a[] + 3 * b[] - - 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) + 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 + + a = dev(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 + + a = dev(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])) + + # Regression test for 0-dimensional in-place broadcasting. + rng = StableRNG(123) + a = dev(BlockSparseArray{elt}(undef)) + @allowscalar a[] = randn(rng, elt) + b = dev(BlockSparseArray{elt}(undef)) + @allowscalar b[] = randn(rng, elt) + c = similar(a) + c .= 2 .* a .+ 3 .* b + @allowscalar @test c[] == 2 * a[] + 3 * b[] + + 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 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 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(2, 1)] == a[Block(1, 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)] - @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 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 - end - function f1() + 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), Block(1)], [Block(2), Block(1)]] - x = randn(elt, 3, 4) - b[Block(1, 1)] .= x - return (; a, b, x) - end - function f2() + 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), 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]) + 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 - @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] + + # 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 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] + @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 index 884ab0d2..cdb17996 100644 --- a/test/test_similartype_unchecked.jl +++ b/test/test_similartype_unchecked.jl @@ -4,23 +4,23 @@ 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 + @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_tensoralgebraext.jl b/test/test_tensoralgebraext.jl index 7caaffd9..019080cb 100644 --- a/test/test_tensoralgebraext.jl +++ b/test/test_tensoralgebraext.jl @@ -5,73 +5,73 @@ using TensorAlgebra: contract using Test: @test, @test_broken, @testset function randn_blockdiagonal(elt::Type, axes::Tuple) - a = BlockSparseArray{elt}(undef, axes) - blockdiaglength = minimum(blocksize(a)) - for i in 1:blockdiaglength - b = Block(ntuple(Returns(i), ndims(a))) - a[b] = randn!(a[b]) - end - return a + a = BlockSparseArray{elt}(undef, axes) + blockdiaglength = minimum(blocksize(a)) + for i in 1:blockdiaglength + b = Block(ntuple(Returns(i), ndims(a))) + a[b] = randn!(a[b]) + end + return a end @testset "Regression test for BlockArrays" begin - # test https://github.com/ITensor/BlockSparseArrays.jl/issues/57 - d = blockedrange([1, 1]) - a = BlockArray(ones((d, d, d, d))) - @test contract((-1, -2, -3, -4), a, (1, -1, 2, -2), a, (2, -3, 1, -4)) isa BlockArray + # test https://github.com/ITensor/BlockSparseArrays.jl/issues/57 + d = blockedrange([1, 1]) + a = BlockArray(ones((d, d, d, d))) + @test contract((-1, -2, -3, -4), a, (1, -1, 2, -2), a, (2, -3, 1, -4)) isa BlockArray end const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @testset "`contract` `BlockSparseArray` (eltype=$elt)" for elt in elts - @testset "BlockedOneTo" begin - d = blockedrange([2, 3]) - a1 = randn_blockdiagonal(elt, (d, d, d, d)) - a2 = randn_blockdiagonal(elt, (d, d, d, d)) - a3 = randn_blockdiagonal(elt, (d, d)) - a1_dense = convert(Array, a1) - a2_dense = convert(Array, a2) - a3_dense = convert(Array, a3) + @testset "BlockedOneTo" begin + d = blockedrange([2, 3]) + a1 = randn_blockdiagonal(elt, (d, d, d, d)) + a2 = randn_blockdiagonal(elt, (d, d, d, d)) + a3 = randn_blockdiagonal(elt, (d, d)) + a1_dense = convert(Array, a1) + a2_dense = convert(Array, a2) + a3_dense = convert(Array, a3) - # matrix matrix - a_dest, dimnames_dest = contract(a1, (1, -1, 2, -2), a2, (2, -3, 1, -4)) - a_dest_dense, dimnames_dest_dense = contract( - a1_dense, (1, -1, 2, -2), a2_dense, (2, -3, 1, -4) - ) - @test dimnames_dest == dimnames_dest_dense - @test size(a_dest) == size(a_dest_dense) - @test a_dest isa BlockSparseArray - @test a_dest ≈ a_dest_dense + # matrix matrix + a_dest, dimnames_dest = contract(a1, (1, -1, 2, -2), a2, (2, -3, 1, -4)) + a_dest_dense, dimnames_dest_dense = contract( + a1_dense, (1, -1, 2, -2), a2_dense, (2, -3, 1, -4) + ) + @test dimnames_dest == dimnames_dest_dense + @test size(a_dest) == size(a_dest_dense) + @test a_dest isa BlockSparseArray + @test a_dest ≈ a_dest_dense - # matrix vector - a_dest, dimnames_dest = contract(a1, (2, -1, -2, 1), a3, (1, 2)) - a_dest_dense, dimnames_dest_dense = contract(a1_dense, (2, -1, -2, 1), a3_dense, (1, 2)) - @test dimnames_dest == dimnames_dest_dense - @test size(a_dest) == size(a_dest_dense) - @test a_dest isa BlockSparseArray - @test a_dest ≈ a_dest_dense + # matrix vector + a_dest, dimnames_dest = contract(a1, (2, -1, -2, 1), a3, (1, 2)) + a_dest_dense, dimnames_dest_dense = contract(a1_dense, (2, -1, -2, 1), a3_dense, (1, 2)) + @test dimnames_dest == dimnames_dest_dense + @test size(a_dest) == size(a_dest_dense) + @test a_dest isa BlockSparseArray + @test a_dest ≈ a_dest_dense - # vector matrix - a_dest, dimnames_dest = contract(a3, (1, 2), a1, (2, -1, -2, 1)) - a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a1_dense, (2, -1, -2, 1)) - @test dimnames_dest == dimnames_dest_dense - @test size(a_dest) == size(a_dest_dense) - @test a_dest isa BlockSparseArray - @test a_dest ≈ a_dest_dense + # vector matrix + a_dest, dimnames_dest = contract(a3, (1, 2), a1, (2, -1, -2, 1)) + a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a1_dense, (2, -1, -2, 1)) + @test dimnames_dest == dimnames_dest_dense + @test size(a_dest) == size(a_dest_dense) + @test a_dest isa BlockSparseArray + @test a_dest ≈ a_dest_dense - # vector vector - a_dest, dimnames_dest = contract(a3, (1, 2), a3, (2, 1)) - a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (2, 1)) - @test dimnames_dest == dimnames_dest_dense - @test size(a_dest) == size(a_dest_dense) - @test a_dest isa BlockSparseArray{elt,0} - @test a_dest ≈ a_dest_dense + # vector vector + a_dest, dimnames_dest = contract(a3, (1, 2), a3, (2, 1)) + a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (2, 1)) + @test dimnames_dest == dimnames_dest_dense + @test size(a_dest) == size(a_dest_dense) + @test a_dest isa BlockSparseArray{elt, 0} + @test a_dest ≈ a_dest_dense - # outer product - a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (3, 4)) - a_dest, dimnames_dest = contract(a3, (1, 2), a3, (3, 4)) - @test dimnames_dest == dimnames_dest_dense - @test size(a_dest) == size(a_dest_dense) - @test a_dest isa BlockSparseArray - @test a_dest ≈ a_dest_dense - end + # outer product + a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (3, 4)) + a_dest, dimnames_dest = contract(a3, (1, 2), a3, (3, 4)) + @test dimnames_dest == dimnames_dest_dense + @test size(a_dest) == size(a_dest_dense) + @test a_dest isa BlockSparseArray + @test a_dest ≈ a_dest_dense + end end diff --git a/test/test_view_bang.jl b/test/test_view_bang.jl index 62eb09d5..329fd23a 100644 --- a/test/test_view_bang.jl +++ b/test/test_view_bang.jl @@ -1,7 +1,7 @@ using Adapt: adapt using BlockArrays: Block using BlockSparseArrays: - BlockSparseArray, @view!, blockstoredlength, eachblockstoredindex, view! + BlockSparseArray, @view!, blockstoredlength, eachblockstoredindex, view! using JLArrays: JLArray using SparseArraysBase: isstored using Test: @test, @testset @@ -9,68 +9,68 @@ 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) + 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 + 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 From f06366f4058a5d7230dd22572ac1a43352a6b9e2 Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Thu, 2 Oct 2025 13:38:18 -0400 Subject: [PATCH 2/2] updates from skeleton --- .JuliaFormatter.toml | 3 --- .github/workflows/CompatHelper.yml | 2 +- .github/workflows/FormatCheck.yml | 13 ++++++++----- .pre-commit-config.yaml | 8 ++++---- README.md | 4 ++-- docs/make.jl | 2 +- examples/README.jl | 4 ++-- 7 files changed, 18 insertions(+), 18 deletions(-) delete mode 100644 .JuliaFormatter.toml diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml deleted file mode 100644 index 4c49a86f..00000000 --- a/.JuliaFormatter.toml +++ /dev/null @@ -1,3 +0,0 @@ -# See https://domluna.github.io/JuliaFormatter.jl/stable/ for a list of options -style = "blue" -indent = 2 diff --git a/.github/workflows/CompatHelper.yml b/.github/workflows/CompatHelper.yml index 456fa05d..0614de9d 100644 --- a/.github/workflows/CompatHelper.yml +++ b/.github/workflows/CompatHelper.yml @@ -2,7 +2,7 @@ name: "CompatHelper" on: schedule: - - cron: 0 0 * * * + - cron: '0 0 * * *' workflow_dispatch: permissions: contents: write diff --git a/.github/workflows/FormatCheck.yml b/.github/workflows/FormatCheck.yml index 3f78afc2..1525861f 100644 --- a/.github/workflows/FormatCheck.yml +++ b/.github/workflows/FormatCheck.yml @@ -1,11 +1,14 @@ name: "Format Check" on: - push: - branches: - - 'main' - tags: '*' - pull_request: + pull_request_target: + paths: ['**/*.jl'] + types: [opened, synchronize, reopened, ready_for_review] + +permissions: + contents: read + actions: write + pull-requests: write jobs: format-check: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 88bc8b49..3fc47436 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,5 +1,5 @@ ci: - skip: [julia-formatter] + skip: [runic] repos: - repo: https://github.com/pre-commit/pre-commit-hooks @@ -11,7 +11,7 @@ repos: - id: end-of-file-fixer exclude_types: [markdown] # incompatible with Literate.jl -- repo: "https://github.com/domluna/JuliaFormatter.jl" - rev: v2.1.6 +- repo: https://github.com/fredrikekre/runic-pre-commit + rev: v2.0.1 hooks: - - id: "julia-formatter" + - id: runic diff --git a/README.md b/README.md index 20e490b7..11ae916a 100644 --- a/README.md +++ b/README.md @@ -27,11 +27,11 @@ This step is only required once. ```julia julia> using Pkg: Pkg -julia> Pkg.Registry.add(url="https://github.com/ITensor/ITensorRegistry") +julia> Pkg.Registry.add(url = "https://github.com/ITensor/ITensorRegistry") ``` or: ```julia -julia> Pkg.Registry.add(url="git@github.com:ITensor/ITensorRegistry.git") +julia> Pkg.Registry.add(url = "git@github.com:ITensor/ITensorRegistry.git") ``` if you want to use SSH credentials, which can make it so you don't have to enter your Github ursername and password when registering packages. diff --git a/docs/make.jl b/docs/make.jl index 4fc83b0e..39aab186 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -15,7 +15,7 @@ makedocs(; authors = "ITensor developers and contributors", sitename = "BlockSparseArrays.jl", format = Documenter.HTML(; - canonical = "https://ITensor.github.io/BlockSparseArrays.jl", + canonical = "https://itensor.github.io/BlockSparseArrays.jl", edit_link = "main", assets = ["assets/favicon.ico", "assets/extras.css"], ), diff --git a/examples/README.jl b/examples/README.jl index f97e3f44..aa895473 100644 --- a/examples/README.jl +++ b/examples/README.jl @@ -24,13 +24,13 @@ ```julia julia> using Pkg: Pkg -julia> Pkg.Registry.add(url="https://github.com/ITensor/ITensorRegistry") +julia> Pkg.Registry.add(url = "https://github.com/ITensor/ITensorRegistry") ``` =# # or: #= ```julia -julia> Pkg.Registry.add(url="git@github.com:ITensor/ITensorRegistry.git") +julia> Pkg.Registry.add(url = "git@github.com:ITensor/ITensorRegistry.git") ``` =# # if you want to use SSH credentials, which can make it so you don't have to enter your Github ursername and password when registering packages.