Make it easier to run example 3 on GPU (#22)

* Make it easier to run example 3 on GPU * Bump to v0.2.4
ITensor · Jun 2, 2024 · 740838b · 740838b · mtfishman · Jun 2, 2024
1 parent 3cc10c1
commit 740838b
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 15 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "ITensorMPS"
 uuid = "0d1a4710-d33b-49a5-8f18-73bdf49b47e2"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org>", "Miles Stoudenmire <mstoudenmire@flatironinstitute.org>"]
-version = "0.2.3"
+version = "0.2.4"
 
 [deps]
 ITensorTDVP = "25707e16-a4db-4a07-99d9-4d67b7af0342"

diff --git a/examples/03_tdvp_time_dependent.jl b/examples/03_tdvp_time_dependent.jl
@@ -6,7 +6,24 @@ using Random: Random
 include("03_models.jl")
 include("03_updaters.jl")
 
-function main()
+"""
+Run the example on CPU:
+```julia
+main()
+```
+
+Run the example on CPU with single precision:
+```julia
+main(; eltype=Float32)
+```
+
+Run the example on GPU:
+```julia
+using CUDA: cu
+main(; eltype=Float32, device=cu)
+```
+"""
+function main(; eltype=Float64, device=identity)
   Random.seed!(1234)
 
   # Time dependent Hamiltonian is:
@@ -24,16 +41,16 @@ function main()
   outputlevel = 3
 
   # Frequency of time dependent terms
-  ω₁ = 0.1
-  ω₂ = 0.2
+  ω₁ = one(eltype) / 10
+  ω₂ = one(eltype) / 5
 
   # Nearest and next-nearest neighbor
   # Heisenberg couplings.
-  J₁ = 1.0
-  J₂ = 1.0
+  J₁ = one(eltype)
+  J₂ = one(eltype)
 
-  time_step = 0.1
-  time_stop = 1.0
+  time_step = one(eltype) / 10
+  time_stop = one(eltype)
 
   # nsite-update TDVP
   nsite = 2
@@ -46,9 +63,9 @@ function main()
 
   # TDVP truncation parameters
   maxdim = 100
-  cutoff = 1e-8
+  cutoff = √(eps(eltype))
 
-  tol = 1e-15
+  tol = 10 * eps(eltype)
 
   @show n
   @show ω₁, ω₂
@@ -61,18 +78,20 @@ function main()
   f⃗ = map(ω -> (t -> cos(ω * t)), ω⃗)
 
   # H₀ = H(0) = H₁(0) + H₂(0) + …
-  ℋ₁₀ = heisenberg(n; J=J₁, J2=0.0)
-  ℋ₂₀ = heisenberg(n; J=0.0, J2=J₂)
+  ℋ₁₀ = heisenberg(n; J=J₁, J2=zero(eltype))
+  ℋ₂₀ = heisenberg(n; J=zero(eltype), J2=J₂)
   ℋ⃗₀ = (ℋ₁₀, ℋ₂₀)
 
   s = siteinds("S=1/2", n)
 
-  H⃗₀ = map(ℋ₀ -> MPO(ℋ₀, s), ℋ⃗₀)
+  H⃗₀ = map(ℋ₀ -> device(MPO(eltype, ℋ₀, s)), ℋ⃗₀)
 
   # Initial state, ψ₀ = ψ(0)
   # Initialize as complex since that is what OrdinaryDiffEq.jl/DifferentialEquations.jl
   # expects.
-  ψ₀ = complex.(random_mps(s, j -> isodd(j) ? "↑" : "↓"; linkdims=start_linkdim))
+  ψ₀ = device(
+    complex.(random_mps(eltype, s, j -> isodd(j) ? "↑" : "↓"; linkdims=start_linkdim))
+  )
 
   @show norm(ψ₀)
 

diff --git a/examples/03_updaters.jl b/examples/03_updaters.jl
@@ -9,7 +9,7 @@ function ode_updater(operator, init; internal_kwargs, alg=Tsit5(), kwargs...)
   time_span = typeof(time_step).((current_time, current_time + time_step))
   init_vec, to_itensor = to_vec(init)
   f(init::ITensor, p, t) = operator(t)(init)
-  f(init_vec::Vector, p, t) = to_vec(f(to_itensor(init_vec), p, t))[1]
+  f(init_vec::AbstractArray, p, t) = to_vec(f(to_itensor(init_vec), p, t))[1]
   prob = ODEProblem(f, init_vec, time_span)
   sol = solve(prob, alg; kwargs...)
   state_vec = sol.u[end]