Fix tests

ITensor · Mar 28, 2024 · 502ac54 · 502ac54
1 parent ab07432
commit 502ac54
Show file tree

Hide file tree

Showing 5 changed files with 62 additions and 16 deletions.
diff --git a/test/Project.toml b/test/Project.toml
@@ -2,7 +2,6 @@
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
-DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 EinExprs = "b1794770-133b-4de1-afb4-526377e9f4c5"
 Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
@@ -19,6 +18,7 @@ Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b"
 NamedGraphs = "678767b0-92e7-4007-89e4-4527a8725b19"
 OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 Observers = "338f10d5-c7f1-4033-a7d1-f9dec39bcaa0"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"

diff --git a/...rnetworks/test_solvers/ITensorNetworksTestSolversUtils/ITensorNetworksTestSolversUtils.jl b/...rnetworks/test_solvers/ITensorNetworksTestSolversUtils/ITensorNetworksTestSolversUtils.jl
@@ -0,0 +1,3 @@
+module ITensorNetworksTestSolversUtils
+include("solvers.jl")
+end
diff --git a/test/test_treetensornetworks/test_solvers/ITensorNetworksTestSolversUtils/solvers.jl b/test/test_treetensornetworks/test_solvers/ITensorNetworksTestSolversUtils/solvers.jl
@@ -0,0 +1,37 @@
+using OrdinaryDiffEq: ODEProblem, Tsit5, solve
+using ITensors: ITensor
+using ITensorNetworks: TimeDependentSum, to_vec
+using KrylovKit: exponentiate
+
+function ode_solver(
+  H::TimeDependentSum,
+  time_step,
+  ψ₀;
+  current_time=0.0,
+  outputlevel=0,
+  solver_alg=Tsit5(),
+  kwargs...,
+)
+  if outputlevel ≥ 3
+    println("    In ODE solver, current_time = $current_time, time_step = $time_step")
+  end
+
+  time_span = (current_time, current_time + time_step)
+  u₀, itensor_from_vec = to_vec(ψ₀)
+  f(ψ::ITensor, p, t) = H(t)(ψ)
+  f(u::Vector, p, t) = to_vec(f(itensor_from_vec(u), p, t))[1]
+  prob = ODEProblem(f, u₀, time_span)
+  sol = solve(prob, solver_alg; kwargs...)
+  uₜ = sol.u[end]
+  return itensor_from_vec(uₜ), nothing
+end
+
+function krylov_solver(
+  H::TimeDependentSum, time_step, ψ₀; current_time=0.0, outputlevel=0, kwargs...
+)
+  if outputlevel ≥ 3
+    println("    In Krylov solver, current_time = $current_time, time_step = $time_step")
+  end
+  ψₜ, info = exponentiate(H(current_time), time_step, ψ₀; kwargs...)
+  return ψₜ, info
+end
diff --git a/test/test_treetensornetworks/test_solvers/Project.toml b/test/test_treetensornetworks/test_solvers/Project.toml
@@ -5,5 +5,6 @@ ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 NamedGraphs = "678767b0-92e7-4007-89e4-4527a8725b19"
 Observers = "338f10d5-c7f1-4033-a7d1-f9dec39bcaa0"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/test_treetensornetworks/test_solvers/test_tdvp_time_dependent.jl b/test/test_treetensornetworks/test_solvers/test_tdvp_time_dependent.jl
@@ -1,16 +1,20 @@
-using DifferentialEquations
-using ITensors
-using ITensorNetworks: NamedGraphs.AbstractNamedEdge
+@eval module $(gensym())
+using ITensors: contract
+using ITensorNetworks: ITensorNetworks, TimeDependentSum, TTN, mpo, mps, siteinds, tdvp
+using OrdinaryDiffEq: Tsit5
 using KrylovKit: exponentiate
-using LinearAlgebra
-using Test
-
-const ttn_solvers_examples_dir = joinpath(
-  pkgdir(ITensorNetworks), "examples", "treetensornetworks", "solvers"
+using LinearAlgebra: norm
+using NamedGraphs: AbstractNamedEdge, named_comb_tree
+using Test: @test, @test_broken, @testset
+
+include(
+  joinpath(
+    @__DIR__, "ITensorNetworksTestSolversUtils", "ITensorNetworksTestSolversUtils.jl"
+  ),
 )
 
-include(joinpath(ttn_solvers_examples_dir, "03_models.jl"))
-include(joinpath(ttn_solvers_examples_dir, "03_solvers.jl"))
+using .ITensorNetworksTestSolversUtils:
+  ITensorNetworksTestSolversUtils, krylov_solver, ode_solver
 
 # Functions need to be defined in global scope (outside
 # of the @testset macro)
@@ -40,7 +44,7 @@ function ode_updater(
 )
   region = first(sweep_plan[which_region_update])
   (; time_step, t) = internal_kwargs
-  t = isa(region, ITensorNetworks.NamedGraphs.AbstractNamedEdge) ? t : t + time_step
+  t = isa(region, AbstractNamedEdge) ? t : t + time_step
 
   H⃗₀ = projected_operator![]
   result, info = ode_solver(
@@ -64,7 +68,9 @@ end
 krylov_kwargs = (; tol=1e-8, krylovdim=15, eager=true)
 krylov_updater_kwargs = (; f=[f⃗], krylov_kwargs)
 
-function krylov_solver(H⃗₀, ψ₀; time_step, ishermitian=false, issymmetric=false, kwargs...)
+function ITensorNetworksTestSolversUtils.krylov_solver(
+  H⃗₀, ψ₀; time_step, ishermitian=false, issymmetric=false, kwargs...
+)
   psi_t, info = krylov_solver(
     -im * TimeDependentSum(f⃗, H⃗₀),
     time_step,
@@ -93,7 +99,7 @@ function krylov_updater(
   (; time_step, t) = internal_kwargs
   H⃗₀ = projected_operator![]
   region = first(sweep_plan[which_region_update])
-  t = isa(region, ITensorNetworks.NamedGraphs.AbstractNamedEdge) ? t : t + time_step
+  t = isa(region, AbstractNamedEdge) ? t : t + time_step
 
   result, info = krylov_solver(
     -im * TimeDependentSum(f, H⃗₀),
@@ -225,5 +231,4 @@ end
   @test ode_err < 1e-2
   @test krylov_err < 1e-2
 end
-
-nothing
+end