Refactor noise model for unitary mixtures

runtime/common/NoiseModel.cpp

@@ -10,8 +10,73 @@
 #include "Logger.h"
 #include "common/CustomOp.h"
 #include "common/EigenDense.h"
+#include <numeric>
+#include <optional>
+
 namespace cudaq {
 
+// Helper to check whether a matrix is a scaled unitary matrix, i.e., `k * U`
+// where U is a unitary matrix. If so, it also returns the `k` factor.
+// Otherwise, return a nullopt.
+static std::optional<double>
+isScaledUnitary(const std::vector<std::complex<double>> &mat, double eps) {
+  typedef Eigen::Matrix<std::complex<double>, 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 0.0;
+  // 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 std::nullopt;
+  // First element
+  std::complex<double> val = kdK(0, 0);
+  if (std::abs(val) > eps && std::abs(val.imag()) < eps) {
+    auto scaledKdK = (std::complex<double>{1.0} / val) * kdK;
+    if (scaledKdK.isIdentity())
+      return val.real();
+  }
+  return std::nullopt;
+}
+
+// Helper to determine if a vector of Kraus ops are actually a unitary mixture.
+// If so, it returns all the unitaries and the probabilities associated with
+// each one of those unitaries.
+static std::optional<std::pair<std::vector<double>,
+                               std::vector<std::vector<std::complex<double>>>>>
+computeUnitaryMixture(
+    const std::vector<std::vector<std::complex<double>>> &krausOps,
+    double tol = 1e-6) {
+  std::vector<double> probs;
+  std::vector<std::vector<std::complex<double>>> mats;
+  const auto scaleMat = [](const std::vector<std::complex<double>> &mat,
+                           double scaleFactor) {
+    std::vector<std::complex<double>> scaledMat = mat;
+    // If scaleFactor is 0, then it means the original matrix was likely all
+    // zeros. In that case, the probability will be 0, so the matrix doesn't
+    // matter, but we don't want NaNs to trickle in anywhere.
+    if (scaleFactor != 0.0)
+      for (auto &x : scaledMat)
+        x /= scaleFactor;
+    return scaledMat;
+  };
+  for (const auto &op : krausOps) {
+    const auto scaledFactor = isScaledUnitary(op, tol);
+    if (!scaledFactor.has_value())
+      return std::nullopt;
+    probs.emplace_back(scaledFactor.value());
+    mats.emplace_back(scaleMat(op, std::sqrt(scaledFactor.value())));
+  }
+
+  if (std::abs(1.0 - std::reduce(probs.begin(), probs.end())) > tol)
+    return std::nullopt;
+
+  return std::make_pair(probs, mats);
+}
+
 template <typename EigenMatTy>
 bool isIdentity(const EigenMatTy &mat, double threshold = 1e-9) {
   EigenMatTy idMat = EigenMatTy::Identity(mat.rows(), mat.cols());
@@ -78,9 +143,51 @@ void validateCompletenessRelation_fp64(const std::vector<kraus_op> &ops) {
         "Provided kraus_ops are not completely positive and trace preserving.");
 }
 
+void generateUnitaryParameters_fp32(
+    const std::vector<kraus_op> &ops,
+    std::vector<std::vector<std::complex<double>>> &unitary_ops,
+    std::vector<double> &probabilities) {
+  std::vector<std::vector<std::complex<double>>> double_kraus_ops;
+  double_kraus_ops.reserve(ops.size());
+  for (auto &op : ops) {
+    // WARNING: danger here. We are intentially treating the incoming op as fp32
+    // type instead of what the compiler thinks it is (fp64). We have to do this
+    // because this file is compiled with cudaq::real = fp64, but the incoming
+    // data for this specific routine is actually fp32.
+    const std::complex<float> *ptr =
+        reinterpret_cast<const std::complex<float> *>(op.data.data());
+    double_kraus_ops.emplace_back(
+        std::vector<std::complex<double>>(ptr, ptr + op.data.size()));
+  }
+
+  auto asUnitaryMixture = computeUnitaryMixture(double_kraus_ops);
+  if (asUnitaryMixture.has_value()) {
+    probabilities = std::move(asUnitaryMixture.value().first);
+    unitary_ops = std::move(asUnitaryMixture.value().second);
+  }
+}
+
+void generateUnitaryParameters_fp64(
+    const std::vector<kraus_op> &ops,
+    std::vector<std::vector<std::complex<double>>> &unitary_ops,
+    std::vector<double> &probabilities) {
+  std::vector<std::vector<std::complex<double>>> double_kraus_ops;
+  double_kraus_ops.reserve(ops.size());
+  for (auto &op : ops)
+    double_kraus_ops.emplace_back(
+        std::vector<std::complex<double>>(op.data.begin(), op.data.end()));
+
+  auto asUnitaryMixture = computeUnitaryMixture(double_kraus_ops);
+  if (asUnitaryMixture.has_value()) {
+    probabilities = std::move(asUnitaryMixture.value().first);
+    unitary_ops = std::move(asUnitaryMixture.value().second);
+  }
+}
+
 kraus_channel::kraus_channel(const kraus_channel &other)
     : ops(other.ops), noise_type(other.noise_type),
-      parameters(other.parameters) {}
+      parameters(other.parameters), unitary_ops(other.unitary_ops),
+      probabilities(other.probabilities) {}
 
 std::size_t kraus_channel::size() const { return ops.size(); }
 
@@ -94,6 +201,8 @@ kraus_channel &kraus_channel::operator=(const kraus_channel &other) {
   ops = other.ops;
   noise_type = other.noise_type;
   parameters = other.parameters;
+  unitary_ops = other.unitary_ops;
+  probabilities = other.probabilities;
   return *this;
 }
 

runtime/common/NoiseModel.h

@@ -102,6 +102,12 @@ struct kraus_op {
 
 void validateCompletenessRelation_fp32(const std::vector<kraus_op> &ops);
 void validateCompletenessRelation_fp64(const std::vector<kraus_op> &ops);
+void generateUnitaryParameters_fp32(
+    const std::vector<kraus_op> &ops,
+    std::vector<std::vector<std::complex<double>>> &, std::vector<double> &);
+void generateUnitaryParameters_fp64(
+    const std::vector<kraus_op> &ops,
+    std::vector<std::vector<std::complex<double>>> &, std::vector<double> &);
 
 /// @brief A kraus_channel represents a quantum noise channel
 /// on specific qubits. The action of the noise channel is
@@ -143,7 +149,17 @@ class kraus_channel {
   // corruption.
   std::vector<double> parameters;
 
-  ~kraus_channel() = default;
+  /// @brief If all Kraus ops are - when scaled - unitary, this holds the
+  /// unitary versions of those ops. These values are always "double" regardless
+  /// of whether cudaq::real is float or double.
+  std::vector<std::vector<std::complex<double>>> unitary_ops;
+
+  /// @brief If all Kraus ops are - when scaled - unitary, this holds the
+  /// probabilities of those ops. These values are always "double" regardless
+  /// of whether cudaq::real is float or double.
+  std::vector<double> probabilities;
+
+  virtual ~kraus_channel() = default;
 
   /// @brief The nullary constructor
   kraus_channel() = default;
@@ -158,12 +174,14 @@ class kraus_channel {
   kraus_channel(std::initializer_list<T> &&...inputLists) {
     (ops.emplace_back(std::move(inputLists)), ...);
     validateCompleteness();
+    generateUnitaryParameters();
   }
 
   /// @brief The constructor, take qubits and channel kraus_ops as lvalue
   /// reference
   kraus_channel(const std::vector<kraus_op> &inOps) : ops(inOps) {
     validateCompleteness();
+    generateUnitaryParameters();
   }
 
   /// @brief The constructor, take qubits and channel kraus_ops as rvalue
@@ -189,6 +207,23 @@ class kraus_channel {
 
   /// @brief Add a kraus_op to this channel.
   void push_back(kraus_op op);
+
+  /// @brief Returns whether or not this is a unitary mixture.
+  bool is_unitary_mixture() const { return !unitary_ops.empty(); }
+
+  /// @brief Checks if Kraus ops have unitary representations and saves them if
+  /// they do. Users should only need to call this if they have modified the
+  /// Kraus ops and want to recompute these values.
+  void generateUnitaryParameters() {
+    unitary_ops.clear();
+    probabilities.clear();
+    if constexpr (std::is_same_v<cudaq::complex::value_type, float>) {
+      generateUnitaryParameters_fp32(ops, this->unitary_ops,
+                                     this->probabilities);
+      return;
+    }
+    generateUnitaryParameters_fp64(ops, this->unitary_ops, this->probabilities);
+  }
 };
 
 /// @brief The noise_model type keeps track of a set of
@@ -381,6 +416,7 @@ class depolarization_channel : public kraus_channel {
     this->parameters.push_back(probability);
     noise_type = noise_model_type::depolarization_channel;
     validateCompleteness();
+    generateUnitaryParameters();
   }
 };
 
@@ -396,6 +432,8 @@ class amplitude_damping_channel : public kraus_channel {
     this->parameters.push_back(probability);
     noise_type = noise_model_type::amplitude_damping_channel;
     validateCompleteness();
+    // Note: amplitude damping is non-unitary, so there is no value in calling
+    // generateUnitaryParameters().
   }
 };
 
@@ -412,6 +450,7 @@ class bit_flip_channel : public kraus_channel {
     this->parameters.push_back(probability);
     noise_type = noise_model_type::bit_flip_channel;
     validateCompleteness();
+    generateUnitaryParameters();
   }
 };
 
@@ -429,6 +468,7 @@ class phase_flip_channel : public kraus_channel {
     this->parameters.push_back(probability);
     noise_type = noise_model_type::phase_flip_channel;
     validateCompleteness();
+    generateUnitaryParameters();
   }
 };
 } // namespace cudaq