[WIP] Add tensor factorizations through MatrixAlgebraKit #23
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,45 +1,83 @@ | ||
| using MatrixAlgebraKit: | ||
| MatrixAlgebraKit, qr_full, qr_compact, svd_full, svd_compact, svd_trunc | ||
| # TODO: consider in-place version | ||
| # TODO: figure out kwargs and document | ||
| # | ||
| """ | ||
| qr(A::AbstractArray, labels_A, labels_codomain, labels_domain; full=true, kwargs...) | ||
| qr(A::AbstractArray, biperm::BlockedPermutation{2}; full=true, kwargs...) | ||
| Compute the QR decomposition of a generic N-dimensional array, by interpreting it as | ||
| a linear map from the domain to the codomain indices. These can be specified either via | ||
| their labels, or directly through a `biperm`. | ||
| """ | ||
| function qr(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) | ||
| biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...) | ||
| return qr(A, biperm) | ||
| end | ||
| function qr(A::AbstractArray, biperm::BlockedPermutation{2}; full::Bool=true, kwargs...) | ||
| # tensor to matrix | ||
| A_mat = fusedims(A, biperm) | ||
|
|
||
| # factorization | ||
| Q, R = full ? qr_full(A_mat; kwargs...) : qr_compact(A_mat; kwargs...) | ||
|
|
||
| # matrix to tensor | ||
| axes_codomain, axes_domain = blockpermute(axes(A), biperm) | ||
| axes_Q = (axes_codomain..., axes(Q, 2)) | ||
| axes_R = (axes(R, 1), axes_domain...) | ||
| return splitdims(Q, axes_Q), splitdims(R, axes_R) | ||
| end | ||
|
|
||
| """ | ||
| svd(A::AbstractArray, labels_A, labels_codomain, labels_domain; full=false, kwargs...) | ||
| svd(A::AbstractArray, biperm::BlockedPermutation{2}; full=false, kwargs...) | ||
| Compute the SVD decomposition of a generic N-dimensional array, by interpreting it as | ||
| a linear map from the domain to the codomain indices. These can be specified either via | ||
| their labels, or directly through a `biperm`. | ||
| """ | ||
| function svd(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) | ||
| biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...) | ||
| return svd(A, biperm; kwargs...) | ||
| end | ||
| function svd(A::AbstractArray, biperm::BlockedPermutation{2}; full::Bool=false, kwargs...) | ||
| # tensor to matrix | ||
| A_mat = fusedims(A, biperm) | ||
|
|
||
| # factorization | ||
| U, S, Vᴴ = full ? svd_full(A_mat; kwargs...) : svd_compact(A_mat; kwargs...) | ||
|
|
||
| # matrix to tensor | ||
| axes_codomain, axes_domain = blockpermute(axes(A), biperm) | ||
| axes_U = (axes_codomain..., axes(U, 2)) | ||
| axes_Vᴴ = (axes(Vᴴ, 1), axes_domain...) | ||
| return splitdims(U, axes_U), S, splitdims(Vᴴ, axes_Vᴴ) | ||
| end | ||
|
|
||
| # TODO: decide on interface | ||
| """ | ||
| tsvd(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) | ||
| tsvd(A::AbstractArray, biperm::BlockedPermutation{2}; kwargs...) | ||
| Compute the truncated SVD decomposition of a generic N-dimensional array, by interpreting it | ||
| as a linear map from the domain to the codomain indices. These can be specified either via | ||
| their labels, or directly through a `biperm`. | ||
| """ | ||
| function tsvd(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) | ||
| biperm = blockedperm_indexin(Tuple.((labels_A, labels_codomain, labels_domain))...) | ||
| return tsvd(A, biperm; kwargs...) | ||
| end | ||
| function tsvd(A::AbstractArray, biperm::BlockedPermutation{2}; kwargs...) | ||
| # tensor to matrix | ||
| A_mat = fusedims(A, biperm) | ||
|
|
||
| # factorization | ||
| U, S, Vᴴ = svd_trunc(A_mat; kwargs...) | ||
|
|
||
| # matrix to tensor | ||
| axes_codomain, axes_domain = blockpermute(axes(A), biperm) | ||
| axes_U = (axes_codomain..., axes(U, 2)) | ||
| axes_Vᴴ = (axes(Vᴴ, 1), axes_domain...) | ||
| return splitdims(U, axes_U), S, splitdims(Vᴴ, axes_Vᴴ) | ||
| end | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,62 @@ | ||
| using Test: @test, @testset, @inferred | ||
| using TestExtras: @constinferred | ||
| using TensorAlgebra: contract, svd, tsvd, TensorAlgebra | ||
| using TensorAlgebra.MatrixAlgebraKit: truncrank | ||
|
There was a problem hiding this comment. Just as a style preference, I prefer There was a problem hiding this comment. To be honest, I added this as a placeholder while we settle on where the keyword arguments are supposed to be handled, since we should either import that into TensorAlgebra, or provide a different interface. I just wanted to run some tests in the meantime, but for sure will adapt this :) |
||
| using LinearAlgebra: norm | ||
|
|
||
| elts = (Float64, ComplexF64) | ||
|
|
||
| @testset "Full SVD ($T)" for T in elts | ||
| A = randn(T, 5, 4, 3, 2) | ||
| A ./= norm(A) | ||
| labels_A = (:a, :b, :c, :d) | ||
| labels_U = (:b, :a) | ||
| labels_Vᴴ = (:d, :c) | ||
|
|
||
| Acopy = deepcopy(A) | ||
| U, S, Vᴴ = @constinferred svd(A, labels_A, labels_U, labels_Vᴴ; full=true) | ||
| @test A == Acopy # should not have altered initial array | ||
| US, labels_US = contract(U, (labels_U..., :u), S, (:u, :v)) | ||
| A′ = contract(labels_A, US, labels_US, Vᴴ, (:v, labels_Vᴴ...)) | ||
| @test A ≈ A′ | ||
| @test size(U, 1) * size(U, 2) == size(U, 3) # U is unitary | ||
| @test size(Vᴴ, 1) == size(Vᴴ, 2) * size(Vᴴ, 3) # V is unitary | ||
| end | ||
|
|
||
| @testset "Compact SVD ($T)" for T in elts | ||
| A = randn(T, 5, 4, 3, 2) | ||
| A ./= norm(A) | ||
| labels_A = (:a, :b, :c, :d) | ||
| labels_U = (:b, :a) | ||
| labels_Vᴴ = (:d, :c) | ||
|
|
||
| Acopy = deepcopy(A) | ||
| U, S, Vᴴ = @constinferred svd(A, labels_A, labels_U, labels_Vᴴ; full=false) | ||
| @test A == Acopy # should not have altered initial array | ||
| US, labels_US = contract(U, (labels_U..., :u), S, (:u, :v)) | ||
| A′ = contract(labels_A, US, labels_US, Vᴴ, (:v, labels_Vᴴ...)) | ||
| @test A ≈ A′ | ||
| k = min(size(S)...) | ||
| @test size(U, 3) == k == size(Vᴴ, 1) | ||
| end | ||
|
|
||
| @testset "Truncated SVD ($T)" for T in elts | ||
| A = randn(T, 5, 4, 3, 2) | ||
| A ./= norm(A) | ||
| labels_A = (:a, :b, :c, :d) | ||
| labels_U = (:b, :a) | ||
| labels_Vᴴ = (:d, :c) | ||
|
|
||
| # test truncated SVD | ||
| Acopy = deepcopy(A) | ||
| _, S_untrunc, _ = svd(A, labels_A, labels_U, labels_Vᴴ) | ||
|
|
||
| trunc = truncrank(size(S_untrunc, 1) - 1) | ||
| U, S, Vᴴ = @constinferred tsvd(A, labels_A, labels_U, labels_Vᴴ; trunc) | ||
|
|
||
| @test A == Acopy # should not have altered initial array | ||
| US, labels_US = contract(U, (labels_U..., :u), S, (:u, :v)) | ||
| A′ = contract(labels_A, US, labels_US, Vᴴ, (:v, labels_Vᴴ...)) | ||
| @test norm(A - A′) ≈ S_untrunc[end] | ||
| @test size(S, 1) == size(S_untrunc, 1) - 1 | ||
| end | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just a suggestion, but what if truncated SVD is implemented as an algorithm backend of a more general
svdfunction (i.e. we havesvd(::Algorithm"truncated", a::AbstractArray, ...)andsvd(::Algorithm"untruncated", a::AbstractArray, ...))?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm definitely also fine with that, behind the scenes this is already what is going on in MatrixAlgebraKit as well, see the definition here. It's just a matter of deciding on where you want to start intercepting the user-input and what you would like the keywords for that to be, and I'm happy to map it to whatever MatrixAlgebraKit has
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah, that is definitely a bit subtle, since ITensor has its own syntax for truncation keyword arguments.
Maybe a strategy we can take for now is starting out by following whatever MatrixAlgebraKit.jl does (say
TensorAlgebra.svdjust forwards kwargs toMatrixAlgebraKit.svd), and then at the ITensor level we can map our truncation arguments likecutoffandmaxdimto the TensorAlgebra.jl/MatrixAlgebraKit.jl names and conventions, and additionally chooseTruncatedAlgorithmas the default SVD algorithm.