[ITensors] HDF5.jl package extension

Project.toml

@@ -11,7 +11,6 @@ Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
-HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 IsApprox = "28f27b66-4bd8-47e7-9110-e2746eb8bed7"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -34,9 +33,11 @@ Zeros = "bd1ec220-6eb4-527a-9b49-e79c3db6233b"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [weakdeps]
+HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 
 [extensions]
+ITensorsHDF5Ext = "HDF5"
 ITensorsVectorInterfaceExt = "VectorInterface"
 
 [compat]
@@ -72,4 +73,5 @@ ZygoteRules = "0.2.2"
 julia = "1.6"
 
 [extras]
+HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"

docs/settings.jl

@@ -1,4 +1,5 @@
-using Documenter, ITensors
+using Documenter
+using ITensors
 
 DocMeta.setdocmeta!(ITensors, :DocTestSetup, :(using ITensors); recursive=true)
 

docs/src/ITensorType.md

@@ -26,7 +26,7 @@ itensor(::Array{<:Number}, ::ITensors.Indices)
 ## QN BlockSparse Constructors
 
 ```@docs
-ITensor(::Type{<:Number}, ::QN, ::ITensors.Indices)
+ITensor(::Type{<:Number}, ::QN, ::ITensors.QNIndices)
 ITensor(::ITensors.AliasStyle, ::Type{<:Number}, ::Array{<:Number}, ::ITensors.QNIndices; tol=0)
 ITensor(::Type{<:Number}, ::UndefInitializer, ::QN, ::ITensors.Indices)
 ```

docs/src/examples/ITensor.md

@@ -501,6 +501,11 @@ UT = dag(C) * CT
 
 ## Write and Read an ITensor to Disk with HDF5
 
+
+
+!!! info 
+    Make sure to install the HDF5 package to use this feature. (Run `julia> ] add HDF5` in the Julia REPL console.)
+
 Saving ITensors to disk can be very useful. For example, you
 might encounter a bug in your own code, and by reading the
 ITensors involved from disk you can shortcut the process of
@@ -520,7 +525,7 @@ from a calculation. To write it to an HDF5 file named "myfile.h5"
 you can use the following pattern:
 
 ```julia
-using ITensors.HDF5
+using HDF5
 f = h5open("myfile.h5","w")
 write(f,"T",T)
 close(f)
@@ -531,16 +536,14 @@ or "Result Tensor" and doesn't have to have the same name as the reference `T`.
 Closing the file `f` is optional and you can also write other objects to the same
 file before closing it.
 
-[*Above we did `using ITensors.HDF5` since HDF5 is already included as a dependency with ITensor. You can also do `using HDF5` but must add the HDF5 package beforehand for that to work.*]
-
 **Reading an ITensor from an HDF5 File**
 
 Say you have an HDF5 file "myfile.h5" which contains an ITensor stored as a dataset with the
 name "T". (Which would be the situation if you wrote it as in the example above.)
 To read this ITensor back from the HDF5 file, use the following pattern:
 
 ```julia
-using ITensors.HDF5
+using HDF5
 f = h5open("myfile.h5","r")
 T = read(f,"T",ITensor)
 close(f)

docs/src/examples/MPSandMPO.md

@@ -384,14 +384,17 @@ can be brought into orthogonal form by calling `psi = orthogonalize(psi, 1)`.
 
 ## Write and Read an MPS or MPO to Disk with HDF5
 
+!!! info 
+    Make sure to install the HDF5 package to use this feature. (Run `julia> ] add HDF5` in the Julia REPL console.)
+
 **Writing an MPS to an HDF5 File**
 
 Let's say you have an MPS `psi` which you have made or obtained
 from a calculation. To write it to an HDF5 file named "myfile.h5"
 you can use the following pattern:
 
 ```julia
-using ITensors.HDF5
+using HDF5
 f = h5open("myfile.h5","w")
 write(f,"psi",psi)
 close(f)
@@ -402,16 +405,14 @@ or "Result MPS" and doesn't have to have the same name as the reference `psi`.
 Closing the file `f` is optional and you can also write other objects to the same
 file before closing it.
 
-[*Above we did `using ITensors.HDF5` since HDF5 is already included as a dependency with ITensor. You can also do `using HDF5` but must add the HDF5 package beforehand for that to work.*]
-
 **Reading an MPS from an HDF5 File**
 
 Say you have an HDF5 file "myfile.h5" which contains an MPS stored as a dataset with the
 name "psi". (Which would be the situation if you wrote it as in the example above.)
 To read this ITensor back from the HDF5 file, use the following pattern:
 
 ```julia
-using ITensors.HDF5
+using HDF5
 f = h5open("myfile.h5","r")
 psi = read(f,"psi",MPS)
 close(f)
