Providing a consistent IO method for all kinds of tensors

[Yue-Zhengyuan]

It would be great if TensorKit provides out-of-box functions to save all kinds of tensors to disk and load it back (just like PyTorch's save and load). The issue I'm currently running into is the IO of a DiagonalTensorMap. For ordinary TensorMaps, I can use the following functions for IO (which, at least for ordinary TensorMaps, actually works for both 0.12.7 and newer versions of TensorKit):

using TensorKit
using JLD2

function save_tensor(filename::AbstractString, t::AbstractTensorMap)
    # the resulting Dict is "dense", the same as an ordinary TensorMap
    return JLD2.save_object(filename, convert(Dict, t))
end

function load_tensor(filename::AbstractString)
    # conversion back to DiagonalTensorMap is not implemented
    return convert(TensorMap, JLD2.load_object(filename))
end

In addition, there is not a convenient conversion from a diagonal TensorMap back to DiagonalTensorMap. Then I cannot easily load a saved tensor back as a DiagonalTensorMap. I guess there is a similar problem for some other special kind of tensors.

[Yue-Zhengyuan]

Currently the converted Dict from a tensor is like follows:

t = randn(ℂ^2 ← ℂ^2)
obj = convert(Dict, t)

Output:

Dict{Symbol, Any} with 3 entries:
  :codomain => "ProductSpace(ℂ^2)"
  :domain   => "ProductSpace(ℂ^2)"
  :data     => Dict{String, Any}("Trivial()"=>[-0.995765 -0.126797; -1.02334 -1…

I suggest to add one more entry :kind to store the tensor kind (TensorMap, DiagonalTensorMap, etc.), and the data is the original data without being converted first to the data of an ordinary tensor. For compatibility, when :kind is missing from old files, just default it to TensorMap. It remains to implement TensorKit.convert from the Dict back to each tensor kind.

[Jutho]

Converting to Dict and then saving that using any of the standard IO packages has been the recommended strategy, though it is probably not super prominent in the manual. The Dict conversion makes sure to only use basic Julia types for the fields, namely Symbol, Strings and Vector{<:Number}, and we commit to at least being able to parse the codomain and domain information in String format back into the proper space structures. Unfortunately, the internal order of the actual numerical data had to change between 0.12 and 0.13, which is why we added a README entry on how to convert tensors between those two versions.

It is definitely also true that the Dict conversion was overlooked with the DiagonalTensorMap implementation. However, the default implementation does work, in the sense that if you convert the DiagonalTensorMap to a Dict, save it and then reload it and convert it back into TensorMap, the resulting tensor is equivalent to the original one, except for the fact that now it is represented as a TensorMap instead of a DiagonalTensorMap.

The question is what kind of behaviour you want. Also an AdjointTensorMap(...) will, upon conversion and back conversion, have changed type and be a regular TensorMap.

Is it really important that a DiagonalTensorMap is stored in the most efficient format with none of the zeros being stored? Disk space is typically quite cheap. For example, would it suffice to have a convert(DiagonalTensorMap, dict) method that just takes the current output of dict = convert(Dict, some_diagonal_tensor_map), that includes all the zeros away from the diagonal, and restores the DiagonalTensorMap, and errors if there are nonzero off-diagonal entries or the spaces do not admit a DiagonalTensorMap? That could easily be implemented without having to change anything to the resulting output of convert(Dict, ::AbstractTensorMap).

I am definitely open to suggestions resulting from actual use cases. I find coming up with the correct interface always the most difficult part of such questions.

[Yue-Zhengyuan]

In the imaginary time evolution algorithms I'm currently working on, there are many operations of absorbing/removing the bond weights into/from the PEPS tensors. My past experience with Python is that if the multiplication with weights is done as usual (with all the zeros in the dense format), the speed can be much slower compared to the code optimized for diagonal matrices. I haven't tested the Julia version yet about whether there is also a such significant performance improvement EDIT: With new TensorKit the simple update algorithm in PEPSKit runs much faster (about 3x speed up for Heisenberg model with spin U(1) symmetry at D = 4). This is the reason I want preserve the DiagonalTensorMaps format, without being affected by the IO process.

In addition, the conversion code on README may need some update.

using JLD2
filename = "choose_some_filename.jld2"
t_dict = jldload(filename)
T = eltype(valtype(t_dict[:data]))
t = TensorMap{T}(undef, t_dict[:space])
for ((f₁, f₂), val) in t_dict[:data]
    t[f₁, f₂] .= val
end

• In the latest version of JLD2 (v0.5.11), there's no jldload. I'm currently replacing it with load_object.
• When trying to convert a real PEPSTensor (ℂ^2 ← (ℂ^4 ⊗ ℂ^4 ⊗ (ℂ^4)' ⊗ (ℂ^4)')) without any symmetry (so there is only one entry in t_dict[:data] with key Tuple{Nothing, Nothing}), I encounter the following error immediately after entering the for loop:

ERROR: LoadError: MethodError: no method matching getindex(::TensorMap{Float64, ComplexSpace, 1, 4, Vector{Float64}}, ::Nothing, ::Nothing)
The function `getindex` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  getindex(::AbstractTensorMap)
   @ TensorKit ~/.julia/packages/TensorKit/IQwop/src/tensors/abstracttensor.jl:356
  getindex(::AbstractTensorMap, ::Union{Colon, AbstractRange{<:Union{Int128, Int16, Int32, Int64, Int8, UInt128, UInt16, UInt32, UInt64, UInt8}}, Int128, Int16, Int32, Int64, Int8, UInt128, UInt16, UInt32, UInt64, UInt8}...)
   @ TensorKit ~/.julia/packages/TensorKit/IQwop/src/tensors/abstracttensor.jl:328
  getindex(::TensorMap{T, S, N₁, N₂, A} where A<:DenseVector{T}, ::FusionTree{Trivial, N₁}, ::FusionTree{Trivial, N₂}) where {T, S, N₁, N₂}
   @ TensorKit ~/.julia/packages/TensorKit/IQwop/src/tensors/tensor.jl:478
  ...

[Jutho]

Thanks; these are valuable bug reports. Do you feel comfortable in starting to prepare a PR to fix those? I am tied up in correcting exams until at least Monday afternoon, so it will probably not be before Tuesday that I have some time to fix this.

[Yue-Zhengyuan]

I did some further testing, and found that save_object(filename, convert(Dict, t)) and convert(TensorMap, load_object(filename)) is actually not affected by the data structure change. For example, I tested tensors with a rather complex $Z_2^f \times U(1) \times SU(2)$ symmetry:

# define the following in both old and new environments
using TensorKit
using JLD2
save_tensor(filename::AbstractString, t::AbstractTensorMap) = save_object(filename, convert(Dict, t))
load_tensor(filename::AbstractString) = convert(TensorMap, load_object(filename))
name_old = "old.jld2"
name_new = "new.jld2"

# in old environment containing TensorKit v0.12.7
stype = FermionParity ⊠ U1Irrep ⊠ SU2Irrep
V1 = ℂ[stype]((0,2,1) => 2, (1,-1,1//2) => 3, (0,-2,0) => 2)
V2 = ℂ[stype]((0,0,1) => 3, (1,3,1//2) => 2, (0,-2,1) => 4)
a = TensorMap(randn, Float64, V1 ⊗ V1', V2' ⊗ V2);
save_tensor(name_old, a)

# in new environment containing TensorKit v0.14.3
save_tensor(name_new, load_tensor(name_old))

# back to old environment
println(a == load_tensor(name_new)) # output `true`

This seems to demonstrate that the (unordered) Dict produced by new TensorKit is the same to the old one. Look at the source code (which is unchanged from v0.12.7 to v0.14.3):

function Base.convert(::Type{Dict}, t::AbstractTensorMap)
    d = Dict{Symbol,Any}()
    d[:codomain] = repr(codomain(t))
    d[:domain] = repr(domain(t))
    data = Dict{String,Any}()
    # the export explicitly saves which block is which
    for (c, b) in blocks(t)
        data[repr(c)] = Array(b)
    end
    d[:data] = data
    return d
end

Although the iteration order of blocks changed, the export explicitly saves which block is which. In addition, the label c of each block is unchanged in newer versions of TensorKit. Thus the exported Dict can be loaded back by both new and old TensorKit (assuming the structure of each block is also unchanged), unless I've missed some edge cases?

On the other hand, the structure of fusion trees has changed from v0.12.7 to newer versions, and the old fusion trees exported by v0.12.7 cannot be rebuild in newer versions. Consider the simplest example of "trivial" bosonic tensors without any symmetry:

# v0.14.3
using TensorKit
a = rand(ℂ^2 ← ℂ^3)
fusiontrees(a)
#= output:
1-element Vector{Tuple{FusionTree{Trivial, 1, 0, 0}, FusionTree{Trivial, 1, 0, 0}}}:
 (FusionTree{Trivial}((Trivial(),), Trivial(), (false,), ()), FusionTree{Trivial}((Trivial(),), Trivial(), (false,), ()))
=# 

# v0.12.7
a = TensorMap(randn, Float64, ℂ^2, ℂ^3)
collect(fusiontrees(a))
#= output:
1-element Vector{Tuple{Nothing, Nothing}}:
 (nothing, nothing)
=#

The new TensorKit cannot build the fusion tree from nothing, causing the error I encountered above. If that's indeed the case I can open a simple PR to fix the upgrade guide in README, without modifying the source code.

[Yue-Zhengyuan]

For the IO of DiagonalTensorMap, I can add a simple function that converts a dense diagonal TensorMap, while all the IO still uses the generic TensorMap format.

[lkdvos]

Just putting this out there: that test does not actually check if the tensor is still compatible, in the sense that there might be basis transformations within the block that change the interpretation of the data, without actually altering the entries. Basically, tensors are basis dependent, so you cannot tell by looking at the entries if they are the same.

A better way to test this would be for example to construct two random tensors in an arbitrary partition, along with their overlap. (Best with some form permutation as well). Since the overlap is a scalar, it is basis independent, and should thus remain the same after loading it in in the new format.

All this being said, I'm actually struggling to come up with an explicit example where the internal ordering of the block has changed. Indeed, the fusiontrees are now iterated in a different order, but it seems like the order within the blocks has remained unchanged. Is that possible @Jutho ?

[Jutho]

The order in the block should be the iteration order, so it should have changed. Let me try to come up with an example.

[Jutho]

On TensorKit 0.12.7

julia> V = Vect[FibonacciAnyon](:I=>1,:τ=>1)
julia> t = TensorMap(randn, V * V * V * V, V * V * V)
julia> sort(collect(pairs(t.rowr[FibonacciAnyon(:I)])); by = first ∘ last)
13-element Vector{Pair{FusionTree{FibonacciAnyon, 4, 2, 3, Nothing}, UnitRange{Int64}}}:
 FusionTree{FibonacciAnyon}((:I, :I, :I, :I), :I, (false, false, false, false), (:I, :I)) => 1:1
 FusionTree{FibonacciAnyon}((:τ, :τ, :I, :I), :I, (false, false, false, false), (:I, :I)) => 2:2
 FusionTree{FibonacciAnyon}((:τ, :I, :τ, :I), :I, (false, false, false, false), (:τ, :I)) => 3:3
 FusionTree{FibonacciAnyon}((:I, :τ, :τ, :I), :I, (false, false, false, false), (:τ, :I)) => 4:4
 FusionTree{FibonacciAnyon}((:τ, :τ, :τ, :I), :I, (false, false, false, false), (:τ, :I)) => 5:5
 FusionTree{FibonacciAnyon}((:τ, :I, :I, :τ), :I, (false, false, false, false), (:τ, :τ)) => 6:6
 FusionTree{FibonacciAnyon}((:I, :τ, :I, :τ), :I, (false, false, false, false), (:τ, :τ)) => 7:7
 FusionTree{FibonacciAnyon}((:τ, :τ, :I, :τ), :I, (false, false, false, false), (:τ, :τ)) => 8:8
 FusionTree{FibonacciAnyon}((:I, :I, :τ, :τ), :I, (false, false, false, false), (:I, :τ)) => 9:9
 FusionTree{FibonacciAnyon}((:τ, :I, :τ, :τ), :I, (false, false, false, false), (:τ, :τ)) => 10:10
 FusionTree{FibonacciAnyon}((:I, :τ, :τ, :τ), :I, (false, false, false, false), (:τ, :τ)) => 11:11
 FusionTree{FibonacciAnyon}((:τ, :τ, :τ, :τ), :I, (false, false, false, false), (:I, :τ)) => 12:12
 FusionTree{FibonacciAnyon}((:τ, :τ, :τ, :τ), :I, (false, false, false, false), (:τ, :τ)) => 13:13

On TensorKit v0.14

V = Vect[FibonacciAnyon](:I=>1,:τ=>1)
julia> collect(enumerate(getindex.(TensorKit.fusionblockstructure(V * V * V *V ← V * V * V).fusiontreelist[1:13], 1)))
13-element Vector{Tuple{Int64, FusionTree{FibonacciAnyon, 4, 2, 3}}}:
 (1, FusionTree{FibonacciAnyon}((:I, :I, :I, :I), :I, (false, false, false, false), (:I, :I)))
 (2, FusionTree{FibonacciAnyon}((:τ, :τ, :I, :I), :I, (false, false, false, false), (:I, :I)))
 (3, FusionTree{FibonacciAnyon}((:τ, :I, :τ, :I), :I, (false, false, false, false), (:τ, :I)))
 (4, FusionTree{FibonacciAnyon}((:I, :τ, :τ, :I), :I, (false, false, false, false), (:τ, :I)))
 (5, FusionTree{FibonacciAnyon}((:τ, :τ, :τ, :I), :I, (false, false, false, false), (:τ, :I)))
 (6, FusionTree{FibonacciAnyon}((:τ, :I, :I, :τ), :I, (false, false, false, false), (:τ, :τ)))
 (7, FusionTree{FibonacciAnyon}((:I, :τ, :I, :τ), :I, (false, false, false, false), (:τ, :τ)))
 (8, FusionTree{FibonacciAnyon}((:τ, :τ, :I, :τ), :I, (false, false, false, false), (:τ, :τ)))
 (9, FusionTree{FibonacciAnyon}((:I, :I, :τ, :τ), :I, (false, false, false, false), (:I, :τ)))
 (10, FusionTree{FibonacciAnyon}((:τ, :τ, :τ, :τ), :I, (false, false, false, false), (:I, :τ)))
 (11, FusionTree{FibonacciAnyon}((:τ, :I, :τ, :τ), :I, (false, false, false, false), (:τ, :τ)))
 (12, FusionTree{FibonacciAnyon}((:I, :τ, :τ, :τ), :I, (false, false, false, false), (:τ, :τ)))
 (13, FusionTree{FibonacciAnyon}((:τ, :τ, :τ, :τ), :I, (false, false, false, false), (:τ, :τ)))

Rows 10 and 12 have changed.

[Yue-Zhengyuan]

I found the following to work for many sector types:

• Save tensors in old TensorKit:

function save_tensor_old(filename::AbstractString, t::AbstractTensorMap)
    t_dict = Dict(
        :codomain => repr(codomain(t)),
        :domain => repr(domain(t)),
        :data => Dict((repr(f₁), repr(f₂)) => Array(t[f₁, f₂]) for (f₁, f₂) in fusiontrees(t))
    )
    save_object(filename, t_dict)
    return nothing
end

• Load tensors in new TensorKit:

function load_tensor_old(filename::AbstractString)
    t_dict = load_object(filename)
    T = eltype(valtype(t_dict[:data]))
    codomain = eval(Meta.parse(t_dict[:codomain]))
    domain = eval(Meta.parse(t_dict[:domain]))
    t = TensorMap{T}(undef, codomain ← domain)
    for ((f₁, f₂), val) in t_dict[:data]
        t[eval(Meta.parse(f₁)), eval(Meta.parse(f₂))] .= val
    end
    return t
end

We then only need to handle the trivial tensors separately.