Abelian symmetric SVD

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.3"
 
 [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.5.2"
 Compat = "4.16.0"
 DerivableInterfaces = "0.4.4"
 FillArrays = "1.13.0"
 HalfIntegers = "1.6.0"
 LinearAlgebra = "1.10.0"
+MatrixAlgebraKit = "0.1.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,75 @@
+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_axis::AbstractUnitRange,
+  alg::BlockPermutedDiagonalAlgorithm,
+)
+  u_axis = s_axis
+  flx = flux(A)
+  axs = eachblockaxis(s_axis)
+  # TODO: Use `gradedrange` constructor.
+  v_axis = mortar_axis(
+    map(axs) do ax
+      return sectorrange(dual(sector(ax)) ⊗ flx, ungrade(ax))
+    end,
+  )
+  U = similar(A, axes(A, 1), dual(u_axis))
+  T = real(eltype(A))
+  S = BlockSparseMatrix{T,Diagonal{T,Vector{T}}}(undef, (u_axis, dual(v_axis)))
+  Vt = similar(A, v_axis, axes(A, 2))
+  return U, S, Vt
+end
+
+function BlockSparseArrays.similar_output(
+  ::typeof(svd_full!),
+  A::GradedMatrix,
+  s_axis::AbstractUnitRange,
+  alg::BlockPermutedDiagonalAlgorithm,
+)
+  U = similar(A, axes(A, 1), dual(s_axis))
+  T = real(eltype(A))
+  S = similar(A, T, (s_axis, axes(A, 2)))
+  Vt = similar(A, dual(axes(A, 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),
+)
+  ax = axes(S, 1)
+  counter = Base.Fix1(count, Base.Fix1(getindex, indexmask))
+  s_lengths = map(counter, blocks(ax))
+  s_sectors = sectors(ax) .=> s_lengths
+  s_sectors_filtered = filter(>(0) ∘ last, s_sectors)
+  s_axis = gradedrange(s_sectors_filtered)
+  u_axis = s_axis
+  flx = flux(S)
+  axs = eachblockaxis(s_axis)
+  # TODO: Use `gradedrange` constructor.
+  v_axis = mortar_axis(
+    map(axs) do ax
+      return sectorrange(dual(sector(ax)) ⊗ flx, ungrade(ax))
+    end,
+  )
+  Ũ = similar(U, axes(U, 1), dual(u_axis))
+  S̃ = similar(S, u_axis, dual(v_axis))
+  Ṽᴴ = similar(Vᴴ, 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,17 @@ 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 +125,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