diff --git a/src/T4AITensorCompat.jl b/src/T4AITensorCompat.jl index 6a8281c..7f6e05e 100644 --- a/src/T4AITensorCompat.jl +++ b/src/T4AITensorCompat.jl @@ -37,6 +37,7 @@ export random_mps, random_mpo # Random tensor train generation export product # Official API name (match ITensorMPS) export apply # Backwards-compatible alias +abstract type AbstractTTN end # Abstract type for tree tensor network include("defaults.jl") include("tensortrain.jl") diff --git a/src/contraction.jl b/src/contraction.jl index e20e3a7..9ab5da1 100644 --- a/src/contraction.jl +++ b/src/contraction.jl @@ -67,50 +67,14 @@ result = contract(M1, M2) # Using default fit algorithm result = contract(M1, M2; alg=Algorithm"densitymatrix"(), maxdim=100) ``` """ -function contract(M1::TensorTrain, M2::TensorTrain; alg=Algorithm"fit"(), cutoff::Real=default_cutoff(), maxdim::Int=default_maxdim(), nsweeps::Int=default_nsweeps(), kwargs...)::TensorTrain - # Detect MPS-like vs MPO-like by counting physical indices per site - is_mps1 = begin - sites_per_tensor = siteinds(M1) - length(M1) > 0 && all(length(s) == 1 for s in sites_per_tensor) - end - is_mps2 = begin - sites_per_tensor = siteinds(M2) - length(M2) > 0 && all(length(s) == 1 for s in sites_per_tensor) - end - - # Convert to ITensorMPS types based on detected type - M1_ = is_mps1 ? ITensorMPS.MPS(M1) : ITensorMPS.MPO(M1) - M2_ = is_mps2 ? ITensorMPS.MPS(M2) : ITensorMPS.MPO(M2) +function contract(M1::TensorTrain, M2::TensorTrain; alg::Union{String, Algorithm}=Algorithm"fit"(), cutoff::Real=default_cutoff(), maxdim::Int=default_maxdim(), nsweeps::Int=default_nsweeps(), kwargs...)::TensorTrain + # Convert both M1 and M2 to MPO format for T4A implementation + M1_ = ITensorMPS.MPO(M1) + M2_ = ITensorMPS.MPO(M2) alg = Algorithm(alg) - # Handle MPO * MPS case - if !is_mps1 && is_mps2 - # MPO * MPS: use ITensorMPS.contract - # Only pass nsweeps for fit algorithm, otherwise remove it from kwargs - if alg == Algorithm"fit"() - result = ITensorMPS.contract(M1_, M2_; alg=alg, cutoff=cutoff, maxdim=maxdim, nsweeps=nsweeps, kwargs...) - else - # Remove nsweeps from kwargs for non-fit algorithms - kwargs_dict = Dict(kwargs) - delete!(kwargs_dict, :nsweeps) - result = ITensorMPS.contract(M1_, M2_; alg=alg, cutoff=cutoff, maxdim=maxdim, kwargs_dict...) - end - return TensorTrain(result) - elseif is_mps1 && !is_mps2 - # MPS * MPO: use ITensorMPS.contract (commutative) - # Only pass nsweeps for fit algorithm, otherwise remove it from kwargs - if alg == Algorithm"fit"() - result = ITensorMPS.contract(M2_, M1_; alg=alg, cutoff=cutoff, maxdim=maxdim, nsweeps=nsweeps, kwargs...) - else - # Remove nsweeps from kwargs for non-fit algorithms - kwargs_dict = Dict(kwargs) - delete!(kwargs_dict, :nsweeps) - result = ITensorMPS.contract(M2_, M1_; alg=alg, cutoff=cutoff, maxdim=maxdim, kwargs_dict...) - end - return TensorTrain(result) - else - # MPO * MPO: use T4AITensorCompat algorithms + # Use T4AITensorCompat algorithms if alg == Algorithm"densitymatrix"() return TensorTrain(ContractionImpl.contract_densitymatrix(M1_, M2_; cutoff, maxdim, kwargs...)) elseif alg == Algorithm"fit"() @@ -121,7 +85,6 @@ function contract(M1::TensorTrain, M2::TensorTrain; alg=Algorithm"fit"(), cutoff return TensorTrain(ITensorMPS.contract(M1_, M2_; alg=Algorithm"naive"(), cutoff, maxdim, kwargs...)) else error("Unknown algorithm: $alg") - end end end @@ -197,7 +160,7 @@ using `contract` internally and adjusting prime levels for compatibility. - `nsweeps::Int`: Number of sweeps for variational algorithms. Defaults to `default_nsweeps()`. - `kwargs...`: Additional keyword arguments passed to contract """ -function product(A::TensorTrain, Ψ::TensorTrain; alg=Algorithm"fit"(), cutoff=default_cutoff(), maxdim=default_maxdim(), nsweeps=default_nsweeps(), kwargs...) +function product(A::TensorTrain, Ψ::TensorTrain; alg::Union{String, Algorithm}=Algorithm"fit"(), cutoff=default_cutoff(), maxdim=default_maxdim(), nsweeps=default_nsweeps(), kwargs...) if :algorithm ∈ keys(kwargs) error("keyword argument :algorithm is not allowed") end @@ -210,6 +173,16 @@ function product(A::TensorTrain, Ψ::TensorTrain; alg=Algorithm"fit"(), cutoff=d @warn "cutoff is too small for densitymatrix algorithm. Use fit algorithm instead." end + # Check that A is MPO-like (has 2 or more physical indices per site) + is_mpo_like_A = begin + sites_per_tensor = siteinds(A) + length(A) > 0 && all(length(s) >= 2 for s in sites_per_tensor) + end + + if !is_mpo_like_A + error("First argument `A` must be MPO-like (have 2 or more physical indices per site), but got a tensor train with $(length(A) > 0 ? length(siteinds(A)[1]) : 0) physical index(es) per site") + end + # Detect MPS-like vs MPO-like by counting physical indices per site: # MPS tensors have 1 physical index per site, MPO tensors have 2 per site. # This is robust to boundary tensors having fewer link indices. diff --git a/src/ext/T4AITensorCompatChainRulesCoreExt.jl b/src/ext/T4AITensorCompatChainRulesCoreExt.jl deleted file mode 100644 index 17c7169..0000000 --- a/src/ext/T4AITensorCompatChainRulesCoreExt.jl +++ /dev/null @@ -1,21 +0,0 @@ -# Extension module to resolve ambiguities with ChainRulesCore and InitialValues -# This module is automatically loaded when ChainRulesCore or InitialValues are available - -module T4AITensorCompatChainRulesCoreExt - -using ChainRulesCore: NoTangent, AbstractThunk, NotImplemented, ZeroTangent -using InitialValues: NonspecificInitialValue, SpecificInitialValue -using ..T4AITensorCompat: TensorTrain - -# Resolve ambiguities with ChainRulesCore types -Base.:+(stt::TensorTrain, ::NoTangent) = error("Cannot add ChainRulesCore.NoTangent to TensorTrain") -Base.:+(stt::TensorTrain, ::AbstractThunk) = error("Cannot add ChainRulesCore.AbstractThunk to TensorTrain") -Base.:+(stt::TensorTrain, ::NotImplemented) = error("Cannot add ChainRulesCore.NotImplemented to TensorTrain") -Base.:+(stt::TensorTrain, ::ZeroTangent) = error("Cannot add ChainRulesCore.ZeroTangent to TensorTrain") - -# Resolve ambiguities with InitialValues types -Base.:+(stt::TensorTrain, ::Union{NonspecificInitialValue, SpecificInitialValue{typeof(+)}}) = - error("Cannot add InitialValues to TensorTrain") - -end - diff --git a/src/tensortrain.jl b/src/tensortrain.jl index 26cc036..447b799 100644 --- a/src/tensortrain.jl +++ b/src/tensortrain.jl @@ -35,7 +35,7 @@ This struct supports conversion to/from ITensorMPS types: - `TensorTrain(mps::MPS)` - Convert from Matrix Product State - `TensorTrain(mpo::MPO)` - Convert from Matrix Product Operator """ -mutable struct TensorTrain +mutable struct TensorTrain <: AbstractTTN data::Vector{ITensor} llim::Int rlim::Int @@ -72,7 +72,7 @@ end Construct a TensorTrain (MPS) from a vector of site indices. -This creates an MPS with empty ITensors initialized with the specified site indices. +This creates an MPS by constructing core arrays with default link dimensions (1) initialized with the specified site indices. # Arguments - `sites::Vector{<:Index}`: Vector of site indices @@ -81,9 +81,8 @@ This creates an MPS with empty ITensors initialized with the specified site indi - `TensorTrain`: A new TensorTrain (MPS) object """ function TensorTrain(sites::Vector{<:Index}) - # Construct an MPS with default link dimensions - mps = ITensorMPS.MPS(Float64, sites) - return TensorTrain(mps) + # Construct an MPS with default link dimensions using the general constructor + return TensorTrain(Float64, sites; linkdims=1) end """ @@ -91,7 +90,7 @@ end Construct a TensorTrain (MPS) from a vector of site indices with specified element type. -This creates an MPS with empty ITensors of type `T` initialized with the specified site indices. +This creates an MPS by constructing core arrays of type `T` with the specified site indices and link dimensions. # Arguments - `T::Type{<:Number}`: Element type (e.g., Float64, ComplexF64) @@ -102,21 +101,53 @@ This creates an MPS with empty ITensors of type `T` initialized with the specifi - `TensorTrain`: A new TensorTrain (MPS) object """ function TensorTrain(::Type{T}, sites::Vector{<:Index}; linkdims::Union{Integer,Vector{<:Integer}}=1) where {T<:Number} - mps = ITensorMPS.MPS(T, sites; linkdims=linkdims) - return TensorTrain(mps) + N = length(sites) + N == 0 && return TensorTrain(Vector{ITensor}()) + + # Normalize linkdims to a vector + if linkdims isa Integer + _linkdims = fill(linkdims, N - 1) + else + _linkdims = linkdims + end + + length(_linkdims) == N - 1 || error("Length mismatch: linkdims ($(length(_linkdims))) must have length $(N - 1)") + + # Create cores with appropriate dimensions + # MPS structure: (left_link, physical, right_link) + cores = Vector{Array{T}}(undef, N) + for n in 1:N + physical_dim = ITensors.dim(sites[n]) + if n == 1 && n == N + # Single site: (physical,) + cores[n] = zeros(T, physical_dim) + elseif n == 1 + # First site: (1, physical, right_link) + cores[n] = zeros(T, 1, physical_dim, _linkdims[n]) + elseif n == N + # Last site: (left_link, physical, 1) + cores[n] = zeros(T, _linkdims[n - 1], physical_dim, 1) + else + # Middle site: (left_link, physical, right_link) + cores[n] = zeros(T, _linkdims[n - 1], physical_dim, _linkdims[n]) + end + end + + # Use the general constructor + return TensorTrain(cores; sites=sites) end """ - TensorTrain(::Type{T}, sites::Vector{Vector{<:Index}}, linkdims::Union{Integer,Vector{<:Integer}}) where {T<:Number} + TensorTrain(::Type{T}, sites::AbstractVector{<:AbstractVector{<:Index}}, linkdims::Union{Integer,Vector{<:Integer}}) where {T<:Number} Construct a TensorTrain (MPO) from a vector of site index pairs with specified element type and link dimensions. -This creates an MPO by manually constructing ITensor objects for each site. Each site is represented by a pair of indices +This creates an MPO by constructing core arrays of type `T` for each site. Each site is represented by a vector of indices (typically the upper and lower indices of the MPO). # Arguments - `T::Type{<:Number}`: Element type (e.g., Float64, ComplexF64) -- `sites::Vector{Vector{<:Index}}`: Vector of site index pairs, where each element is a vector of indices for that site +- `sites::AbstractVector{<:AbstractVector{<:Index}}`: Vector of site index vectors, where each element is a vector of indices for that site - `linkdims::Union{Integer,Vector{<:Integer}}`: Link dimension(s) between sites # Returns @@ -141,57 +172,48 @@ function TensorTrain(::Type{T}, sites::AbstractVector{<:AbstractVector{<:Index}} length(_linkdims) == N - 1 || error("Length mismatch: linkdims ($(length(_linkdims))) must have length $(N - 1)") - # Create internal link indices (no boundary links) - links = N > 1 ? [Index(_linkdims[n], "Link,l=$n") for n in 1:(N - 1)] : Index[] - - # Create ITensors for each site - tensors = Vector{ITensor}(undef, N) + # Create cores with appropriate dimensions + # MPO structure: (left_link, upper_physical, lower_physical, right_link) + cores = Vector{Array{T}}(undef, N) for n in 1:N site_inds = sites[n] if length(site_inds) != 2 error("Each site must have exactly 2 indices (upper and lower), but site $n has $(length(site_inds))") end - # For MPO: sites[n] = [lower_index, upper_index] - # Lower index is unprimed, upper index is primed - lower_ind = site_inds[1] # unprimed - upper_ind = ITensors.prime(site_inds[2]) # primed + upper_dim = ITensors.dim(site_inds[2]) + lower_dim = ITensors.dim(site_inds[1]) - # MPO structure: - # - First site: (upper_site, lower_site, right_link) - # - Last site: (left_link, upper_site, lower_site) - # - Middle: (left_link, upper_site, lower_site, right_link) if n == 1 && n == N - inds_tuple = (upper_ind, lower_ind) + # Single site: (upper_physical, lower_physical) + cores[n] = zeros(T, upper_dim, lower_dim) elseif n == 1 - inds_tuple = (upper_ind, lower_ind, links[n]) + # First site: (1, upper_physical, lower_physical, right_link) + cores[n] = zeros(T, 1, upper_dim, lower_dim, _linkdims[n]) elseif n == N - inds_tuple = (links[n - 1], upper_ind, lower_ind) + # Last site: (left_link, upper_physical, lower_physical, 1) + cores[n] = zeros(T, _linkdims[n - 1], upper_dim, lower_dim, 1) else - inds_tuple = (links[n - 1], upper_ind, lower_ind, links[n]) + # Middle site: (left_link, upper_physical, lower_physical, right_link) + cores[n] = zeros(T, _linkdims[n - 1], upper_dim, lower_dim, _linkdims[n]) end - - # Create zero ITensor with appropriate dimensions - dims = map(ITensors.dim, inds_tuple) - data = zeros(T, dims...) - tensors[n] = ITensor(data, inds_tuple...) end - return TensorTrain(tensors) + # Use the general constructor + return TensorTrain(cores; sites=sites) end """ - TensorTrain(A::ITensor, sites; kwargs...) + TensorTrain(A::ITensor, sites::AbstractVector{<:Index}; kwargs...) Construct a TensorTrain from an ITensor by decomposing it according to site indices. -This function creates an MPS or MPO by decomposing the ITensor `A` site by site -according to the site indices `sites`. The `sites` can be either `Vector{Index}` -(for MPS) or `Vector{Vector{Index}}` (for MPO). +This function creates a TensorTrain by decomposing the ITensor `A` site by site according to the site indices `sites`. +If `sites` is `Vector{Index}`, it is automatically converted to `Vector{Vector{Index}}` format. # Arguments - `A::ITensor`: The ITensor to decompose -- `sites`: Site indices - either `Vector{Index}` (for MPS) or `Vector{Vector{Index}}` (for MPO) +- `sites::AbstractVector{<:Index}`: Vector of site indices (will be converted to `Vector{Vector{Index}}` format) # Keyword Arguments - `leftinds=nothing`: Optional left dangling indices @@ -202,7 +224,7 @@ according to the site indices `sites`. The `sites` can be either `Vector{Index}` - `kwargs...`: Additional keyword arguments passed to the decomposition # Returns -- `TensorTrain`: A new TensorTrain object (MPS or MPO depending on `sites`) +- `TensorTrain`: A new TensorTrain object # Examples ```julia @@ -211,47 +233,85 @@ A = randomITensor(sites...) tt = TensorTrain(A, sites) ``` """ -function TensorTrain(A::ITensor, sites; kwargs...) - # Detect if sites is Vector{Index} (MPS) or Vector{Vector{Index}} (MPO) - if length(sites) > 0 && sites[1] isa Index - # MPS case: sites is Vector{Index} - mps = ITensorMPS.MPS(A, sites; kwargs...) - return TensorTrain(mps) - else - # MPO case: sites is Vector{Vector{Index}} or similar - mpo = ITensorMPS.MPO(A, sites; kwargs...) - return TensorTrain(mpo) - end +function TensorTrain(A::ITensor, sites::AbstractVector{<:Index}; kwargs...) + # Convert Vector{Index} to Vector{Vector{Index}} and delegate + sites_normalized = [[s] for s in sites] + return TensorTrain(A, sites_normalized; kwargs...) +end + +""" + TensorTrain(A::ITensor, sites::AbstractVector{<:AbstractVector{<:Index}}; kwargs...) + +Construct a TensorTrain from an ITensor by decomposing it according to site indices. + +This function creates a TensorTrain by decomposing the ITensor `A` site by site according to the site indices `sites`. + +# Arguments +- `A::ITensor`: The ITensor to decompose +- `sites::AbstractVector{<:AbstractVector{<:Index}}`: Vector of site index vectors, where each element is a vector of indices for that site + +# Keyword Arguments +- `leftinds=nothing`: Optional left dangling indices +- `orthocenter::Integer=length(sites)`: Desired orthogonality center +- `tags`: Tags for link indices +- `cutoff`: Truncation error at each link +- `maxdim`: Maximum link dimension +- `kwargs...`: Additional keyword arguments passed to the decomposition + +# Returns +- `TensorTrain`: A new TensorTrain object +""" +function TensorTrain(A::ITensor, sites::AbstractVector{<:AbstractVector{<:Index}}; kwargs...) + # Always use MPO decomposition + mpo = ITensorMPS.MPO(A, sites; kwargs...) + return TensorTrain(mpo) end """ - TensorTrain(A::AbstractArray, sites; kwargs...) + TensorTrain(A::AbstractArray, sites::AbstractVector{<:Index}; kwargs...) Construct a TensorTrain from an AbstractArray by converting it to an ITensor first. +This function creates a TensorTrain by converting the array `A` to an ITensor and then decomposing it. +If `sites` is `Vector{Index}`, it is automatically converted to `Vector{Vector{Index}}` format. + # Arguments - `A::AbstractArray`: The array to convert -- `sites`: Site indices - either `Vector{Index}` (for MPS) or `Vector{Vector{Index}}` (for MPO) +- `sites::AbstractVector{<:Index}`: Vector of site indices (will be converted to `Vector{Vector{Index}}` format) - `kwargs...`: Keyword arguments passed to the decomposition # Returns - `TensorTrain`: A new TensorTrain object """ -function TensorTrain(A::AbstractArray, sites; kwargs...) - # Convert array to ITensor - if length(sites) > 0 && sites[1] isa Index - # MPS case: sites is Vector{Index} - A_itensor = ITensor(A, sites...) - mps = ITensorMPS.MPS(A_itensor, sites; kwargs...) - return TensorTrain(mps) - else - # MPO case: sites is Vector{Vector{Index}} - # Flatten sites for ITensor construction - sites_flat = collect(Iterators.flatten(sites)) - A_itensor = ITensor(A, sites_flat...) - mpo = ITensorMPS.MPO(A_itensor, sites; kwargs...) - return TensorTrain(mpo) - end +function TensorTrain(A::AbstractArray, sites::AbstractVector{<:Index}; kwargs...) + # Convert Vector{Index} to Vector{Vector{Index}} and delegate + sites_normalized = [[s] for s in sites] + return TensorTrain(A, sites_normalized; kwargs...) +end + +""" + TensorTrain(A::AbstractArray, sites::AbstractVector{<:AbstractVector{<:Index}}; kwargs...) + +Construct a TensorTrain from an AbstractArray by converting it to an ITensor first. + +This function creates a TensorTrain by converting the array `A` to an ITensor and then decomposing it. + +# Arguments +- `A::AbstractArray`: The array to convert +- `sites::AbstractVector{<:AbstractVector{<:Index}}`: Vector of site index vectors, where each element is a vector of indices for that site +- `kwargs...`: Keyword arguments passed to the decomposition + +# Returns +- `TensorTrain`: A new TensorTrain object +""" +function TensorTrain(A::AbstractArray, sites::AbstractVector{<:AbstractVector{<:Index}}; kwargs...) + # Flatten sites for ITensor construction + sites_flat = collect(Iterators.flatten(sites)) + A_itensor = ITensor(A, sites_flat...) + + # Always use MPO decomposition + mpo = ITensorMPS.MPO(A_itensor, sites; kwargs...) + return TensorTrain(mpo) end """