[WIP] Start using TensorAlgebra.factorize

Author: mtfishman

Project.toml

@@ -10,6 +10,7 @@ FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MapBroadcast = "ebd9b9da-f48d-417c-9660-449667d60261"
 NamedDimsArrays = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
+TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 UnallocatedArrays = "43c9e47c-e622-40fb-bf18-a09fc8c466b6"
 UnspecifiedTypes = "42b3faec-625b-4613-8ddc-352bf9672b8d"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
@@ -34,6 +35,7 @@ LinearAlgebra = "1.10"
 MapBroadcast = "0.1.5"
 NamedDimsArrays = "0.6"
 SparseArraysBase = "0.5"
+TensorAlgebra = "0.2.10"
 UnallocatedArrays = "0.1.1"
 UnspecifiedTypes = "0.1.1"
 VectorInterface = "0.5"

src/quirks.jl

@@ -20,46 +20,30 @@ hasqns(a::AbstractITensor) = all(hasqns, inds(a))
 # TODO: Investigate this and see if we can get rid of it.
 Base.Broadcast.extrude(a::AbstractITensor) = a
 
-# TODO: This is just a stand-in for truncated SVD
-# that only makes use of `maxdim`, just to get some
-# functionality running in `ITensorMPS.jl`.
-# Define a proper truncated SVD in
-# `MatrixAlgebra.jl`/`TensorAlgebra.jl`.
-function svd_truncated(a::AbstractITensor, codomain_inds; maxdim)
-  U, S, V = svd(a, codomain_inds)
-  r = Base.OneTo(min(maxdim, minimum(Int.(size(S)))))
-  u = commonind(U, S)
-  v = commonind(V, S)
-  us = uniqueinds(U, S)
-  vs = uniqueinds(V, S)
-  U′ = U[(us .=> :)..., u => r]
-  S′ = S[u => r, v => r]
-  V′ = V[v => r, (vs .=> :)...]
-  return U′, S′, V′
+function translate_factorize_kwargs(;
+  # ITensors.jl kwargs.
+  ortho=nothing,
+  cutoff=nothing,
+  maxdim=nothing,
+  # MatrixAlgebraKit.jl/TensorAlgebra.jl kwargs.
+  orth=nothing,
+  trunc=nothing,
+  kwargs...,
+)
+  @show ortho, cutoff, maxdim
+  @show orth, trunc
+  @show kwargs
+  return error()
 end
 
-using LinearAlgebra: qr, svd
-# TODO: Define this in `MatrixAlgebra.jl`/`TensorAlgebra.jl`.
-function factorize(
-  a::AbstractITensor, codomain_inds; maxdim=nothing, cutoff=nothing, ortho="left", kwargs...
-)
-  # TODO: Perform this intersection in `TensorAlgebra.qr`/`TensorAlgebra.svd`?
-  # See https://github.com/ITensor/NamedDimsArrays.jl/issues/22.
-  codomain_inds′ = if ortho == "left"
-    intersect(inds(a), codomain_inds)
-  elseif ortho == "right"
-    setdiff(inds(a), codomain_inds)
-  else
-    error("Bad `ortho` input.")
-  end
-  F1, F2 = if isnothing(maxdim) && isnothing(cutoff)
-    qr(a, codomain_inds′)
-  else
-    U, S, V = svd_truncated(a, codomain_inds′; maxdim)
-    U, S * V
-  end
-  if ortho == "right"
-    F2, F1 = F1, F2
-  end
-  return F1, F2, (; truncerr=zero(Bool),)
+using TensorAlgebra: TensorAlgebra, factorize
+function TensorAlgebra.factorize(a::AbstractITensor, codomain_inds, domain_inds; kwargs...)
+  return invoke(
+    factorize,
+    Tuple{AbstractNamedDimsArray,Any,Any},
+    a,
+    codomain_inds,
+    domain_inds;
+    translate_factorize_kwargs(; kwargs...)...,
+  )
 end

test/Project.toml

@@ -6,11 +6,13 @@ DiagonalArrays = "74fd4be6-21e2-4f6f-823a-4360d37c7a77"
 GradedArrays = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 ITensorBase = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 NamedDimsArrays = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseArraysBase = "0d5efcca-f356-4864-8770-e1ed8d78f208"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 SymmetrySectors = "f8a8ad64-adbc-4fce-92f7-ffe2bb36a86e"
+TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]

test/test_basics.jl

@@ -19,6 +19,7 @@ using ITensorBase:
 using NamedDimsArrays: dename, name, named
 using SparseArraysBase: oneelement
 using SymmetrySectors: U1
+using LinearAlgebra: factorize
 using Test: @test, @test_broken, @test_throws, @testset
 
 @testset "ITensorBase" begin
@@ -164,4 +165,10 @@ using Test: @test, @test_broken, @test_throws, @testset
     @test hasqns(j)
     @test hasqns(a)
   end
+  @testset "factorize" begin
+    i = Index(2)
+    j = Index(2)
+    a = randn(i, j)
+    x, y = factorize(a, (i,))
+  end
 end