diff --git a/Project.toml b/Project.toml index 7f0208a3..1e1903ec 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.8.1" +version = "0.9.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" @@ -22,9 +22,11 @@ TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138" [weakdeps] TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" +TensorProducts = "decf83d6-1968-43f4-96dc-fdb3fe15fc6d" [extensions] BlockSparseArraysTensorAlgebraExt = "TensorAlgebra" +BlockSparseArraysTensorProductsExt = "TensorProducts" [compat] Adapt = "4.1.1" @@ -43,6 +45,7 @@ MatrixAlgebraKit = "0.2.2" SparseArraysBase = "0.7.1" SplitApplyCombine = "1.2.3" TensorAlgebra = "0.3.2" +TensorProducts = "0.1.7" Test = "1.10" TypeParameterAccessors = "0.4.1" julia = "1.10" diff --git a/docs/Project.toml b/docs/Project.toml index 7885be60..edbaa0bc 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -6,6 +6,6 @@ Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306" [compat] BlockArrays = "1" -BlockSparseArrays = "0.8" +BlockSparseArrays = "0.9" Documenter = "1" Literate = "2" diff --git a/examples/Project.toml b/examples/Project.toml index 6c17bd70..7e02e0de 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -5,5 +5,5 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] BlockArrays = "1" -BlockSparseArrays = "0.8" +BlockSparseArrays = "0.9" Test = "1" diff --git a/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl b/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl index d5551f43..7912ef51 100644 --- a/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl +++ b/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl @@ -15,19 +15,51 @@ function TensorAlgebra.FusionStyle(::Type{<:AbstractBlockSparseArray}) return BlockReshapeFusion() end +using BlockArrays: Block, blocklength, blocks +using BlockSparseArrays: blocksparse +using SparseArraysBase: eachstoredindex +using TensorAlgebra: TensorAlgebra, matricize, unmatricize function TensorAlgebra.matricize( ::BlockReshapeFusion, a::AbstractArray, biperm::BlockedTrivialPermutation{2} ) - new_axes = fuseaxes(axes(a), biperm) - return blockreshape(a, new_axes) + 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_axes::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}}, + blocked_ax::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}}, ) - return blockreshape(m, Tuple(blocked_axes)...) + 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 new file mode 100644 index 00000000..d3733d89 --- /dev/null +++ b/ext/BlockSparseArraysTensorProductsExt/BlockSparseArraysTensorProductsExt.jl @@ -0,0 +1,14 @@ +module BlockSparseArraysTensorProductsExt + +using BlockSparseArrays: BlockUnitRange, blockrange, eachblockaxis +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) +end + +end diff --git a/src/BlockArraysExtensions/BlockArraysExtensions.jl b/src/BlockArraysExtensions/BlockArraysExtensions.jl index 58155403..4a86a980 100644 --- a/src/BlockArraysExtensions/BlockArraysExtensions.jl +++ b/src/BlockArraysExtensions/BlockArraysExtensions.jl @@ -180,7 +180,10 @@ end # ``` # but includes `BlockIndices`, where the blocks aren't contiguous. const BlockSliceCollection = Union{ - Base.Slice,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}}} @@ -267,13 +270,15 @@ tuple_oneto(n) = ntuple(identity, n) function _blockreshape(a::AbstractArray, axes::Tuple{Vararg{AbstractUnitRange}}) reshaped_blocks_a = reshape(blocks(a), blocklength.(axes)) - reshaped_a = similar(a, axes) - for I in eachstoredindex(reshaped_blocks_a) - block_size_I = map(i -> length(axes[i][Block(I[i])]), tuple_oneto(length(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. - reshaped_a[Block(Tuple(I))] = reshape(reshaped_blocks_a[I], block_size_I) + return reshape(reshaped_blocks_a[I], block_axes_I) end - return reshaped_a + bs = Dict(Block(Tuple(I)) => f(I) for I in eachstoredindex(reshaped_blocks_a)) + return blocksparse(bs, axes) end function blockreshape( diff --git a/src/BlockArraysExtensions/blockrange.jl b/src/BlockArraysExtensions/blockrange.jl index a1f291eb..ea991e13 100644 --- a/src/BlockArraysExtensions/blockrange.jl +++ b/src/BlockArraysExtensions/blockrange.jl @@ -8,6 +8,9 @@ end function blockrange(eachblockaxis) return BlockUnitRange(blockedrange(length.(eachblockaxis)), eachblockaxis) end +function blockrange(first::Integer, 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) @@ -15,6 +18,11 @@ eachblockaxis(r::BlockUnitRange) = r.eachblockaxis function Base.getindex(r::BlockUnitRange, I::Block{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)]) +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} diff --git a/src/abstractblocksparsearray/abstractblocksparsearray.jl b/src/abstractblocksparsearray/abstractblocksparsearray.jl index 3d688e45..5d5805cd 100644 --- a/src/abstractblocksparsearray/abstractblocksparsearray.jl +++ b/src/abstractblocksparsearray/abstractblocksparsearray.jl @@ -87,12 +87,19 @@ function Base.setindex!( ), ) end - # Custom `_convert` works around the issue that - # `convert(::Type{<:Diagonal}, ::AbstractMatrix)` isnt' defined - # in Julia v1.10 (https://github.com/JuliaLang/julia/pull/48895, - # https://github.com/JuliaLang/julia/pull/52487). - # TODO: Delete once we drop support for Julia v1.10. - blocks(a)[Int.(I)...] = _convert(blocktype(a), value) + 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 diff --git a/src/blocksparsearray/blocksparsearray.jl b/src/blocksparsearray/blocksparsearray.jl index 15af2f77..f708f345 100644 --- a/src/blocksparsearray/blocksparsearray.jl +++ b/src/blocksparsearray/blocksparsearray.jl @@ -262,13 +262,18 @@ function blocksparsezeros(::BlockType{A}, axes...) where {A<:AbstractArray} # to make a bit more generic. return BlockSparseArray{eltype(A),ndims(A),A}(undef, axes...) end -function blocksparse(d::Dict{<:Block,<:AbstractArray}, axes...) - a = blocksparsezeros(BlockType(valtype(d)), axes...) +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)) +end # Base `AbstractArray` interface Base.axes(a::BlockSparseArray) = a.axes diff --git a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index 30d75baf..e017e286 100644 --- a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -537,7 +537,11 @@ function SparseArraysBase.getunstoredindex( return error("Not implemented.") end +# Convert a blockwise slice on a block sparse array +# to an elementwise slice on the corresponding sparse array +# of blocks. to_blocks_indices(I::BlockSlice{<:BlockRange{1}}) = Int.(I.block) +to_blocks_indices(I::BlockSlice{<:Block{1}}) = Int(I.block):Int(I.block) to_blocks_indices(I::BlockIndices{<:Vector{<:Block{1}}}) = Int.(I.blocks) to_blocks_indices(I::Base.Slice) = Base.OneTo(blocklength(I.indices)) diff --git a/src/blocksparsearrayinterface/getunstoredblock.jl b/src/blocksparsearrayinterface/getunstoredblock.jl index d704f05a..0669096f 100644 --- a/src/blocksparsearrayinterface/getunstoredblock.jl +++ b/src/blocksparsearrayinterface/getunstoredblock.jl @@ -18,10 +18,17 @@ function Base.AbstractArray{A}(a::ZeroBlocks{N}) where {N,A} end @inline function Base.getindex(a::ZeroBlocks{N,A}, I::Vararg{Int,N}) where {N,A} + + #foreach(display, a.parentaxes) + #foreach(display ∘ eachblockaxis, a.parentaxes) + # TODO: Use `BlockArrays.eachblockaxes`. ax = ntuple(N) do d - return only(axes(a.parentaxes[d][Block(I[d])])) + return eachblockaxis(a.parentaxes[d])[I[d]] end + + #@show ax + !isconcretetype(A) && return zero!(similar(Array{eltype(A),N}, ax)) return zero!(similar(A, ax)) end diff --git a/test/Project.toml b/test/Project.toml index 52ef8620..12fb6f03 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -25,7 +25,7 @@ Adapt = "4" Aqua = "0.8" ArrayLayouts = "1" BlockArrays = "1" -BlockSparseArrays = "0.8" +BlockSparseArrays = "0.9" DerivableInterfaces = "0.5" DiagonalArrays = "0.3" GPUArraysCore = "0.2"