Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions Project.toml
Original file line number Diff line number Diff line change
@@ -1,14 +1,16 @@
name = "TensorKitManifolds"
uuid = "11fa318c-39cb-4a83-b1ed-cdc7ba1e3684"
authors = ["Jutho Haegeman <jutho.haegeman@ugent.be>", "Markus Hauru <markus@mhauru.org>"]
version = "0.7.2"
version = "0.7.3"

[deps]
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"

[compat]
TensorKit = "0.13,0.14"
MatrixAlgebraKit = "0.5.0"
TensorKit = "0.15"
julia = "1.10"

[extras]
Expand Down
61 changes: 15 additions & 46 deletions src/TensorKitManifolds.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@ export Grassmann, Stiefel, Unitary
export inner, retract, transport, transport!

using TensorKit
using MatrixAlgebraKit: MatrixAlgebraKit, AbstractAlgorithm, Algorithm, PolarViaSVD,
LAPACK_DivideAndConquer, diagview
import MatrixAlgebraKit as MAK

# Every submodule -- Grassmann, Stiefel, and Unitary -- implements their own methods for
# these. The signatures should be
Expand All @@ -28,23 +31,8 @@ checkbase(x, y, z, args...) = checkbase(checkbase(x, y), z, args...)
# the machine epsilon for the elements of an object X, name inspired from eltype
scalareps(X) = eps(real(scalartype(X)))

# default SVD algorithm used in the algorithms
default_svd_alg(::AbstractTensorMap) = TensorKit.SVD()

function isisometry(W::AbstractTensorMap; tol=10 * scalareps(W))
WdW = W' * W
s = zero(float(real(scalartype(W))))
for (c, b) in blocks(WdW)
_subtractone!(b)
s += dim(c) * length(b)
end
return norm(WdW) <= tol * sqrt(s)
end

