[NDTensors] HDF5 package extension

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org>"]
-version = "0.3.0"
+version = "0.4.0"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
@@ -15,7 +15,6 @@ FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Folds = "41a02a25-b8f0-4f67-bc48-60067656b558"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
-HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 HalfIntegers = "f0d1745a-41c9-11e9-1dd9-e5d34d218721"
 InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -34,13 +33,15 @@ VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
 TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9"
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 
 [extensions]
 NDTensorsCUDAExt = "CUDA"
+NDTensorsHDF5Ext = "HDF5"
 NDTensorsMetalExt = "Metal"
 NDTensorsOctavianExt = "Octavian"
 NDTensorsTBLISExt = "TBLIS"
@@ -75,3 +76,6 @@ TimerOutputs = "0.5.5"
 TupleTools = "1.2.0"
 VectorInterface = "0.4.2"
 julia = "1.6"
+
+[extras]
+HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"

NDTensors/ext/NDTensorsHDF5Ext/NDTensorsHDF5Ext.jl

@@ -0,0 +1,8 @@
+module NDTensorsHDF5Ext
+
+include("blocksparse.jl")
+include("dense.jl")
+include("diag.jl")
+include("empty.jl")
+
+end # module NDTensorsHDF5Ext

NDTensors/ext/NDTensorsHDF5Ext/blocksparse.jl

@@ -0,0 +1,68 @@
+using HDF5: HDF5, attributes, create_group, open_group, read, write
+using NDTensors: data, Block, blockoffsets, BlockOffsets, BlockSparse
+
+# Helper function for HDF5 write/read of BlockSparse
+function offsets_to_array(boff::BlockOffsets{N}) where {N}
+  nblocks = length(boff)
+  asize = (N + 1) * nblocks
+  n = 1
+  a = Vector{Int}(undef, asize)
+  for bo in pairs(boff)
+    for j in 1:N
+      a[n] = bo[1][j]
+      n += 1
+    end
+    a[n] = bo[2]
+    n += 1
+  end
+  return a
+end
+
+# Helper function for HDF5 write/read of BlockSparse
+function array_to_offsets(a, N::Int)
+  asize = length(a)
+  nblocks = div(asize, N + 1)
+  boff = BlockOffsets{N}()
+  j = 0
+  for b in 1:nblocks
+    insert!(boff, Block(ntuple(i -> (a[j + i]), N)), a[j + N + 1])
+    j += (N + 1)
+  end
+  return boff
+end
+
+function HDF5.write(parent::Union{HDF5.File,HDF5.Group}, name::String, B::BlockSparse)
+  g = create_group(parent, name)
+  attributes(g)["type"] = "BlockSparse{$(eltype(B))}"
+  attributes(g)["version"] = 1
+  if eltype(B) != Nothing
+    write(g, "ndims", ndims(B))
+    write(g, "data", data(B))
+    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}
+) where {Store<:BlockSparse}
+  g = open_group(parent, name)
+  ElT = eltype(Store)
+  typestr = "BlockSparse{$ElT}"
+  if read(attributes(g)["type"]) != typestr
+    error("HDF5 group or file does not contain $typestr data")
+  end
+  N = read(g, "ndims")
+  off_array = read(g, "offsets")
+  boff = array_to_offsets(off_array, N)
+  # 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")
+    nelt = size(M, 1) * size(M, 2)
+    data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
+  else
+    data = read(g, "data")
+  end
+  return BlockSparse(data, boff)
+end

NDTensors/ext/NDTensorsHDF5Ext/dense.jl

@@ -0,0 +1,37 @@
+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
+) where {Store<:Dense}
+  g = create_group(parent, name)
+  attributes(g)["type"] = "Dense{$(eltype(Store))}"
+  attributes(g)["version"] = 1
+  if eltype(D) != Nothing
+    write(g, "data", D.data)
+  end
+end
+
+function HDF5.read(
+  parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{Store}
+) where {Store<:Dense}
+  g = open_group(parent, name)
+  ElT = eltype(Store)
+  typestr = "Dense{$ElT}"
+  if read(attributes(g)["type"]) != typestr
+    error("HDF5 group or file does not contain $typestr data")
+  end
+  if ElT == Nothing
+    return Dense{Nothing}()
+  end
+  # 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")
+    nelt = size(M, 1) * size(M, 2)
+    data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
+  else
+    data = read(g, "data")
+  end
+  return Dense{ElT}(data)
+end

