From 06b484724e5a7bf9925a688fbff7a5665f61c5ab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Tue, 8 Apr 2025 14:33:27 -0400 Subject: [PATCH 01/10] WIP trivial_axis --- Project.toml | 6 +- .../GradedArraysTensorAlgebraExt.jl | 111 ++++++++---------- test/Project.toml | 4 +- test/test_tensoralgebraext.jl | 75 +++++++----- 4 files changed, 100 insertions(+), 96 deletions(-) diff --git a/Project.toml b/Project.toml index 56b1b8c..a0bf91e 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" @@ -24,7 +24,7 @@ GradedArraysTensorAlgebraExt = "TensorAlgebra" [compat] BlockArrays = "1.5.0" -BlockSparseArrays = "0.4.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.1" 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..5bf7205 100644 --- a/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl +++ b/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl @@ -9,7 +9,8 @@ using BlockSparseArrays: BlockSparseArrayInterface, BlockSparseMatrix, BlockSparseVector, - block_merge + block_merge, + blockreshape using DerivableInterfaces: @interface using GradedArrays.GradedUnitRanges: GradedUnitRanges, @@ -17,93 +18,79 @@ using GradedArrays.GradedUnitRanges: blockmergesortperm, blocksortperm, dual, + flip, invblockperm, nondual, unmerged_tensor_product +using GradedArrays.SymmetrySectors: trivial using LinearAlgebra: Adjoint, Transpose using TensorAlgebra: - TensorAlgebra, FusionStyle, BlockReshapeFusion, SectorFusion, fusedims, splitdims -using TensorProducts: OneToOne + TensorAlgebra, + AbstractBlockPermutation, + BlockedTuple, + FusionStyle, + matricize, + 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) +TensorAlgebra.FusionStyle(::AbstractGradedUnitRange) = SectorFusion() -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 +function TensorAlgebra.FusionStyle(::AbstractBlockSparseArray, ::SectorFusion) + return SectorFusion() end -TensorAlgebra.FusionStyle(::AbstractGradedUnitRange) = SectorFusion() +# TODO 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...] -end +maybe_trivial_axis(axes::Tuple, ::Tuple) = unmerged_tensor_product(axes...) +maybe_trivial_axis(::Tuple{}, axes::Tuple) = trivial_axis(axes) -function TensorAlgebra.fusedims( - ::SectorFusion, a::AbstractArray, merged_axes::AbstractUnitRange... +function row_and_column_axes( + blocked_axes::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}} ) - # 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 + codomain_axes, domain_axes = blocks(blocked_axes) + @assert !(isempty(codomain_axes) && isempty(domain_axes)) + row_axis = maybe_trivial_axis(codomain_axes, domain_axes) + unflipped_col_axis = maybe_trivial_axis(domain_axes, codomain_axes) + return row_axis, flip(unflipped_col_axis) +end - a_reshaped = fusedims(BlockReshapeFusion(), a, unmerged_axes...) +function TensorAlgebra.matricize( + ::SectorFusion, a::AbstractArray, biperm::AbstractBlockPermutation{2} +) + a_perm = permutedims(a, Tuple(biperm)) + row_axis, col_axis = row_and_column_axes(axes(a)[biperm]) + a_reshaped = blockreshape(a_perm, (row_axis, col_axis)) # Sort the blocks by sector and merge the equivalent sectors. return block_mergesort(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) + row_col_axes = row_and_column_axes(blocked_axes) + + blockperms = blocksortperm.(row_col_axes) + sorted_axes = map((r, I) -> only(axes(r[I])), row_col_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(m, ()), m_blockpermed, blocked_axes) end +# Sort the blocks by sector and then merge the common sectors. +function block_mergesort(a::AbstractArray) + I = blockmergesortperm.(axes(a)) + return a[I...] +end end diff --git a/test/Project.toml b/test/Project.toml index ed52e81..b890ed0 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -16,14 +16,14 @@ TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a" [compat] Aqua = "0.8.11" BlockArrays = "1.5.0" -BlockSparseArrays = "0.4.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.1" TensorProducts = "0.1.3" Test = "1.10.0" TestExtras = "0.3.1" diff --git a/test/test_tensoralgebraext.jl b/test/test_tensoralgebraext.jl index f0cf4ee..f8a9c7d 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,6 +18,49 @@ end const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @testset "`contract` `GradedArray` (eltype=$elt)" for elt in elts + @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 + + 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]))) + @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, U1(1) => 2, U1(2) => 1])) + @test space_isequal(axes(m, 2), flip(gradedrange([U1(0) => 1]))) + @test a == unmatricize(m, (), (d1, d2, dual(d1), dual(d2))) + end + @testset "GradedOneTo with U(1)" begin d = gradedrange([U1(0) => 2, U1(1) => 3]) a1 = randn_blockdiagonal(elt, (d, d, dual(d), dual(d))) @@ -71,30 +114,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 From b5c9ecd16f86ffb00b5e7aaad75106973b66bc5e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Wed, 9 Apr 2025 11:57:15 -0400 Subject: [PATCH 02/10] mostly working --- .../GradedArraysTensorAlgebraExt.jl | 36 +++++-------------- src/gradedarray.jl | 10 ++++++ test/Project.toml | 2 +- test/test_tensoralgebraext.jl | 17 ++++----- 4 files changed, 29 insertions(+), 36 deletions(-) diff --git a/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl b/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl index 5bf7205..8834f44 100644 --- a/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl +++ b/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl @@ -1,59 +1,41 @@ module GradedArraysTensorAlgebraExt -using BlockArrays: Block, BlockIndexRange, blockedrange, blocks -using BlockSparseArrays: - BlockSparseArrays, - AbstractBlockSparseArray, - AbstractBlockSparseArrayInterface, - BlockSparseArray, - BlockSparseArrayInterface, - BlockSparseMatrix, - BlockSparseVector, - block_merge, - blockreshape -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 GradedArrays.SymmetrySectors: trivial -using LinearAlgebra: Adjoint, Transpose using TensorAlgebra: TensorAlgebra, AbstractBlockPermutation, BlockedTuple, FusionStyle, - matricize, trivial_axis, unmatricize struct SectorFusion <: FusionStyle end -TensorAlgebra.FusionStyle(::AbstractGradedUnitRange) = SectorFusion() - -function TensorAlgebra.FusionStyle(::AbstractBlockSparseArray, ::SectorFusion) - return SectorFusion() -end +TensorAlgebra.FusionStyle(::Type{<:GradedArray}) = SectorFusion() # TODO consider heterogeneous sectors? TensorAlgebra.trivial_axis(t::Tuple{Vararg{AbstractGradedUnitRange}}) = trivial(first(t)) -maybe_trivial_axis(axes::Tuple, ::Tuple) = unmerged_tensor_product(axes...) -maybe_trivial_axis(::Tuple{}, axes::Tuple) = trivial_axis(axes) +fuse_or_trivial_axis(axes::Tuple, ::Tuple) = unmerged_tensor_product(axes...) +fuse_or_trivial_axis(::Tuple{}, axes::Tuple) = trivial_axis(axes) function row_and_column_axes( blocked_axes::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}} ) codomain_axes, domain_axes = blocks(blocked_axes) @assert !(isempty(codomain_axes) && isempty(domain_axes)) - row_axis = maybe_trivial_axis(codomain_axes, domain_axes) - unflipped_col_axis = maybe_trivial_axis(domain_axes, codomain_axes) + row_axis = fuse_or_trivial_axis(codomain_axes, domain_axes) + unflipped_col_axis = fuse_or_trivial_axis(domain_axes, codomain_axes) return row_axis, flip(unflipped_col_axis) end @@ -85,7 +67,7 @@ function TensorAlgebra.unmatricize( # for this combination of slicing. m_unblocked = m[sorted_axes...] m_blockpermed = m_unblocked[invblockperm.(blockperms)...] - return unmatricize(FusionStyle(m, ()), m_blockpermed, blocked_axes) + return unmatricize(FusionStyle(BlockSparseArray), m_blockpermed, blocked_axes) end # Sort the blocks by sector and then merge the common sectors. diff --git a/src/gradedarray.jl b/src/gradedarray.jl index 4d23d72..7652716 100644 --- a/src/gradedarray.jl +++ b/src/gradedarray.jl @@ -4,11 +4,21 @@ 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} = + BlockSparseMatrix{T,A,Blocks,Axes} where {Axes<:Tuple{Vararg{AbstractGradedUnitRange}}} +const GradedVector{T,A,Blocks,Axes} = + BlockSparseVector{T,A,Blocks,Axes} where {Axes<:Tuple{Vararg{AbstractGradedUnitRange}}} + # 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 b890ed0..0cd72ca 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -23,7 +23,7 @@ Random = "1.10.0" SafeTestsets = "0.1.0" SparseArraysBase = "0.5.4" Suppressor = "0.2.8" -TensorAlgebra = "0.3.1" +TensorAlgebra = "0.3.2" TensorProducts = "0.1.3" Test = "1.10.0" TestExtras = "0.3.1" diff --git a/test/test_tensoralgebraext.jl b/test/test_tensoralgebraext.jl index f8a9c7d..41db7e8 100644 --- a/test/test_tensoralgebraext.jl +++ b/test/test_tensoralgebraext.jl @@ -48,20 +48,21 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) 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), gradedrange([U1(0) => 1, U1(1) => 2, U1(2) => 1])) + @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, U1(1) => 2, U1(2) => 1])) - @test space_isequal(axes(m, 2), flip(gradedrange([U1(0) => 1]))) + @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 "GradedOneTo with U(1)" begin + @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))) @@ -81,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)) @@ -99,12 +98,14 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test a_dest ≈ a_dest_dense # vector vector - a_dest, dimnames_dest = contract(a3, (1, 2), a3, (2, 1)) + @test_broken 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, dimnames_dest = contract(a3, (1, 2), a3, (3, 4)) From fcc862ba54b1626eada3c47e2876e469c39e690b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Wed, 9 Apr 2025 12:01:06 -0400 Subject: [PATCH 03/10] bump BlockArrays compat --- Project.toml | 4 ++-- test/Project.toml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index a0bf91e..7b7e77c 100644 --- a/Project.toml +++ b/Project.toml @@ -23,7 +23,7 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" GradedArraysTensorAlgebraExt = "TensorAlgebra" [compat] -BlockArrays = "1.5.0" +BlockArrays = "1.6.0" BlockSparseArrays = "0.4.2" Compat = "4.16.0" DerivableInterfaces = "0.4.4" @@ -32,7 +32,7 @@ HalfIntegers = "1.6.0" LinearAlgebra = "1.10.0" Random = "1.10.0" SplitApplyCombine = "1.2.3" -TensorAlgebra = "0.3.1" +TensorAlgebra = "0.3.2" TensorProducts = "0.1.3" TypeParameterAccessors = "0.3.9" julia = "1.10" diff --git a/test/Project.toml b/test/Project.toml index 0cd72ca..03fb5eb 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -15,7 +15,7 @@ TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a" [compat] Aqua = "0.8.11" -BlockArrays = "1.5.0" +BlockArrays = "1.6.0" BlockSparseArrays = "0.4.2" GradedArrays = "0.2.0" LinearAlgebra = "1.10.0" From ea7f0aa8600d7947700daec7725d152fdfe5fd7c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Thu, 10 Apr 2025 10:42:58 -0400 Subject: [PATCH 04/10] test GradedArray --- test/test_gradedarray.jl | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) 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) From 3f8b70956a6f193e7ccb85b66f58142d1ed40a8f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Thu, 10 Apr 2025 11:00:52 -0400 Subject: [PATCH 05/10] matricize_axes and sectormergesort --- .../GradedArraysTensorAlgebraExt.jl | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl b/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl index 8834f44..2be8ffd 100644 --- a/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl +++ b/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl @@ -26,16 +26,13 @@ TensorAlgebra.FusionStyle(::Type{<:GradedArray}) = SectorFusion() # TODO consider heterogeneous sectors? TensorAlgebra.trivial_axis(t::Tuple{Vararg{AbstractGradedUnitRange}}) = trivial(first(t)) -fuse_or_trivial_axis(axes::Tuple, ::Tuple) = unmerged_tensor_product(axes...) -fuse_or_trivial_axis(::Tuple{}, axes::Tuple) = trivial_axis(axes) - -function row_and_column_axes( +function matricize_axes( blocked_axes::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}} ) codomain_axes, domain_axes = blocks(blocked_axes) @assert !(isempty(codomain_axes) && isempty(domain_axes)) - row_axis = fuse_or_trivial_axis(codomain_axes, domain_axes) - unflipped_col_axis = fuse_or_trivial_axis(domain_axes, codomain_axes) + row_axis = unmerged_tensor_product(trivial_axes(domain_axes), codomain_axes...) + unflipped_col_axis = unmerged_tensor_product(trivial_axes(codomain_axes), domain_axes...) return row_axis, flip(unflipped_col_axis) end @@ -43,10 +40,10 @@ function TensorAlgebra.matricize( ::SectorFusion, a::AbstractArray, biperm::AbstractBlockPermutation{2} ) a_perm = permutedims(a, Tuple(biperm)) - row_axis, col_axis = row_and_column_axes(axes(a)[biperm]) + row_axis, col_axis = matricize_axes(axes(a)[biperm]) a_reshaped = blockreshape(a_perm, (row_axis, col_axis)) # Sort the blocks by sector and merge the equivalent sectors. - return block_mergesort(a_reshaped) + return sectormergesort(a_reshaped) end function TensorAlgebra.unmatricize( @@ -71,7 +68,7 @@ function TensorAlgebra.unmatricize( end # Sort the blocks by sector and then merge the common sectors. -function block_mergesort(a::AbstractArray) +function sectormergesort(a::AbstractArray) I = blockmergesortperm.(axes(a)) return a[I...] end From 448079e394e907fefe5aa9379e1659c235c05149 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Thu, 10 Apr 2025 11:07:34 -0400 Subject: [PATCH 06/10] fix trivial_axis --- .../GradedArraysTensorAlgebraExt.jl | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl b/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl index 2be8ffd..d95b4ed 100644 --- a/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl +++ b/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl @@ -23,16 +23,17 @@ struct SectorFusion <: FusionStyle end TensorAlgebra.FusionStyle(::Type{<:GradedArray}) = SectorFusion() -# TODO consider heterogeneous sectors? +# TBD consider heterogeneous sectors? TensorAlgebra.trivial_axis(t::Tuple{Vararg{AbstractGradedUnitRange}}) = trivial(first(t)) 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) - @assert !(isempty(codomain_axes) && isempty(domain_axes)) - row_axis = unmerged_tensor_product(trivial_axes(domain_axes), codomain_axes...) - unflipped_col_axis = unmerged_tensor_product(trivial_axes(codomain_axes), domain_axes...) + row_axis = unmerged_tensor_product(default_axis, codomain_axes...) + unflipped_col_axis = unmerged_tensor_product(default_axis, domain_axes...) return row_axis, flip(unflipped_col_axis) end @@ -53,7 +54,7 @@ function TensorAlgebra.unmatricize( ) # First, fuse axes to get `blockmergesortperm`. # Then unpermute the blocks. - row_col_axes = row_and_column_axes(blocked_axes) + row_col_axes = matricize_axes(blocked_axes) blockperms = blocksortperm.(row_col_axes) sorted_axes = map((r, I) -> only(axes(r[I])), row_col_axes, blockperms) From 8b292c177ce347cb881d0307409d0d2ce406ffff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Thu, 10 Apr 2025 11:46:28 -0400 Subject: [PATCH 07/10] uncomment test --- test/test_tensoralgebraext.jl | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/test/test_tensoralgebraext.jl b/test/test_tensoralgebraext.jl index 41db7e8..80282f5 100644 --- a/test/test_tensoralgebraext.jl +++ b/test/test_tensoralgebraext.jl @@ -98,14 +98,12 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test a_dest ≈ a_dest_dense # vector vector - @test_broken a_dest, dimnames_dest = contract(a3, (1, 2), a3, (2, 1)) - #= + @test 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, dimnames_dest = contract(a3, (1, 2), a3, (3, 4)) From 163c7157d57b1f4c76b042afdc0927522c18eb53 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Thu, 10 Apr 2025 12:13:54 -0400 Subject: [PATCH 08/10] row/col axis -> domain --- .../GradedArraysTensorAlgebraExt.jl | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl b/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl index d95b4ed..48b64f6 100644 --- a/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl +++ b/ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl @@ -32,17 +32,17 @@ function matricize_axes( @assert !isempty(blocked_axes) default_axis = trivial_axis(Tuple(blocked_axes)) codomain_axes, domain_axes = blocks(blocked_axes) - row_axis = unmerged_tensor_product(default_axis, codomain_axes...) - unflipped_col_axis = unmerged_tensor_product(default_axis, domain_axes...) - return row_axis, flip(unflipped_col_axis) + 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.matricize( ::SectorFusion, a::AbstractArray, biperm::AbstractBlockPermutation{2} ) a_perm = permutedims(a, Tuple(biperm)) - row_axis, col_axis = matricize_axes(axes(a)[biperm]) - a_reshaped = blockreshape(a_perm, (row_axis, col_axis)) + 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 sectormergesort(a_reshaped) end @@ -54,10 +54,10 @@ function TensorAlgebra.unmatricize( ) # First, fuse axes to get `blockmergesortperm`. # Then unpermute the blocks. - row_col_axes = matricize_axes(blocked_axes) + fused_axes = matricize_axes(blocked_axes) - blockperms = blocksortperm.(row_col_axes) - sorted_axes = map((r, I) -> only(axes(r[I])), row_col_axes, blockperms) + 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. From 750ce3f01cf37ad3a81ac1fcdf5340e458c30d9d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Thu, 10 Apr 2025 12:14:49 -0400 Subject: [PATCH 09/10] remove @test --- test/test_tensoralgebraext.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_tensoralgebraext.jl b/test/test_tensoralgebraext.jl index 80282f5..b2b1b29 100644 --- a/test/test_tensoralgebraext.jl +++ b/test/test_tensoralgebraext.jl @@ -98,7 +98,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test a_dest ≈ a_dest_dense # vector vector - @test a_dest, dimnames_dest = contract(a3, (1, 2), a3, (2, 1)) + 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) From 1c372d123e921393eb7ff74fdb19c93d7978ff73 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Thu, 10 Apr 2025 15:56:04 -0400 Subject: [PATCH 10/10] simpler definitions --- src/gradedarray.jl | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/gradedarray.jl b/src/gradedarray.jl index 7652716..918b528 100644 --- a/src/gradedarray.jl +++ b/src/gradedarray.jl @@ -14,10 +14,8 @@ 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} = - BlockSparseMatrix{T,A,Blocks,Axes} where {Axes<:Tuple{Vararg{AbstractGradedUnitRange}}} -const GradedVector{T,A,Blocks,Axes} = - BlockSparseVector{T,A,Blocks,Axes} where {Axes<:Tuple{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