Abelian symmetric SVD (#33)

Co-authored-by: Lukas Devos <[email protected]>

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "GradedArrays"
 uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.4.2"
+version = "0.4.4"
 
 [deps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
@@ -11,6 +11,7 @@ DerivableInterfaces = "6c5e35bf-e59e-4898-b73c-732dcc4ba65f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 HalfIntegers = "f0d1745a-41c9-11e9-1dd9-e5d34d218721"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
 TensorProducts = "decf83d6-1968-43f4-96dc-fdb3fe15fc6d"
@@ -24,12 +25,13 @@ GradedArraysTensorAlgebraExt = "TensorAlgebra"
 
 [compat]
 BlockArrays = "1.6.0"
-BlockSparseArrays = "0.5"
+BlockSparseArrays = "0.6.1"
 Compat = "4.16.0"
 DerivableInterfaces = "0.4.4"
 FillArrays = "1.13.0"
 HalfIntegers = "1.6.0"
 LinearAlgebra = "1.10.0"
+MatrixAlgebraKit = "0.2"
 Random = "1.10.0"
 SplitApplyCombine = "1.2.3"
 TensorAlgebra = "0.3.2"

src/GradedArrays.jl

@@ -20,6 +20,7 @@ include("sector_product.jl")
 
 include("fusion.jl")
 include("gradedarray.jl")
+include("factorizations.jl")
 
 export SU2,
   U1,

src/factorizations.jl

@@ -0,0 +1,58 @@
+using BlockArrays: blocks
+using BlockSparseArrays:
+  BlockSparseArrays,
+  BlockSparseMatrix,
+  BlockPermutedDiagonalAlgorithm,
+  BlockPermutedDiagonalTruncationStrategy,
+  diagview,
+  eachblockaxis,
+  mortar_axis
+using LinearAlgebra: Diagonal
+using MatrixAlgebraKit: MatrixAlgebraKit, svd_compact!, svd_full!, svd_trunc!
+
+function BlockSparseArrays.similar_output(
+  ::typeof(svd_compact!), A::GradedMatrix, S_axes, alg::BlockPermutedDiagonalAlgorithm
+)
+  u_axis, v_axis = S_axes
+  U = similar(A, axes(A, 1), dual(u_axis))
+  T = real(eltype(A))
+  S = BlockSparseMatrix{T,Diagonal{T,Vector{T}}}(undef, (u_axis, v_axis))
+  Vt = similar(A, dual(v_axis), axes(A, 2))
+  return U, S, Vt
+end
+
+function BlockSparseArrays.similar_output(
+  ::typeof(svd_full!), A::GradedMatrix, S_axes, alg::BlockPermutedDiagonalAlgorithm
+)
+  u_axis, s_axis = S_axes
+  U = similar(A, axes(A, 1), dual(u_axis))
+  T = real(eltype(A))
+  S = similar(A, T, S_axes)
+  Vt = similar(A, dual(S_axes[2]), axes(A, 2))
+  return U, S, Vt
+end
+
+const TGradedUSVᴴ = Tuple{<:GradedMatrix,<:GradedMatrix,<:GradedMatrix}
+
+function BlockSparseArrays.similar_truncate(
+  ::typeof(svd_trunc!),
+  (U, S, Vᴴ)::TGradedUSVᴴ,
+  strategy::BlockPermutedDiagonalTruncationStrategy,
+  indexmask=MatrixAlgebraKit.findtruncated(diagview(S), strategy),
+)
+  u_axis, v_axis = axes(S)
+  counter = Base.Fix1(count, Base.Fix1(getindex, indexmask))
+  s_lengths = map(counter, blocks(u_axis))
+  u_sectors = sectors(u_axis) .=> s_lengths
+  v_sectors = sectors(v_axis) .=> s_lengths
+  u_sectors_filtered = filter(>(0) ∘ last, u_sectors)
+  v_sectors_filtered = filter(>(0) ∘ last, v_sectors)
+  u_axis′ = gradedrange(u_sectors_filtered)
+  u_axis = isdual(u_axis) ? dual(u_axis′) : u_axis′
+  v_axis′ = gradedrange(v_sectors_filtered)
+  v_axis = isdual(v_axis) ? dual(v_axis′) : v_axis′
+  Ũ = similar(U, axes(U, 1), dual(u_axis))
+  S̃ = similar(S, u_axis, v_axis)
+  Ṽᴴ = similar(Vᴴ, dual(v_axis), axes(Vᴴ, 2))
+  return Ũ, S̃, Ṽᴴ
+end

src/gradedarray.jl

@@ -5,6 +5,7 @@ using BlockSparseArrays:
   AnyAbstractBlockSparseArray,
   BlockSparseArray,
   blocktype,
+  eachblockstoredindex,
   sparsemortar
 using LinearAlgebra: Adjoint
 using TypeParameterAccessors: similartype, unwrap_array_type
@@ -41,13 +42,19 @@ function similar_blocksparse(
   elt::Type,
   axes::Tuple{AbstractGradedUnitRange,Vararg{AbstractGradedUnitRange}},
 )
-  # TODO: Probably need to unwrap the type of `a` in certain cases
-  # to make a proper block type.
-  return BlockSparseArray{
-    elt,length(axes),similartype(unwrap_array_type(blocktype(a)), elt, axes)
-  }(
-    undef, axes
+  blockaxistypes = map(axes) do axis
+    return eltype(Base.promote_op(eachblockaxis, typeof(axis)))
+  end
+  similar_blocktype = Base.promote_op(
+    similar, blocktype(a), Type{elt}, Tuple{blockaxistypes...}
   )
+  return BlockSparseArray{elt,length(axes),similar_blocktype}(undef, axes)
+end
+
+function Base.similar(
+  a::AbstractArray, elt::Type, axes::Tuple{SectorOneTo,Vararg{SectorOneTo}}
+)
+  return similar(a, elt, Base.OneTo.(length.(axes)))
 end
 
 function Base.similar(
@@ -120,6 +127,25 @@ end
 
 ungrade(a::GradedArray) = sparsemortar(blocks(a), ungrade.(axes(a)))
 
+function flux(a::GradedArray{<:Any,N}, I::Vararg{Block{1},N}) where {N}
+  sects = ntuple(N) do d
+    return flux(axes(a, d), I[d])
+  end
+  return ⊗(sects...)
+end
+function flux(a::GradedArray{<:Any,N}, I::Block{N}) where {N}
+  return flux(a, Tuple(I)...)
+end
+function flux(a::GradedArray)
+  sect = nothing
+  for I in eachblockstoredindex(a)
+    sect_I = flux(a, I)
+    isnothing(sect) || sect_I == sect || throw(ArgumentError("Inconsistent flux."))
+    sect = sect_I
+  end
+  return sect
+end
+
 # Copy of `Base.dims2string` defined in `show.jl`.
 function dims_to_string(d)
   isempty(d) && return "0-dimensional"

src/gradedunitrange.jl

@@ -18,7 +18,8 @@ using BlockArrays:
   combine_blockaxes,
   findblock,
   mortar
-using BlockSparseArrays: BlockSparseArrays, blockedunitrange_getindices
+using BlockSparseArrays:
+  BlockSparseArrays, blockedunitrange_getindices, eachblockaxis, mortar_axis
 using Compat: allequal
 
 # ====================================  Definitions  =======================================
@@ -60,7 +61,7 @@ end
 
 # =====================================  Accessors  ========================================
 
-eachblockaxis(g::GradedUnitRange) = g.eachblockaxis
+BlockSparseArrays.eachblockaxis(g::GradedUnitRange) = g.eachblockaxis
 ungrade(g::GradedUnitRange) = g.full_range
 
 sector_multiplicities(g::GradedUnitRange) = sector_multiplicity.(eachblockaxis(g))
@@ -69,12 +70,12 @@ sector_type(::Type{<:GradedUnitRange{<:Any,SUR}}) where {SUR} = sector_type(SUR)
 
 # ====================================  Constructors  ======================================
 
-function mortar_axis(geachblockaxis::AbstractVector{<:SectorOneTo})
+function BlockSparseArrays.mortar_axis(geachblockaxis::AbstractVector{<:SectorOneTo})
   brange = blockedrange(length.(geachblockaxis))
   return GradedUnitRange(geachblockaxis, brange)
 end
 
-function mortar_axis(gaxes::AbstractVector{<:GradedOneTo})
+function BlockSparseArrays.mortar_axis(gaxes::AbstractVector{<:GradedOneTo})
   return mortar_axis(mapreduce(eachblockaxis, vcat, gaxes))
 end
 
@@ -102,6 +103,11 @@ function sectors(g::AbstractGradedUnitRange)
   return sector.(eachblockaxis(g))
 end
 
+function flux(a::AbstractGradedUnitRange, I::Block{1})
+  sect = sector(a[I])
+  return isdual(a) ? dual(sect) : sect
+end
+
 function map_sectors(f, g::GradedUnitRange)
   return GradedUnitRange(map_sectors.(f, eachblockaxis(g)), ungrade(g))
 end

test/Project.toml

@@ -4,6 +4,7 @@ BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 BlockSparseArrays = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 GradedArrays = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseArraysBase = "0d5efcca-f356-4864-8770-e1ed8d78f208"
@@ -16,9 +17,10 @@ TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
 [compat]
 Aqua = "0.8.11"
 BlockArrays = "1.6.0"
-BlockSparseArrays = "0.5"
+BlockSparseArrays = "0.6"
 GradedArrays = "0.4"
 LinearAlgebra = "1.10.0"
+MatrixAlgebraKit = "0.2"
 Random = "1.10.0"
 SafeTestsets = "0.1.0"
 SparseArraysBase = "0.5.4"

test/test_factorizations.jl

@@ -0,0 +1,108 @@
+using BlockArrays: Block, blocksizes
+using GradedArrays: U1, dual, flux, gradedrange
+using LinearAlgebra: I, diag, svdvals
+using MatrixAlgebraKit: svd_compact, svd_full, svd_trunc
+using Test: @test, @testset
+
+const elts = (Float32, Float64, ComplexF32, ComplexF64)
+@testset "svd_compact (eltype=$elt)" for elt in elts
+  for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3]
+    r1 = gradedrange([U1(0) => i, U1(1) => j])
+    r2 = gradedrange([U1(0) => k, U1(1) => l])
+    a = zeros(elt, r1, dual(r2))
+    a[Block(2, 2)] = randn(elt, blocksizes(a)[2, 2])
+    @test flux(a) == U1(0)
+    u, s, vᴴ = svd_compact(a)
+    @test sort(diag(Matrix(s)); rev=true) ≈ svdvals(Matrix(a))[1:size(s, 1)]
+    @test u * s * vᴴ ≈ a
+    @test Array(u'u) ≈ I
+    @test Array(vᴴ * vᴴ') ≈ I
+    @test flux(u) == U1(0)
+    @test flux(s) == flux(a)
+    @test flux(vᴴ) == U1(0)
+
+    r1 = gradedrange([U1(0) => i, U1(1) => j])
+    r2 = gradedrange([U1(0) => k, U1(1) => l])
+    a = zeros(elt, r1, dual(r2))
+    a[Block(1, 2)] = randn(elt, blocksizes(a)[1, 2])
+    @test flux(a) == U1(-1)
+    u, s, vᴴ = svd_compact(a)
+    @test sort(diag(Matrix(s)); rev=true) ≈ svdvals(Matrix(a))[1:size(s, 1)]
+    @test u * s * vᴴ ≈ a
+    @test Array(u'u) ≈ I
+    @test Array(vᴴ * vᴴ') ≈ I
+    @test flux(u) == U1(0)
+    @test flux(s) == flux(a)
+    @test flux(vᴴ) == U1(0)
+  end
+end
+
+@testset "svd_full (eltype=$elt)" for elt in elts
+  for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3]
+    r1 = gradedrange([U1(0) => i, U1(1) => j])
+    r2 = gradedrange([U1(0) => k, U1(1) => l])
+    a = zeros(elt, r1, dual(r2))
+    a[Block(2, 2)] = randn(elt, blocksizes(a)[2, 2])
+    @test flux(a) == U1(0)
+    u, s, vᴴ = svd_full(a)
+    @test u * s * vᴴ ≈ a
+    @test Array(u'u) ≈ I
+    @test Array(u * u') ≈ I
+    @test Array(vᴴ * vᴴ') ≈ I
+    @test Array(vᴴ'vᴴ) ≈ I
+    @test flux(u) == U1(0)
+    @test flux(s) == flux(a)
+    @test flux(vᴴ) == U1(0)
+
+    r1 = gradedrange([U1(0) => i, U1(1) => j])
+    r2 = gradedrange([U1(0) => k, U1(1) => l])
+    a = zeros(elt, r1, dual(r2))
+    a[Block(1, 2)] = randn(elt, blocksizes(a)[1, 2])
+    @test flux(a) == U1(-1)
+    u, s, vᴴ = svd_full(a)
+    @test u * s * vᴴ ≈ a
+    @test Array(u'u) ≈ I
+    @test Array(u * u') ≈ I
+    @test Array(vᴴ * vᴴ') ≈ I
+    @test Array(vᴴ'vᴴ) ≈ I
+    @test flux(u) == U1(0)
+    @test flux(s) == flux(a)
+    @test flux(vᴴ) == U1(0)
+  end
+end
+
+@testset "svd_trunc (eltype=$elt)" for elt in elts
+  for i in [2, 3], j in [2, 3], k in [2, 3], l in [2, 3]
+    r1 = gradedrange([U1(0) => i, U1(1) => j])
+    r2 = gradedrange([U1(0) => k, U1(1) => l])
+    a = zeros(elt, r1, dual(r2))
+    a[Block(2, 2)] = randn(elt, blocksizes(a)[2, 2])
+    @test flux(a) == U1(0)
+    u, s, vᴴ = svd_trunc(a; trunc=(; maxrank=1))
+    @test sort(diag(Matrix(s)); rev=true) ≈ svdvals(Matrix(a))[1:size(s, 1)]
+    @test size(u) == (size(a, 1), 1)
+    @test size(s) == (1, 1)
+    @test size(vᴴ) == (1, size(a, 2))
+    @test Array(u'u) ≈ I
+    @test Array(vᴴ * vᴴ') ≈ I
+    @test flux(u) == U1(0)
+    @test flux(s) == flux(a)
+    @test flux(vᴴ) == U1(0)
+
+    r1 = gradedrange([U1(0) => i, U1(1) => j])
+    r2 = gradedrange([U1(0) => k, U1(1) => l])
+    a = zeros(elt, r1, dual(r2))
+    a[Block(1, 2)] = randn(elt, blocksizes(a)[1, 2])
+    @test flux(a) == U1(-1)
+    u, s, vᴴ = svd_trunc(a; trunc=(; maxrank=1))
+    @test sort(diag(Matrix(s)); rev=true) ≈ svdvals(Matrix(a))[1:size(s, 1)]
+    @test size(u) == (size(a, 1), 1)
+    @test size(s) == (1, 1)
+    @test size(vᴴ) == (1, size(a, 2))
+    @test Array(u'u) ≈ I
+    @test Array(vᴴ * vᴴ') ≈ I
+    @test flux(u) == U1(0)
+    @test flux(s) == flux(a)
+    @test flux(vᴴ) == U1(0)
+  end
+end