[WIP] Matrix factorizations, *, and one! for AbelianGradedMatrix

Project.toml

@@ -1,6 +1,6 @@
 name = "GradedArrays"
 uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
-version = "0.13.8"
+version = "0.13.9"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -26,6 +26,10 @@ TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f"
 SUNRepresentations = "1a50b95c-7aac-476d-a9ce-2bfc675fc617"
 TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 
+[sources.TensorAlgebra]
+rev = "mf/nextgen-followups-round1"
+url = "https://github.com/ITensor/TensorAlgebra.jl"
+
 [extensions]
 GradedArraysSUNRepresentationsExt = "SUNRepresentations"
 GradedArraysTensorKitExt = "TensorKit"

src/abeliangradedarray.jl

@@ -822,3 +822,15 @@ end
 function Base.getindex(a::Array, ax1::GradedOneTo)
     return invoke(getindex, Tuple{AbstractArray, GradedOneTo, Vararg{GradedOneTo}}, a, ax1)
 end
+
+# ---------------------------------------------------------------------------
+#  matrix multiplication
+# ---------------------------------------------------------------------------
+
+# Matrix-matrix multiply via `TensorAlgebra.contract` (which matricizes internally). The
+# contracted label `2` pairs `a`'s column axis with `b`'s row axis; the explicit output
+# labels `(1, 3)` fix the result as `a`'s row axis (codomain) × `b`'s column axis (domain),
+# so `(a * b)[i, j] == sum_k a[i, k] * b[k, j]`.
+function Base.:*(a::AbelianGradedMatrix, b::AbelianGradedMatrix)
+    return TensorAlgebra.contract((1, 3), a, (1, 2), b, (2, 3))
+end

src/matrixalgebrakit.jl

