Subspace expansion

Project.toml

@@ -54,12 +54,12 @@ DocStringExtensions = "0.8, 0.9"
 EinExprs = "0.6.4"
 Graphs = "1.8"
 GraphsFlows = "0.1.1"
-ITensors = "0.3.58"
+ITensors = "0.4"
 IsApprox = "0.1"
 IterTools = "1.4.0"
 KrylovKit = "0.6, 0.7"
-NamedGraphs = "0.1.23"
 NDTensors = "0.2, 0.3"
+NamedGraphs = "0.1.23"
 Observers = "0.2"
 PackageExtensionCompat = "1"
 Requires = "1.3"

src/ITensorNetworks.jl

@@ -42,6 +42,7 @@ include("solvers/local_solvers/exponentiate.jl")
 include("solvers/local_solvers/dmrg_x.jl")
 include("solvers/local_solvers/contract.jl")
 include("solvers/local_solvers/linsolve.jl")
+include("solvers/local_solvers/expand.jl")
 include("treetensornetworks/abstracttreetensornetwork.jl")
 include("treetensornetworks/ttn.jl")
 include("treetensornetworks/opsum_to_ttn.jl")
@@ -53,6 +54,9 @@ include("solvers/solver_utils.jl")
 include("solvers/defaults.jl")
 include("solvers/insert/insert.jl")
 include("solvers/extract/extract.jl")
+include("solvers/subspace_expansions/linalg/standard_svd.jl")
+#include("solvers/subspace_expansions/linalg/rsvd_aux.jl")
+#include("solvers/subspace_expansions/linalg/rsvd_linalg.jl")
 include("solvers/alternating_update/alternating_update.jl")
 include("solvers/alternating_update/region_update.jl")
 include("solvers/tdvp.jl")

src/solvers/defaults.jl

@@ -7,6 +7,8 @@ default_extracter() = default_extracter
 default_inserter() = default_inserter
 default_checkdone() = (; kws...) -> false
 default_transform_operator() = nothing
