From 3cd522243e10bdba0a5c2f40a7d8cbeea64f65d2 Mon Sep 17 00:00:00 2001
From: Joey <joeytindall94@googlemail.com>
Date: Thu, 8 Feb 2024 16:52:12 -0500
Subject: [PATCH] Ket <-> Bra, Added Link Map to QuadraticForm

---
 src/Forms/construct_form.jl | 12 ++++++++----
 src/Forms/quadraticform.jl  | 27 ++++++++++++++++++---------
 2 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/src/Forms/construct_form.jl b/src/Forms/construct_form.jl
index 6b84b5ca..e399137d 100644
--- a/src/Forms/construct_form.jl
+++ b/src/Forms/construct_form.jl
@@ -5,12 +5,13 @@ function Form(
   operator::ITensorNetwork,
   ket::ITensorNetwork;
   dual_space_map=default_index_map,
+  link_space_map=default_index_map,
   bra_vertex_map=default_bra_vertex_map,
   ket_vertex_map=default_ket_vertex_map,
   operator_vertex_map=default_operator_vertex_map,
   form_type=default_form_type,
 )
-  dual_map = map(dual_space_map, Indices(externalinds(bra)))
+  dual_map = map(dual_space_map, Indices(externalinds(ket)))
 
   @assert Set(externalinds(bra)) == Set(externalinds(ket))
   @assert Set(externalinds(operator)) ==
@@ -26,14 +27,17 @@ function Form(
   bra_renamed = rename_vertices_itn(bra, bra_vs_map)
   ket_renamed = rename_vertices_itn(ket, ket_vs_map)
   operator_renamed = rename_vertices_itn(operator, operator_vs_map)
-  ket_renamed_dual = map_inds(dual_space_map, ket_renamed; links=[])
+  bra_renamed_dual = map_inds(dual_space_map, bra_renamed; links=[])
 
-  tn = union(union(bra_renamed, operator_renamed), ket_renamed_dual)
+  tn = union(union(bra_renamed_dual, operator_renamed), ket_renamed)
 
   if form_type == "Bilinear"
     return BilinearForm(tn, bra_vs_map, ket_vs_map, operator_vs_map, dual_map)
   elseif form_type == "Quadratic"
-    return QuadraticForm(tn, bra_vs_map, ket_vs_map, operator_vs_map, dual_map)
+    link_map = map(link_space_map, Indices(internalinds(ket)))
+    return QuadraticForm(
+      tn, bra_vs_map, ket_vs_map, operator_vs_map, dual_map, link_space_map
+    )
   else
     return error("Form type not supported")
   end
diff --git a/src/Forms/quadraticform.jl b/src/Forms/quadraticform.jl
index b25babf5..0589ee84 100644
--- a/src/Forms/quadraticform.jl
+++ b/src/Forms/quadraticform.jl
@@ -1,10 +1,11 @@
-struct QuadraticForm{V,KetMap,BraMap,OperatorMap,SpaceMap} <:
-       AbstractForm{V,KetMap,BraMap,OperatorMap,SpaceMap}
+struct QuadraticForm{V,KetMap,BraMap,OperatorMap,VectorSpaceMap,LinkSpaceMap} <:
+       AbstractForm{V,KetMap,BraMap,OperatorMap,VectorSpaceMap}
   tn::ITensorNetwork{V}
   bra_vs_map::BraMap
   ket_vs_map::KetMap
   operator_vs_map::OperatorMap
-  dual_space_map::SpaceMap
+  dual_space_map::VectorSpaceMap
+  link_space_map::LinkSpaceMap
 end
 
 #Needed for implementation
@@ -17,22 +18,22 @@ data_graph_type(::Type{<:QuadraticForm}) = data_graph_type(tensornetwork(qf))
 data_graph(qf::QuadraticForm) = data_graph(tensornetwork(qf))
 
 function QuadraticForm(
-  bra::ITensorNetwork, operator::ITensorNetwork; link_space_map=default_index_map, kwargs...
+  ket::ITensorNetwork, operator::ITensorNetwork; link_space_map=default_index_map, kwargs...
 )
-  ket = link_space_map(dag(bra); sites=[])
-  return Form(bra, operator, ket; form_type="Quadratic", kwargs...)
+  bra = link_space_map(dag(ket); sites=[])
+  return Form(bra, operator, ket; form_type="Quadratic", link_space_map, kwargs...)
 end
 
 function QuadraticForm(
-  bra::ITensorNetwork;
+  ket::ITensorNetwork;
   dual_space_map=default_index_map,
   tno_constructor=default_tno_constructor,
   kwargs...,
 )
-  s = siteinds(bra)
+  s = siteinds(ket)
   operator_space = union_all_inds(s, dual_space_map(s; links=[]))
   operator = tno_constructor(operator_space)
-  return QuadraticForm(bra, operator; kwargs...)
+  return QuadraticForm(ket, operator; kwargs...)
 end
 
 function bra_ket_vertices(qf::QuadraticForm, state_vertices::Vector)
@@ -59,3 +60,11 @@ function gradient(
   reduced_tn, _ = induced_subgraph(tensornetwork(qf), tn_vertices)
   return contract(reduced_tn, contractor_kwargs...)
 end
+
+function update_qf!(qf::QuadraticForm, state_vertex, state::ITensor)
+  setindex_preserve_graph!(tensornetwork(qf)[ket_map(state_vertex)], state, state_vertex)
+  #Change to be consistent with dual space map and link space map. How can we do this elegantly??
+  return setindex_preserve_graph!(
+    tensornetwork(qf)[bra_map(state_vertex)], prime(dag(state)), state_vertex
+  )
+end