Add versioned TensorMap save and load

Project.toml

@@ -9,6 +9,7 @@ projects = ["test", "docs"]
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
@@ -52,6 +53,7 @@ Enzyme = "0.13.195"
 EnzymeTestUtils = "0.2.8"
 FiniteDifferences = "0.12"
 GPUArrays = "11.4.1"
+JLD2 = "0.6"
 LRUCache = "1.6"
 LinearAlgebra = "1"
 MatrixAlgebraKit = "0.6.9"

docs/src/Changelog.md

@@ -22,6 +22,8 @@ When releasing a new version, move the "Unreleased" changes to a new version sec
 
 ### Added
 
+- Versioned `save` and `load` support for `TensorMap`, `DiagonalTensorMap`, and `BraidingTensor` objects.
+
 ### Changed
 
 ### Deprecated

docs/src/lib/tensors.md

@@ -20,6 +20,12 @@ AdjointTensorMap
 BraidingTensor
 ```
 
+Tensor maps can be stored and restored with:
+```@docs
+save
+load
+```
+
 Of those, `TensorMap` provides the generic instantiation of our tensor concept. It supports various constructors, which are discussed in the next subsection.
 
 Furthermore, some aliases are provided for convenience:

docs/src/man/tensors.md

@@ -427,12 +427,21 @@ f1, f2 = first(fusiontrees(t))
 t[f1,f2]
 ```
 
-## [Reading and writing tensors: `Dict` conversion](@id ss_tensor_readwrite)
+## [Reading and writing tensors](@id ss_tensor_readwrite)
 
-There are no custom or dedicated methods for reading, writing or storing `TensorMap`s, however, there is the possibility to convert a `t::AbstractTensorMap` into a `Dict`, simply as `convert(Dict, t)`.
-The backward conversion `convert(TensorMap, dict)` will return a tensor that is equal to `t`, i.e. `t == convert(TensorMap, convert(Dict, t))`.
+TensorKit provides [`save`](@ref) and [`load`](@ref) for storing one tensor map in a versioned JLD2 file.
 