+default_scale_cutoff_by_timestep(cutoff;time_step, kwargs...) = isinf(time_step) ? cutoff : cutoff / abs(time_step)
+default_scale_maxdim(maxdim;kwargs...) = Int(ceil(2^(1.0/3.0) * maxdim))
 function default_region_printer(;
   cutoff,
   maxdim,

src/solvers/extract/extract.jl

@@ -24,3 +24,29 @@ function default_extracter(state, projected_operator, region, ortho; internal_kw
   projected_operator = position(projected_operator, state, region)
   return state, projected_operator, local_tensor
 end
+
+function extract_and_truncate(state,projected_operator, region, ortho; maxdim=nothing, cutoff=nothing,internal_kwargs)
+  #ToDo: move default into function signature?
+  #svd_kwargs = NamedTuple()
+  #svd_kwargs = isnothing(maxdim) ? svd_kwargs : merge(svd_kwargs,(;maxdim))
+  #svd_kwargs = isnothing(cutoff) ? svd_kwargs : merge(svd_kwargs,(;cutoff))
+  svd_kwargs= (;
+  (isnothing(maxdim) ? (;) : (;maxdim))...,
+  (isnothing(cutoff) ? (;) : (;cutoff))...
+  )
+  state = orthogonalize(state, ortho)
+  if isa(region, AbstractEdge)
+    other_vertex = only(setdiff(support(region), [ortho]))
+    left_inds = uniqueinds(state[ortho], state[other_vertex])
+    #ToDo: replace with call to factorize
+    U, S, V = svd(
+      state[ortho], left_inds; lefttags=tags(state, region), righttags=tags(state, region), svd_kwargs...
+    )
+    state[ortho] = U
+    local_tensor = S * V
+  else
+    local_tensor = prod(state[v] for v in region)
+  end
+  projected_operator = position(projected_operator, state, region)
+  return state, projected_operator, local_tensor
+end

src/solvers/local_solvers/expand.jl

@@ -0,0 +1,207 @@
+#ToDo: is this the correct scaling for infinite timestep? DMRG will not have infinite timestep, unless we set it explicitly
+#ToDo: implement this as a closure, to be constructed in sweep_plan (where time_step is known)
+#scale_cutoff_by_timestep(cutoff;time_step) = isinf(time_step) ? cutoff : cutoff / abs(time_step)
+
+function two_site_expansion_updater(
+    init;
+    state!,
+    projected_operator!,
+    outputlevel,
+    which_sweep,
+    sweep_plan,
+    which_region_update,
+    internal_kwargs,
+    svd_func_expand=ITensorNetworks.rsvd_iterative,
+    maxdim,
+    maxdim_func= (arg;kwargs...) -> identity(arg),
+    cutoff,
+    cutoff_func= (arg;kwargs...) -> identity(arg),
+    use_relative_cutoff=false,
+    use_absolute_cutoff=true,
+  )
+    maxdim=maxdim_func(maxdim;internal_kwargs...)
+    cutoff=cutoff_func(cutoff;internal_kwargs...)
+    region = first(sweep_plan[which_region_update])
+    typeof(region) <: NamedEdge && return init, (;)
+    region = only(region)
+    # figure out next region, since we want to expand there
+    # ToDo: account for non-integer indices into which_region_update
+    next_region = if which_region_update == length(sweep_plan)
+      nothing
+    else
+      first(sweep_plan[which_region_update + 1])
+    end
+    previous_region =
+      which_region_update == 1 ? nothing : first(sweep_plan[which_region_update - 1])
+    isnothing(next_region) && return init, (;)
+    !(typeof(next_region) <: NamedEdge) && return init, (;)
+    !(region == src(next_region) || region == dst(next_region)) && return init, (;)
+    next_vertex = src(next_region) == region ? dst(next_region) : src(next_region)
+
+    phi, has_changed = _two_site_expand_core(
+      init, region, region => next_vertex; 
+      projected_operator!,
+      state!,
+      svd_func_expand,
+      maxdim,
+      cutoff,
+      use_relative_cutoff,
+      use_absolute_cutoff,
+    )
+    !has_changed && return init, (;)
+    return phi, (;)
+  end
+
+  """
+Returns a function which applies Nullspace projector to an ITensor with matching indices without
+explicitly constructing the projector as an ITensor or Network.
+"""
+function implicit_nullspace(A, linkind)
+  # only works when applied in the direction of the environment tensor, not the other way (due to use of ITensorNetworkMap which is not reversible)
+  outind = uniqueinds(A, linkind)
+  inind = outind  #?
+  # ToDo: implement without ITensorNetworkMap
+  P = ITensors.ITensorNetworkMaps.ITensorNetworkMap(
+    [prime(dag(A), linkind), prime(A, linkind)], inind, outind
+  )
+  return x::ITensor -> x - P(x)
+end
+
+"""
+Performs a local expansion using the two-site projected variance.
+The input state should have orthogonality center on a single vertex (region), and phi0 is that site_tensors.
+Expansion is performed on vertex that would be visited next (if next vertex is a neighbor of region). 
+"""
+function _two_site_expand_core(
+  phi0,
+  region,
+  vertexpair::Pair;
+  projected_operator!,
+  state!,
+  svd_func_expand,
+  cutoff,
+  maxdim,
+  use_relative_cutoff,
+  use_absolute_cutoff,
+)
+  # preliminaries
+  theflux = flux(phi0)
+  v1 = first(vertexpair)
+  v2 = last(vertexpair)
+  verts = [v1, v2]
+  PH = copy(projected_operator![])
+  state = copy(state![])
+  old_linkdim = dim(commonind(state[v1],state[v2]))
+  (old_linkdim >= maxdim) && return phi0, false
+  # orthogonalize on the edge
+  #update of environment not strictly necessary here
+  state, PH, phi = default_extracter(state, PH, edgetype(state)(vertexpair), v1;
+                                  internal_kwargs=(;))
+  psis = map(n -> state[n], verts)  # extract local site tensors
+  linkinds = map(psi -> commonind(psi, phi), psis)
+
+  # compute nullspace to the left and right 
+  #@timeit_debug timer "nullvector" begin
+    nullVecs = map(zip(psis, linkinds)) do (psi, linkind)
+      #return nullspace(psi, linkind; atol=atol)
+      return ITensorNetworks.implicit_nullspace(psi, linkind)
+    end
+  #end
+
+  # update the projected operator
+  #ToDo: may not be necessary if we use the extracter anyway
+  PH = set_nsite(PH, 2)
+  PH = position(PH, state, verts)
+
+  # build environments
+  g = underlying_graph(PH)
+  #@timeit_debug timer "build environments" begin
+    envs = map(zip(verts, psis)) do (n, psi)
+      return noprime(
+        mapreduce(*, [v => n for v in neighbors(g, n) if !(v ∈ verts)]; init=psi) do e
+          return PH.environments[NamedEdge(e)]
+        end * PH.operator[n],
+      )
+    end
+  #end
+
+  # apply the projection into nullspace,
+  # FIXME: think through again whether we really want to evaluate null space projector already
+  envs=[ first(nullVecs)(first(envs)),last(nullVecs)(last(envs))]
+  ininds = uniqueinds(last(psis), phi)
+  outinds = uniqueinds(first(psis), phi)
+  cin = combiner(ininds)
+  cout = combiner(outinds)
+  envMap=[cout * envs[1], phi* (cin * envs[2])]
+
+  # factorize
+  #FIXME: remove specialization on svd_func, define interface for svd_funcs
+  #@timeit_debug timer "svd_func" begin
+  U, S, V = svd_func_expand(
+        envMap,
+        uniqueinds(inds(cout), outinds);
+        maxdim=maxdim - old_linkdim,
+        cutoff=cutoff,
+        use_relative_cutoff,
+        use_absolute_cutoff,
+        )
+  #end
+
+  # catch cases when we decompose a map of norm==0.0
+  (isnothing(U) || iszero(norm(U))) && return phi0, false
+  #FIXME: somehow the svd funcs sometimes return empty ITensors instead of nothing, that should be caught in the SVD routines instead...
+  all(isempty.([U, S, V])) && return phi0, false #ToDo: do not rely on isempty here
+
+  # uncombine indices on the non-link-indices
+  U *= dag(cout)
+  V *= dag(cin)
+
+  # direct sum the site tensors
+  #@timeit_debug timer "direct sum" begin
+    new_psis = map(zip(psis, [U, V])) do (psi, exp_basis)
+      return ITensors.directsum(
+        psi => commonind(psi, phi),
+        exp_basis => uniqueind(exp_basis, psi);
+        tags=tags(commonind(psi, phi)),
+      )
+    end
+  #end
+  new_inds = [last(x) for x in new_psis]
+  new_psis = [first(x) for x in new_psis]
+
+  # extract the expanded linkinds from expanded site tensors and create a zero-tensor
+  phi_indices = replace(
+    inds(phi), (commonind(phi, psis[n]) => dag(new_inds[n]) for n in 1:2)...
+  )
+  if hasqns(phi)
+    new_phi = ITensor(eltype(phi), flux(phi), phi_indices...)
+    #ToDo: Remove? This used to be necessary, but may have been fixed with bugfixes in ITensor
+    #fill!(new_phi, 0.0)
+  else
+    new_phi = ITensor(eltype(phi), phi_indices...)
+  end
+
+  # set the new bond tensor elements from the old bond tensor
+  map(eachindex(phi)) do I
+    v = phi[I]
+    !iszero(v) && (return new_phi[I] = v) # I think this line errors without the fill! with zeros above
+  end
+
+  # apply combiners on linkinds #ToDo: figure out why this is strictly necessary 
+  combiners = map(
+    new_ind -> combiner(new_ind; tags=tags(new_ind), dir=dir(new_ind)), new_inds
+  )
+  for (v, new_psi, C) in zip(verts, new_psis, combiners)
+    state[v] = noprime(new_psi * C)
+  end
+  new_phi = dag(first(combiners)) * new_phi * dag(last(combiners))
+
+  # apply expanded bond tensor to site tensor and reset projected operator to region
+  state[v1] *= new_phi
+  new_phi = state[v1]
+  PH = set_nsite(PH, 1)
+  PH = position(PH, state, [region])
+  projected_operator![] = PH
+  state![] = state
+  return new_phi, true
+end

src/solvers/solver_utils.jl

@@ -87,3 +87,27 @@ function cache_operator_to_disk(
   end
   return operator
 end
+
+#ToDo: Move? This belongs more into local_solvers
+function compose_updaters(;kwargs...)
+  funcs=values(kwargs)
+  kwarg_symbols=Symbol.(keys(kwargs),"_kwargs")
+  info_symbols=Symbol.(keys(kwargs),"_info")
+
+  function composed_updater(init; kwargs...)
+    info=(;)
+    kwargs_for_updaters=map(x -> kwargs[x], kwarg_symbols)
+    other_kwargs=Base.structdiff((;kwargs...),NamedTuple(kwarg_symbols .=> kwargs_for_updaters))
+
+    for (func,kwargs_for_updater,info_symbol) in zip(funcs,kwargs_for_updaters,info_symbols)
+      init, new_info=func(init;
+      other_kwargs...,
+      kwargs_for_updater...
+      )
+      #figure out a better way to handle composing the info?
+      info=(;info...,NamedTuple((info_symbol=>new_info,))...)
+    end
+    return init, info
+  end
+  return composed_updater
+end