@@ -436,7 +437,6 @@ H = MPO(os,sites)
 energy_psi = inner(psi',H,psi)
 ```
 
-
 Note the `MPS` argument to the read function, which tells Julia which read function
 to call and how to interpret the data stored in the HDF5 dataset named "psi". In the
 future we might lift the requirement of providing the type and have it be detected

ext/ITensorsHDF5Ext/ITensorMPS/mpo.jl

@@ -0,0 +1,27 @@
+using ITensors: ITensor
+using ITensors.ITensorMPS: MPO
+
+function HDF5.write(parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, M::MPO)
+  g = create_group(parent, name)
+  attributes(g)["type"] = "MPO"
+  attributes(g)["version"] = 1
+  N = length(M)
+  write(g, "rlim", M.rlim)
+  write(g, "llim", M.llim)
+  write(g, "length", N)
+  for n in 1:N
+    write(g, "MPO[$(n)]", M[n])
+  end
+end
+
+function HDF5.read(parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{MPO})
+  g = open_group(parent, name)
+  if read(attributes(g)["type"]) != "MPO"
+    error("HDF5 group or file does not contain MPO data")
+  end
+  N = read(g, "length")
+  rlim = read(g, "rlim")
+  llim = read(g, "llim")
+  v = [read(g, "MPO[$(i)]", ITensor) for i in 1:N]
+  return MPO(v, llim, rlim)
+end

ext/ITensorsHDF5Ext/ITensorMPS/mps.jl

@@ -0,0 +1,27 @@
+using ITensors: ITensor
+using ITensors.ITensorMPS: MPS
+
+function HDF5.write(parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, M::MPS)
+  g = create_group(parent, name)
+  attributes(g)["type"] = "MPS"
+  attributes(g)["version"] = 1
+  N = length(M)
+  write(g, "length", N)
+  write(g, "rlim", M.rlim)
+  write(g, "llim", M.llim)
+  for n in 1:N
+    write(g, "MPS[$(n)]", M[n])
+  end
+end
+
+function HDF5.read(parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{MPS})
+  g = open_group(parent, name)
+  if read(attributes(g)["type"]) != "MPS"
+    error("HDF5 group or file does not contain MPS data")
+  end
+  N = read(g, "length")
+  rlim = read(g, "rlim")
+  llim = read(g, "llim")
+  v = [read(g, "MPS[$(i)]", ITensor) for i in 1:N]
+  return MPS(v, llim, rlim)
+end

ext/ITensorsHDF5Ext/ITensorsHDF5Ext.jl

@@ -0,0 +1,14 @@
+module ITensorsHDF5Ext
+
+using HDF5: HDF5, attributes, create_group, open_group, read, write
+
+include("index.jl")
+include("itensor.jl")
+include("qnindex.jl")
+include("indexset.jl")
+include("qn.jl")
+include("tagset.jl")
+include("ITensorMPS/mps.jl")
+include("ITensorMPS/mpo.jl")
+
+end

ext/ITensorsHDF5Ext/index.jl

@@ -0,0 +1,42 @@
+using ITensors: Arrow, dim, dir, id, Index, plev, QNBlocks, space, tags, TagSet
+
+function HDF5.write(parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, I::Index)
+  g = create_group(parent, name)
+  attributes(g)["type"] = "Index"
+  attributes(g)["version"] = 1
+  write(g, "id", id(I))
+  write(g, "dim", dim(I))
+  write(g, "dir", Int(dir(I)))
+  write(g, "tags", tags(I))
+  write(g, "plev", plev(I))
+  if typeof(space(I)) == Int
+    attributes(g)["space_type"] = "Int"
+  elseif typeof(space(I)) == QNBlocks
+    attributes(g)["space_type"] = "QNBlocks"
+    write(g, "space", space(I))
+  else
+    error("Index space type not recognized")
+  end
+end
+
+function HDF5.read(parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{Index})
+  g = open_group(parent, name)
+  if read(attributes(g)["type"]) != "Index"
+    error("HDF5 group or file does not contain Index data")
+  end
+  id = read(g, "id")
+  dim = read(g, "dim")
+  dir = Arrow(read(g, "dir"))
+  tags = read(g, "tags", TagSet)
+  plev = read(g, "plev")
+  space_type = "Int"
+  if haskey(attributes(g), "space_type")
+    space_type = read(attributes(g)["space_type"])
+  end
+  if space_type == "Int"
+    space = dim
+  elseif space_type == "QNBlocks"
+    space = read(g, "space", QNBlocks)
+  end
+  return Index(id, space, dir, tags, plev)
+end

ext/ITensorsHDF5Ext/indexset.jl

@@ -0,0 +1,23 @@
+using ITensors: Index, Indices
+
+function HDF5.write(parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, is::Indices)
+  g = create_group(parent, name)
+  attributes(g)["type"] = "IndexSet"
+  attributes(g)["version"] = 1
+  N = length(is)
+  write(g, "length", N)
+  for n in 1:N
+    write(g, "index_$n", is[n])
+  end
+end
+
+function HDF5.read(
+  parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, T::Type{<:Indices}
+)
+  g = open_group(parent, name)
+  if read(attributes(g)["type"]) != "IndexSet"
+    error("HDF5 group or file does not contain IndexSet data")
+  end
+  n = read(g, "length")
+  return T(Index[read(g, "index_$j", Index) for j in 1:n])
+end

ext/ITensorsHDF5Ext/itensor.jl

@@ -0,0 +1,36 @@
+using ITensors: IndexSet, inds, itensor, ITensor, storage
+using NDTensors:
+  NDTensors, BlockSparse, Combiner, Dense, Diag, DiagBlockSparse, EmptyStorage
+
+function HDF5.write(parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, T::ITensor)
+  g = create_group(parent, name)
+  attributes(g)["type"] = "ITensor"
+  attributes(g)["version"] = 1
+  write(g, "inds", inds(T))
+  return write(g, "storage", storage(T))
+end
+
+function HDF5.read(
+  parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{ITensor}
+)
+  g = open_group(parent, name)
+  if read(attributes(g)["type"]) != "ITensor"
+    error("HDF5 group or file does not contain ITensor data")
+  end
+  # TODO: use Vector{Index} here?
+  inds = read(g, "inds", IndexSet)
+
+  # check input file for key name of ITensor data
+  # ITensors.jl <= v0.1.x uses `store` as key
+  # whereas ITensors.jl >= v0.2.x uses `storage` as key
+  for key in ["storage", "store"]
+    if haskey(g, key)
+      stypestr = read(attributes(open_group(g, key))["type"])
+      stype = eval(Meta.parse(stypestr))
+      storage = read(g, key, stype)
+      return itensor(storage, inds)
+    end
+  end
+  return error("HDF5 file: $(g) does not contain correct ITensor data.\nNeither key
+               `store` nor `storage` could be found.")
+end

ext/ITensorsHDF5Ext/qn.jl

@@ -0,0 +1,25 @@
+using ITensors: maxQNs, modulus, name, QN, QNVal, val
+
+function HDF5.write(parent::Union{HDF5.File,HDF5.Group}, gname::AbstractString, q::QN)
+  g = create_group(parent, gname)
+  attributes(g)["type"] = "QN"
+  attributes(g)["version"] = 1
+  names = [String(name(q[n])) for n in 1:maxQNs]
+  vals = [val(q[n]) for n in 1:maxQNs]
+  mods = [modulus(q[n]) for n in 1:maxQNs]
+  write(g, "names", names)
+  write(g, "vals", vals)
+  return write(g, "mods", mods)
+end
+
+function HDF5.read(parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{QN})
+  g = open_group(parent, name)
+  if read(attributes(g)["type"]) != "QN"
+    error("HDF5 group or file does not contain QN data")
+  end
+  names = read(g, "names")
+  vals = read(g, "vals")
+  mods = read(g, "mods")
+  mqn = ntuple(n -> QNVal(names[n], vals[n], mods[n]), maxQNs)
+  return QN(mqn)
+end

ext/ITensorsHDF5Ext/qnindex.jl

@@ -0,0 +1,29 @@
+using ITensors: dims, QNBlock, QNBlocks
+
+function HDF5.write(parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, B::QNBlocks)
+  g = create_group(parent, name)
+  attributes(g)["type"] = "QNBlocks"
+  attributes(g)["version"] = 1
+  write(g, "length", length(B))
+  dims = [block[2] for block in B]
+  write(g, "dims", dims)
+  for n in 1:length(B)
+    write(g, "QN[$n]", B[n][1])
+  end
+end
+
+function HDF5.read(
+  parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{QNBlocks}
+)
+  g = open_group(parent, name)
+  if read(attributes(g)["type"]) != "QNBlocks"
+    error("HDF5 group or file does not contain QNBlocks data")
+  end
+  N = read(g, "length")
+  dims = read(g, "dims")
+  B = QNBlocks(undef, N)
+  for n in 1:length(B)
+    B[n] = QNBlock(read(g, "QN[$n]", QN), dims[n])
+  end
+  return B
+end

ext/ITensorsHDF5Ext/tagset.jl

@@ -0,0 +1,19 @@
+using ITensors: TagSet, tagstring
+
+function HDF5.write(parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, T::TagSet)
+  g = create_group(parent, name)
+  attributes(g)["type"] = "TagSet"
+  attributes(g)["version"] = 1
+  return write(g, "tags", tagstring(T))
+end
+
+function HDF5.read(
+  parent::Union{HDF5.File,HDF5.Group}, name::AbstractString, ::Type{TagSet}
+)
+  g = open_group(parent, name)
+  if read(attributes(g)["type"]) != "TagSet"
+    error("HDF5 group '$name' does not contain TagSet data")
+  end
+  tstring = read(g, "tags")
+  return TagSet(tstring)
+end