@@ -34,6 +34,30 @@ for f in [
     end
 end
 
+# Bare-matrix factorizations
+# ---------------------------
+# There is no in-place block algorithm for an unfused `AbelianGradedMatrix`, so the plain
+# matrix forms (`MAK.svd_compact(m)`, etc.) route through the matricizing `TensorAlgebra`
+# factorizations: matricize to a `FusedGradedMatrix`, run the block factorization, then
+# unmatricize back. The factors are returned as graded matrices. Only the functions with an
+# identically-named `TensorAlgebra` perm-form are delegated here (`qr_null`/`lq_null` are
+# spelled `left_null`/`right_null` there, and `project_antihermitian`/`project_isometric`
+# have no perm-form).
+for f in (
+        :svd_compact, :svd_full, :svd_vals, :qr_compact, :qr_full, :lq_compact,
+        :lq_full, :eig_full, :eig_vals, :eigh_full, :eigh_vals, :left_polar,
+        :right_polar, :project_hermitian,
+    )
+    @eval function MAK.$f(m::AbelianGradedMatrix; kwargs...)
+        return TensorAlgebra.$f(m, (1,), (2,); kwargs...)
+    end
+end
+
+# In-place graded identity fill. Filling the unfused data blocks with identities is not the
+# graded identity map in general, so route through the fused path: `TensorAlgebra.one!`
+# matricizes, fills the fused matrix with `MAK.one!`, and scatters it back into `a`.
+MAK.one!(a::AbelianGradedMatrix) = TensorAlgebra.one!(a, Val(1))
+
 # Generic Implementations
 # -----------------------
 # utility function to do something with each block

test/test_factorizations.jl

@@ -1,7 +1,9 @@
 import MatrixAlgebraKit as MAK
-using GradedArrays: FusedGradedMatrix, FusedGradedVector, GradedBlockAlgorithm, U1, Z2
+using GradedArrays: AbelianGradedMatrix, FusedGradedMatrix, FusedGradedVector,
+    GradedBlockAlgorithm, U1, Z2, dual, gradedrange
 using LinearAlgebra: Diagonal, I, eigvals, isposdef, istril, istriu, norm
 using MatrixAlgebraKit: isisometric, isunitary
+using TensorAlgebra: TensorAlgebra
 using Test: @test, @testset
 
 # ---------------------------------------------------------------------------
@@ -408,4 +410,112 @@ end
             end
         end
     end
+
+    # -----------------------------------------------------------------------
+    @testset "matrix multiplication" begin
+        g = gradedrange([U1(0) => 2, U1(1) => 3, U1(2) => 2])
+        h = gradedrange([U1(0) => 3, U1(1) => 2, U1(2) => 4])
+        a = randn(Float64, (g, dual(g)))
+        b = randn(Float64, (g, dual(h)))
+        c = a * b
+        @test c isa AbelianGradedMatrix
+        # `(a * b)[i, j] == sum_k a[i, k] * b[k, j]`.
+        @test Array(c) ≈ Array(a) * Array(b)
+        # Result axes: codomain from `a`, domain from `b`.
+        @test axes(c, 1) == axes(a, 1)
+        @test axes(c, 2) == axes(b, 2)
+    end
+
+    # -----------------------------------------------------------------------
+    # Bare-matrix factorizations delegate to the matricizing `TensorAlgebra` forms.
+    @testset "factorizations on a bare AbelianGradedMatrix" begin
+        g = gradedrange([U1(0) => 2, U1(1) => 3, U1(2) => 2])
+        h = gradedrange([U1(0) => 3, U1(1) => 2, U1(2) => 4])
+        m_rect = randn(Float64, (g, dual(h)))
+        m_sq = randn(Float64, (g, dual(g)))
+        m_herm = MAK.project_hermitian(randn(Float64, (g, dual(g))))
+
+        # helper: compare two `FusedGradedVector`s block-by-block (the broadcasting `-`
+        # path they would otherwise take is not supported).
+        fgv_approx(x, y) =
+            keys(x.blocks) == keys(y.blocks) &&
+            all(x.blocks[k] ≈ y.blocks[k] for k in keys(x.blocks))
+
+        @testset "svd_compact" begin
+            U, S, Vᴴ = MAK.svd_compact(m_rect)
+            @test all(x -> x isa AbelianGradedMatrix, (U, S, Vᴴ))
+            @test axes(U, 1) == axes(m_rect, 1)
+            @test axes(Vᴴ, 2) == axes(m_rect, 2)
+            @test U * S * Vᴴ ≈ m_rect
+            @test Array(U) * Array(S) * Array(Vᴴ) ≈ Array(m_rect)
+        end
+
+        @testset "svd_full" begin
+            U, S, Vᴴ = MAK.svd_full(m_rect)
+            @test all(x -> x isa AbelianGradedMatrix, (U, S, Vᴴ))
+            @test Array(U) * Array(S) * Array(Vᴴ) ≈ Array(m_rect)
+        end
+
+        @testset "svd_vals" begin
+            @test fgv_approx(MAK.svd_vals(m_rect), MAK.svd_vals(FusedGradedMatrix(m_rect)))
+        end
+
+        @testset "qr_compact / qr_full" begin
+            Q, R = MAK.qr_compact(m_rect)
+            @test Array(Q) * Array(R) ≈ Array(m_rect)
+            Q, R = MAK.qr_full(m_rect)
+            @test Array(Q) * Array(R) ≈ Array(m_rect)
+        end
+
+        @testset "lq_compact / lq_full" begin
+            L, Q = MAK.lq_compact(m_rect)
+            @test Array(L) * Array(Q) ≈ Array(m_rect)
+            L, Q = MAK.lq_full(m_rect)
+            @test Array(L) * Array(Q) ≈ Array(m_rect)
+        end
+
+        @testset "eig_full / eig_vals" begin
+            D, V = MAK.eig_full(m_sq)
+            @test Array(m_sq) * Array(V) ≈ Array(V) * Array(D)
+            @test fgv_approx(MAK.eig_vals(m_sq), MAK.eig_vals(FusedGradedMatrix(m_sq)))
+        end
+
+        @testset "eigh_full / eigh_vals" begin
+            D, V = MAK.eigh_full(m_herm)
+            @test Array(m_herm) ≈ Array(V) * Array(D) * Array(V)'
+            @test fgv_approx(
+                MAK.eigh_vals(m_herm),
+                MAK.eigh_vals(FusedGradedMatrix(m_herm))
+            )
+        end
+
+        @testset "left_polar / right_polar" begin
+            W, P = MAK.left_polar(m_sq)
+            @test Array(W) * Array(P) ≈ Array(m_sq)
+            P, W = MAK.right_polar(m_sq)
+            @test Array(P) * Array(W) ≈ Array(m_sq)
+        end
+
+        @testset "project_hermitian" begin
+            m = randn(Float64, (g, dual(g)))
+            @test Array(MAK.project_hermitian(m)) ≈ (Array(m) + Array(m)') / 2
+        end
+    end
+
+    # -----------------------------------------------------------------------
+    @testset "one! on an AbelianGradedMatrix" begin
+        g = gradedrange([U1(0) => 2, U1(1) => 3, U1(2) => 2])
+        a = randn(Float64, (g, dual(g)))
+        a_before = Array(copy(a))
+
+        b = MAK.one!(a)
+        # In place: returns `a` and mutates its contents.
+        @test b === a
+        @test Array(a) != a_before
+
+        # Same as the existing graded identity constructor and the dense identity.
+        id = TensorAlgebra.one(randn(Float64, (g, dual(g))), (1,), (2,))
+        @test Array(a) ≈ Array(id)
+        @test Array(a) ≈ Matrix(1.0I, size(a)...)
+    end
 end  # @testset "Factorizations"

test/test_tensoralgebra.jl

@@ -449,14 +449,11 @@ end
     A = AbelianGradedArray{Float64}(undef, s, dual(s))
     randn!(A)
     U, S, Vᴴ = TensorAlgebra.svd_compact(A, (1,), (2,))
-    # The natural `U * S * Vᴴ` form falls into LinearAlgebra's `_tri_matmul`,
-    # which scalar-indexes on `AbstractGradedArray`. `contract` is the
-    # block-wise route, but the chain form should also work once a block-aware
-    # matmul lands on `AbstractGradedMatrix`.
     US = contract((:a, :r), U, (:a, :i), S, (:i, :r))
     USV = contract((:a, :b), US, (:a, :r), Vᴴ, (:r, :b))
     @test A ≈ USV
-    @test_broken A ≈ U * S * Vᴴ
+    # `*` on `AbelianGradedMatrix` routes through the block-wise `contract`.
+    @test A ≈ U * S * Vᴴ
 end
 
 @testset "TA.gram_eigh_full_with_pinv on AbelianGradedMatrix (axes_Y regression)" begin
@@ -467,14 +464,13 @@ end
     # so we stay on the graded matmul path; the natural `*` form is broken
     # against the same scalar-indexing path as the SVD round-trip above.
     A = contract((:a, :b), B, (:a, :r), conj(B), (:b, :r))
+    # `*` on two `AbelianGradedMatrix` works, but the adjoint forms (`B * B'`,
+    # `X * X'` below) still need a block-aware `adjoint`; `B'` is an `Adjoint`
+    # wrapper that falls through to LinearAlgebra's scalar-indexing path.
     @test_broken A ≈ B * B'
     X, Y = TensorAlgebra.gram_eigh_full_with_pinv(A, (1,), (2,))
     # X · conj(X) ≈ A on the rank subspace.
     @test A ≈ contract((:a, :b), X, (:a, :r), conj(X), (:b, :r))
-    # Matmul (`*`) on `AbelianGradedMatrix` is unimplemented, so any `X * Y`
-    # falls through to LinearAlgebra's scalar-indexing path and throws. The
-    # adjoint forms (`X * X'`, `B * B'`) additionally need a block-aware
-    # `adjoint`. Both will pass once those land.
     @test_broken A ≈ X * X'
     # Y is a left inverse of X on the rank subspace.
     YX = contract((:r, :s), Y, (:r, :a), X, (:a, :s))