NDTensors/ext/NDTensorsHDF5Ext/diag.jl

@@ -0,0 +1,38 @@
+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
+) where {Store<:Diag}
+  g = create_group(parent, name)
+  attributes(g)["type"] = "Diag{$(eltype(Store)),$(datatype(Store))}"
+  attributes(g)["version"] = 1
+  if eltype(D) != Nothing
+    write(g, "data", D.data)
+  end
+end
+
+function HDF5.read(
+  parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{Store}
+) where {Store<:Diag}
+  g = open_group(parent, name)
+  ElT = eltype(Store)
+  DataT = datatype(Store)
+  typestr = "Diag{$ElT,$DataT}"
+  if read(attributes(g)["type"]) != typestr
+    error("HDF5 group or file does not contain $typestr data")
+  end
+  if ElT == Nothing
+    return Dense{Nothing}()
+  end
+  # 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")
+    nelt = size(M, 1) * size(M, 2)
+    data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
+  else
+    data = read(g, "data")
+  end
+  return Diag{ElT,DataT}(data)
+end

NDTensors/ext/NDTensorsHDF5Ext/empty.jl

@@ -0,0 +1,23 @@
+using HDF5: HDF5, attributes, create_group, open_group, read, write
+using NDTensors: EmptyStorage
+
+# XXX: this seems a bit strange and fragile?
+# Takes the type very literally.
+function HDF5.read(
+  parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{StoreT}
+) where {StoreT<:EmptyStorage}
+  g = open_group(parent, name)
+  typestr = string(StoreT)
+  if read(attributes(g)["type"]) != typestr
+    error("HDF5 group or file does not contain $typestr data")
+  end
+  return StoreT()
+end
+
+function HDF5.write(
+  parent::Union{HDF5.File,HDF5.Group}, name::String, ::StoreT
+) where {StoreT<:EmptyStorage}
+  g = create_group(parent, name)
+  attributes(g)["type"] = string(StoreT)
+  return attributes(g)["version"] = 1
+end

NDTensors/src/blocksparse/blocksparse.jl

@@ -205,69 +205,3 @@ end
 #  end
 #  return R
 #end
-
-# Helper function for HDF5 write/read of BlockSparse
-function offsets_to_array(boff::BlockOffsets{N}) where {N}
-  nblocks = length(boff)
-  asize = (N + 1) * nblocks
-  n = 1
-  a = Vector{Int}(undef, asize)
-  for bo in pairs(boff)
-    for j in 1:N
-      a[n] = bo[1][j]
-      n += 1
-    end
-    a[n] = bo[2]
-    n += 1
-  end
-  return a
-end
-
-# Helper function for HDF5 write/read of BlockSparse
-function array_to_offsets(a, N::Int)
-  asize = length(a)
-  nblocks = div(asize, N + 1)
-  boff = BlockOffsets{N}()
-  j = 0
-  for b in 1:nblocks
-    insert!(boff, Block(ntuple(i -> (a[j + i]), N)), a[j + N + 1])
-    j += (N + 1)
-  end
-  return boff
-end
-
-function HDF5.write(parent::Union{HDF5.File,HDF5.Group}, name::String, B::BlockSparse)
-  g = create_group(parent, name)
-  attributes(g)["type"] = "BlockSparse{$(eltype(B))}"
-  attributes(g)["version"] = 1
-  if eltype(B) != Nothing
-    write(g, "ndims", ndims(B))
-    write(g, "data", data(B))
-    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}
-) where {Store<:BlockSparse}
-  g = open_group(parent, name)
-  ElT = eltype(Store)
-  typestr = "BlockSparse{$ElT}"
-  if read(attributes(g)["type"]) != typestr
-    error("HDF5 group or file does not contain $typestr data")
-  end
-  N = read(g, "ndims")
-  off_array = read(g, "offsets")
-  boff = array_to_offsets(off_array, N)
-  # 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")
-    nelt = size(M, 1) * size(M, 2)
-    data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
-  else
-    data = read(g, "data")
-  end
-  return BlockSparse(data, boff)
-end

