From afb85bb8f88f0f25a7b21bfdae1f5bfb9342bd9b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jofre=20Vall=C3=A8s=20Muns?=
 <61060572+jofrevalles@users.noreply.github.com>
Date: Tue, 1 Oct 2024 09:18:30 +0200
Subject: [PATCH] Fix `Yao` extension by properly transposing matrices (#212)

* Transpose the matrix from Yao Blocks

* Enhance Yao extension tests

* Format code from tests

* Add two-qubit gate test

* Change permutation of indices instead of array

* Format tests
---
 ext/TenetYaoExt.jl                 |  2 +-
 test/integration/YaoBlocks_test.jl | 42 ++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)
 create mode 100644 test/integration/YaoBlocks_test.jl

diff --git a/ext/TenetYaoExt.jl b/ext/TenetYaoExt.jl
index 03a0a946a..0b68da28b 100644
--- a/ext/TenetYaoExt.jl
+++ b/ext/TenetYaoExt.jl
@@ -33,7 +33,7 @@ function Tenet.Quantum(circuit::AbstractBlock)
             map(occupied_locs(gate)) do l
                 from, to = last(wire[l]), Tenet.nextindex!(gen)
                 push!(wire[l], to)
-                (from, to)
+                (to, from)
             end,
         )
 
diff --git a/test/integration/YaoBlocks_test.jl b/test/integration/YaoBlocks_test.jl
new file mode 100644
index 000000000..2b9e11ca8
--- /dev/null
+++ b/test/integration/YaoBlocks_test.jl
@@ -0,0 +1,42 @@
+@testset "YaoBlocks" begin
+    using YaoBlocks
+
+    # NOTE qubit #3 left empty on purpose
+    circuit = chain(3, put(1 => X), cnot(1, 2))
+    tn = Quantum(circuit)
+
+    @test issetequal(sites(tn), [site"1", site"2", site"1'", site"2'"])
+    @test Tenet.ntensors(tn) == 2
+
+    @testset "GHZ Circuit" begin
+        circuit_GHZ = chain(n_qubits, put(1 => Yao.H), Yao.control(1, 2 => Yao.X), Yao.control(2, 3 => Yao.X))
+
+        quantum_circuit = Quantum(circuit_GHZ)
+
+        zeros = Quantum(Product(fill([1, 0], n_qubits))) #|000>
+        ones = Quantum(Product(fill([0, 1], n_qubits))) #|111>
+
+        expected_value = Tenet.contract(merge(zeros, quantum_circuit, ones')) # <111|circuit|000>
+        @test only(expected_value) ≈ 1 / √2
+
+        SV_Yao = apply!(zero_state(n_qubits), circuit_GHZ) # circuit|000>
+        @test only(statevec(ArrayReg(bit"111"))' * statevec(SV_Yao)) ≈ 1 / √2
+    end
+
+    @testset "two-qubit gate" begin
+        U = matblock(rand(ComplexF64, 4, 4); tag="U")
+        circuit = chain(2, put((1, 2) => U))
+        psi = zero_state(2)
+        apply!(psi, circuit)
+
+        quantum_circuit = Quantum(circuit)
+        zeros = Quantum(Product(fill([1, 0], 2))) #|00>
+        ones = Quantum(Product(fill([0, 1], 2))) #|11>
+
+        expected_value = Tenet.contract(merge(zeros, quantum_circuit, ones')) # <11|circuit|00>
+
+        SV_Yao = apply!(zero_state(2), circuit) # circuit|00>
+
+        @test only(expected_value) ≈ only(statevec(ArrayReg(bit"11"))' * statevec(SV_Yao))
+    end
+end