-This conversion relies on that the string representation of objects such as `VectorSpace`, `FusionTree` or `Sector` should be such that it represents valid code to recreate the object.
-Hence, we store information about the domain and codomain of the tensor, and the sector associated with each data block, as a `String` obtained with `repr`.
-This provides the flexibility to still change the internal structure of such objects, without this breaking the ability to load older data files.
-The resulting dictionary can then be stored using any of the provided Julia packages such as [JLD.jl](https://github.com/JuliaIO/JLD.jl), [JLD2.jl](https://github.com/JuliaIO/JLD2.jl), [BSON.jl](https://github.com/JuliaIO/BSON.jl), [JSON.jl](https://github.com/JuliaIO/JSON.jl), ...
+```julia
+filename = "tensor.jld2"
+save(filename, t)
+t′ = load(filename)
+```
+
+`TensorMap`, `DiagonalTensorMap`, and `BraidingTensor` retain their semantic types, while numerical storage is copied to a CPU `Vector` when saving and loading.
+The compact data of `DiagonalTensorMap` and the structural description of `BraidingTensor` are stored without materializing dense blocks.
+A lazy `AdjointTensorMap` must be materialized explicitly before saving, for example with `save(filename, convert(TensorMap, t'))`.
+TensorKit does not add a filename extension and replaces an existing file at the requested path.
+When another loaded package exports functions with the same names, use `TensorKit.save` and `TensorKit.load` explicitly.
+
+The older `convert(Dict, t)` and `convert(TensorMap, dict)` workflow remains available for compatibility.
+That representation stores spaces and block sectors as strings and does not preserve specialized tensor-map types, so it is no longer recommended for new files.

src/TensorKit.jl

@@ -29,6 +29,7 @@ export FusionTree
 export IndexSpace, HomSpace, TensorSpace, TensorMapSpace
 export AbstractTensorMap, AbstractTensor, TensorMap, Tensor # tensors and tensor properties
 export DiagonalTensorMap, BraidingTensor
+export save, load
 export SpaceMismatch, SectorMismatch, IndexError # error types
 
 # Export general vector space methods
@@ -121,6 +122,7 @@ using MatrixAlgebraKit
 
 using Dictionaries: Dictionaries, Dictionary, Indices, gettoken, gettokenvalue
 using LRUCache
+import JLD2
 using OhMyThreads
 using ScopedValues
 
@@ -265,6 +267,7 @@ include("tensors/treetransformers.jl")
 include("tensors/indexmanipulations.jl")
 include("tensors/diagonal.jl")
 include("tensors/braidingtensor.jl")
+include("tensors/io.jl")
 
 include("factorizations/factorizations.jl")
 using .Factorizations

src/tensors/io.jl

@@ -0,0 +1,198 @@
+# TensorMap IO
+#=============#
+
+const TENSORMAP_FILE_FORMAT = "TensorKit.AbstractTensorMap"
+const TENSORMAP_FILE_VERSION = UInt16(1)
+
+abstract type AbstractTensorMapRecordV1 end
+
+struct DenseTensorMapRecordV1{S, I, T} <: AbstractTensorMapRecordV1
+    space::S
+    sectors::Vector{I}
+    blockshapes::Vector{Tuple{Int, Int}}
+    data::Vector{T}
+end
+
+struct DiagonalTensorMapRecordV1{S, I, T} <: AbstractTensorMapRecordV1
+    domain::S
+    sectors::Vector{I}
+    blocklengths::Vector{Int}
+    data::Vector{T}
+end
+
+struct BraidingTensorRecordV1{S, T} <: AbstractTensorMapRecordV1
+    V1::S
+    V2::S
+    adjoint::Bool
+    scalartype::Type{T}
+end
+
+"""Pack a dense tensor map into the portable version-one representation."""
+function _pack_tensormap(t::TensorMap{T}) where {T}
+    I = sectortype(t)
+    sectors = I[]
+    blockshapes = Tuple{Int, Int}[]
+    data = Vector{T}(undef, dim(t))
+    offset = 0
+    for (c, b) in blocks(t)
+        push!(sectors, c)
+        push!(blockshapes, size(b))
+        blockdata = vec(Array(b))
+        copyto!(data, offset + 1, blockdata, 1, length(blockdata))
+        offset += length(blockdata)
+    end
+    offset == length(data) || error("inconsistent TensorMap block storage")
+    return DenseTensorMapRecordV1(space(t), sectors, blockshapes, data)
+end
+
+"""Pack a diagonal tensor map without expanding its zero off-diagonal entries."""
+function _pack_tensormap(t::DiagonalTensorMap{T}) where {T}
+    I = sectortype(t)
+    sectors = I[]
+    blocklengths = Int[]
+    data = Vector{T}(undef, length(t.data))
+    offset = 0
+    for (c, b) in blocks(t)
+        diagonal = Array(b.diag)
+        push!(sectors, c)
+        push!(blocklengths, length(diagonal))
+        copyto!(data, offset + 1, diagonal, 1, length(diagonal))
+        offset += length(diagonal)
+    end
+    offset == length(data) || error("inconsistent DiagonalTensorMap block storage")
+    return DiagonalTensorMapRecordV1(only(domain(t)), sectors, blocklengths, data)
+end
+
+"""Pack a braiding tensor using only the spaces and orientation that define it."""
+function _pack_tensormap(t::BraidingTensor{T}) where {T}
+    return BraidingTensorRecordV1(t.V1, t.V2, t.adjoint, T)
+end
+
+"""Reject lazy adjoints so that saving never hides an implicit materialization choice."""
+function _pack_tensormap(::AdjointTensorMap)
+    throw(ArgumentError("AdjointTensorMap must be materialized with `convert(TensorMap, tensor)` before saving"))
+end
+
+"""Reject tensor-map implementations without an explicit stable serialization record."""
+function _pack_tensormap(t::AbstractTensorMap)
+    throw(ArgumentError("saving $(typeof(t)) is not supported; materialize it as a built-in TensorMap type first"))
+end
+
+"""Check that serialized block labels are unique."""
+function _check_unique_sectors(sectors)
+    length(unique(sectors)) == length(sectors) ||
+        throw(ArgumentError("serialized tensor contains duplicate block sectors"))
+    return nothing
+end
+
+"""Reconstruct a dense tensor map from a validated version-one record."""
+function _unpack_tensormap(record::DenseTensorMapRecordV1{S, I, T}) where {S, I, T}
+    length(record.sectors) == length(record.blockshapes) ||
+        throw(ArgumentError("serialized TensorMap has inconsistent block metadata"))
+    _check_unique_sectors(record.sectors)
+
+    tensor = TensorMap{T}(undef, record.space)
+    expected_sectors = collect(blocksectors(tensor))
+    length(record.sectors) == length(expected_sectors) &&
+        all(c -> c in expected_sectors, record.sectors) ||
+        throw(ArgumentError("serialized TensorMap block sectors do not match its space"))
+
+    offset = 0
+    for (c, shape) in zip(record.sectors, record.blockshapes)
+        destination = block(tensor, c)
+        size(destination) == shape ||
+            throw(DimensionMismatch("serialized TensorMap block for sector $c has shape $shape, expected $(size(destination))"))
+        blocklength = prod(shape)
+        offset + blocklength <= length(record.data) ||
+            throw(DimensionMismatch("serialized TensorMap data is shorter than its block metadata"))
+        copyto!(destination, reshape(view(record.data, (offset + 1):(offset + blocklength)), shape))
+        offset += blocklength
+    end
+    offset == length(record.data) ||
+        throw(DimensionMismatch("serialized TensorMap data is longer than its block metadata"))
+    return tensor
+end
+
+"""Reconstruct a diagonal tensor map from a validated compact record."""
+function _unpack_tensormap(record::DiagonalTensorMapRecordV1{S, I, T}) where {S, I, T}
+    length(record.sectors) == length(record.blocklengths) ||
+        throw(ArgumentError("serialized DiagonalTensorMap has inconsistent block metadata"))
+    _check_unique_sectors(record.sectors)
+
+    tensor = DiagonalTensorMap{T}(undef, record.domain)
+    expected_sectors = collect(blocksectors(tensor))
+    length(record.sectors) == length(expected_sectors) &&
+        all(c -> c in expected_sectors, record.sectors) ||
+        throw(ArgumentError("serialized DiagonalTensorMap block sectors do not match its space"))
+
+    offset = 0
+    for (c, blocklength) in zip(record.sectors, record.blocklengths)
+        blocklength >= 0 || throw(ArgumentError("serialized diagonal block length is negative"))
+        destination = block(tensor, c).diag
+        length(destination) == blocklength ||
+            throw(DimensionMismatch("serialized DiagonalTensorMap block for sector $c has length $blocklength, expected $(length(destination))"))
+        offset + blocklength <= length(record.data) ||
+            throw(DimensionMismatch("serialized DiagonalTensorMap data is shorter than its block metadata"))
+        copyto!(destination, view(record.data, (offset + 1):(offset + blocklength)))
+        offset += blocklength
+    end
+    offset == length(record.data) ||
+        throw(DimensionMismatch("serialized DiagonalTensorMap data is longer than its block metadata"))
+    return tensor
+end
+
+"""Reconstruct a braiding tensor from its structural version-one record."""
+function _unpack_tensormap(record::BraidingTensorRecordV1{S, T}) where {S, T}
+    record.scalartype === T || throw(ArgumentError("serialized BraidingTensor has an inconsistent scalar type"))
+    return BraidingTensor{T}(record.V1, record.V2, record.adjoint)
+end
+
+"""Reject unknown serialization record types."""
+function _unpack_tensormap(record)
+    throw(ArgumentError("unsupported TensorKit tensor record $(typeof(record))"))
+end
+
+"""
+    save(path::AbstractString, tensor::AbstractTensorMap)
+
+Save one materialized tensor map to `path` using TensorKit's versioned JLD2 format.
+Numerical data is copied to CPU storage, and an existing file is replaced.
+"""
+function save(path::AbstractString, tensor::AbstractTensorMap)
+    record = _pack_tensormap(tensor)
+    destination = abspath(path)
+    temporary, io = mktemp(dirname(destination))
+    close(io)
+    committed = false
+    try
+        JLD2.jldopen(temporary, "w") do file
+            file["format"] = TENSORMAP_FILE_FORMAT
+            file["version"] = TENSORMAP_FILE_VERSION
+            file["tensor"] = record
+        end
+        mv(temporary, destination; force = true)
+        committed = true
+    finally
+        !committed && isfile(temporary) && rm(temporary)
+    end
+    return nothing
+end
+
+"""
+    load(path::AbstractString) -> AbstractTensorMap
+
+Load one tensor map saved with [`save`](@ref), using CPU storage for numerical data.
+"""
+function load(path::AbstractString)
+    record = JLD2.jldopen(path, "r") do file
+        all(key -> haskey(file, key), ("format", "version", "tensor")) ||
+            throw(ArgumentError("file is not a TensorKit tensor-map file"))
+        file["format"] == TENSORMAP_FILE_FORMAT ||
+            throw(ArgumentError("file has an invalid TensorKit tensor-map format marker"))
+        version = file["version"]
+        version == TENSORMAP_FILE_VERSION ||
+            throw(ArgumentError("unsupported TensorKit tensor-map file version $version"))
+        return file["tensor"]
+    end
+    return _unpack_tensormap(record)
+end

test/Project.toml

@@ -14,6 +14,7 @@ EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
@@ -43,4 +44,3 @@ ParallelTestRunner = "2"
 Test = "1"
 TestExtras = "0.2,0.3"
 Zygote = "0.7"
-