Add write-to-disk during DMRG with SerializedElementArrays.jl (ITensor#648)

mtfishman · achuchmala · commit ab3dfcedcc1d · 2022-04-11T13:36:38.000+02:00
* Add write-to-disk during DMRG with SerializedElementArrays. Introduces `AbstractProjMPO` supertype for `ProjMPO` and new type `DiskProjMPO` that saves the environment tensors on disk. New keyword argument for `dmrg` called `write_when_maxdim_exceeds` to enable.
diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml
@@ -14,7 +14,7 @@ jobs:
       - name: Install dependencies
         run: julia -e 'using Pkg; pkg"add PkgBenchmark BenchmarkCI@0.1"'
       - name: Run benchmarks
-        run: julia -e 'using BenchmarkCI; BenchmarkCI.judge()'
+        run: julia -e 'using BenchmarkCI; BenchmarkCI.judge(; baseline="origin/main")'
       - name: Post results
         run: julia -e 'using BenchmarkCI; BenchmarkCI.postjudge()'
         env:
diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml
@@ -3,7 +3,7 @@ name: Documentation
 on:
   push:
     branches:
-      - master
+      - main
     tags: '*'
   pull_request:
 
diff --git a/.github/workflows/format_check.yml b/.github/workflows/format_check.yml
@@ -3,7 +3,7 @@ name: format-check
 on:
   push:
     branches:
-      - 'master'
+      - 'main'
       - 'release-'
     tags: '*'
   pull_request:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -2,7 +2,7 @@ name: Tests
 on:
   push:
     branches:
-      - master
+      - main
     tags: '*'
   pull_request:
 jobs:
diff --git a/.gitignore b/.gitignore
@@ -1,5 +1,6 @@
 *.o
 *.swp
+.tmp
 benchmark/mult
 Manifest.toml
 test/data.h5
diff --git a/Project.toml b/Project.toml
@@ -14,6 +14,7 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+SerializedElementArrays = "d3ce8812-9567-47e9-a7b5-65a6d70a3065"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
@@ -26,6 +27,7 @@ HDF5 = "0.14,0.15"
 KrylovKit = "0.4.2, 0.5"
 PackageCompiler = "1.0.0"
 Requires = "1.1"
+SerializedElementArrays = "0.1"
 StaticArrays = "0.12, 1.0"
 Strided = "1.1"
 TimerOutputs = "0.5.5"
diff --git a/docs/src/AdvancedUsageGuide.md b/docs/src/AdvancedUsageGuide.md
@@ -35,8 +35,8 @@ because new releases may be breaking.
 
 To try the "development branch" of ITensors.jl (for example, if 
 there is a feature or fix we added that hasn't been released yet), 
-you can do `add ITensors#master`. You can switch back to the latest
-released version with `add ITensors`. Using the development/master
+you can do `add ITensors#main`. You can switch back to the latest
+released version with `add ITensors`. Using the development/main
 branch is generally not encouraged unless you know what you are doing.
 
 ## Using ITensors.jl in the REPL
@@ -732,7 +732,7 @@ process more automated in the future.
 
 Julia has great built-in tools for benchmarking and profiling.
 For benchmarking fast code at the command line, you can use
-[BenchmarkTools](https://github.com/JuliaCI/BenchmarkTools.jl/blob/master/doc/manual.md):
+[BenchmarkTools](https://github.com/JuliaCI/BenchmarkTools.jl/blob/main/doc/manual.md):
 ```julia
 julia> using ITensors;
 
diff --git a/docs/src/examples/DMRG.md b/docs/src/examples/DMRG.md
@@ -167,7 +167,7 @@ some helper functions which
 return an array of bonds. Each bond object has an
 "s1" field and an "s2" field which are the integers numbering
 the two sites the bond connects.
-(You can view the source for these functions at [this link](https://github.com/ITensor/ITensors.jl/blob/master/src/physics/lattices.jl).)
+(You can view the source for these functions at [this link](https://github.com/ITensor/ITensors.jl/blob/main/src/physics/lattices.jl).)
 
 The two provided functions currently are `square_lattice` and 
 `triangular_lattice`. It is not hard to write your own similar lattice
diff --git a/docs/src/examples/Physics.md b/docs/src/examples/Physics.md
@@ -226,11 +226,11 @@ operator names into AutoMPO and it will know how to use these operators.
 **Further Steps**
 
 See how the built-in site types are defined inside the ITensor library:
-* [S=1/2 sites](https://github.com/ITensor/ITensors.jl/blob/master/src/physics/site_types/spinhalf.jl)
-* [S=1 sites](https://github.com/ITensor/ITensors.jl/blob/master/src/physics/site_types/spinone.jl)
-* [Fermion sites](https://github.com/ITensor/ITensors.jl/blob/master/src/physics/site_types/fermion.jl)
-* [Electron sites](https://github.com/ITensor/ITensors.jl/blob/master/src/physics/site_types/electron.jl)
-* [tJ sites](https://github.com/ITensor/ITensors.jl/blob/master/src/physics/site_types/tj.jl)
+* [S=1/2 sites](https://github.com/ITensor/ITensors.jl/blob/main/src/physics/site_types/spinhalf.jl)
+* [S=1 sites](https://github.com/ITensor/ITensors.jl/blob/main/src/physics/site_types/spinone.jl)
+* [Fermion sites](https://github.com/ITensor/ITensors.jl/blob/main/src/physics/site_types/fermion.jl)
+* [Electron sites](https://github.com/ITensor/ITensors.jl/blob/main/src/physics/site_types/electron.jl)
+* [tJ sites](https://github.com/ITensor/ITensors.jl/blob/main/src/physics/site_types/tj.jl)
 
 
 
diff --git a/docs/src/index.md b/docs/src/index.md
@@ -6,7 +6,7 @@
 
 |**Build Status**                                                                                |
 :-----------------------------------------------------------------------------------------------:|
-| [![Tests](https://github.com/ITensor/ITensors.jl/workflows/Tests/badge.svg)](https://github.com/ITensor/ITensors.jl/actions?query=workflow%3ATests) [![codecov](https://codecov.io/gh/ITensor/ITensors.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/ITensor/ITensors.jl) |
+| [![Tests](https://github.com/ITensor/ITensors.jl/workflows/Tests/badge.svg)](https://github.com/ITensor/ITensors.jl/actions?query=workflow%3ATests) [![codecov](https://codecov.io/gh/ITensor/ITensors.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/ITensor/ITensors.jl) |
 
 |**Citation**                                                                    |
 |:-------------------------------------------------------------------------------:|
@@ -84,7 +84,7 @@ If you use ITensors.jl in your work, for now please cite the [arXiv preprint](ht
 ## Full Example Codes
 
 The ITensors.jl package contains a directory of examples, which we
-will continue to add to. You can find them online [here](https://github.com/ITensor/ITensors.jl/tree/master/examples).
+will continue to add to. You can find them online [here](https://github.com/ITensor/ITensors.jl/tree/main/examples).
 Additionally, once you have installed ITensors.jl you can find a local version
 of the examples in the directory `ITensors.examples_dir()`, and you can run them
 as follows from the Julia REPL:
diff --git a/examples/contraction_sequence_optimization/main.jl b/examples/contraction_sequence_optimization/main.jl
@@ -256,15 +256,7 @@ function mera_2d_4_to_1(dχ; indlabels=indlabels)
   sequence_inds, sequence_cost = contraction_cost(network, sequence)
   @show hassameinds(sequence_inds, noncommoninds(network...))
   @show sequence_cost
-  @show 4dχ^26 +
-        2dχ^25 +
-        2dχ^23 +
-        3dχ^22 +
-        3dχ^20 +
-        dχ^16 +
-        dχ^14 +
-        dχ^13 +
-        dχ^12 +
+  @show 4dχ^26 + 2dχ^25 + 2dχ^23 + 3dχ^22 + 3dχ^20 + dχ^16 + dχ^14 + dχ^13 + dχ^12 +
         4dχ^8 +
         4dχ^7
   return (sequence=sequence, sequence_labels=sequence_labels, sequence_cost=sequence_cost)
diff --git a/examples/dmrg/write_to_disk/1d_heisenberg.jl b/examples/dmrg/write_to_disk/1d_heisenberg.jl
@@ -0,0 +1,43 @@
+using ITensors
+using ITensors.Strided
+using LinearAlgebra
+using Printf
+using Random
+
+Random.seed!(1234)
+Strided.set_num_threads(1)
+
+let
+  N = 100
+
+  sites = siteinds("S=1", N; conserve_qns=true)
+
+  ampo = AutoMPO()
+  for j in 1:(N - 1)
+    ampo .+= 0.5, "S+", j, "S-", j + 1
+    ampo .+= 0.5, "S-", j, "S+", j + 1
+    ampo .+= "Sz", j, "Sz", j + 1
+  end
+  H = MPO(ampo, sites)
+
+  # This step makes the MPO more sparse.
+  # It generally improves DMRG performance
+  # at large bond dimensions but makes DMRG slower at
+  # small bond dimensions.
+  H = splitblocks(linkinds, H)
+
+  state = [isodd(n) ? "Up" : "Dn" for n in 1:N]
+  psi0 = randomMPS(sites, state, 10)
+
+  # Plan to do 5 DMRG sweeps:
+  sweeps = Sweeps(5)
+  # Set maximum MPS bond dimensions for each sweep
+  setmaxdim!(sweeps, 10, 20, 100, 100, 200)
+  # Set maximum truncation error allowed when adapting bond dimensions
+  setcutoff!(sweeps, 1E-10)
+  @show sweeps
+
+  # Run the DMRG algorithm, returning energy and optimized MPS
+  energy, psi = dmrg(H, psi0, sweeps; write_when_maxdim_exceeds=25)
+  @printf("Final energy = %.12f\n", energy)
+end
diff --git a/examples/src/ctmrg_anisotropic.jl b/examples/src/ctmrg_anisotropic.jl
@@ -65,11 +65,7 @@ function calc_κ(iy, ix, T, (Clu, Cru, Cld, Crd), (Al, Ar, Au, Ad); dir="left")
     Clu[iy, ix] * Cld[iy, ix] * Au[iy, ix] * Ad[iy, ix] * Cru[iy, ixp] * Crd[iy, ixp]
   )
   normT = scalar(
-    Clu[iy, ix] *
-    Al[iy, ix] *
-    Cld[iyp, ix] *
-    Au[iy, ix] *
-    T[iy, ix] *
+    Clu[iy, ix] * Al[iy, ix] * Cld[iyp, ix] * Au[iy, ix] * T[iy, ix] *
     Ad[iyp, ix] *
     Cru[iy, ixp] *
     Ar[iy, ixp] *
@@ -148,24 +144,14 @@ function leftright_move!(T, (Clu, Cru, Cld, Crd), (Al, Ar, Au, Ad); dir="left",
       iyp = per(iy + 1, ny)
 
       Cu =
-        Al[iym, ix] *
-        Clu[iym, ix] *
-        Au[iym, ix] *
-        T[iym, ix] *
-        Au[iym, ixp] *
+        Al[iym, ix] * Clu[iym, ix] * Au[iym, ix] * T[iym, ix] * Au[iym, ixp] *
         T[iym, ixp] *
-        Cru[iym, ixpp] *
-        Ar[iym, ixpp]
+        Cru[iym, ixpp] * Ar[iym, ixpp]
       @assert order(Cu) == 4
       Cd =
-        Al[iy, ix] *
-        Cld[iyp, ix] *
-        Ad[iyp, ix] *
-        T[iy, ix] *
-        Ad[iyp, ixp] *
+        Al[iy, ix] * Cld[iyp, ix] * Ad[iyp, ix] * T[iy, ix] * Ad[iyp, ixp] *
         T[iy, ixp] *
-        Crd[iyp, ixpp] *
-        Ar[iy, ixpp]
+        Crd[iyp, ixpp] * Ar[iy, ixpp]
       @assert order(Cd) == 4
       if dir == "left" || dir == "up"
         li = commonindex(Cru[iy, ixpp], Crd[iy, ixpp])
diff --git a/examples/src/trg.jl b/examples/src/trg.jl
@@ -42,9 +42,7 @@ function trg(T::ITensor; χmax::Int, nsteps::Int, cutoff=0.0, svd_alg="divide_an
     Fᵥ′ *= δ(dag(s̃ᵥ), s̃ᵥ')
 
     T =
-      (Fₕ * δ(dag(sₕ'), sₕ)) *
-      (Fᵥ * δ(dag(sᵥ'), sᵥ)) *
-      (Fₕ′ * δ(dag(sₕ), sₕ')) *
+      (Fₕ * δ(dag(sₕ'), sₕ)) * (Fᵥ * δ(dag(sᵥ'), sᵥ)) * (Fₕ′ * δ(dag(sₕ), sₕ')) *
       (Fᵥ′ * δ(dag(sᵥ), sᵥ'))
 
     sₕ, sᵥ = s̃ₕ, s̃ᵥ
diff --git a/src/ITensors.jl b/src/ITensors.jl
@@ -18,6 +18,7 @@ using PackageCompiler
 using Pkg
 using Printf
 using Random
+using SerializedElementArrays
 using StaticArrays
 using TimerOutputs
 
@@ -102,7 +103,9 @@ include("mps/deprecated.jl")
 include("mps/mps.jl")
 include("mps/mpo.jl")
 include("mps/sweeps.jl")
+include("mps/abstractprojmpo.jl")
 include("mps/projmpo.jl")
+include("mps/diskprojmpo.jl")
 include("mps/projmposum.jl")
 include("mps/projmps.jl")
 include("mps/projmpo_mps.jl")
diff --git a/src/imports.jl b/src/imports.jl
@@ -145,3 +145,9 @@ import ITensors.NDTensors:
   store
 
 import Random: randn!
+
+using SerializedElementArrays: SerializedElementVector
+
+const DiskVector{T} = SerializedElementVector{T}
+
+import SerializedElementArrays: disk
diff --git a/src/mps/abstractprojmpo.jl b/src/mps/abstractprojmpo.jl
diff --git a/src/mps/diskprojmpo.jl b/src/mps/diskprojmpo.jl
diff --git a/src/mps/dmrg.jl b/src/mps/dmrg.jl
diff --git a/src/mps/projmpo.jl b/src/mps/projmpo.jl
diff --git a/src/physics/sitetype.jl b/src/physics/sitetype.jl
diff --git a/test/dmrg.jl b/test/dmrg.jl
