Adapt new solver codes into ITensorNetworks module

Author: emstoudenmire

src/ITensorNetworks.jl

@@ -35,6 +35,23 @@ include("gauging.jl")
 include("utils.jl")
 include("update_observer.jl")
 
+include("solvers/local_solvers/eigsolve.jl")
+include("solvers/local_solvers/exponentiate.jl")
+include("solvers/local_solvers/runge_kutta.jl")
+include("solvers/truncation_parameters.jl")
+include("solvers/sweep_solve.jl")
+include("solvers/iterators.jl")
+include("solvers/region_plans/dfs_plans.jl")
+include("solvers/region_plans/euler_tour.jl")
+include("solvers/region_plans/euler_plans.jl")
+include("solvers/region_plans/tdvp_region_plans.jl")
+include("solvers/extracter.jl")
+include("solvers/inserter.jl")
+include("solvers/subspace/subspace.jl")
+include("solvers/subspace/densitymatrix.jl")
+include("solvers/eigsolve.jl")
+include("solvers/applyexp.jl")
+
 include("treetensornetworks/abstracttreetensornetwork.jl")
 include("treetensornetworks/treetensornetwork.jl")
 include("treetensornetworks/opsum_to_ttn/matelem.jl")

src/solvers/applyexp.jl

@@ -1,13 +1,13 @@
-import ITensorNetworks as itn
-using Printf
+using Printf: @printf
+import ConstructionBase: setproperties
 
 @kwdef mutable struct ApplyExpProblem{State}
   state::State
   operator
   current_time::Number = 0.0
 end
 
-state(tdvp::ApplyExpProblem) = tdvp.state
+ITensorNetworks.state(tdvp::ApplyExpProblem) = tdvp.state
 operator(tdvp::ApplyExpProblem) = tdvp.operator
 current_time(tdvp::ApplyExpProblem) = tdvp.current_time
 
