* Adding helpers functionlity

* Aggregating parameters * Extracting documentation * Extracting positional and keyword arguments * Generating all quantum states for given degrees and dimensions * Permuting a given matrix * Canonicalizing degrees * Adding test cases for the above functions * Formatting Signed-off-by: Sachin Pisal <[email protected]>
NVIDIA · Jan 14, 2025 · 143827f · 143827f
1 parent 760dd9c
commit 143827f
Show file tree

Hide file tree

Showing 10 changed files with 549 additions and 177 deletions.
diff --git a/runtime/cudaq/dynamics/CMakeLists.txt b/runtime/cudaq/dynamics/CMakeLists.txt
@@ -1,5 +1,5 @@
 # ============================================================================ #
-# Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates.                   #
+# Copyright (c) 2022 - 2025 NVIDIA Corporation & Affiliates.                   #
 # All rights reserved.                                                         #
 #                                                                              #
 # This source code and the accompanying materials are made available under     #
@@ -11,7 +11,7 @@ set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-ctad-maybe-unsupported")
 set(INTERFACE_POSITION_INDEPENDENT_CODE ON)
 
 set(CUDAQ_OPS_SRC
-  scalar_operators.cpp elementary_operators.cpp product_operators.cpp operator_sum.cpp schedule.cpp definition.cpp
+  scalar_operators.cpp elementary_operators.cpp product_operators.cpp operator_sum.cpp schedule.cpp definition.cpp helpers.cpp
 )
 
 add_library(${LIBRARY_NAME} SHARED ${CUDAQ_OPS_SRC})

