Move code to ITensors.ITensorMPS

Project.toml

@@ -5,11 +5,10 @@ version = "0.2.1"
 
 [deps]
 ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5"
-Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 
 [compat]
 ITensors = "0.3.60"
-Reexport = "1"
 julia = "1.6"
 
 [extras]

src/ITensorTDVP.jl

@@ -1,9 +1,25 @@
 module ITensorTDVP
 
-using Reexport: @reexport
-@reexport using ITensors.ITensorMPS: tdvp, dmrg_x, to_vec, TimeDependentSum, linsolve
-
+using ITensors: ITensors, MPO
+using ITensors.NDTensors: @Algorithm_str
 using ITensors.ITensorMPS: ITensorMPS
-dmrg(args...; kwargs...) = ITensorMPS.alternating_update_dmrg(args...; kwargs...)
+using KrylovKit: KrylovKit
+
+tdvp(args...; kwargs...) = ITensorMPS.itensortdvp_tdvp(args...; kwargs...)
+
+function KrylovKit.linsolve(A::MPO, args...; kwargs...)
+  return ITensorMPS.itensortdvp_linsolve(A, args...; kwargs...)
+end
+
+dmrg(args...; kwargs...) = ITensorMPS.itensortdvp_dmrg(args...; kwargs...)
+
+dmrg_x(args...; kwargs...) = ITensorMPS.itensortdvp_dmrg_x(args...; kwargs...)
+
+function ITensors.contract(alg::Algorithm"fit", args...; kwargs...)
+  return ITensorMPS.itensortdvp_contract(alg, args...; kwargs...)
+end
+
+include("solver_utils.jl")
+export TimeDependentSum, to_vec
 
 end

src/solver_utils.jl

@@ -0,0 +1,69 @@
+using ITensors: ITensor, apply, array, inds, itensor, permute, ProjMPOSum
+
+# Utilities for making it easier
+# to define solvers (like ODE solvers)
+# for TDVP
+
+"""
+    to_vec(x)
+Transform `x` into a `Vector`, and return the vector, and a closure which inverts the
+transformation.
+
+Modeled after FiniteDifferences.to_vec:
+
+https://github.com/JuliaDiff/FiniteDifferences.jl/blob/main/src/to_vec.jl
+"""
+to_vec(x) = error("Not implemented")
+
+function to_vec(x::ITensor)
+  function to_itensor(x_vec)
+    return itensor(x_vec, inds(x))
+  end
+  return vec(array(x)), to_itensor
+end
+
+# Represents a time-dependent sum of terms:
+#
+# H(t) = f[1](t) * H0[1] + f[2](t) * H0[2] + …
+#
+struct TimeDependentSum{S,T}
+  f::Vector{S}
+  H0::T
+end
+TimeDependentSum(f::Vector, H0::ProjMPOSum) = TimeDependentSum(f, ITensors.terms(H0))
+Base.length(H::TimeDependentSum) = length(H.f)
+
+function Base.:*(c::Number, H::TimeDependentSum)
+  return TimeDependentSum([t -> c * fₙ(t) for fₙ in H.f], H.H0)
+end
+Base.:*(H::TimeDependentSum, c::Number) = c * H
+
+# Calling a `TimeDependentOpSum` at a certain time like:
+#
+# H(t)
+#
+# Returns a `ScaledSum` at that time.
+(H::TimeDependentSum)(t::Number) = ScaledSum([fₙ(t) for fₙ in H.f], H.H0)
+
+# Represents the sum of scaled terms:
+#
+# H = coefficient[1] * H[1] + coefficient * H[2] + …
+#
+struct ScaledSum{S,T}
+  coefficients::Vector{S}
+  H::T
+end
+Base.length(H::ScaledSum) = length(H.coefficients)
+
+# Apply the scaled sum of terms:
+#
+# H(ψ₀) = coefficient[1] * H[1](ψ₀) + coefficient[2] * H[2](ψ₀) + …
+#
+# onto ψ₀.
+function (H::ScaledSum)(ψ₀)
+  ψ = ITensor(inds(ψ₀))
+  for n in 1:length(H)
+    ψ += H.coefficients[n] * apply(H.H[n], ψ₀)
+  end
+  return permute(ψ, inds(ψ₀))
+end

test/test_examples.jl

@@ -4,11 +4,7 @@ using ITensorTDVP: ITensorTDVP
 using Test: @testset
 @testset "Run examples" begin
   examples_files = [
-    "01_tdvp.jl",
-    "02_dmrg-x.jl",
-    "03_tdvp_time_dependent.jl",
-    "04_tdvp_observers.jl",
-    "05_tdvp_nonuniform_timesteps.jl",
+    "01_tdvp.jl", "02_dmrg-x.jl", "03_tdvp_time_dependent.jl", "04_tdvp_observers.jl"
   ]
   examples_path = joinpath(pkgdir(ITensorTDVP), "examples")
   @testset "Running example file $f" for f in examples_files