From 1ce884db39daa2e5271220571b0bf2bed8ea8c88 Mon Sep 17 00:00:00 2001
From: Benedikt Kloss <bene.kloss@gmail.com>
Date: Fri, 2 Feb 2024 11:48:18 -0500
Subject: [PATCH] Fix factors and edit tests.

---
 .../projttns/abstractprojttn.jl               |  2 +-
 src/treetensornetworks/projttns/projttnsum.jl | 25 ++++++++++++-------
 .../test_solvers/test_contract.jl             | 21 ++++++++++------
 3 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/src/treetensornetworks/projttns/abstractprojttn.jl b/src/treetensornetworks/projttns/abstractprojttn.jl
index 7ac54106..fc8eeb7e 100644
--- a/src/treetensornetworks/projttns/abstractprojttn.jl
+++ b/src/treetensornetworks/projttns/abstractprojttn.jl
@@ -102,7 +102,7 @@ function Base.eltype(P::AbstractProjTTN)::Type
 end
 
 vertextype(::Type{<:AbstractProjTTN{V}}) where {V} = V
-vertextype(p::AbstractProjTTN) = typeof(p)
+vertextype(p::AbstractProjTTN) = vertextype(typeof(p))
 
 function Base.size(P::AbstractProjTTN)::Tuple{Int,Int}
   d = 1
diff --git a/src/treetensornetworks/projttns/projttnsum.jl b/src/treetensornetworks/projttns/projttnsum.jl
index bbe93759..b731b989 100644
--- a/src/treetensornetworks/projttns/projttnsum.jl
+++ b/src/treetensornetworks/projttns/projttnsum.jl
@@ -12,13 +12,13 @@ end
 terms(P::ProjTTNSum) = P.terms
 factors(P::ProjTTNSum) = P.factors
 
-copy(P::ProjTTNSum) = ProjTTNSum(copy.(terms(P)))
+copy(P::ProjTTNSum) = ProjTTNSum(copy.(terms(P)), copy(factors(P)))
 
 function ProjTTNSum(operators::Vector{<:AbstractProjTTN})
-  return ProjTTNSum(operators, fill(1, length(operators)))
+  return ProjTTNSum(operators, fill(one(Bool), length(operators)))
 end
 function ProjTTNSum(operators::Vector{<:AbstractTTN})
-  return ProjTTNSum(ProjTTN.(operators), fill(1, length(operators)))
+  return ProjTTNSum(ProjTTN.(operators))
 end
 
 on_edge(P::ProjTTNSum) = on_edge(terms(P)[1])
@@ -26,7 +26,7 @@ on_edge(P::ProjTTNSum) = on_edge(terms(P)[1])
 nsite(P::ProjTTNSum) = nsite(terms(P)[1])
 
 function set_nsite(Ps::ProjTTNSum, nsite)
-  return ProjTTNSum(map(p -> set_nsite(p, nsite), terms(Ps)))
+  return ProjTTNSum(map(p -> set_nsite(p, nsite), terms(Ps)), factors(Ps))
 end
 
 underlying_graph(P::ProjTTNSum) = underlying_graph(terms(P)[1])
@@ -41,15 +41,19 @@ internal_edges(P::ProjTTNSum) = internal_edges(terms(P)[1])
 
 product(P::ProjTTNSum, v::ITensor) = noprime(contract(P, v))
 
-contract(P::ProjTTNSum, v::ITensor) =
-  mapreduce(+, zip(factors(P), terms(P))) do (f, p)
+function contract(P::ProjTTNSum, v::ITensor)
+  res = mapreduce(+, zip(factors(P), terms(P))) do (f, p)
     f * contract(p, v)
   end
+  return res
+end
 
-contract_ket(P::ProjTTNSum, v::ITensor) =
-  mapreduce(+, zip(factors(P), terms(P))) do (f, p)
+function contract_ket(P::ProjTTNSum, v::ITensor)
+  res = mapreduce(+, zip(factors(P), terms(P))) do (f, p)
     f * contract_ket(p, v)
   end
