diff --git a/Project.toml b/Project.toml index 56b1b8c..7b7e77c 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "GradedArrays" uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2" authors = ["ITensor developers and contributors"] -version = "0.2.3" +version = "0.2.4" [deps] BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" @@ -23,8 +23,8 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" GradedArraysTensorAlgebraExt = "TensorAlgebra" [compat] -BlockArrays = "1.5.0" -BlockSparseArrays = "0.4.0" +BlockArrays = "1.6.0" +BlockSparseArrays = "0.4.2" Compat = "4.16.0" DerivableInterfaces = "0.4.4" FillArrays = "1.13.0" @@ -32,7 +32,7 @@ HalfIntegers = "1.6.0" LinearAlgebra = "1.10.0" Random = "1.10.0" SplitApplyCombine = "1.2.3" -TensorAlgebra = "0.2.7" +TensorAlgebra = "0.3.2" TensorProducts = "0.1.3" TypeParameterAccessors = "0.3.9" julia = "1.10" diff --git a/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl b/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl index 9b11d25..48b64f6 100644 --- a/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl +++ b/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl @@ -1,109 +1,76 @@ module GradedArraysTensorAlgebraExt -using BlockArrays: Block, BlockIndexRange, blockedrange, blocks -using BlockSparseArrays: - BlockSparseArrays, - AbstractBlockSparseArray, - AbstractBlockSparseArrayInterface, - BlockSparseArray, - BlockSparseArrayInterface, - BlockSparseMatrix, - BlockSparseVector, - block_merge -using DerivableInterfaces: @interface +using BlockArrays: blocks +using BlockSparseArrays: BlockSparseArray, blockreshape +using GradedArrays: GradedArray using GradedArrays.GradedUnitRanges: - GradedUnitRanges, AbstractGradedUnitRange, blockmergesortperm, blocksortperm, - dual, + flip, invblockperm, - nondual, unmerged_tensor_product -using LinearAlgebra: Adjoint, Transpose +using GradedArrays.SymmetrySectors: trivial using TensorAlgebra: - TensorAlgebra, FusionStyle, BlockReshapeFusion, SectorFusion, fusedims, splitdims -using TensorProducts: OneToOne + TensorAlgebra, + AbstractBlockPermutation, + BlockedTuple, + FusionStyle, + trivial_axis, + unmatricize -#= - reducewhile(f, op, collection, state) +struct SectorFusion <: FusionStyle end -reducewhile(x -> length(x) < 3, vcat, ["a", "b", "c", "d"], 2; init=String[]) == - (["b", "c"], 4) -=# -function reducewhile(f, op, collection, state; init) - prev_result = init - prev_state = state - result = prev_result - while f(result) - prev_result = result - prev_state = state - value_and_state = iterate(collection, state) - isnothing(value_and_state) && break - value, state = value_and_state - result = op(result, value) - end - return prev_result, prev_state -end - -#= - groupreducewhile(f, op, collection, ngroups) - -groupreducewhile((i, x) -> length(x) ≤ i, vcat, ["a", "b", "c", "d", "e", "f"], 3; init=String[]) == - (["a"], ["b", "c"], ["d", "e", "f"]) -=# -function groupreducewhile(f, op, collection, ngroups; init) - state = firstindex(collection) - return ntuple(ngroups) do group_number - result, state = reducewhile(x -> f(group_number, x), op, collection, state; init) - return result - end -end +TensorAlgebra.FusionStyle(::Type{<:GradedArray}) = SectorFusion() -TensorAlgebra.FusionStyle(::AbstractGradedUnitRange) = SectorFusion() +# TBD consider heterogeneous sectors? +TensorAlgebra.trivial_axis(t::Tuple{Vararg{AbstractGradedUnitRange}}) = trivial(first(t)) -# Sort the blocks by sector and then merge the common sectors. -function block_mergesort(a::AbstractArray) - I = blockmergesortperm.(axes(a)) - return a[I...] +function matricize_axes( + blocked_axes::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}} +) + @assert !isempty(blocked_axes) + default_axis = trivial_axis(Tuple(blocked_axes)) + codomain_axes, domain_axes = blocks(blocked_axes) + codomain_axis = unmerged_tensor_product(default_axis, codomain_axes...) + unflipped_domain_axis = unmerged_tensor_product(default_axis, domain_axes...) + return codomain_axis, flip(unflipped_domain_axis) end -function TensorAlgebra.fusedims( - ::SectorFusion, a::AbstractArray, merged_axes::AbstractUnitRange... +function TensorAlgebra.matricize( + ::SectorFusion, a::AbstractArray, biperm::AbstractBlockPermutation{2} ) - # First perform a fusion using a block reshape. - # TODO avoid groupreducewhile. Require refactor of fusedims. - unmerged_axes = groupreducewhile( - unmerged_tensor_product, axes(a), length(merged_axes); init=OneToOne() - ) do i, axis - return length(axis) ≤ length(merged_axes[i]) - end - - a_reshaped = fusedims(BlockReshapeFusion(), a, unmerged_axes...) + a_perm = permutedims(a, Tuple(biperm)) + codomain_axis, domain_axis = matricize_axes(axes(a)[biperm]) + a_reshaped = blockreshape(a_perm, (codomain_axis, domain_axis)) # Sort the blocks by sector and merge the equivalent sectors. - return block_mergesort(a_reshaped) + return sectormergesort(a_reshaped) end -function TensorAlgebra.splitdims( - ::SectorFusion, a::AbstractArray, split_axes::AbstractUnitRange... +function TensorAlgebra.unmatricize( + ::SectorFusion, + m::AbstractMatrix, + blocked_axes::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}}, ) # First, fuse axes to get `blockmergesortperm`. # Then unpermute the blocks. - axes_prod = groupreducewhile( - unmerged_tensor_product, split_axes, ndims(a); init=OneToOne() - ) do i, axis - return length(axis) ≤ length(axes(a, i)) - end - blockperms = blocksortperm.(axes_prod) - sorted_axes = map((r, I) -> only(axes(r[I])), axes_prod, blockperms) + fused_axes = matricize_axes(blocked_axes) + + blockperms = blocksortperm.(fused_axes) + sorted_axes = map((r, I) -> only(axes(r[I])), fused_axes, blockperms) # TODO: This is doing extra copies of the blocks, # use `@view a[axes_prod...]` instead. # That will require implementing some reindexing logic # for this combination of slicing. - a_unblocked = a[sorted_axes...] - a_blockpermed = a_unblocked[invblockperm.(blockperms)...] - return splitdims(BlockReshapeFusion(), a_blockpermed, split_axes...) + m_unblocked = m[sorted_axes...] + m_blockpermed = m_unblocked[invblockperm.(blockperms)...] + return unmatricize(FusionStyle(BlockSparseArray), m_blockpermed, blocked_axes) end +# Sort the blocks by sector and then merge the common sectors. +function sectormergesort(a::AbstractArray) + I = blockmergesortperm.(axes(a)) + return a[I...] +end end diff --git a/src/gradedarray.jl b/src/gradedarray.jl index 4d23d72..918b528 100644 --- a/src/gradedarray.jl +++ b/src/gradedarray.jl @@ -4,11 +4,19 @@ using BlockSparseArrays: AbstractBlockSparseMatrix, AnyAbstractBlockSparseArray, BlockSparseArray, + BlockSparseMatrix, + BlockSparseVector, blocktype using ..GradedUnitRanges: AbstractGradedUnitRange, dual using LinearAlgebra: Adjoint using TypeParameterAccessors: similartype, unwrap_array_type +const GradedArray{T,M,A,Blocks,Axes} = BlockSparseArray{ + T,M,A,Blocks,Axes +} where {Axes<:Tuple{AbstractGradedUnitRange,Vararg{AbstractGradedUnitRange}}} +const GradedMatrix{T,A,Blocks,Axes} = GradedArray{T,2,A,Blocks,Axes} +const GradedVector{T,A,Blocks,Axes} = GradedArray{T,1,A,Blocks,Axes} + # TODO: Handle this through some kind of trait dispatch, maybe # a `SymmetryStyle`-like trait to check if the block sparse # matrix has graded axes. diff --git a/test/Project.toml b/test/Project.toml index ed52e81..03fb5eb 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -15,15 +15,15 @@ TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a" [compat] Aqua = "0.8.11" -BlockArrays = "1.5.0" -BlockSparseArrays = "0.4.0" +BlockArrays = "1.6.0" +BlockSparseArrays = "0.4.2" GradedArrays = "0.2.0" LinearAlgebra = "1.10.0" Random = "1.10.0" SafeTestsets = "0.1.0" SparseArraysBase = "0.5.4" Suppressor = "0.2.8" -TensorAlgebra = "0.2.7" +TensorAlgebra = "0.3.2" TensorProducts = "0.1.3" Test = "1.10.0" TestExtras = "0.3.1" diff --git a/test/test_gradedarray.jl b/test/test_gradedarray.jl index 233c386..f23d095 100644 --- a/test/test_gradedarray.jl +++ b/test/test_gradedarray.jl @@ -2,12 +2,12 @@ using BlockArrays: AbstractBlockArray, Block, BlockedOneTo, blockedrange, blocklengths, blocksize using BlockSparseArrays: BlockSparseArray, BlockSparseMatrix, BlockSparseVector, blockstoredlength -using GradedArrays: +using GradedArrays: GradedArray, GradedMatrix, GradedVector +using GradedArrays.GradedUnitRanges: GradedUnitRanges, GradedOneTo, GradedUnitRange, GradedUnitRangeDual, - blocklabels, dag, dual, gradedrange, @@ -31,6 +31,19 @@ end const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @testset "GradedArray (eltype=$elt)" for elt in elts + @testset "definitions" begin + r = gradedrange([U1(0) => 2, U1(1) => 2]) + a = BlockSparseArray{elt}(undef) + @test !(a isa GradedArray) # no type piracy + v = BlockSparseArray{elt}(undef, r) + @test v isa GradedArray + @test v isa GradedVector + m = BlockSparseArray{elt}(undef, r, r) + @test m isa GradedArray + @test m isa GradedMatrix + a = BlockSparseArray{elt}(undef, r, r, r) + @test a isa GradedArray + end @testset "map" begin d1 = gradedrange([U1(0) => 2, U1(1) => 2]) d2 = gradedrange([U1(0) => 2, U1(1) => 2]) @@ -60,6 +73,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) r = gradedrange([U1(0) => 2, U1(1) => 2]) a = zeros(r, r, r, r) @test a isa BlockSparseArray{Float64} + @test a isa GradedArray @test eltype(a) === Float64 @test size(a) == (4, 4, 4, 4) @test iszero(a) diff --git a/test/test_tensoralgebraext.jl b/test/test_tensoralgebraext.jl index f0cf4ee..b2b1b29 100644 --- a/test/test_tensoralgebraext.jl +++ b/test/test_tensoralgebraext.jl @@ -1,9 +1,9 @@ using BlockArrays: Block, blocksize -using BlockSparseArrays: BlockSparseArray -using GradedArrays: GradedOneTo, blocklabels, dual, gradedrange +using BlockSparseArrays: BlockSparseArray, BlockSparseMatrix +using GradedArrays: GradedOneTo, blocklabels, dual, flip, gradedrange, space_isequal using GradedArrays.SymmetrySectors: U1 using Random: randn! -using TensorAlgebra: contract, fusedims, splitdims +using TensorAlgebra: contract, matricize, unmatricize using Test: @test, @test_broken, @testset function randn_blockdiagonal(elt::Type, axes::Tuple) @@ -18,7 +18,51 @@ end const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @testset "`contract` `GradedArray` (eltype=$elt)" for elt in elts - @testset "GradedOneTo with U(1)" begin + @testset "matricize" begin + d1 = gradedrange([U1(0) => 1, U1(1) => 1]) + d2 = gradedrange([U1(0) => 1, U1(1) => 1]) + a = randn_blockdiagonal(elt, (d1, d2, dual(d1), dual(d2))) + m = matricize(a, (1, 2), (3, 4)) + @test m isa BlockSparseMatrix + @test space_isequal(axes(m, 1), gradedrange([U1(0) => 1, U1(1) => 2, U1(2) => 1])) + @test space_isequal( + axes(m, 2), flip(gradedrange([U1(0) => 1, U1(-1) => 2, U1(-2) => 1])) + ) + + for I in CartesianIndices(m) + if I ∈ CartesianIndex.([(1, 1), (4, 4)]) + @test !iszero(m[I]) + else + @test iszero(m[I]) + end + end + @test a[1, 1, 1, 1] == m[1, 1] + @test a[2, 2, 2, 2] == m[4, 4] + @test blocksize(m) == (3, 3) + @test a == unmatricize(m, (d1, d2), (dual(d1), dual(d2))) + + # check block fusing and splitting + d = gradedrange([U1(0) => 2, U1(1) => 1]) + b = randn_blockdiagonal(elt, (d, d, dual(d), dual(d))) + @test unmatricize( + matricize(b, (1, 2), (3, 4)), (axes(b, 1), axes(b, 2)), (axes(b, 3), axes(b, 4)) + ) == b + + d1234 = gradedrange([U1(-2) => 1, U1(-1) => 4, U1(0) => 6, U1(1) => 4, U1(2) => 1]) + m = matricize(a, (1, 2, 3, 4), ()) + @test m isa BlockSparseMatrix + @test space_isequal(axes(m, 1), d1234) + @test space_isequal(axes(m, 2), flip(gradedrange([U1(0) => 1]))) + @test a == unmatricize(m, (d1, d2, dual(d1), dual(d2)), ()) + + m = matricize(a, (), (1, 2, 3, 4)) + @test m isa BlockSparseMatrix + @test space_isequal(axes(m, 1), gradedrange([U1(0) => 1])) + @test space_isequal(axes(m, 2), dual(d1234)) + @test a == unmatricize(m, (), (d1, d2, dual(d1), dual(d2))) + end + + @testset "contract with U(1)" begin d = gradedrange([U1(0) => 2, U1(1) => 3]) a1 = randn_blockdiagonal(elt, (d, d, dual(d), dual(d))) a2 = randn_blockdiagonal(elt, (d, d, dual(d), dual(d))) @@ -38,14 +82,12 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test a_dest ≈ a_dest_dense # matrix vector - @test_broken a_dest, dimnames_dest = contract(a1, (2, -1, -2, 1), a3, (1, 2)) - #= + 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)) @@ -71,30 +113,4 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test a_dest isa BlockSparseArray @test a_dest ≈ a_dest_dense end - @testset "fusedims" begin - d1 = gradedrange([U1(0) => 1, U1(1) => 1]) - d2 = gradedrange([U1(0) => 1, U1(1) => 1]) - a = randn_blockdiagonal(elt, (d1, d2, d1, d2)) - m = fusedims(a, (1, 2), (3, 4)) - for ax in axes(m) - @test ax isa GradedOneTo - @test blocklabels(ax) == [U1(0), U1(1), U1(2)] - end - for I in CartesianIndices(m) - if I ∈ CartesianIndex.([(1, 1), (4, 4)]) - @test !iszero(m[I]) - else - @test iszero(m[I]) - end - end - @test a[1, 1, 1, 1] == m[1, 1] - @test a[2, 2, 2, 2] == m[4, 4] - @test blocksize(m) == (3, 3) - @test a == splitdims(m, (d1, d2), (d1, d2)) - - # check block fusing and splitting - d = gradedrange([U1(0) => 2, U1(1) => 1]) - a = randn_blockdiagonal(elt, (d, d, dual(d), dual(d))) - @test splitdims(fusedims(a, (1, 2), (3, 4)), axes(a)...) == a - end end