function isunitary(W::AbstractTensorMap; tol=10 * scalareps(W))
return isisometry(W; tol=tol) && isisometry(W'; tol=tol)
end

# TODO: these functions should be replaced by MAK functions
projecthermitian(W::AbstractTensorMap) = projecthermitian!(copy(W))
function projecthermitian!(W::AbstractTensorMap)
codomain(W) == domain(W) ||
throw(DomainError("Tensor with distinct domain and codomain cannot be hermitian."))
Expand All @@ -53,6 +41,8 @@ function projecthermitian!(W::AbstractTensorMap)
end
return W
end

projectantihermitian(W::AbstractTensorMap) = projectantihermitian!(copy(W))
function projectantihermitian!(W::AbstractTensorMap)
codomain(W) == domain(W) ||
throw(DomainError("Tensor with distinct domain and codomain cannot be anithermitian."))
Expand All @@ -62,27 +52,18 @@ function projectantihermitian!(W::AbstractTensorMap)
return W
end

struct PolarNewton <: TensorKit.OrthogonalFactorizationAlgorithm
end
function projectisometric!(W::AbstractTensorMap; alg=default_svd_alg(W))
if alg isa TensorKit.Polar || alg isa TensorKit.SDD
foreach(blocks(W)) do (c, b)
return _polarsdd!(b)
end
elseif alg isa TensorKit.SVD
foreach(blocks(W)) do (c, b)
return _polarsvd!(b)
end
elseif alg isa PolarNewton
foreach(blocks(W)) do (c, b)
return _polarnewton!(b)
end
else
throw(ArgumentError("unkown algorithm for projectisometric!: alg = $alg"))
projectisometric(W::AbstractTensorMap; kwargs...) = projectisometric!(copy(W); kwargs...)
function projectisometric!(W::AbstractTensorMap;
alg::AbstractAlgorithm=MAK.select_algorithm(left_polar!, W))
TensorKit.foreachblock(W) do c, (b,)
return _left_polar!(b, alg)
end
return W
end

function projectcomplement(X::AbstractTensorMap, W::AbstractTensorMap, kwargs...)
return projectcomplement!(copy(X), W; kwargs...)
end
function projectcomplement!(X::AbstractTensorMap, W::AbstractTensorMap;
tol=10 * scalareps(X))
P = W' * X
Expand All @@ -97,18 +78,6 @@ function projectcomplement!(X::AbstractTensorMap, W::AbstractTensorMap;
return X
end

projecthermitian(W::AbstractTensorMap) = projecthermitian!(copy(W))
projectantihermitian(W::AbstractTensorMap) = projectantihermitian!(copy(W))

function projectisometric(W::AbstractTensorMap;
alg=default_svd_alg(W))
return projectisometric!(copy(W); alg=alg)
end
function projectcomplement(X::AbstractTensorMap, W::AbstractTensorMap,
tol=10 * scalareps(X))
return projectcomplement!(copy(X), W; tol=tol)
end

include("auxiliary.jl")
include("grassmann.jl")
include("stiefel.jl")
Expand Down
52 changes: 26 additions & 26 deletions src/auxiliary.jl
Original file line number Diff line number Diff line change
Expand Up @@ -70,37 +70,43 @@ function _subtractone!(a::AbstractMatrix)
view(a, diagind(a)) .= view(a, diagind(a)) .- 1
return a
end
function _polarsdd!(A::StridedMatrix)
U, S, V = svd!(A; alg=LinearAlgebra.DivideAndConquer())
return mul!(A, U, V')
end
function _polarsvd!(A::StridedMatrix)
U, S, V = svd!(A; alg=LinearAlgebra.QRIteration())
return mul!(A, U, V')

# TODO: _left_polar! is more or less the same as MAK.left_polar! but doesn't compute the P
# which is not needed here. Can we unify this?
function _left_polar!(A::StridedMatrix,
alg::PolarViaSVD=PolarViaSVD(LAPACK_DivideAndConquer()))
U, _, Vᴴ = svd_compact!(A, alg.svdalg)
return mul!(A, U, Vᴴ)
end

# TODO: can we move this to a dedicated MAK algorithm?
MatrixAlgebraKit.@algdef PolarNewton

_left_polar!(A::StridedMatrix, alg::PolarNewton) = _polarnewton!(A; alg.kwargs...)
function _polarnewton!(A::StridedMatrix; tol=10 * scalareps(A), maxiter=5)
m, n = size(A)
@assert m >= n
A2 = copy(A)
Q, R = qr!(A2)
Ri = ldiv!(UpperTriangular(R)', TensorKit.MatrixAlgebra.one!(similar(R)))
Q, R = LinearAlgebra.qr!(A2)
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not use MAK.qr_compact! for this?

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same for the other qr! calls below.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Because this is the compact form that is immediately in-place multiplied into a different matrix below, so technically that's doing an additional multiplication lmul! below. I'm not sure if it really matters, but it felt a bit weird to change this into something slower if it was working before :)

Ri = ldiv!(UpperTriangular(R)', MatrixAlgebraKit.one!(similar(R)))
R, Ri = _avgdiff!(R, Ri)
i = 1
R2 = view(A, 1:n, 1:n)
fill!(view(A, (n + 1):m, 1:n), zero(eltype(A)))
copyto!(R2, R)
while maximum(abs, Ri) > tol
if i == maxiter # if not converged by now, fall back to sdd
_polarsdd!(Ri)
_left_polar!(Ri)
break
end
Ri = ldiv!(lu!(R2)', TensorKit.MatrixAlgebra.one!(Ri))
Ri = ldiv!(lu!(R2)', MatrixAlgebraKit.one!(Ri))
R, Ri = _avgdiff!(R, Ri)
copyto!(R2, R)
i += 1
end
return lmul!(Q, A)
end

# in place computation of the average and difference of two arrays
function _avgdiff!(A::AbstractArray, B::AbstractArray)
axes(A) == axes(B) || throw(DimensionMismatch())
Expand All @@ -124,7 +130,7 @@ end
function _stiefelexp(W::StridedMatrix, A::StridedMatrix, Z::StridedMatrix, α)
n, p = size(W)
r = min(2 * p, n)
QQ, _ = qr!([W Z])
QQ, _ = LinearAlgebra.qr!([W Z])
Q = similar(W, n, r - p)
@inbounds for j in Base.OneTo(r - p)
for i in Base.OneTo(n)
Expand All @@ -139,7 +145,7 @@ function _stiefelexp(W::StridedMatrix, A::StridedMatrix, Z::StridedMatrix, α)
A2[1:p, (p + 1):end] .= (-α) .* (R')
A2[(p + 1):end, (p + 1):end] .= 0
U = [W Q] * exp(A2)
U = _polarnewton!(U)
U = _left_polar!(U, PolarNewton())
W′ = U[:, 1:p]
Q′ = U[:, (p + 1):end]
R′ = R
Expand All @@ -152,7 +158,7 @@ function _stiefellog(Wold::StridedMatrix, Wnew::StridedMatrix;
r = min(2 * p, n)
P = Wold' * Wnew
dW = Wnew - Wold * P
QQ, _ = qr!([Wold dW])
QQ, _ = LinearAlgebra.qr!([Wold dW])
Q = similar(Wold, n, r - p)
@inbounds for j in Base.OneTo(r - p)
for i in Base.OneTo(n)
Expand All @@ -161,23 +167,17 @@ function _stiefellog(Wold::StridedMatrix, Wnew::StridedMatrix;
end
Q = lmul!(QQ, Q)
R = Q' * dW
Wext = [Wold Q]
F = qr!([P; R])
U = lmul!(F.Q, TensorKit.MatrixAlgebra.one!(similar(P, r, r)))
F = LinearAlgebra.qr!([P; R])
U = lmul!(F.Q, MatrixAlgebraKit.one!(similar(P, r, r)))
U[1:p, 1:p] .= P
U[(p + 1):r, 1:p] .= R
X = view(U, 1:p, (p + 1):r)
Y = view(U, (p + 1):r, (p + 1):r)
if p < n
YSVD = svd!(Y)
mul!(X, X * (YSVD.V), (YSVD.U)')
UsqrtS = YSVD.U
@inbounds for j in 1:size(UsqrtS, 2)
s = sqrt(YSVD.S[j])
@simd for i in 1:size(UsqrtS, 1)
UsqrtS[i, j] *= s
end
end
USVᴴ = svd_compact!(Y)
mul!(X, X * USVᴴ[3]', USVᴴ[1]')
diagview(USVᴴ[2]) .= sqrt.(diagview(USVᴴ[2]))
UsqrtS = rmul!(USVᴴ[1], USVᴴ[2])
mul!(Y, UsqrtS, UsqrtS')
end
logU = _projectantihermitian!(log(U))
Expand Down
11 changes: 5 additions & 6 deletions src/grassmann.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,7 @@ module Grassmann
using TensorKit
using TensorKit: similarstoragetype, SectorDict
using ..TensorKitManifolds: projecthermitian!, projectantihermitian!,
projectisometric!, projectcomplement!, PolarNewton,
default_svd_alg
projectisometric!, projectcomplement!, PolarNewton
import ..TensorKitManifolds: base, checkbase, inner, retract, transport, transport!

# special type to store tangent vectors using Z
Expand All @@ -32,8 +31,8 @@ end

# output type of U, S, V in tsvd
function _tsvd_types(Z::AbstractTensorMap)
TUSV = Core.Compiler.return_type(tsvd, Tuple{typeof(Z)})
TU, TS, TV, = TUSV.types
TUSV = Core.Compiler.return_type(svd_compact, Tuple{typeof(Z)})
TU, TS, TV = TUSV.types
return TU, TS, TV
end

Expand All @@ -60,7 +59,7 @@ function Base.getproperty(Δ::GrassmannTangent, sym::Symbol)
elseif sym ∈ (:U, :S, :V)
v = Base.getfield(Δ, sym)
v !== nothing && return v
U, S, V, = tsvd(Δ.Z; alg=default_svd_alg(Δ.Z))
U, S, V = svd_compact(Δ.Z)
Base.setfield!(Δ, :U, U)
Base.setfield!(Δ, :S, S)
Base.setfield!(Δ, :V, V)
Expand Down Expand Up @@ -198,7 +197,7 @@ function invretract(Wold::AbstractTensorMap, Wnew::AbstractTensorMap; alg=nothin
space(Wold) == space(Wnew) || throw(SpaceMismatch())
WodWn = Wold' * Wnew # V' * cos(S) * V * Y
Wneworth = Wnew - Wold * WodWn
Vd, cS, VY = tsvd!(WodWn; alg=default_svd_alg(WodWn))
Vd, cS, VY = svd_compact!(WodWn)
Scmplx = acos(cS)
# acos always returns a complex TensorMap. We cast back to real if possible.
S = scalartype(WodWn) <: Real && isreal(sectortype(Scmplx)) ? real(Scmplx) : Scmplx
Expand Down
8 changes: 5 additions & 3 deletions src/unitary.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ using TensorKit
import TensorKit: similarstoragetype, SectorDict
using ..TensorKitManifolds: projectantihermitian!, projectisometric!, PolarNewton
import ..TensorKitManifolds: base, checkbase, inner, retract, transport, transport!
import MatrixAlgebraKit as MAK

struct UnitaryTangent{T<:AbstractTensorMap,TA<:AbstractTensorMap}
W::T
Expand Down Expand Up @@ -82,10 +83,11 @@ end
project(X, W; metric=:euclidean) = project!(copy(X), W; metric=:euclidean)

# geodesic retraction, coincides with Stiefel retraction (which is not geodesic for p < n)
function retract(W::AbstractTensorMap, Δ::UnitaryTangent, α; alg=nothing)
function retract(W::AbstractTensorMap, Δ::UnitaryTangent, α;
alg=MAK.select_algorithm(left_polar!, W))
W == base(Δ) || throw(ArgumentError("not a valid tangent vector at base point"))
E = exp(α * Δ.A)
W′ = projectisometric!(W * E; alg=SDD())
W′ = projectisometric!(W * E; alg)
A′ = Δ.A
return W′, UnitaryTangent(W′, A′)
end
Expand All @@ -104,7 +106,7 @@ function transport!(Θ::UnitaryTangent, W::AbstractTensorMap, Δ::UnitaryTangent
end
function transport(Θ::UnitaryTangent, W::AbstractTensorMap, Δ::UnitaryTangent, α::Real, W′;
alg=:stiefel)
return transport!(copy(Θ), W, Δ, α, W′; alg=alg)
return transport!(copy(Θ), W, Δ, α, W′; alg)
end

# transport_parallel correspondings to the torsion-free Levi-Civita connection
Expand Down
4 changes: 2 additions & 2 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ const α = 0.75

@testset "Grassmann with space $V" for V in spaces
for T in (Float64,)
W, = leftorth(randn(T, V * V * V, V * V); alg=Polar())
W, = left_polar(randn(T, V * V * V, V * V))
X = randn(T, space(W))
Y = randn(T, space(W))
Δ = @inferred Grassmann.project(X, W)
Expand Down Expand Up @@ -124,7 +124,7 @@ end

@testset "Unitary with space $V" for V in spaces
for T in (Float64, ComplexF64)
W, = leftorth(randn(T, V * V * V, V * V); alg=Polar())
W, = left_polar(randn(T, V * V * V, V * V))
X = randn(T, space(W))
Y = randn(T, space(W))
Δ = @inferred Unitary.project(X, W)
Expand Down
Loading