diff --git a/runtime/cudaq/dynamics/helpers.cpp b/runtime/cudaq/dynamics/helpers.cpp
@@ -0,0 +1,135 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2025 NVIDIA Corporation & Affiliates.                  *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+
+#include "cudaq/helpers.h"
+#include <sstream>
+#include <map>
+
+namespace cudaq {
+// Aggregate parameters from multiple mappings.
+std::map<std::string, std::string> OperatorHelpers::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) {
+        for (const auto &[key, new_desc] : descriptions) {
+            if (!parameter_descriptions[key].empty() && !new_desc.empty()) {
+                parameter_descriptions[key] += "\n---\n" + new_desc;
+            } else {
+                parameter_descriptions[key] = new_desc;
+            }
+        }
+    }
+
+    return parameter_descriptions;
+}
+
+// Extract documentation for a specific parameter from docstring.
+std::string OperatorHelpers::parameter_docs(const std::string &param_name, 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 param_pattern(R"(^\s*)" + param_name + R"(\s*(\(.*\))?:\s*(.*)$)", std::regex::multiline);
+
+        std::smatch match;
+        std::sregex_iterator it(docs.begin(), docs.end(), keyword_pattern);
+        std::sregex_iterator end;
+
+        if (it != end) {
+            std::string params_section = docs.substr(it->position() + it->length());
+            if(std::regex_search(params_section, match, param_pattern)) {
+                std::string param_docs = match.str(2);
+                return std::regex_replace(param_docs, std::regex(R"(\s+)"), " ");
+            }
+        }
+    } catch (...) {
+        return "";
+    }
+
+    return "";
+}
+
+// Extract positional arguments and keyword-only arguments.
+std::pair<std::vector<std::string>, std::map<std::string, std::string>> OperatorHelpers::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;
+
+    for (const auto &arg : required_args) {
+        if (kwargs.count(arg)) {
+            extracted_args.push_back(kwargs.at(arg));
+        } else {
+            throw std::invalid_argument("Missing required argument: " + arg);
+        }
+    }
+
+    for (const auto &arg : kwonly_args) {
+        if (kwargs.count(arg)) {
+            kwonly_dict[arg] = kwargs.at(arg);
+        }
+    }
+
+    return {extracted_args, kwonly_dict};
+}
+
+// Generate all possible quantum states for given degrees and dimensions
+std::vector<std::string> OperatorHelpers::generate_all_states(const std::vector<int> &degrees, const std::map<int, int> &dimensions) {
+    if (degrees.empty()) {
+        return {};
+    }
+
+    std::vector<std::vector<std::string>> states;
+    for (int state = 0; state < dimensions.at(degrees[0]); state++) {
+        states.push_back({std::to_string(state)});
+    }
+
+    for (size_t i = 1; i < degrees.size(); i++) {
+        std::vector<std::vector<std::string>> new_states;
+        for (const auto &current : states) {
+            for (int state = 0; state < dimensions.at(degrees[i]); state++) {
+                auto new_entry = current;
+                new_entry.push_back(std::to_string(state));
+                new_states.push_back(new_entry);
+            }
+        }
+        states = new_states;
+    }
+
+    std::vector<std::string> result;
+    for (const auto &state : states) {
+        std::ostringstream joined;
+        for (const auto &s : state) {
+            joined << s;
+        }
+        result.push_back(joined.str());
+    }
+    return result;
+}
+
+// Permute a given eigen matrix
+void OperatorHelpers::permute_matrix(Eigen::MatrixXcd &matrix, const std::vector<int> &permutation) {
+    Eigen::MatrixXcd permuted_matrix(matrix.rows(), matrix.cols());
+
+    for (size_t i = 0; i < permutation.size(); i++) {
+        for (size_t j = 0; j < permutation.size(); j++) {
+            permuted_matrix(i, j) = matrix(permutation[i], permutation[j]);
+        }
+    }
+
+    matrix = permuted_matrix;
+}
+
+// Canonicalize degrees by sorting in descending order
+std::vector<int> OperatorHelpers::canonicalize_degrees(const std::vector<int> &degrees) {
+    std::vector<int> sorted_degrees = degrees;
+    std::sort(sorted_degrees.rbegin(), sorted_degrees.rend());
+    return sorted_degrees;
+}
+}
diff --git a/runtime/cudaq/helpers.h b/runtime/cudaq/helpers.h
@@ -0,0 +1,42 @@
+/****************************************************************-*- C++ -*-****
+ * Copyright (c) 2022 - 2025 NVIDIA Corporation & Affiliates.                  *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+
+#pragma once
+
+#include <map>
+#include <vector>
+#include <string>
+#include <complex>
+#include <functional>
+#include <regex>
+#include <numeric>
+#include <Eigen/Dense>
+
+namespace cudaq {
+class OperatorHelpers {
+public:
+    // Aggregate parameters from multiple mappings.
+    static std::map<std::string, std::string> aggregate_parameters(const std::vector<std::map<std::string, std::string>> &parameter_mappings);
+
+    // Extract documentation for a specific parameter from docstring.
+    static std::string parameter_docs(const std::string &param_name, const std::string &docs);
+
+    // Extract positional arguments and keyword-only arguments.
+    static 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);
+
+    // Generate all possible quantum states for given degrees and dimensions.
+    static std::vector<std::string> generate_all_states(const std::vector<int> &degrees, const std::map<int, int> &dimensions);
+
+    // Permute a given Eigen matrix.
+    static void permute_matrix(Eigen::MatrixXcd &matrix, const std::vector<int> &permutation);
+
+    // Canonicalize degrees by sorting in descending order.
+    static std::vector<int> canonicalize_degrees(const std::vector<int> &degrees);
+};
+}
diff --git a/runtime/cudaq/runge_kutta_integrator.h b/runtime/cudaq/runge_kutta_integrator.h
@@ -16,31 +16,32 @@ namespace cudaq {
 template <typename TState>
 class RungeKuttaIntegrator : public BaseIntegrator<TState> {
 public:
-    using DerivativeFunction = std::function<TState(const TState &, double)>;
+  using DerivativeFunction = std::function<TState(const TState &, double)>;
 
-    explicit RungeKuttaIntegrator(DerivativeFunction f) : stepper(std::make_shared<RungeKuttaTimeStepper<TState>>(f)) {}
+  explicit RungeKuttaIntegrator(DerivativeFunction f)
+      : stepper(std::make_shared<RungeKuttaTimeStepper<TState>>(f)) {}
 
-    // Initializes the integrator
-    void post_init() override {
-        if (!this->stepper) {
-            throw std::runtime_error("Time stepper is not set");
-        }
+  // Initializes the integrator
+  void post_init() override {
+    if (!this->stepper) {
+      throw std::runtime_error("Time stepper is not set");
     }
+  }
 
-    // Advances the system's state from current time to `t`
-    void integrate(double target_t) override {
-        if (!this->schedule || !this->hamiltonian) {
-            throw std::runtime_error("System is not properly set!");
-        }
-
-        while (this->t < target_t) {
-            stepper->compute(this->state, this->t);
-            // Time step size
-            this->t += 0.01;
-        }
+  // Advances the system's state from current time to `t`
+  void integrate(double target_t) override {
+    if (!this->schedule || !this->hamiltonian) {
+      throw std::runtime_error("System is not properly set!");
     }
 
+    while (this->t < target_t) {
+      stepper->compute(this->state, this->t);
+      // Time step size
+      this->t += 0.01;
+    }
+  }
+
 private:
-    std::shared_ptr<RungeKuttaTimeStepper<TState>> stepper;
+  std::shared_ptr<RungeKuttaTimeStepper<TState>> stepper;
 };
-}
+} // namespace cudaq
diff --git a/runtime/cudaq/runge_kutta_time_stepper.h b/runtime/cudaq/runge_kutta_time_stepper.h
@@ -13,21 +13,21 @@ namespace cudaq {
 template <typename TState>
 class RungeKuttaTimeStepper : public BaseTimeStepper<TState> {
 public:
-    using DerivativeFunction = std::function<TState(const TState &, double)>;
+  using DerivativeFunction = std::function<TState(const TState &, double)>;
 
-    RungeKuttaTimeStepper(DerivativeFunction f) : derivativeFunc(f) {}
+  RungeKuttaTimeStepper(DerivativeFunction f) : derivativeFunc(f) {}
 
-    void compute(TState &state, double t, double dt = 0.01) override {
-        // 4th order Runge-Kutta method
-        TState k1 = derivativeFunc(state, t);
-        TState k2 = derivativeFunc(state + (dt / 2.0) * k1, t + dt / 2.0);
-        TState k3 = derivativeFunc(state + (dt / 2.0) * k2, t + dt / 2.0);
-        TState k4 = derivativeFunc(state + dt * k3, t + dt);
+  void compute(TState &state, double t, double dt = 0.01) override {
+    // 4th order Runge-Kutta method
+    TState k1 = derivativeFunc(state, t);
+    TState k2 = derivativeFunc(state + (dt / 2.0) * k1, t + dt / 2.0);
+    TState k3 = derivativeFunc(state + (dt / 2.0) * k2, t + dt / 2.0);
+    TState k4 = derivativeFunc(state + dt * k3, t + dt);
 
-        state = state + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4);
-    }
+    state = state + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4);
+  }
 
 private:
-    DerivativeFunction derivativeFunc;
+  DerivativeFunction derivativeFunc;
 };
-}
+} // namespace cudaq
diff --git a/unittests/CMakeLists.txt b/unittests/CMakeLists.txt
@@ -51,6 +51,7 @@ set(CUDAQ_RUNTIME_TEST_SOURCES
   dynamics/product_operators_arithmetic.cpp
   dynamics/test_runge_kutta_time_stepper.cpp
   dynamics/test_runge_kutta_integrator.cpp
+  dynamics/test_helpers.cpp
 )
 
 # Make it so we can get function symbols

diff --git a/unittests/dynamics/runge_kutta_test_helpers.h b/unittests/dynamics/runge_kutta_test_helpers.h
@@ -15,10 +15,10 @@ using TestState = double;
 
 // Simple derivative function: dx/dt = -x (exponential decay)
 inline TestState simple_derivative(const TestState &state, double t) {
-    return -state;
+  return -state;
 }
 
 // A complex function: dx/dt = sin(t)
 inline TestState sine_derivative(const TestState &state, double t) {
-    return std::sin(t);
+  return std::sin(t);
 }