trim excess stuff, add docs

src/UniformTensorTrains/UniformTensorTrains.jl

@@ -10,8 +10,7 @@ using Tullio: @tullio
 
 export AbstractPeriodicTensorTrain, PeriodicTensorTrain, flat_periodic_tt, rand_periodic_tt,
        AbstractUniformTensorTrain, UniformTensorTrain, periodic_tensor_train,
-       symmetrized_uniform_tensor_train, InfiniteUniformTensorTrain,
-       transfer_operator, infinite_transfer_operator, leading_eig
+       symmetrized_uniform_tensor_train, InfiniteUniformTensorTrain
 
 
 include("uniform_tensor_train.jl")

src/UniformTensorTrains/transfer_operator.jl

@@ -1,16 +1,36 @@
-abstract type AbstractTransferOperator{F1<:Number,F2<:Number} end
-abstract type AbstractFiniteTransferOperator{F1<:Number,F2<:Number} <: AbstractTransferOperator{F1,F2} end
+"""
+    AbstractTransferOperator{TA,TM}
 
-struct TransferOperator{F1<:Number,F2<:Number} <: AbstractFiniteTransferOperator{F1,F2}
-    A :: Array{F1,3}
-    M :: Array{F2,3}
+A type to represent a transfer operator.
+"""
+abstract type AbstractTransferOperator{TA<:Number,TM<:Number} end
+
+@doc raw"""
+    TransferOperator{TA<:Number,TM<:Number} <: AbstractTransferOperator{TA,TM}
+
+A type to represent a transfer operator $G$ obtained from variable-dependent matrices $A(x)$, $M(x)$ as
+```math
+G_{i,j,k,l} = \sum_x A_{i,k}(x) M^*_{j,l}(x)
+    ```
+"""
+struct TransferOperator{TA<:Number,TM<:Number} <: AbstractTransferOperator{TA,TM}
+    A :: Array{TA,3}
+    M :: Array{TM,3}
 end
 
-struct HomogeneousTransferOperator{F<:Number} <: AbstractFiniteTransferOperator{F,F}
-    A :: Array{F,3}
+@doc raw"""
+    HomogeneousTransferOperator{T<:Number} <: AbstractTransferOperator{T,T}
+
+A type to represent a transfer operator $G$ obtained from variable-dependent matrix $A(x)$ as
+```math
+G_{i,j,k,l} = \sum_x A_{i,k}(x) A^*_{j,l}(x)
+    ```
+"""
+struct HomogeneousTransferOperator{T<:Number} <: AbstractTransferOperator{T,T}
+    A :: Array{T,3}
 end
 
-function sizes(G::AbstractFiniteTransferOperator)
+function sizes(G::AbstractTransferOperator)
     A, M = get_tensors(G)
     size(A, 1), size(M, 1), size(A, 2), size(M, 2)
 end
@@ -23,25 +43,24 @@ function Base.convert(::Type{TransferOperator}, G::HomogeneousTransferOperator)
 end
 TransferOperator(G::HomogeneousTransferOperator) = convert(TransferOperator, G)
 
-# the first argument `p` is the one with `A` matrices
 function transfer_operator(q::AbstractUniformTensorTrain, p::AbstractUniformTensorTrain)
     return TransferOperator(_reshape1(q.tensor), _reshape1(p.tensor))
 end
 function transfer_operator(q::AbstractUniformTensorTrain)
     return HomogeneousTransferOperator(_reshape1(q.tensor))
 end
 
-function Base.collect(G::AbstractFiniteTransferOperator)
+function Base.collect(G::AbstractTransferOperator)
     A, M = get_tensors(G)
     return @tullio B[i,j,k,l] := A[i,k,x] * conj(M[j,l,x])
 end
 
-function Base.:(*)(G::AbstractFiniteTransferOperator, B::AbstractMatrix)
+function Base.:(*)(G::AbstractTransferOperator, B::AbstractMatrix)
     A, M = get_tensors(G)
     return @tullio C[i,j] := A[i,k,x] * conj(M[j,l,x]) * B[k,l]
 end
 
-function Base.:(*)(B::AbstractMatrix, G::AbstractFiniteTransferOperator)
+function Base.:(*)(B::AbstractMatrix, G::AbstractTransferOperator)
     A, M = get_tensors(G)
     return @tullio C[k,l] := B[i,j] * A[i,k,x] * conj(M[j,l,x])
 end
@@ -59,44 +78,22 @@ function leading_eig(G::AbstractTransferOperator)
     r = reshape(R, d[1], d[2])
     l = reshape(L, d[1], d[2])
     l ./= dot(l, r)
-    return (; l, r, λ)
-end
-
-
-struct InfiniteTransferOperator{F<:Number,M<:AbstractMatrix{F}} <: AbstractTransferOperator{F,F}
-    l :: M
-    r :: M
-    λ :: F
-end
-
-function infinite_transfer_operator(G::AbstractTransferOperator; lambda1::Bool=false)
-    l, r, λ_ = leading_eig(G)
-    λ = lambda1 ? one(λ_) : λ_
-    λ = convert(eltype(r), λ)
-    InfiniteTransferOperator(l, r, λ)
-end
-
-function Base.collect(G::InfiniteTransferOperator)
-    (; l, r, λ) = G
-    return @tullio B[i,j,k,m] := r[i,j] * l[k,m]
+    return λ, l, r
 end
 
-function sizes(G::InfiniteTransferOperator)
-    (; l, r) = G
-    return tuple(size(r)..., size(l)...)
-end
+"""
+    dot(q::InfiniteUniformTensorTrain, p::InfiniteUniformTensorTrain)
 
-function infinite_transfer_operator(q::AbstractUniformTensorTrain, p::AbstractUniformTensorTrain)
-    return infinite_transfer_operator(transfer_operator(q, p))
-end
-
-function infinite_transfer_operator(q::AbstractUniformTensorTrain)
-    return infinite_transfer_operator(transfer_operator(q))
-end
+Return the "dot product" between the infinite tensor trains `q` and `p`.
 
-function LinearAlgebra.dot(q::InfiniteUniformTensorTrain, p::InfiniteUniformTensorTrain;
-        G = infinite_transfer_operator(q, p),
-        Ep = infinite_transfer_operator(p),
-        Eq = infinite_transfer_operator(q))
-    return G.λ / sqrt(abs(Ep.λ * Eq.λ))
+Since two infinite tensor trains are either equal or orthogonal (orthogonality catastrophe), what is actually returned here is the leading eigenvalue of the transfer operator obtained from the matrices of `q` and `p`
+"""
+function LinearAlgebra.dot(q::InfiniteUniformTensorTrain, p::InfiniteUniformTensorTrain)
+    G = transfer_operator(q, p)
+    Eq = transfer_operator(q)
+    Ep = transfer_operator(p)
+    λG, = leading_eig(G)
+    λq, = leading_eig(Eq)
+    λp, = leading_eig(Ep)
+    return λG / sqrt(abs(λp * λq))
 end

test/runtests.jl

@@ -1,5 +1,6 @@
 using TensorTrains
 using TensorTrains.UniformTensorTrains
+using TensorTrains.UniformTensorTrains: transfer_operator, leading_eig
 using Test
 using Random, Suppressor, InvertedIndices
 using Aqua

test/uniform_tensor_train.jl

@@ -103,7 +103,7 @@ end
 
     G = transfer_operator(q, p)
 
-    (; l, r, λ) = leading_eig(transfer_operator(q, p))
+    λ, l, r = leading_eig(transfer_operator(q, p))
     @test l * G ≈ l * λ
     @test G * r ≈ λ * r