diff --git a/.gitignore b/.gitignore index 22b0fc821..018bdba07 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ __pycache__ .ipynb* *Manifest.toml +LocalPreferences.toml .vscode experimental refs diff --git a/Project.toml b/Project.toml index 8a93b4358..af155c4d2 100644 --- a/Project.toml +++ b/Project.toml @@ -33,6 +33,7 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" +JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" [extensions] @@ -43,6 +44,7 @@ TensorKitEnzymeExt = "Enzyme" TensorKitEnzymeTestUtilsExt = "EnzymeTestUtils" TensorKitFiniteDifferencesExt = "FiniteDifferences" TensorKitGPUArraysExt = "GPUArrays" +TensorKitJLD2Ext = "JLD2" TensorKitMooncakeExt = "Mooncake" [compat] @@ -55,6 +57,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 282799375..a59893f77 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, fusion-tree-based `save_tensor` and `load_tensor` support for `TensorMap`, `DiagonalTensorMap`, and `BraidingTensor` objects through the optional JLD2-based `TensorKitJLD2Ext` extension. + ### Changed - For sector types with `GenericUnit` such that colorings are not unique, `GradedSpace`, `ProductSpace` and `HomSpace` now check for this compatibility. In particular, this prevents the construction of `TensorMap`s with incompatible colorings, which previously either errored or produced empty tensors inconsistently. ([#515](https://github.com/QuantumKitHub/TensorKit.jl/pull/515)) diff --git a/docs/src/lib/tensors.md b/docs/src/lib/tensors.md index 8594534f7..d25a39c22 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_tensor +load_tensor +``` + 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..d574343ee 100644 --- a/docs/src/man/tensors.md +++ b/docs/src/man/tensors.md @@ -427,12 +427,24 @@ 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_tensor`](@ref) and [`load_tensor`](@ref) for storing one tensor map in a versioned JLD2 file. +Install JLD2 and load it with `using JLD2` to activate these functions through the `TensorKitJLD2Ext` extension. -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 +using JLD2 + +filename = "tensor.jld2" +save_tensor(filename, t) +t′ = load_tensor(filename) +``` + +`TensorMap`, `DiagonalTensorMap`, and `BraidingTensor` retain their semantic types, while numerical storage is copied to a CPU `Vector` when saving and loading. +Dense numerical segments are labeled by the semantic fields of their codomain and domain fusion trees, so loading does not depend on fusion-tree or block iteration order. +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_tensor(filename, convert(TensorMap, t'))`. +TensorKit does not add a filename extension and replaces an existing file at the requested path. + +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/ext/TensorKitJLD2Ext/TensorKitJLD2Ext.jl b/ext/TensorKitJLD2Ext/TensorKitJLD2Ext.jl new file mode 100644 index 000000000..40d96d327 --- /dev/null +++ b/ext/TensorKitJLD2Ext/TensorKitJLD2Ext.jl @@ -0,0 +1,18 @@ +module TensorKitJLD2Ext + +using TensorKit +using TensorKit: AdjointTensorMap +import TensorKit: save_tensor, load_tensor +import JLD2 + +# TensorMap IO +#=============# + +const TENSORMAP_FILE_FORMAT = "TensorKit.AbstractTensorMap" +const TENSORMAP_FILE_VERSION = UInt16(1) + +include("fusiontrees.jl") +include("records.jl") +include("io.jl") + +end diff --git a/ext/TensorKitJLD2Ext/fusiontrees.jl b/ext/TensorKitJLD2Ext/fusiontrees.jl new file mode 100644 index 000000000..7f8b6f68f --- /dev/null +++ b/ext/TensorKitJLD2Ext/fusiontrees.jl @@ -0,0 +1,118 @@ +const _FUSIONTREE_TABLE_FIELDS = (:uncoupled, :coupled, :isdual, :innerlines, :vertices) + +"""Return the position of a fusion tree, adding it to the table when necessary.""" +function _intern_fusiontree!(trees::AbstractVector, tree::FusionTree) + index = findfirst(==(tree), trees) + if isnothing(index) + push!(trees, tree) + return length(trees) + end + return index +end + +"""Encode fusion trees as columnar arrays of their semantic fields.""" +function _encode_fusiontrees(trees::AbstractVector, ::Type{I}, numlegs::Int) where {I <: Sector} + numtrees = length(trees) + numinner = max(0, numlegs - 2) + numvertices = max(0, numlegs - 1) + uncoupled = Matrix{I}(undef, numlegs, numtrees) + coupled = Vector{I}(undef, numtrees) + isdual = falses(numlegs, numtrees) + innerlines = Matrix{I}(undef, numinner, numtrees) + vertices = Matrix{Int}(undef, numvertices, numtrees) + for (column, tree) in enumerate(trees) + length(tree.uncoupled) == numlegs || + error("inconsistent fusion-tree leg count while saving") + length(tree.innerlines) == numinner || + error("inconsistent fusion-tree inner-line count while saving") + length(tree.vertices) == numvertices || + error("inconsistent fusion-tree vertex count while saving") + uncoupled[:, column] .= tree.uncoupled + coupled[column] = tree.coupled + isdual[:, column] .= tree.isdual + innerlines[:, column] .= tree.innerlines + vertices[:, column] .= tree.vertices + end + return (; uncoupled, coupled, isdual, innerlines, vertices) +end + +"""Decode a columnar fusion-tree table and validate its basic representation.""" +function _decode_fusiontrees(table, ::Type{I}, numlegs::Int, description::AbstractString) where {I <: Sector} + _require_record_fields(table, _FUSIONTREE_TABLE_FIELDS, description) + table.uncoupled isa Matrix{I} || + throw(ArgumentError("serialized $description has invalid uncoupled sectors")) + table.coupled isa Vector{I} || + throw(ArgumentError("serialized $description has invalid coupled sectors")) + table.isdual isa BitMatrix || + throw(ArgumentError("serialized $description has invalid duality flags")) + table.innerlines isa Matrix{I} || + throw(ArgumentError("serialized $description has invalid inner lines")) + table.vertices isa Matrix{Int} || + throw(ArgumentError("serialized $description has invalid vertices")) + + numtrees = length(table.coupled) + expected_sizes = ( + (numlegs, numtrees), + (numlegs, numtrees), + (max(0, numlegs - 2), numtrees), + (max(0, numlegs - 1), numtrees), + ) + actual_sizes = ( + size(table.uncoupled), size(table.isdual), + size(table.innerlines), size(table.vertices), + ) + actual_sizes == expected_sizes || + throw(DimensionMismatch("serialized $description has inconsistent table dimensions")) + + trees = Vector{FusionTree{I, numlegs}}(undef, numtrees) + for column in 1:numtrees + uncoupled = ntuple(row -> table.uncoupled[row, column], numlegs) + isdual = ntuple(row -> table.isdual[row, column], numlegs) + innerlines = ntuple(row -> table.innerlines[row, column], max(0, numlegs - 2)) + vertices = ntuple(row -> table.vertices[row, column], max(0, numlegs - 1)) + trees[column] = try + FusionTree{I}(uncoupled, table.coupled[column], isdual, innerlines, vertices) + catch error + message = sprint(showerror, error) + throw(ArgumentError("serialized $description contains an invalid fusion tree: $message")) + end + end + _check_unique_values(trees, description) + return trees +end + +"""Decode explicit fusion-tree pair identifiers and validate them against a tensor-map space.""" +function _decode_fusiontree_pairs(pair_ids, codomain_trees, domain_trees, tensor_space::TensorMapSpace) + pair_ids isa Matrix{Int} || + throw(ArgumentError("serialized tensor has invalid fusion-tree pair identifiers")) + size(pair_ids, 1) == 2 || + throw(DimensionMismatch("serialized fusion-tree pair identifiers must have two rows")) + + numpairs = size(pair_ids, 2) + pairs = Vector{Tuple{eltype(codomain_trees), eltype(domain_trees)}}(undef, numpairs) + used_codomain = falses(length(codomain_trees)) + used_domain = falses(length(domain_trees)) + for column in 1:numpairs + codomain_id = pair_ids[1, column] + domain_id = pair_ids[2, column] + checkbounds(Bool, codomain_trees, codomain_id) || + throw(ArgumentError("serialized tensor has an out-of-range codomain fusion-tree identifier")) + checkbounds(Bool, domain_trees, domain_id) || + throw(ArgumentError("serialized tensor has an out-of-range domain fusion-tree identifier")) + pairs[column] = (codomain_trees[codomain_id], domain_trees[domain_id]) + used_codomain[codomain_id] = true + used_domain[domain_id] = true + end + _check_unique_values(pairs, "fusion-tree pairs") + + all(used_codomain) || + throw(ArgumentError("serialized tensor contains an unused codomain fusion tree")) + all(used_domain) || + throw(ArgumentError("serialized tensor contains an unused domain fusion tree")) + + expected_pairs = collect(fusiontrees(tensor_space)) + length(pairs) == length(expected_pairs) && + all(pair -> any(==(pair), expected_pairs), pairs) || + throw(ArgumentError("serialized fusion-tree pairs do not match the tensor-map space")) + return pairs +end diff --git a/ext/TensorKitJLD2Ext/io.jl b/ext/TensorKitJLD2Ext/io.jl new file mode 100644 index 000000000..705b0c43f --- /dev/null +++ b/ext/TensorKitJLD2Ext/io.jl @@ -0,0 +1,27 @@ +function save_tensor(path::AbstractString, tensor::AbstractTensorMap) + record = _pack_tensormap(tensor) + destination = abspath(path) + mktemp(dirname(destination)) do temporary, io + close(io) + JLD2.jldsave( + temporary; format = TENSORMAP_FILE_FORMAT, + version = TENSORMAP_FILE_VERSION, tensor = record + ) + mv(temporary, destination; force = true) + end + return nothing +end + +function load_tensor(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/ext/TensorKitJLD2Ext/records.jl b/ext/TensorKitJLD2Ext/records.jl new file mode 100644 index 000000000..912be6806 --- /dev/null +++ b/ext/TensorKitJLD2Ext/records.jl @@ -0,0 +1,197 @@ +"""Require a named tuple to contain all fields used by a serialized record.""" +function _require_record_fields(record, fields::Tuple, description::AbstractString) + record isa NamedTuple || + throw(ArgumentError("serialized $description must be a NamedTuple")) + missing = filter(field -> !hasproperty(record, field), fields) + isempty(missing) || + throw(ArgumentError("serialized $description is missing fields $(join(missing, ", "))")) + return nothing +end + +"""Check that values are unique using equality without imposing an ordering or hash contract.""" +function _check_unique_values(values, description::AbstractString) + for index in eachindex(values) + any(previous -> previous == values[index], @view(values[firstindex(values):(index - 1)])) && + throw(ArgumentError("serialized tensor contains duplicate $description")) + end + return nothing +end + +"""Pack a dense tensor map using explicit fusion-tree pairs and CPU subblock elements.""" +function _pack_tensormap(t::TensorMap{T}) where {T} + I = sectortype(t) + Nout = numout(t) + Nin = numin(t) + tree_pairs = fusiontrees(t) + numpairs = length(tree_pairs) + codomain_trees = FusionTree{I, Nout}[] + domain_trees = FusionTree{I, Nin}[] + pair_ids = Matrix{Int}(undef, 2, numpairs) + shapes = Matrix{Int}(undef, numind(t), numpairs) + data = sizehint!(T[], length(t.data)) + for (column, (codomain_tree, domain_tree)) in enumerate(tree_pairs) + pair_ids[1, column] = _intern_fusiontree!(codomain_trees, codomain_tree) + pair_ids[2, column] = _intern_fusiontree!(domain_trees, domain_tree) + source = subblock(t, (codomain_tree, domain_tree)) + shapes[:, column] .= size(source) + append!(data, vec(Array(source))) + end + length(data) == length(t.data) || error("inconsistent TensorMap subblock storage") + codomain_table = _encode_fusiontrees(codomain_trees, I, Nout) + domain_table = _encode_fusiontrees(domain_trees, I, Nin) + return (; + kind = :dense, space = space(t), codomain_trees = codomain_table, + domain_trees = domain_table, pair_ids, shapes, data, + ) +end + +"""Pack a diagonal tensor map using fusion-tree labels and compact diagonal elements.""" +function _pack_tensormap(t::DiagonalTensorMap{T}) where {T} + I = sectortype(t) + tree_pairs = fusiontrees(t) + numpairs = length(tree_pairs) + trees = FusionTree{I, 1}[] + pair_ids = Matrix{Int}(undef, 2, numpairs) + lengths = Vector{Int}(undef, numpairs) + data = sizehint!(T[], length(t.data)) + for (column, (codomain_tree, domain_tree)) in enumerate(tree_pairs) + pair_ids[1, column] = _intern_fusiontree!(trees, codomain_tree) + pair_ids[2, column] = _intern_fusiontree!(trees, domain_tree) + elements = Vector(subblock(t, (codomain_tree, domain_tree)).diag) + lengths[column] = length(elements) + append!(data, elements) + end + length(data) == length(t.data) || error("inconsistent DiagonalTensorMap subblock storage") + tree_table = _encode_fusiontrees(trees, I, 1) + return (; kind = :diagonal, domain = only(domain(t)), trees = tree_table, pair_ids, lengths, data) +end + +"""Pack a braiding tensor using only the spaces and orientation that define it.""" +function _pack_tensormap(t::BraidingTensor{T}) where {T} + return (; kind = :braiding, V1 = t.V1, V2 = t.V2, adjoint = t.adjoint, scalartype = 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 + +"""Reconstruct a dense tensor map from an order-independent version-one record.""" +function _unpack_dense_tensormap(record) + fields = (:space, :codomain_trees, :domain_trees, :pair_ids, :shapes, :data) + _require_record_fields(record, fields, "dense TensorMap record") + record.space isa TensorMapSpace || + throw(ArgumentError("serialized TensorMap has an invalid tensor-map space")) + record.data isa Vector || + throw(ArgumentError("serialized TensorMap data must be a Vector")) + eltype(record.data) <: Number || + throw(ArgumentError("serialized TensorMap has an invalid scalar type")) + + I = sectortype(record.space) + codomain_trees = _decode_fusiontrees( + record.codomain_trees, I, numout(record.space), "codomain fusion trees" + ) + domain_trees = _decode_fusiontrees( + record.domain_trees, I, numin(record.space), "domain fusion trees" + ) + pairs = _decode_fusiontree_pairs( + record.pair_ids, codomain_trees, domain_trees, record.space + ) + record.shapes isa Matrix{Int} || + throw(ArgumentError("serialized TensorMap has invalid subblock shapes")) + size(record.shapes) == (numind(record.space), length(pairs)) || + throw(DimensionMismatch("serialized TensorMap has inconsistent subblock shape metadata")) + + T = eltype(record.data) + tensor = TensorMap{T}(undef, record.space) + offset = 0 + for (column, pair) in enumerate(pairs) + shape = Tuple(@view record.shapes[:, column]) + all(>=(0), shape) || + throw(ArgumentError("serialized TensorMap contains a negative subblock dimension")) + destination = subblock(tensor, pair) + size(destination) == shape || + throw(DimensionMismatch("serialized TensorMap subblock has shape $shape, expected $(size(destination))")) + blocklength = length(destination) + blocklength <= length(record.data) - offset || + throw(DimensionMismatch("serialized TensorMap data is shorter than its subblock metadata")) + source = reshape(@view(record.data[(offset + 1):(offset + blocklength)]), shape) + copyto!(destination, source) + offset += blocklength + end + offset == length(record.data) || + throw(DimensionMismatch("serialized TensorMap data is longer than its subblock metadata")) + return tensor +end + +"""Reconstruct a compact diagonal tensor map from an order-independent version-one record.""" +function _unpack_diagonal_tensormap(record) + fields = (:domain, :trees, :pair_ids, :lengths, :data) + _require_record_fields(record, fields, "DiagonalTensorMap record") + record.domain isa IndexSpace || + throw(ArgumentError("serialized DiagonalTensorMap has an invalid domain")) + record.data isa Vector || + throw(ArgumentError("serialized DiagonalTensorMap data must be a Vector")) + eltype(record.data) <: Number || + throw(ArgumentError("serialized DiagonalTensorMap has an invalid scalar type")) + record.lengths isa Vector{Int} || + throw(ArgumentError("serialized DiagonalTensorMap has invalid segment lengths")) + + I = sectortype(record.domain) + trees = _decode_fusiontrees(record.trees, I, 1, "diagonal fusion trees") + tensor_space = record.domain ← record.domain + pairs = _decode_fusiontree_pairs(record.pair_ids, trees, trees, tensor_space) + length(record.lengths) == length(pairs) || + throw(ArgumentError("serialized DiagonalTensorMap has inconsistent segment metadata")) + + T = eltype(record.data) + tensor = DiagonalTensorMap{T}(undef, record.domain) + offset = 0 + for (column, pair) in enumerate(pairs) + blocklength = record.lengths[column] + blocklength >= 0 || + throw(ArgumentError("serialized DiagonalTensorMap contains a negative segment length")) + destination = subblock(tensor, pair).diag + length(destination) == blocklength || + throw(DimensionMismatch("serialized diagonal segment has length $blocklength, expected $(length(destination))")) + blocklength <= length(record.data) - offset || + throw(DimensionMismatch("serialized DiagonalTensorMap data is shorter than its segment 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 segment metadata")) + return tensor +end + +"""Reconstruct a braiding tensor from its structural version-one record.""" +function _unpack_braiding_tensormap(record) + fields = (:V1, :V2, :adjoint, :scalartype) + _require_record_fields(record, fields, "BraidingTensor record") + record.V1 isa IndexSpace && record.V2 isa IndexSpace || + throw(ArgumentError("serialized BraidingTensor has invalid spaces")) + record.adjoint isa Bool || + throw(ArgumentError("serialized BraidingTensor has an invalid orientation flag")) + record.scalartype isa Type && record.scalartype <: Number || + throw(ArgumentError("serialized BraidingTensor has an invalid scalar type")) + return try + BraidingTensor{record.scalartype}(record.V1, record.V2, record.adjoint) + catch error + message = sprint(showerror, error) + throw(ArgumentError("serialized BraidingTensor is inconsistent: $message")) + end +end + +"""Dispatch a version-one basic Julia record to its tensor-map decoder.""" +function _unpack_tensormap(record) + _require_record_fields(record, (:kind,), "tensor record") + record.kind === :dense && return _unpack_dense_tensormap(record) + record.kind === :diagonal && return _unpack_diagonal_tensormap(record) + record.kind === :braiding && return _unpack_braiding_tensormap(record) + throw(ArgumentError("unsupported TensorKit tensor record kind $(repr(record.kind))")) +end diff --git a/src/TensorKit.jl b/src/TensorKit.jl index 41a7d8ae1..866e8b5d7 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_tensor, load_tensor export SpaceMismatch, SectorMismatch, IndexError # error types # Export general vector space methods @@ -267,6 +268,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..e19f30fb8 --- /dev/null +++ b/src/tensors/io.jl @@ -0,0 +1,21 @@ +# TensorMap IO interface +#=======================# + +""" + save_tensor(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. + +Requires loading JLD2 with `using JLD2` to activate the `TensorKitJLD2Ext` extension. +""" +function save_tensor end + +""" + load_tensor(path::AbstractString) -> AbstractTensorMap + +Load one tensor map saved with [`save_tensor`](@ref), using CPU storage for numerical data. + +Requires loading JLD2 with `using JLD2` to activate the `TensorKitJLD2Ext` extension. +""" +function load_tensor end diff --git a/test/Project.toml b/test/Project.toml index ce51b70a3..39951cfd5 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" @@ -45,4 +46,3 @@ Test = "1" TestExtras = "0.2,0.3" TimerOutputs = "1" Zygote = "0.7" - diff --git a/test/tensors/io.jl b/test/tensors/io.jl new file mode 100644 index 000000000..da3ebc870 --- /dev/null +++ b/test/tensors/io.jl @@ -0,0 +1,320 @@ +using Test +import JLD2 +using TensorKit + +const TensorKitJLD2Ext = Base.get_extension(TensorKit, :TensorKitJLD2Ext) + +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 and permutation tests.""" +function write_record(path, record; format = TensorKitJLD2Ext.TENSORMAP_FILE_FORMAT, version = TensorKitJLD2Ext.TENSORMAP_FILE_VERSION) + return JLD2.jldopen(path, "w") do file + file["format"] = format + file["version"] = version + file["tensor"] = record + end +end + +"""Reverse every order-bearing part of a tensor record while preserving its semantics.""" +function reverse_record(record) + permute_table = (table, permutation) -> ( + uncoupled = table.uncoupled[:, permutation], + coupled = table.coupled[permutation], + isdual = table.isdual[:, permutation], + innerlines = table.innerlines[:, permutation], + vertices = table.vertices[:, permutation], + ) + pair_permutation = reverse(axes(record.pair_ids, 2)) + lengths = record.kind === :dense ? + [prod(@view(record.shapes[:, column])) for column in axes(record.shapes, 2)] : + record.lengths + offsets = cumsum(vcat(0, lengths)) + data = similar(record.data) + destination = 0 + for source_column in pair_permutation + source_range = (offsets[source_column] + 1):offsets[source_column + 1] + copyto!(data, destination + 1, record.data, first(source_range), length(source_range)) + destination += length(source_range) + end + + pair_ids = copy(record.pair_ids) + if record.kind === :dense + codomain_permutation = reverse(eachindex(record.codomain_trees.coupled)) + domain_permutation = reverse(eachindex(record.domain_trees.coupled)) + pair_ids[1, :] .= invperm(codomain_permutation)[pair_ids[1, :]] + pair_ids[2, :] .= invperm(domain_permutation)[pair_ids[2, :]] + return merge( + record, ( + codomain_trees = permute_table(record.codomain_trees, codomain_permutation), + domain_trees = permute_table(record.domain_trees, domain_permutation), + pair_ids = pair_ids[:, pair_permutation], + shapes = record.shapes[:, pair_permutation], + data, + ) + ) + end + tree_permutation = reverse(eachindex(record.trees.coupled)) + remapping = invperm(tree_permutation) + pair_ids .= remapping[pair_ids] + return merge( + record, ( + trees = permute_table(record.trees, tree_permutation), + pair_ids = pair_ids[:, pair_permutation], + lengths = record.lengths[pair_permutation], + data, + ) + ) +end + +@testset "TensorMap save_tensor and load_tensor" begin + @test !isnothing(TensorKitJLD2Ext) + @test :save_tensor in names(TensorKit) + @test :load_tensor in names(TensorKit) + @test :save ∉ names(TensorKit) + @test :load ∉ names(TensorKit) + + 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_tensor(path, tensor) === nothing + restored = load_tensor(path) + @test restored isa TensorMap + @test storagetype(restored) === Vector{ComplexF64} + @test space(restored) == space(tensor) + @test restored == tensor + end + + tensor = randn(Float64, ℂ^2 ⊗ ℂ^3) + path = joinpath(directory, "plain-tensor.jld2") + save_tensor(path, tensor) + restored = load_tensor(path) + @test restored isa Tensor + @test restored == tensor + + save_tensor(path, 2 * tensor) + @test load_tensor(path) == 2 * tensor + + scalar_space = one(ℂ^1) + scalar = randn(ComplexF64, scalar_space ← scalar_space) + path = joinpath(directory, "scalar.jld2") + save_tensor(path, scalar) + restored_scalar = load_tensor(path) + @test numout(restored_scalar) == numin(restored_scalar) == 0 + @test restored_scalar == scalar + + empty_tensor = randn(Float64, zero(ℂ^2) ← ℂ^2) + path = joinpath(directory, "empty.jld2") + save_tensor(path, empty_tensor) + restored_empty = load_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} + path = joinpath(directory, "custom.jld2") + save_tensor(path, custom) + restored_custom = load_tensor(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) + empty_diagonal = DiagonalTensorMap(ComplexF64[], zero(diagonal_space)) + for (index, value) in enumerate((diagonal, diagonal', empty_diagonal)) + path = joinpath(directory, "diagonal-$index.jld2") + save_tensor(path, value) + restored_diagonal = load_tensor(path) + @test restored_diagonal isa DiagonalTensorMap + @test storagetype(restored_diagonal) === Vector{ComplexF64} + @test restored_diagonal == value + end + + 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_tensor(path, braiding) + restored_braiding = load_tensor(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 convert(TensorMap, convert(Dict, source)) == source + @test_throws ArgumentError save_tensor(joinpath(directory, "adjoint.jld2"), source') + @test_throws ArgumentError save_tensor(joinpath(directory, "unsupported.jld2"), UnregisteredTensorMap()) + end +end + +@testset "TensorMap record representation" begin + V1, V2, V3, V4, V5 = TestSetup.VRepA4 + tensor = randn(ComplexF64, V1 ⊗ V2 ← (V3 ⊗ V4 ⊗ V5)') + record = TensorKitJLD2Ext._pack_tensormap(tensor) + @test record isa NamedTuple + @test record.kind === :dense + for table in (record.codomain_trees, record.domain_trees) + @test table isa NamedTuple + @test table.uncoupled isa Matrix{sectortype(tensor)} + @test table.coupled isa Vector{sectortype(tensor)} + @test table.isdual isa BitMatrix + @test table.innerlines isa Matrix{sectortype(tensor)} + @test table.vertices isa Matrix{Int} + @test all(!(value isa FusionTree) for value in values(table)) + @test all(!(value isa AbstractString) for value in values(table)) + end + @test record.pair_ids isa Matrix{Int} + @test record.shapes isa Matrix{Int} + @test record.data isa Vector{ComplexF64} +end + +@testset "Fusion-tree iteration-order independence" begin + V1, V2, V3, V4, V5 = TestSetup.VRepA4 + tensor = randn(ComplexF64, V1 ⊗ V2 ← (V3 ⊗ V4 ⊗ V5)') + @test TensorKitJLD2Ext._unpack_tensormap(reverse_record(TensorKitJLD2Ext._pack_tensormap(tensor))) == tensor + + V = Vect[SU2Irrep](0 => 3, 1 // 2 => 2, 1 => 1)' + diagonal = DiagonalTensorMap(randn(ComplexF64, reduceddim(V)), V) + @test TensorKitJLD2Ext._unpack_tensormap(reverse_record(TensorKitJLD2Ext._pack_tensormap(diagonal))) == diagonal +end + +@testset "TensorMap file validation" begin + V = Vect[Z2Irrep](0 => 2, 1 => 3) + tensor = randn(Float64, V ⊗ V ← V ⊗ V) + record = TensorKitJLD2Ext._pack_tensormap(tensor) + mktempdir() do directory + path = joinpath(directory, "invalid.jld2") + + write_record(path, record; format = "not TensorKit") + @test_throws ArgumentError load_tensor(path) + + write_record(path, record; version = TensorKitJLD2Ext.TENSORMAP_FILE_VERSION + 1) + @test_throws ArgumentError load_tensor(path) + + JLD2.jldsave(path; unrelated = tensor.data) + @test_throws ArgumentError load_tensor(path) + + write_record(path, merge(record, (kind = :unknown,))) + @test_throws ArgumentError load_tensor(path) + + duplicate_ids = copy(record.pair_ids) + duplicate_ids[:, 2] .= duplicate_ids[:, 1] + write_record(path, merge(record, (pair_ids = duplicate_ids,))) + @test_throws ArgumentError load_tensor(path) + + write_record( + path, merge( + record, ( + pair_ids = record.pair_ids[:, 1:(end - 1)], + shapes = record.shapes[:, 1:(end - 1)], + ) + ) + ) + @test_throws ArgumentError load_tensor(path) + + invalid_ids = copy(record.pair_ids) + invalid_ids[1, 1] = size(record.codomain_trees.coupled, 1) + 1 + write_record(path, merge(record, (pair_ids = invalid_ids,))) + @test_throws ArgumentError load_tensor(path) + + used_pairs = record.pair_ids[:, record.pair_ids[1, :] .!= 1] + write_record(path, merge(record, (pair_ids = used_pairs,))) + @test_throws "unused codomain fusion tree" load_tensor(path) + + duplicate_table = merge( + record.codomain_trees, ( + uncoupled = hcat(record.codomain_trees.uncoupled, record.codomain_trees.uncoupled[:, 1]), + coupled = vcat(record.codomain_trees.coupled, record.codomain_trees.coupled[1]), + isdual = hcat(record.codomain_trees.isdual, record.codomain_trees.isdual[:, 1]), + innerlines = hcat(record.codomain_trees.innerlines, record.codomain_trees.innerlines[:, 1]), + vertices = hcat(record.codomain_trees.vertices, record.codomain_trees.vertices[:, 1]), + ) + ) + write_record(path, merge(record, (codomain_trees = duplicate_table,))) + @test_throws ArgumentError load_tensor(path) + + bad_dimensions = merge( + record.codomain_trees, ( + uncoupled = record.codomain_trees.uncoupled[1:(end - 1), :], + ) + ) + write_record(path, merge(record, (codomain_trees = bad_dimensions,))) + @test_throws DimensionMismatch load_tensor(path) + + bad_vertices = copy(record.codomain_trees.vertices) + bad_vertices[1, 1] = 2 + invalid_tree = merge(record.codomain_trees, (vertices = bad_vertices,)) + write_record(path, merge(record, (codomain_trees = invalid_tree,))) + @test_throws ArgumentError load_tensor(path) + + incompatible_space = Vect[Z2Irrep](0 => 2) ⊗ Vect[Z2Irrep](0 => 2) ← + Vect[Z2Irrep](0 => 2) ⊗ Vect[Z2Irrep](0 => 2) + write_record(path, merge(record, (space = incompatible_space,))) + @test_throws ArgumentError load_tensor(path) + + badshape = copy(record.shapes) + badshape[1, 1] += 1 + write_record(path, merge(record, (shapes = badshape,))) + @test_throws DimensionMismatch load_tensor(path) + + write_record(path, merge(record, (data = record.data[1:(end - 1)],))) + @test_throws DimensionMismatch load_tensor(path) + + write_record(path, merge(record, (data = vcat(record.data, zero(eltype(record.data))),))) + @test_throws DimensionMismatch load_tensor(path) + + write_record(path, 1) + @test_throws ArgumentError load_tensor(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) + tensor_path = joinpath(directory, "tensor.jld2") + dict_path = joinpath(directory, "tensor-dict.jld2") + save_tensor(tensor_path, tensor) + JLD2.jldsave(dict_path; tensor = convert(Dict, tensor)) + tensor_size = filesize(tensor_path) + dict_size = filesize(dict_path) + @info "TensorMap IO size comparison" tensor_size dict_size + @test tensor_size > 0 + @test dict_size > 0 + + 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_tensor(diagonal_path, diagonal) + save_tensor(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_tensor(braiding_path, braiding) + save_tensor(dense_braiding_path, TensorMap(braiding)) + @test filesize(braiding_path) < filesize(dense_braiding_path) + end +end