[NDTensors] Allow passing keyword arguments to HDF5 read and write

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org>"]
-version = "0.4.16"
+version = "0.4.17"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

NDTensors/ext/NDTensorsHDF5Ext/blocksparse.jl

@@ -31,20 +31,24 @@ function array_to_offsets(a, N::Int)
     return boff
 end
 
-function HDF5.write(parent::Union{HDF5.File, HDF5.Group}, name::String, B::BlockSparse)
+function HDF5.write(
+        parent::Union{HDF5.File, HDF5.Group}, name::String, B::BlockSparse; kwargs...
+    )
     g = create_group(parent, name)
     attributes(g)["type"] = "BlockSparse{$(eltype(B))}"
     attributes(g)["version"] = 1
     return if eltype(B) != Nothing
         write(g, "ndims", ndims(B))
-        write(g, "data", data(B))
+        # Use `setindex!` so that chunking happens automatically
+        # when compression is used, see: https://github.com/JuliaIO/HDF5.jl/issues/822
+        setindex!(g, data(B), "data"; kwargs...)
         off_array = offsets_to_array(blockoffsets(B))
         write(g, "offsets", off_array)
     end
 end
 
 function HDF5.read(
-        parent::Union{HDF5.File, HDF5.Group}, name::AbstractString, ::Type{Store}
+        parent::Union{HDF5.File, HDF5.Group}, name::AbstractString, ::Type{Store}; kwargs...
     ) where {Store <: BlockSparse}
     g = open_group(parent, name)
     ElT = eltype(Store)
@@ -58,11 +62,11 @@ function HDF5.read(
     # Attribute __complex__ is attached to the "data" dataset
     # by the h5 library used by C++ version of ITensor:
     if haskey(attributes(g["data"]), "__complex__")
-        M = read(g, "data")
+        M = read(g, "data"; kwargs...)
         nelt = size(M, 1) * size(M, 2)
         data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
     else
-        data = read(g, "data")
+        data = read(g, "data"; kwargs...)
     end
     return BlockSparse(data, boff)
 end

NDTensors/ext/NDTensorsHDF5Ext/dense.jl

@@ -2,18 +2,20 @@ using HDF5: HDF5, attributes, create_group, open_group, read, write
 using NDTensors: Dense
 
 function HDF5.write(
-        parent::Union{HDF5.File, HDF5.Group}, name::String, D::Store
+        parent::Union{HDF5.File, HDF5.Group}, name::String, D::Store; kwargs...
     ) where {Store <: Dense}
     g = create_group(parent, name)
     attributes(g)["type"] = "Dense{$(eltype(Store))}"
     attributes(g)["version"] = 1
     return if eltype(D) != Nothing
-        write(g, "data", D.data)
+        # Use `setindex!` so that chunking happens automatically
+        # when compression is used, see: https://github.com/JuliaIO/HDF5.jl/issues/822
+        setindex!(g, D.data, "data"; kwargs...)
     end
 end
 
 function HDF5.read(
-        parent::Union{HDF5.File, HDF5.Group}, name::AbstractString, ::Type{Store}
+        parent::Union{HDF5.File, HDF5.Group}, name::AbstractString, ::Type{Store}; kwargs...
     ) where {Store <: Dense}
     g = open_group(parent, name)
     ElT = eltype(Store)
@@ -31,7 +33,7 @@ function HDF5.read(
         nelt = size(M, 1) * size(M, 2)
         data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
     else
-        data = read(g, "data")
+        data = read(g, "data"; kwargs...)
     end
     return Dense{ElT}(data)
 end

NDTensors/ext/NDTensorsHDF5Ext/diag.jl

@@ -2,18 +2,20 @@ using HDF5: HDF5, attributes, create_group, open_group, read, write
 using NDTensors: datatype, Dense, Diag
 
 function HDF5.write(
-        parent::Union{HDF5.File, HDF5.Group}, name::String, D::Store
+        parent::Union{HDF5.File, HDF5.Group}, name::String, D::Store; kwargs...
     ) where {Store <: Diag}
     g = create_group(parent, name)
     attributes(g)["type"] = "Diag{$(eltype(Store)),$(datatype(Store))}"
     attributes(g)["version"] = 1
     return if eltype(D) != Nothing
-        write(g, "data", D.data)
+        # Use `setindex!` so that chunking happens automatically
+        # when compression is used, see: https://github.com/JuliaIO/HDF5.jl/issues/822
+        setindex!(g, D.data, "data"; kwargs...)
     end
 end
 
 function HDF5.read(
-        parent::Union{HDF5.File, HDF5.Group}, name::AbstractString, ::Type{Store}
+        parent::Union{HDF5.File, HDF5.Group}, name::AbstractString, ::Type{Store}; kwargs...
     ) where {Store <: Diag}
     g = open_group(parent, name)
     ElT = eltype(Store)
@@ -28,11 +30,11 @@ function HDF5.read(
     # Attribute __complex__ is attached to the "data" dataset
     # by the h5 library used by C++ version of ITensor:
     if haskey(attributes(g["data"]), "__complex__")
-        M = read(g, "data")
+        M = read(g, "data"; kwargs...)
         nelt = size(M, 1) * size(M, 2)
         data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
     else
-        data = read(g, "data")
+        data = read(g, "data"; kwargs...)
     end
     return Diag{ElT, DataT}(data)
 end

NDTensors/ext/NDTensorsHDF5Ext/empty.jl

@@ -3,8 +3,10 @@ using NDTensors: EmptyStorage
 
 # XXX: this seems a bit strange and fragile?
 # Takes the type very literally.
+# Trailing `kwargs` are used to capture chunking/compression options,
+# which are ignored for EmptyStorage.
 function HDF5.read(
-        parent::Union{HDF5.File, HDF5.Group}, name::AbstractString, ::Type{StoreT}
+        parent::Union{HDF5.File, HDF5.Group}, name::AbstractString, ::Type{StoreT}; kwargs...
     ) where {StoreT <: EmptyStorage}
     g = open_group(parent, name)
     typestr = string(StoreT)
@@ -14,8 +16,10 @@ function HDF5.read(
     return StoreT()
 end
 
+# Trailing `kwargs` are used to capture chunking/compression options,
+# which are ignored for EmptyStorage.
 function HDF5.write(
-        parent::Union{HDF5.File, HDF5.Group}, name::String, ::StoreT
+        parent::Union{HDF5.File, HDF5.Group}, name::String, ::StoreT; kwargs...
     ) where {StoreT <: EmptyStorage}
     g = create_group(parent, name)
     attributes(g)["type"] = string(StoreT)