NDTensors/src/dense/dense.jl

@@ -141,38 +141,3 @@ end
 function convert(::Type{<:Dense{ElR,DataT}}, D::Dense) where {ElR,DataT}
   return Dense(convert(DataT, data(D)))
 end
-
-function HDF5.write(
-  parent::Union{HDF5.File,HDF5.Group}, name::String, D::Store
-) where {Store<:Dense}
-  g = create_group(parent, name)
-  attributes(g)["type"] = "Dense{$(eltype(Store))}"
-  attributes(g)["version"] = 1
-  if eltype(D) != Nothing
-    write(g, "data", D.data)
-  end
-end
-
-function HDF5.read(
-  parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{Store}
-) where {Store<:Dense}
-  g = open_group(parent, name)
-  ElT = eltype(Store)
-  typestr = "Dense{$ElT}"
-  if read(attributes(g)["type"]) != typestr
-    error("HDF5 group or file does not contain $typestr data")
-  end
-  if ElT == Nothing
-    return Dense{Nothing}()
-  end
-  # 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")
-    nelt = size(M, 1) * size(M, 2)
-    data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
-  else
-    data = read(g, "data")
-  end
-  return Dense{ElT}(data)
-end

NDTensors/src/diag/diag.jl

@@ -132,39 +132,3 @@ end
 # Convert a Diag storage type to the closest Dense storage type
 dense(::Type{<:NonuniformDiag{ElT,DataT}}) where {ElT,DataT} = Dense{ElT,DataT}
 dense(::Type{<:UniformDiag{ElT}}) where {ElT} = Dense{ElT,default_datatype(ElT)}
-
-function HDF5.write(
-  parent::Union{HDF5.File,HDF5.Group}, name::String, D::Store
-) where {Store<:Diag}
-  g = create_group(parent, name)
-  attributes(g)["type"] = "Diag{$(eltype(Store)),$(datatype(Store))}"
-  attributes(g)["version"] = 1
-  if eltype(D) != Nothing
-    write(g, "data", D.data)
-  end
-end
-
-function HDF5.read(
-  parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{Store}
-) where {Store<:Diag}
-  g = open_group(parent, name)
-  ElT = eltype(Store)
-  DataT = datatype(Store)
-  typestr = "Diag{$ElT,$DataT}"
-  if read(attributes(g)["type"]) != typestr
-    error("HDF5 group or file does not contain $typestr data")
-  end
-  if ElT == Nothing
-    return Dense{Nothing}()
-  end
-  # 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")
-    nelt = size(M, 1) * size(M, 2)
-    data = Vector(reinterpret(ComplexF64, reshape(M, nelt)))
-  else
-    data = read(g, "data")
-  end
-  return Diag{ElT,DataT}(data)
-end

NDTensors/src/empty/empty.jl

@@ -94,23 +94,3 @@ end
 using .TypeParameterAccessors: TypeParameterAccessors
 TypeParameterAccessors.parenttype(empty::Type{<:EmptyStorage}) = storagetype(empty)
 zero(empty::EmptyStorage) = empty
-# XXX: this seems a bit strange and fragile?
-# Takes the type very literally.
-function HDF5.read(
-  parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{StoreT}
-) where {StoreT<:EmptyStorage}
-  g = open_group(parent, name)
-  typestr = string(StoreT)
-  if read(attributes(g)["type"]) != typestr
-    error("HDF5 group or file does not contain $typestr data")
-  end
-  return StoreT()
-end
-
-function HDF5.write(
-  parent::Union{HDF5.File,HDF5.Group}, name::String, ::StoreT
-) where {StoreT<:EmptyStorage}
-  g = create_group(parent, name)
-  attributes(g)["type"] = string(StoreT)
-  return attributes(g)["version"] = 1
-end

NDTensors/src/imports.jl

@@ -15,7 +15,6 @@ using Random
 using LinearAlgebra
 using StaticArrays
 using Functors
-using HDF5
 using SimpleTraits
 using SplitApplyCombine
 using Strided