From dee8c3bd0e73c1b8257a2d10e9d6331992db9701 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?=
 <15837247+mofeing@users.noreply.github.com>
Date: Fri, 31 Jan 2025 06:46:46 +0100
Subject: [PATCH] Add conversion between `ITensorNetwork` and `Ansatz` (#305)

* Add conversion between `ITensorNetwork` and `Ansatz`

* Fix typo
---
 Project.toml                             |  5 +++-
 ext/TenetITensorNetworksExt.jl           | 37 +++++++++++++++++++++---
 test/integration/ITensorNetworks_test.jl | 20 +++++++++++--
 3 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/Project.toml b/Project.toml
index f1b02d304..5d3b6e343 100644
--- a/Project.toml
+++ b/Project.toml
@@ -25,6 +25,7 @@ ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
+DataGraphs = "b5a273c3-7e6c-41f6-98bd-8d7f1525a36a"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 GraphMakie = "1ecd5474-83a3-4783-bb4f-06765db800d2"
 ITensorMPS = "0d1a4710-d33b-49a5-8f18-73bdf49b47e2"
@@ -32,6 +33,7 @@ ITensorNetworks = "2919e153-833c-4bdc-8836-1ea460a35fc7"
 ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+NamedGraphs = "678767b0-92e7-4007-89e4-4527a8725b19"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Quac = "b9105292-1415-45cf-bff1-d6ccf71e6143"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
@@ -46,7 +48,7 @@ TenetDaggerExt = "Dagger"
 TenetFiniteDifferencesExt = "FiniteDifferences"
 TenetGraphMakieExt = ["GraphMakie", "Makie"]
 TenetITensorMPSExt = ["ITensors", "ITensorMPS"]
-TenetITensorNetworksExt = "ITensorNetworks"
+TenetITensorNetworksExt = ["ITensorNetworks", "ITensors", "NamedGraphs", "DataGraphs"]
 TenetITensorsExt = "ITensors"
 TenetKrylovKitExt = ["KrylovKit"]
 TenetPythonCallExt = "PythonCall"
@@ -75,6 +77,7 @@ ITensors = "0.6"
 KrylovKit = "0.7, 0.8, 0.9"
 LinearAlgebra = "1.10"
 Makie = "0.18,0.19,0.20, 0.21, 0.22"
+NamedGraphs = "0.6.4"
 OMEinsum = "0.7, 0.8"
 PythonCall = "0.9"
 Quac = "0.3"
diff --git a/ext/TenetITensorNetworksExt.jl b/ext/TenetITensorNetworksExt.jl
index 48dcd855d..3a0ed82ab 100644
--- a/ext/TenetITensorNetworksExt.jl
+++ b/ext/TenetITensorNetworksExt.jl
@@ -1,9 +1,12 @@
 module TenetITensorNetworksExt
 
 using Tenet
-using ITensorNetworks: ITensorNetworks, ITensorNetwork, ITensor, Index, siteinds, plev, vertices, rename_vertices
-const ITensors = ITensorNetworks.ITensors
-const DataGraphs = ITensorNetworks.DataGraphs
+using ITensorNetworks: ITensorNetworks, ITensorNetwork, IndsNetwork, plev, rename_vertices
+using ITensors: ITensors, ITensor, Index, siteinds
+using BijectiveDicts: BijectiveDict
+using Graphs: vertices
+using DataGraphs: DataGraphs, underlying_graph
+using NamedGraphs: NamedGraph, position_graph
 const TenetITensorsExt = Base.get_extension(Tenet, :TenetITensorsExt)
 
 Base.convert(::Type{TensorNetwork}, tn::ITensorNetwork) = TensorNetwork([convert(Tensor, tn[v]) for v in vertices(tn)])
@@ -35,8 +38,34 @@ function Base.convert(::Type{ITensorNetwork}, tn::Tenet.AbstractQuantum)
         indices = map(x -> Symbol(replace(x, "\"" => "")), string.(ITensors.tags.(ITensors.inds(itensor))))
         tensor = only(tensors(tn; contains=indices))
         physical_index = only(inds(tn; set=:physical) ∩ inds(tensor))
-        return sites(tn; at=physical_index).id
+        return Tenet.id(sites(tn; at=physical_index))
     end
 end
 
+function Base.convert(::Type{Lattice}, s::IndsNetwork)
+    # NOTE they don't suffer from the issue of removing a vertex from SimpleGraph (i.e. vertex renaming), because their vertice list keeps ordered
+    # TODO maybe we can use Vector for the mapping
+    ng = underlying_graph(s) # namedgraph
+    g = position_graph(ng) # simplegraph
+
+    mapping = BijectiveDict{Lane,Int}()
+    for (i, v) in enumerate(vertices(ng))
+        mapping[Lane(v)] = i
+    end
+
+    return Lattice(mapping, copy(g))
+end
+
+function Base.convert(::Type{IndsNetwork}, l::Lattice)
+    return NamedGraph(copy(l.graph), Tenet.id.(vertices(l)))
+end
+
+function Base.convert(::Type{Ansatz}, tn::ITensorNetwork)
+    return Ansatz(convert(Quantum, tn), convert(Lattice, siteinds(tn)))
+end
+
+function Base.convert(::Type{ITensorNetwork}, tn::Tenet.AbstractAnsatz)
+    return @invoke convert(ITensorNetwork, tn::Tenet.AbstractQuantum)
+end
+
 end
diff --git a/test/integration/ITensorNetworks_test.jl b/test/integration/ITensorNetworks_test.jl
index e86a14ffb..9f41d631b 100644
--- a/test/integration/ITensorNetworks_test.jl
+++ b/test/integration/ITensorNetworks_test.jl
@@ -1,6 +1,8 @@
 @testset "ITensorNetworks" begin
+    using Tenet
+    using Graphs
     using ITensors: ITensors, ITensor, Index, array
-    using ITensorNetworks: ITensorNetwork, vertices
+    using ITensorNetworks: ITensorNetwork
 
     i = Index(2, "i")
     j = Index(3, "j")
@@ -21,5 +23,19 @@
     @test itn isa ITensorNetwork
     @test issetequal(map(v -> array(itn[v]), vertices(itn)), array.([a, b, c]))
 
-    # TODO test Quantum
+    tn = convert(Quantum, itn)
+    @test tn isa Quantum
+    @test issetequal(arrays(tn), array.([a, b, c]))
+
+    itn = convert(ITensorNetwork, tn)
+    @test itn isa ITensorNetwork
+    @test issetequal(map(v -> array(itn[v]), vertices(itn)), array.([a, b, c]))
+
+    tn = convert(Ansatz, itn)
+    @test tn isa Ansatz
+    @test issetequal(arrays(tn), array.([a, b, c]))
+
+    itn = convert(ITensorNetwork, tn)
+    @test itn isa ITensorNetwork
+    @test issetequal(map(v -> array(itn[v]), vertices(itn)), array.([a, b, c]))
 end