From e41ddf82a4e359f6cad3dbd3b147fe08ac25a15c Mon Sep 17 00:00:00 2001
From: Sergei Isakov <iserge@google.com>
Date: Fri, 30 Oct 2020 20:23:28 +0100
Subject: [PATCH] Fix unitary calculator.

---
 lib/unitary_calculator_basic.h         |  3 ++
 tests/BUILD                            |  3 ++
 tests/unitary_calculator_basic_test.cc |  4 +++
 tests/unitary_calculator_testfixture.h | 47 +++++++++++++++++++++++---
 4 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/lib/unitary_calculator_basic.h b/lib/unitary_calculator_basic.h
index 47ee61c5..013dfb54 100644
--- a/lib/unitary_calculator_basic.h
+++ b/lib/unitary_calculator_basic.h
@@ -74,6 +74,9 @@ class UnitaryCalculatorBasic final {
   using Unitary = typename UnitarySpace::Unitary;
   using fp_type = typename UnitarySpace::fp_type;
 
+  using StateSpace = UnitarySpace;
+  using State = Unitary;
+
   template <typename... ForArgs>
   explicit UnitaryCalculatorBasic(unsigned num_qubits, ForArgs&&... args)
       : for_(args...), num_qubits_(num_qubits) {}
diff --git a/tests/BUILD b/tests/BUILD
index 44007663..ed8d4a79 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -303,6 +303,9 @@ cc_library(
     }),
     deps = [
         "@com_google_googletest//:gtest_main",
+        "//lib:fuser",
+        "//lib:gate_appl",
+        "//lib:gates_cirq",
     ],
     testonly = 1,
 )
diff --git a/tests/unitary_calculator_basic_test.cc b/tests/unitary_calculator_basic_test.cc
index f0eed5b3..5494a161 100644
--- a/tests/unitary_calculator_basic_test.cc
+++ b/tests/unitary_calculator_basic_test.cc
@@ -33,6 +33,10 @@ TEST(UnitaryCalculatorTest, ApplyGate2) {
   TestApplyGate2<UnitaryCalculatorBasic<For, float>>();
 }
 
+TEST(UnitaryCalculatorTest, ApplyFusedGate) {
+  TestApplyFusedGate<UnitaryCalculatorBasic<For, float>>();
+}
+
 }  // namspace
 }  // namespace unitary
 }  // namespace qsim
diff --git a/tests/unitary_calculator_testfixture.h b/tests/unitary_calculator_testfixture.h
index c4e86b55..1a79b352 100644
--- a/tests/unitary_calculator_testfixture.h
+++ b/tests/unitary_calculator_testfixture.h
@@ -19,6 +19,10 @@
 
 #include "gtest/gtest.h"
 
+#include "../lib/fuser.h"
+#include "../lib/gate_appl.h"
+#include "../lib/gates_cirq.h"
+
 namespace qsim {
 
 namespace unitary {
@@ -28,8 +32,8 @@ namespace {
 template <typename UnitarySpace, typename Unitary>
 void FillMatrix(UnitarySpace& us, Unitary& u, int n) {
   // Intentionally create non-unitary matrix with ascending elements.
-  for(int i =0; i < (1 << n); i++){
-    for(int j =0;j < (1 << n); j++) {
+  for(int i = 0; i < (1 << n); i++) {
+    for(int j = 0; j < (1 << n); j++) {
       us.SetEntry(u, i, j, 2 * i * (1 << n) + 2 * j,
         2 * i * (1 << n) + 2 * j + 1);
     }
@@ -38,8 +42,6 @@ void FillMatrix(UnitarySpace& us, Unitary& u, int n) {
 
 }  // namespace
 
-constexpr char provider[] = "unitary_calculator_test";
-
 template <typename UC>
 void TestApplyGate1() {
   const int n_qubits = 3;
@@ -222,6 +224,43 @@ void TestApplyGate2() {
   }
 }
 
+template <typename UnitaryCalculator>
+void TestApplyFusedGate() {
+  using UnitarySpace = typename UnitaryCalculator::UnitarySpace;
+  using Unitary = typename UnitaryCalculator::Unitary;
+  using fp_type = typename UnitaryCalculator::fp_type;
+  using Gate = Cirq::GateCirq<fp_type>;
+
+  unsigned num_qubits = 1;
+
+  UnitaryCalculator uc(num_qubits, 1);
+  UnitarySpace us(num_qubits, 1);
+
+  Unitary u = us.CreateUnitary();
+  us.SetIdentity(u);
+
+  std::vector<Gate> gates = {Cirq::H<fp_type>::Create(0, 0),
+                             Cirq::H<fp_type>::Create(1, 0)};
+
+  GateFused<Gate> fgate{Cirq::kH, 0, {0}, &gates[0], {&gates[0], &gates[1]}};
+
+  ApplyFusedGate(uc, fgate, u);
+
+  unsigned size = 1 << num_qubits;
+  for (unsigned i = 0; i < size; ++i) {
+    for (unsigned j = 0; j < size; ++j) {
+      auto a = us.GetEntry(u, i, j);
+      if (i == j) {
+        EXPECT_NEAR(std::real(a), 1, 1e-6);
+      } else {
+        EXPECT_NEAR(std::real(a), 0, 1e-6);
+      }
+
+      EXPECT_NEAR(std::imag(a), 0, 1e-6);
+    }
+  }
+}
+
 }  // namespace unitary
 }  // namespace qsim