From b40d4d1f02bea695f0ae462161f13ef966eee2e1 Mon Sep 17 00:00:00 2001
From: Thien Nguyen <thiennguyen@nvidia.com>
Date: Tue, 13 Aug 2024 02:53:28 +0000
Subject: [PATCH] Add a property flag in kraus_op to identify whether it is a
 'scaled' unitary matrix

---
 runtime/common/NoiseModel.cpp          | 31 +++++++++++++++++++++++++
 runtime/common/NoiseModel.h            | 11 +++++++++
 unittests/integration/noise_tester.cpp | 32 ++++++++++++++++++++++++++
 3 files changed, 74 insertions(+)

diff --git a/runtime/common/NoiseModel.cpp b/runtime/common/NoiseModel.cpp
index 5cb37cfaf8..fbdf5ccf13 100644
--- a/runtime/common/NoiseModel.cpp
+++ b/runtime/common/NoiseModel.cpp
@@ -149,4 +149,35 @@ noise_model::get_channels(const std::string &quantumOp,
   return iter->second;
 }
 
+template <typename T>
+bool isScaledUnitaryImpl(const std::vector<std::complex<T>> &mat) {
+  typedef Eigen::Matrix<std::complex<T>, Eigen::Dynamic, Eigen::Dynamic,
+                        Eigen::RowMajor>
+      RowMajorMatTy;
+  const int dim = std::log2(mat.size());
+  Eigen::Map<const RowMajorMatTy> kMat(mat.data(), dim, dim);
+  if (kMat.isZero())
+    return true;
+  // Check that (K_dag * K) is a scaled identity matrix
+  // i.e., the K matrix is a scaled unitary.
+  auto kdK = kMat.adjoint() * kMat;
+  if (!kdK.isDiagonal())
+    return false;
+  // First element
+  std::complex<T> val = kdK(0, 0);
+  constexpr T eps = 1e-12;
+  if (std::abs(val) > eps) {
+    auto scaledKdK = (std::complex<T>{1.0} / val) * kdK;
+    return scaledKdK.isIdentity();
+  }
+  return false;
+}
+
+bool kraus_op::isScaledUnitary(const std::vector<std::complex<double>> &mat) {
+  return isScaledUnitaryImpl(mat);
+}
+
+bool kraus_op::isScaledUnitary(const std::vector<std::complex<float>> &mat) {
+  return isScaledUnitaryImpl(mat);
+}
 } // namespace cudaq
diff --git a/runtime/common/NoiseModel.h b/runtime/common/NoiseModel.h
index 4823134faf..37e1ba7db2 100644
--- a/runtime/common/NoiseModel.h
+++ b/runtime/common/NoiseModel.h
@@ -34,6 +34,9 @@ struct kraus_op {
   /// NOTE we currently assume nRows == nCols
   std::size_t nCols = 0;
 
+  /// @brief If true, the Kraus operator is a unitary operator times a constant.
+  bool unitary = false;
+
   /// @brief Copy constructor
   kraus_op(const kraus_op &) = default;
 
@@ -47,6 +50,7 @@ struct kraus_op {
 
     nRows = (std::size_t)std::round(sqrtNEl);
     nCols = nRows;
+    unitary = isScaledUnitary(data);
   }
 
   /// @brief Constructor, initialize from initializer_list
@@ -61,11 +65,13 @@ struct kraus_op {
 
     nRows = (std::size_t)std::round(sqrtNEl);
     nCols = nRows;
+    unitary = isScaledUnitary(data);
   }
 
   /// @brief Set this kraus_op equal to the other
   kraus_op &operator=(const kraus_op &other) {
     data = other.data;
+    unitary = other.unitary;
     return *this;
   }
 
@@ -78,6 +84,11 @@ struct kraus_op {
         newData[i * nRows + j] = std::conj(data[j * nCols + i]);
     return kraus_op(newData);
   }
+
+private:
+  /// @brief Return true if the input matrix is a scaled unitary matrix.
+  static bool isScaledUnitary(const std::vector<std::complex<double>> &mat);
+  static bool isScaledUnitary(const std::vector<std::complex<float>> &mat);
 };
 
 void validateCompletenessRelation_fp32(const std::vector<kraus_op> &ops);
diff --git a/unittests/integration/noise_tester.cpp b/unittests/integration/noise_tester.cpp
index 6d67613950..236d272b20 100644
--- a/unittests/integration/noise_tester.cpp
+++ b/unittests/integration/noise_tester.cpp
@@ -290,4 +290,36 @@ CUDAQ_TEST(NoiseTest, checkPhaseFlipType) {
   cudaq::unset_noise(); // clear for subsequent tests
 }
 
+CUDAQ_TEST(NoiseTest, checkNoiseModelUtils) {
+  {
+    cudaq::depolarization_channel depol(0.1);
+    for (const auto &op : depol.get_ops())
+      EXPECT_TRUE(op.unitary);
+  }
+  {
+    cudaq::depolarization_channel depol(1.0);
+    for (const auto &op : depol.get_ops())
+      EXPECT_TRUE(op.unitary);
+  }
+  {
+    cudaq::bit_flip_channel bf(0.1);
+    for (const auto &op : bf.get_ops())
+      EXPECT_TRUE(op.unitary);
+  }
+  {
+    cudaq::phase_flip_channel pf(0.1);
+    for (const auto &op : pf.get_ops())
+      EXPECT_TRUE(op.unitary);
+  }
+  {
+    // Amplitude damping has a non-unitary Kraus matrix.
+    cudaq::amplitude_damping_channel ad(0.1);
+    const auto krausOps = ad.get_ops();
+    const bool hasNonUnitaryOp =
+        std::any_of(krausOps.begin(), krausOps.end(),
+                    [](const auto &op) { return !op.unitary; });
+    EXPECT_TRUE(hasNonUnitaryOp);
+  }
+}
+
 #endif