+  return res
+end
 
 function Base.eltype(P::ProjTTNSum)
   return mapreduce(eltype, promote_type, terms(P))
@@ -60,5 +64,8 @@ end
 Base.size(P::ProjTTNSum) = size(terms(P)[1])
 
 function position(P::ProjTTNSum, psi::AbstractTTN, pos)
-  return ProjTTNSum(map(M -> position(M, psi, pos), terms(P)))
+  theterms = map(M -> position(M, psi, pos), terms(P))
+  #@show typeof(theterms)
+  #@show factors(P)
+  return ProjTTNSum(theterms, factors(P))
 end
diff --git a/test/test_treetensornetworks/test_solvers/test_contract.jl b/test/test_treetensornetworks/test_solvers/test_contract.jl
index 42e0f068..49f79e57 100644
--- a/test/test_treetensornetworks/test_solvers/test_contract.jl
+++ b/test/test_treetensornetworks/test_solvers/test_contract.jl
@@ -29,8 +29,8 @@ using Test
   Hpsi_via_dmrg = dmrg(Hfit, psi; updater_kwargs=(; which_eigval=:LR,), nsweeps=1)
   @test abs(inner(Hpsi_via_dmrg, Hpsi / norm(Hpsi))) ≈ 1 atol = 1E-4
   # Test whether the interface works for ProjTTNSum with factors
-  Hfit = ProjTTNSum([ProjOuterProdTTN(psi', H), ProjOuterProdTTN(psi', H)], [0.5, 0.5])
-  Hpsi_via_dmrg = dmrg(Hfit, psi; nsweeps=1, updater_kwargs=(; which_eigval=:LR,))
+  Hfit = ProjTTNSum([ProjOuterProdTTN(psi', H), ProjOuterProdTTN(psi', H)], [-0.2, -0.8])
+  Hpsi_via_dmrg = dmrg(Hfit, psi; nsweeps=1, updater_kwargs=(; which_eigval=:SR,))
   @test abs(inner(Hpsi_via_dmrg, Hpsi / norm(Hpsi))) ≈ 1 atol = 1E-4
 
   # Test basic usage for use with multiple ProjOuterProdTTN with default parameters
@@ -38,15 +38,20 @@ using Test
   os_id = OpSum()
   os_id += -1, "Id", 1, "Id", 2
   minus_identity = mpo(os_id, s)
+  os_id = OpSum()
+  os_id += +1, "Id", 1, "Id", 2
+  identity = mpo(os_id, s)
   Hpsi = ITensorNetworks.sum_apply(
-    [(H, psi), (minus_identity, psi)]; alg="fit", init=psi, nsweeps=1
+    [(H, psi), (minus_identity, psi)]; alg="fit", init=psi, nsweeps=3
   )
   @test inner(psi, Hpsi) ≈ (inner(psi', H, psi) - norm(psi)^2) atol = 1E-5
-
-  #Hpsi = ITensorNetworks.dmrg(
-  #  [(H, psi), (minus_identity, psi)]; alg="fit", init=psi, nsweeps=1
-  #)
-  #@test inner(psi, Hpsi) ≈ (inner(psi', H, psi) - norm(psi)^2) atol = 1E-5
+  # Test the above via DMRG
+  # ToDo: Investigate why this is broken
+  Hfit = ProjTTNSum([ProjOuterProdTTN(psi', H), ProjOuterProdTTN(psi', identity)], [-1, 1])
+  Hpsi_normalized = ITensorNetworks.dmrg(
+    Hfit, psi; nsweeps=3, updater_kwargs=(; which_eigval=:SR)
+  )
+  @test_broken abs(inner(Hpsi, (Hpsi_normalized) / norm(Hpsi))) ≈ 1 atol = 1E-5
 
   #
   # Change "top" indices of MPO to be a different set