Adding template parameter HandlerTy for operator_sum and formatting

Signed-off-by: Sachin Pisal <spisal@nvidia.com>
NVIDIA · Feb 4, 2025 · 49d7823 · 49d7823
1 parent 37cca6d
commit 49d7823
Show file tree

Hide file tree

Showing 18 changed files with 1,556 additions and 1,262 deletions.
diff --git a/runtime/cudaq/base_integrator.h b/runtime/cudaq/base_integrator.h
@@ -16,17 +16,17 @@
 #include <vector>
 
 namespace cudaq {
-template <typename TState>
+template <typename TState, typename HandlerTy>
 class BaseIntegrator {
 protected:
   std::map<std::string, double> integrator_options;
   TState state;
   double t;
   std::map<int, int> dimensions;
   std::shared_ptr<Schedule> schedule;
-  std::shared_ptr<operator_sum> hamiltonian;
+  std::shared_ptr<operator_sum<HandlerTy>> hamiltonian;
   std::shared_ptr<BaseTimeStepper<TState>> stepper;
-  std::vector<std::shared_ptr<operator_sum>> collapse_operators;
+  std::vector<std::shared_ptr<operator_sum<HandlerTy>>> collapse_operators;
 
   virtual void post_init() = 0;
 
@@ -61,10 +61,11 @@ class BaseIntegrator {
   }
 
   /// @brief Set the system parameters (dimensions, schedule, and operators)
-  void set_system(
-      const std::map<int, int> &dimensions, std::shared_ptr<Schedule> schedule,
-      std::shared_ptr<operator_sum> hamiltonian,
-      std::vector<std::shared_ptr<operator_sum>> collapse_operators = {}) {
+  void set_system(const std::map<int, int> &dimensions,
+                  std::shared_ptr<Schedule> schedule,
+                  std::shared_ptr<operator_sum<HandlerTy>> hamiltonian,
+                  std::vector<std::shared_ptr<operator_sum<HandlerTy>>>
+                      collapse_operators = {}) {
     this->dimensions = dimensions;
     this->schedule = schedule;
     this->hamiltonian = hamiltonian;

diff --git a/runtime/cudaq/cudm_helpers.h b/runtime/cudaq/cudm_helpers.h
@@ -25,10 +25,12 @@ compute_lindblad_operator(cudensitymatHandle_t handle,
                           const std::vector<matrix_2> &c_ops,
                           const std::vector<int64_t> &mode_extents);
 
+template <typename HandlerTy>
 cudensitymatOperator_t convert_to_cudensitymat_operator(
     cudensitymatHandle_t handle,
     const std::map<std::string, std::complex<double>> &parameters,
-    const operator_sum &op, const std::vector<int64_t> &mode_extents);
+    const operator_sum<HandlerTy> &op,
+    const std::vector<int64_t> &mode_extents);
 
 cudensitymatOperator_t construct_liovillian(
     cudensitymatHandle_t handle, const cudensitymatOperator_t &hamiltonian,

diff --git a/runtime/cudaq/definition.h b/runtime/cudaq/definition.h
@@ -25,8 +25,8 @@ using Func = std::function<matrix_2(
 
 class CallbackFunction {
 private:
-  // The user provided callback function that takes a vector defining the 
-  // dimension for each degree of freedom it acts on, and a map of complex 
+  // The user provided callback function that takes a vector defining the
+  // dimension for each degree of freedom it acts on, and a map of complex
   // parameters.
   Func _callback_func;
 
@@ -56,15 +56,15 @@ class CallbackFunction {
   }
 
   // assignment operator
-  CallbackFunction& operator=(const CallbackFunction &other) {
+  CallbackFunction &operator=(const CallbackFunction &other) {
     if (this != &other) {
       _callback_func = other._callback_func;
     }
     return *this;
   }
 
   // move assignment operator
-  CallbackFunction& operator=(CallbackFunction &&other) {
+  CallbackFunction &operator=(CallbackFunction &&other) {
     if (this != &other) {
       _callback_func = std::move(other._callback_func);
     }
@@ -76,7 +76,8 @@ class CallbackFunction {
   matrix_2
   operator()(std::vector<int> relevant_dimensions,
              std::map<std::string, std::complex<double>> parameters) const {
-    return _callback_func(std::move(relevant_dimensions), std::move(parameters));
+    return _callback_func(std::move(relevant_dimensions),
+                          std::move(parameters));
   }
 };
 
@@ -114,15 +115,15 @@ class ScalarCallbackFunction : CallbackFunction {
   }
 
   // assignment operator
-  ScalarCallbackFunction& operator=(const ScalarCallbackFunction &other) {
+  ScalarCallbackFunction &operator=(const ScalarCallbackFunction &other) {
     if (this != &other) {
       _callback_func = other._callback_func;
     }
     return *this;
   }
 
   // move assignment operator
-  ScalarCallbackFunction& operator=(ScalarCallbackFunction &&other) {
+  ScalarCallbackFunction &operator=(ScalarCallbackFunction &&other) {
     if (this != &other) {
       _callback_func = std::move(other._callback_func);
     }
@@ -141,14 +142,14 @@ class ScalarCallbackFunction : CallbackFunction {
 /// or scalar operator is instantiated by other means than the `define`
 /// class method.
 class Definition {
-private: 
+private:
   std::string id;
   CallbackFunction generator;
   std::vector<int> m_expected_dimensions;
 
 public:
-
-  Definition(const std::string &operator_id, std::vector<int> expected_dimensions, CallbackFunction &&create);
+  Definition(const std::string &operator_id,
+             std::vector<int> expected_dimensions, CallbackFunction &&create);
   Definition(Definition &&def);
   ~Definition();
 

diff --git a/runtime/cudaq/dynamics/definition.cpp b/runtime/cudaq/dynamics/definition.cpp
@@ -16,11 +16,15 @@
 
 namespace cudaq {
 
-Definition::Definition(const std::string &operator_id, std::vector<int> expected_dimensions, CallbackFunction &&create) 
-  : id(operator_id), generator(std::move(create)), m_expected_dimensions(std::move(expected_dimensions)) {}
+Definition::Definition(const std::string &operator_id,
+                       std::vector<int> expected_dimensions,
+                       CallbackFunction &&create)
+    : id(operator_id), generator(std::move(create)),
+      m_expected_dimensions(std::move(expected_dimensions)) {}
 
-Definition::Definition(Definition &&def) 
-  : id(def.id), generator(std::move(def.generator)), m_expected_dimensions(std::move(def.m_expected_dimensions)) {}
+Definition::Definition(Definition &&def)
+    : id(def.id), generator(std::move(def.generator)),
+      m_expected_dimensions(std::move(def.m_expected_dimensions)) {}
 
 matrix_2 Definition::generate_matrix(
     const std::vector<int> &relevant_dimensions,

diff --git a/runtime/cudaq/dynamics/helpers.cpp b/runtime/cudaq/dynamics/helpers.cpp
@@ -6,12 +6,12 @@
  * the terms of the Apache License 2.0 which accompanies this distribution.    *
  ******************************************************************************/
 
-#include "cudaq/operators.h"
-#include <ranges>
 #include "cudaq/helpers.h"
 #include "cudaq/cudm_error_handling.h"
+#include "cudaq/operators.h"
 #include <iostream>
 #include <map>
+#include <ranges>
 #include <sstream>
 
 namespace cudaq {
@@ -22,8 +22,9 @@ class _OperatorHelpers {
   _OperatorHelpers() = default;
 
   // Aggregate parameters from multiple mappings.
-  std::map<std::string, std::string> aggregate_parameters(
-      const std::vector<std::map<std::string, std::string>> &parameter_mappings) {
+  std::map<std::string, std::string>
+  aggregate_parameters(const std::vector<std::map<std::string, std::string>>
+                           &parameter_mappings) {
     std::map<std::string, std::string> parameter_descriptions;
 
     for (const auto &descriptions : parameter_mappings) {
@@ -41,17 +42,17 @@ class _OperatorHelpers {
 
   // Extract documentation for a specific parameter from docstring.
   std::string parameter_docs(const std::string &param_name,
-                                              const std::string &docs) {
+                             const std::string &docs) {
     if (param_name.empty() || docs.empty()) {
       return "";
     }
 
     try {
       std::regex keyword_pattern(R"(^\s*(Arguments|Args):\s*$)",
-                                std::regex::multiline);
+                                 std::regex::multiline);
       std::regex param_pattern(R"(^\s*)" + param_name +
-                                  R"(\s*(\(.*\))?:\s*(.*)$)",
-                              std::regex::multiline);
+                                   R"(\s*(\(.*\))?:\s*(.*)$)",
+                               std::regex::multiline);
 
       std::smatch match;
       std::sregex_iterator it(docs.begin(), docs.end(), keyword_pattern);
@@ -73,10 +74,9 @@ class _OperatorHelpers {
 
   // Extract positional arguments and keyword-only arguments.
   std::pair<std::vector<std::string>, std::map<std::string, std::string>>
-  args_from_kwargs(
-      const std::map<std::string, std::string> &kwargs,
-      const std::vector<std::string> &required_args,
-      const std::vector<std::string> &kwonly_args) {
+  args_from_kwargs(const std::map<std::string, std::string> &kwargs,
+                   const std::vector<std::string> &required_args,
+                   const std::vector<std::string> &kwonly_args) {
     std::vector<std::string> extracted_args;
     std::map<std::string, std::string> kwonly_dict;
 
@@ -99,8 +99,8 @@ class _OperatorHelpers {
 
   /// Generates all possible states for the given dimensions ordered according
   /// to the sequence of degrees (ordering is relevant if dimensions differ).
-  std::vector<std::string>
-  generate_all_states(std::vector<int> degrees, std::map<int, int> dimensions) {
+  std::vector<std::string> generate_all_states(std::vector<int> degrees,
+                                               std::map<int, int> dimensions) {
     if (degrees.size() == 0)
       return {};
 
@@ -124,7 +124,7 @@ class _OperatorHelpers {
   }
 
   cudaq::matrix_2 permute_matrix(cudaq::matrix_2 matrix,
-                                        std::vector<int> permutation) {
+                                 std::vector<int> permutation) {
     auto result = cudaq::matrix_2(matrix.get_rows(), matrix.get_columns());
     std::vector<std::complex<double>> sorted_values;
     for (std::size_t permuted : permutation) {
@@ -146,7 +146,6 @@ class _OperatorHelpers {
     std::sort(degrees.begin(), degrees.end(), std::greater<int>());
     return degrees;
   }
-
 };
-}
-}
+} // namespace detail
+} // namespace cudaq
diff --git a/runtime/cudaq/dynamics/helpers.h b/runtime/cudaq/dynamics/helpers.h
@@ -6,27 +6,27 @@
  * the terms of the Apache License 2.0 which accompanies this distribution.    *
  ******************************************************************************/
 
+#include "cudaq/utils/tensor.h"
 #include <map>
 #include <vector>
-#include "cudaq/utils/tensor.h"
 
 namespace cudaq {
 namespace detail {
 
-  /// Generates all possible states for the given dimensions ordered according
-  /// to the sequence of degrees (ordering is relevant if dimensions differ).
-  std::vector<std::string> generate_all_states(std::vector<int> degrees, std::map<int, int> dimensions);
-
-  // Permutes the given matrix according to the given permutation.
-  // If states is the current order of vector entries on which the given matrix
-  // acts, and permuted_states is the desired order of an array on which the
-  // permuted matrix should act, then the permutation is defined such that
-  // [states[i] for i in permutation] produces permuted_states.
-  cudaq::matrix_2 permute_matrix(cudaq::matrix_2 matrix,
-                                        std::vector<int> permutation);
+/// Generates all possible states for the given dimensions ordered according
+/// to the sequence of degrees (ordering is relevant if dimensions differ).
+std::vector<std::string> generate_all_states(std::vector<int> degrees,
+                                             std::map<int, int> dimensions);
 
-  // Returns the degrees sorted in canonical order.
-  std::vector<int> canonicalize_degrees(std::vector<int> degrees);
-}
-}
+// Permutes the given matrix according to the given permutation.
+// If states is the current order of vector entries on which the given matrix
+// acts, and permuted_states is the desired order of an array on which the
+// permuted matrix should act, then the permutation is defined such that
+// [states[i] for i in permutation] produces permuted_states.
+cudaq::matrix_2 permute_matrix(cudaq::matrix_2 matrix,
+                               std::vector<int> permutation);
 
+// Returns the degrees sorted in canonical order.
+std::vector<int> canonicalize_degrees(std::vector<int> degrees);
+} // namespace detail
+} // namespace cudaq
diff --git a/runtime/cudaq/dynamics/manipulation.cpp b/runtime/cudaq/dynamics/manipulation.cpp
@@ -14,16 +14,15 @@ namespace cudaq {
 std::vector<int>
 MatrixArithmetics::_compute_permutation(std::vector<int> op_degrees,
                                         std::vector<int> canon_degrees) {
-  auto states =
-      cudaq::detail::generate_all_states(canon_degrees, m_dimensions);
+  auto states = cudaq::detail::generate_all_states(canon_degrees, m_dimensions);
 
   std::vector<int> reordering;
   for (auto degree : op_degrees) {
     auto it = std::find(canon_degrees.begin(), canon_degrees.end(), degree);
     reordering.push_back(it - canon_degrees.begin());
   }
 
-  std::vector<std::string> op_states = 
+  std::vector<std::string> op_states =
       cudaq::detail::generate_all_states(op_degrees, m_dimensions);
 
   std::vector<int> permutation;
@@ -99,8 +98,10 @@ EvaluatedMatrix MatrixArithmetics::add(EvaluatedMatrix op1,
   return EvaluatedMatrix(op1.m_degrees, (op1.m_matrix + op2.m_matrix));
 }
 
-EvaluatedMatrix MatrixArithmetics::evaluate(
-    std::variant<scalar_operator, matrix_operator, product_operator<matrix_operator>> op) {
+EvaluatedMatrix
+MatrixArithmetics::evaluate(std::variant<scalar_operator, matrix_operator,
+                                         product_operator<matrix_operator>>
+                                op) {
   // auto getDegrees = [](auto &&t) { return t.degrees; };
   // auto toMatrix = [&](auto &&t) {
   //   return t.to_matrix(this->m_dimensions, this->m_parameters);