@@ -31,11 +31,11 @@ function updater(
     curr_reg = current_region(region_iterator)
     next_reg = next_region(region_iterator)
     if !isnothing(next_reg) && next_reg != curr_reg
-      next_edge = first(itn.edge_sequence_between_regions(state(T), curr_reg, next_reg))
-      v1, v2 = itn.src(next_edge), itn.dst(next_edge)
+      next_edge = first(edge_sequence_between_regions(state(T), curr_reg, next_reg))
+      v1, v2 = src(next_edge), dst(next_edge)
       psi = copy(state(T))
       psi[v1], R = qr(local_state, uniqueinds(local_state, psi[v2]))
-      shifted_operator = itn.position(operator(T), psi, itn.NamedEdge(v1=>v2))
+      shifted_operator = position(operator(T), psi, NamedEdge(v1=>v2))
       R_t, _ = solver(x->optimal_map(shifted_operator, x), -time_step, R; kws...)
       local_state = psi[v1]*R_t
     end
@@ -53,7 +53,7 @@ function applyexp_sweep_printer(
   if outputlevel >= 1
     T = problem(region_iterator)
     @printf("  Current time = %s, ", process_time(current_time(T)))
-    @printf("maxlinkdim=%d", itn.maxlinkdim(state(T)))
+    @printf("maxlinkdim=%d", maxlinkdim(state(T)))
     println()
     flush(stdout)
   end
@@ -83,7 +83,7 @@ end
 
 function applyexp(H, init_state, exponents; kws...)
   init_prob = ApplyExpProblem(;
-    state=permute_indices(init_state), operator=itn.ProjTTN(permute_indices(H))
+    state=permute_indices(init_state), operator=ProjTTN(permute_indices(H))
   )
   return applyexp(init_prob, exponents; kws...)
 end

src/solvers/eigsolve.jl

@@ -1,5 +1,5 @@
-import ITensorNetworks as itn
-using Printf
+using Printf: @printf
+import ConstructionBase: setproperties
 
 @kwdef mutable struct EigsolveProblem{State,Operator}
   state::State
@@ -8,7 +8,7 @@ using Printf
 end
 
 eigenvalue(E::EigsolveProblem) = E.eigenvalue
-state(E::EigsolveProblem) = E.state
+ITensorNetworks.state(E::EigsolveProblem) = E.state
 operator(E::EigsolveProblem) = E.operator
 
 function updater(
@@ -36,7 +36,7 @@ function eigsolve_sweep_printer(region_iterator; outputlevel, sweep, nsweeps, kw
     end
     E = problem(region_iterator)
     @printf("eigenvalue=%.12f ", eigenvalue(E))
-    @printf("maxlinkdim=%d", itn.maxlinkdim(state(E)))
+    @printf("maxlinkdim=%d", maxlinkdim(state(E)))
     println()
     flush(stdout)
   end
@@ -68,7 +68,7 @@ end
 
 function eigsolve(H, init_state; kws...)
   init_prob = EigsolveProblem(;
-    state=permute_indices(init_state), operator=itn.ProjTTN(permute_indices(H))
+    state=permute_indices(init_state), operator=ProjTTN(permute_indices(H))
   )
   return eigsolve(init_prob; kws...)
 end

src/solvers/extracter.jl

@@ -3,15 +3,15 @@ import ConstructionBase: setproperties
 function extracter(problem, region_iterator; sweep, trunc=(;), kws...)
   trunc = truncation_parameters(sweep; trunc...)
   region = current_region(region_iterator)
-  psi = itn.orthogonalize(state(problem), region)
+  psi = orthogonalize(state(problem), region)
   local_state = prod(psi[v] for v in region)
   problem = setproperties(problem; state=psi)
 
   problem, local_state = subspace_expand(
     problem, local_state, region_iterator; sweep, trunc, kws...
   )
 
-  shifted_operator = itn.position(operator(problem), state(problem), region)
+  shifted_operator = position(operator(problem), state(problem), region)
 
   return setproperties(problem; operator=shifted_operator), local_state
 end

src/solvers/fitting.jl

@@ -21,13 +21,13 @@ function inserter(
     indsTe = it.inds(psi[first(region)])
     tags = it.tags(psi, e)
     U, C, _ = it.factorize(local_tensor, indsTe; tags, trunc...)
-    itn.@preserve_graph psi[first(region)] = U
+    @preserve_graph psi[first(region)] = U
   else
     error("Region of length $(length(region)) not currently supported")
   end
   v = last(region)
-  itn.@preserve_graph psi[v] = C
-  psi = set_orthogonal_region ? itn.set_ortho_region(psi, [v]) : psi
-  normalize && itn.@preserve_graph psi[v] = psi[v] / norm(psi[v])
+  @preserve_graph psi[v] = C
+  psi = set_orthogonal_region ? set_ortho_region(psi, [v]) : psi
+  normalize && @preserve_graph psi[v] = psi[v] / norm(psi[v])
   return setproperties(problem; state=psi)
 end

src/solvers/inserter.jl

@@ -1,42 +1,9 @@
-import ITensorNetworks as itn
 
-function optimal_map(P::itn.ProjTTN, ψ)
-  envs = [itn.environment(P, e) for e in itn.incident_edges(P)]
-  site_ops = [itn.operator(P)[s] for s in itn.sites(P)]
+function optimal_map(P::ProjTTN, ψ)
+  envs = [environment(P, e) for e in incident_edges(P)]
+  site_ops = [operator(P)[s] for s in sites(P)]
   contract_list = [envs..., site_ops..., ψ]
-  sequence = itn.contraction_sequence(contract_list; alg="optimal")
-  Pψ = itn.contract(contract_list; sequence)
+  sequence = contraction_sequence(contract_list; alg="optimal")
+  Pψ = contract(contract_list; sequence)
   return noprime(Pψ)
 end
-
-# This function is a workaround for the slow contraction order
-# heuristic in ITensorNetworks/src/treetensornetworks/projttns/projttn.jl
-# in the projected_operator_tensors(P::ProjTTN) function (line 97 or so)
-function operator_map(P::itn.ProjTTN, ψ)
-  ψ = copy(ψ)
-  if itn.on_edge(P)
-    for edge in itn.incident_edges(P)
-      ψ *= itn.environment(P, edge)
-    end
-  else
-    region = itn.sites(P)
-    ie = itn.incident_edges(P)
-    # TODO: improvement ideas
-    # - check which vertex (first(region) vs. last(region)
-    #   has more incident edges and contract those environments first
-    for edge in ie
-      if itn.dst(edge) == first(region)
-        ψ *= itn.environment(P, edge)
-      end
-    end
-    for s in itn.sites(P)
-      ψ *= itn.operator(P)[s]
-    end
-    for edge in ie
-      if itn.dst(edge) != first(region)
-        ψ *= itn.environment(P, edge)
-      end
-    end
-  end
-  return noprime(ψ)
-end

src/solvers/operator_map.jl

@@ -1,10 +1,8 @@
-import ITensorNetworks as itn
-using ITensors
 
 function permute_indices(tn)
-  si = itn.siteinds(tn)
+  si = siteinds(tn)
   ptn = copy(tn)
-  for v in itn.vertices(tn)
+  for v in vertices(tn)
     is = inds(tn[v])
     ls = setdiff(is, si[v])
     isempty(ls) && continue

src/solvers/permute_indices.jl

@@ -6,7 +6,9 @@ import NamedGraphs: GraphsExtensions
 #  return [fwd_sweep..., reverse(fwd_sweep)...]
 #end
 
-function tdvp_regions(g::AbstractGraph, time_step; nsites=1, updater_kwargs, sweep_kwargs...)
+function tdvp_regions(
+  g::AbstractGraph, time_step; nsites=1, updater_kwargs, sweep_kwargs...
+)
   @assert nsites==1
   fwd_up_args = (; time=(time_step / 2), updater_kwargs...)
   rev_up_args = (; time=(-time_step / 2), updater_kwargs...)
@@ -25,7 +27,7 @@ function tdvp_regions(g::AbstractGraph, time_step; nsites=1, updater_kwargs, swe
 end
 
 function overlap(ea::AbstractEdge, eb::AbstractEdge)
-  return intersect([src(ea),dst(ea)], [src(eb),dst(eb)])
+  return intersect([src(ea), dst(ea)], [src(eb), dst(eb)])
 end
 
 function forward_region(edges, which_edge; nsites=1, region_kwargs=(;))

src/solvers/region_plans.jl

@@ -1,9 +1,21 @@
-using ITensors: commonind, dag, dim, directsum, dot, hascommoninds, Index, norm, onehot, uniqueinds, random_itensor
+using ITensors:
+  commonind,
+  dag,
+  dim,
+  directsum,
+  dot,
+  hascommoninds,
+  Index,
+  norm,
+  onehot,
+  uniqueinds,
+  random_itensor
 
 # TODO: hoist num_expand default value out to a function or similar
-function subspace_expand!(problem::EigsolveProblem, local_tensor, region; prev_region, num_expand=4, kws...)
-
-  if isnothing(prev_region) || isa(region, AbstractEdge) 
+function subspace_expand!(
+  problem::EigsolveProblem, local_tensor, region; prev_region, num_expand=4, kws...
+)
+  if isnothing(prev_region) || isa(region, AbstractEdge)
     return local_tensor
   end
 
@@ -49,4 +61,4 @@ function subspace_expand!(problem::EigsolveProblem, local_tensor, region; prev_r
   local_tensor = prod(psi[v] for v in region)
 
   return local_tensor
-end 
+end