diff --git a/Project.toml b/Project.toml index 5edabec65..36bd1a819 100644 --- a/Project.toml +++ b/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" diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index 74a238fc2..3500fa07e 100644 --- a/docs/src/Changelog.md +++ b/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 diff --git a/docs/src/lib/tensors.md b/docs/src/lib/tensors.md index 217c6640a..e92c360e8 100644 --- a/docs/src/lib/tensors.md +++ b/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: diff --git a/docs/src/man/tensors.md b/docs/src/man/tensors.md index 7aa61f6fe..c68b7ffba 100644 --- a/docs/src/man/tensors.md +++ b/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. diff --git a/src/TensorKit.jl b/src/TensorKit.jl index 87a3a2380..bbbdd4c32 100644 --- a/src/TensorKit.jl +++ b/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 diff --git a/src/tensors/io.jl b/src/tensors/io.jl new file mode 100644 index 000000000..bc7dce190 --- /dev/null +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index ca00312e2..eb583dae4 100644 --- a/test/Project.toml +++ b/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" - diff --git a/test/tensors/io.jl b/test/tensors/io.jl new file mode 100644 index 000000000..b07d569f4 --- /dev/null +++ b/test/tensors/io.jl @@ -0,0 +1,155 @@ +using Test +import JLD2 +using TensorKit + +struct UnregisteredTensorMap <: AbstractTensorMap{Float64, ComplexSpace, 1, 0} end + +struct TestDenseVector{T} <: DenseVector{T} + data::Vector{T} +end +Base.size(vector::TestDenseVector) = size(vector.data) +Base.IndexStyle(::Type{<:TestDenseVector}) = IndexLinear() +Base.getindex(vector::TestDenseVector, index::Int) = vector.data[index] +Base.setindex!(vector::TestDenseVector, value, index::Int) = (vector.data[index] = value) + +"""Write a raw TensorKit tensor record for malformed-file tests.""" +function write_record(path, record; format = TensorKit.TENSORMAP_FILE_FORMAT, version = TensorKit.TENSORMAP_FILE_VERSION) + return JLD2.jldopen(path, "w") do file + file["format"] = format + file["version"] = version + file["tensor"] = record + end +end + +@testset "TensorMap save and load" begin + spacelists = ( + TestSetup.Vtr, + TestSetup.VRepℤ₂, + TestSetup.VRepSU₂, + TestSetup.VRepA4, + TestSetup.VIBM, + ) + mktempdir() do directory + for (index, spaces) in enumerate(spacelists) + V1, V2, V3, V4, V5 = spaces + tensor = randn(ComplexF64, V1 ⊗ V2 ← (V3 ⊗ V4 ⊗ V5)') + path = joinpath(directory, "tensor-$index.jld2") + @test save(path, tensor) === nothing + restored = load(path) + @test restored isa TensorMap + @test storagetype(restored) === Vector{ComplexF64} + @test space(restored) == space(tensor) + @test restored == tensor + end + + tensor = randn(Float64, ℂ^2 ⊗ ℂ^3) + restored = load((path = joinpath(directory, "plain-tensor.jld2"); save(path, tensor); path)) + @test restored isa Tensor + @test restored == tensor + + empty_tensor = randn(Float64, zero(ℂ^2) ← ℂ^2) + restored_empty = load((path = joinpath(directory, "empty.jld2"); save(path, empty_tensor); path)) + @test restored_empty == empty_tensor + @test isempty(restored_empty.data) + + source = randn(Float64, ℂ^3 ← ℂ^2) + custom_data = TestDenseVector(copy(source.data)) + custom = TensorMap{Float64, ComplexSpace, 1, 1, typeof(custom_data)}(custom_data, space(source)) + @test storagetype(custom) === TestDenseVector{Float64} + restored_custom = load((path = joinpath(directory, "custom.jld2"); save(path, custom); path)) + @test storagetype(restored_custom) === Vector{Float64} + @test restored_custom == custom + + diagonal_space = Vect[SU2Irrep](0 => 3, 1 // 2 => 2, 1 => 1)' + diagonal = DiagonalTensorMap(randn(ComplexF64, reduceddim(diagonal_space)), diagonal_space) + restored_diagonal = load((path = joinpath(directory, "diagonal.jld2"); save(path, diagonal); path)) + @test restored_diagonal isa DiagonalTensorMap + @test storagetype(restored_diagonal) === Vector{ComplexF64} + @test restored_diagonal == diagonal + + braid_space = Vect[FibonacciAnyon](:I => 3, :τ => 2) + for braiding in (BraidingTensor(braid_space, braid_space'), BraidingTensor(braid_space, braid_space')') + path = joinpath(directory, "braiding-$(braiding.adjoint).jld2") + save(path, braiding) + restored_braiding = load(path) + @test restored_braiding isa BraidingTensor + @test storagetype(restored_braiding) === Vector{eltype(braiding)} + @test restored_braiding.V1 == braiding.V1 + @test restored_braiding.V2 == braiding.V2 + @test restored_braiding.adjoint == braiding.adjoint + @test TensorMap(restored_braiding) == TensorMap(braiding) + end + + @test_throws ArgumentError save(joinpath(directory, "adjoint.jld2"), source') + @test_throws ArgumentError save(joinpath(directory, "unsupported.jld2"), UnregisteredTensorMap()) + end +end + +@testset "TensorMap file validation" begin + tensor = randn(Float64, Vect[Z2Irrep](0 => 2, 1 => 3) ← Vect[Z2Irrep](0 => 3, 1 => 2)) + record = TensorKit._pack_tensormap(tensor) + mktempdir() do directory + path = joinpath(directory, "invalid.jld2") + + write_record(path, record; format = "not TensorKit") + @test_throws ArgumentError load(path) + + write_record(path, record; version = TensorKit.TENSORMAP_FILE_VERSION + 1) + @test_throws ArgumentError load(path) + + JLD2.jldsave(path; unrelated = tensor.data) + @test_throws ArgumentError load(path) + + duplicate = TensorKit.DenseTensorMapRecordV1( + record.space, + [record.sectors[1], record.sectors[1]], + [record.blockshapes[1], record.blockshapes[1]], + vcat(record.data[1:prod(record.blockshapes[1])], record.data[1:prod(record.blockshapes[1])]), + ) + write_record(path, duplicate) + @test_throws ArgumentError load(path) + + badshape = copy(record.blockshapes) + badshape[1] = (badshape[1][1] + 1, badshape[1][2]) + write_record(path, TensorKit.DenseTensorMapRecordV1(record.space, record.sectors, badshape, record.data)) + @test_throws DimensionMismatch load(path) + + write_record( + path, + TensorKit.DenseTensorMapRecordV1(record.space, record.sectors, record.blockshapes, record.data[1:(end - 1)]), + ) + @test_throws DimensionMismatch load(path) + + write_record(path, 1) + @test_throws ArgumentError load(path) + end +end + +@testset "TensorMap file compactness" begin + mktempdir() do directory + V = Vect[U1Irrep](i => 3 for i in -3:3) + tensor = randn(ComplexF64, V ⊗ V ← V ⊗ V) + compact_path = joinpath(directory, "tensor.jld2") + dict_path = joinpath(directory, "tensor-dict.jld2") + save(compact_path, tensor) + JLD2.jldsave(dict_path; tensor = convert(Dict, tensor)) + compact_size = filesize(compact_path) + dict_size = filesize(dict_path) + @test compact_size < dict_size + + Vd = Vect[Z2Irrep](0 => 20, 1 => 20) + diagonal = DiagonalTensorMap(randn(Float64, reduceddim(Vd)), Vd) + diagonal_path = joinpath(directory, "diagonal.jld2") + dense_diagonal_path = joinpath(directory, "dense-diagonal.jld2") + save(diagonal_path, diagonal) + save(dense_diagonal_path, TensorMap(diagonal)) + @test filesize(diagonal_path) < filesize(dense_diagonal_path) + + braiding = BraidingTensor(Vd, Vd) + braiding_path = joinpath(directory, "braiding.jld2") + dense_braiding_path = joinpath(directory, "dense-braiding.jld2") + save(braiding_path, braiding) + save(dense_braiding_path, TensorMap(braiding)) + @test filesize(braiding_path) < filesize(dense_braiding_path) + end +end