Fix RK 1st order and add a test

Signed-off-by: Thien Nguyen <[email protected]>

runtime/cudaq/cudm_state.h

@@ -60,6 +60,9 @@ class cudm_state {
   /// @return String representation of the state data.
   std::string dump() const;
 
+  /// @brief Dump the state data to the console for debugging purposes.
+  void dumpDeviceData() const;
+
   /// @brief Convert the state vector to a density matrix.
   /// @return A new cudm_state representing the density matrix.
   cudm_state to_density_matrix() const;
@@ -94,7 +97,7 @@ class cudm_state {
 
   /// @brief Scalar multiplication operator
   /// @return The new state after multiplying scalar with the current state.
-  cudm_state operator*(double scalar) const;
+  cudm_state &operator*=(const std::complex<double> &scalar);
 
 private:
   std::vector<std::complex<double>> rawData_;

runtime/cudaq/dynamics/cudm_helpers.cpp

@@ -218,14 +218,22 @@ cudensitymatOperator_t convert_to_cudensitymat_operator(
         } else if (const auto *scalar_op =
                        dynamic_cast<const cudaq::scalar_operator *>(
                            &component)) {
+          // FIXME: do we need this code path?
+          // The product_op already has get_coefficient method.
           auto coeff = scalar_op->evaluate(parameters);
           append_scalar_to_term(handle, term, coeff);
+        } else {
+          // Catch anything that we don't know
+          throw std::runtime_error("Unhandled type!");
         }
       }
 
+      // Handle the coefficient
+      auto coeff = product_op.get_coefficient().evaluate(parameters);
       // Append the product operator term to the top-level operator
       HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(
-          handle, operator_handle, term, 0, make_cuDoubleComplex(1.0, 0.0),
+          handle, operator_handle, term, 0,
+          make_cuDoubleComplex(coeff.real(), coeff.imag()),
           {nullptr, nullptr}));
 
       // FIXME: leak

runtime/cudaq/dynamics/cudm_state.cpp

@@ -167,20 +167,18 @@ cudm_state &cudm_state::operator+=(const cudm_state &other) {
 
   return *this;
 }
+cudm_state &cudm_state::operator*=(const std::complex<double> &scalar) {
+  void *gpuScalar;
+  HANDLE_CUDA_ERROR(cudaMalloc(&gpuScalar, sizeof(std::complex<double>)));
+  HANDLE_CUDA_ERROR(cudaMemcpy(gpuScalar, &scalar, sizeof(std::complex<double>),
+                               cudaMemcpyHostToDevice));
 
-cudm_state cudm_state::operator*(double scalar) const {
-  cudm_state result = cudm_state(handle_, rawData_, hilbertSpaceDims_);
-
-  double *gpuScalar;
-  cudaMalloc(reinterpret_cast<void **>(&gpuScalar), sizeof(double));
-  cudaMemcpy(gpuScalar, &scalar, sizeof(double), cudaMemcpyHostToDevice);
-
-  HANDLE_CUDM_ERROR(cudensitymatStateComputeScaling(handle_, result.get_impl(),
-                                                    gpuScalar, 0));
+  HANDLE_CUDM_ERROR(
+      cudensitymatStateComputeScaling(handle_, state_, gpuScalar, 0));
 
-  cudaFree(gpuScalar);
+  HANDLE_CUDA_ERROR(cudaFree(gpuScalar));
 
-  return result;
+  return *this;
 }
 
 std::string cudm_state::dump() const {
@@ -200,6 +198,25 @@ std::string cudm_state::dump() const {
   return oss.str();
 }
 
+void cudm_state::dumpDeviceData() const {
+  if (!is_initialized()) {
+    throw std::runtime_error("State is not initialized.");
+  }
+
+  std::vector<std::complex<double>> hostBuffer(rawData_.size());
+  HANDLE_CUDA_ERROR(cudaMemcpy(hostBuffer.data(), get_device_pointer(),
+                               hostBuffer.size() * sizeof(std::complex<double>),
+                               cudaMemcpyDefault));
+  std::cout << "State data: [";
+  for (size_t i = 0; i < hostBuffer.size(); i++) {
+    std::cout << hostBuffer[i];
+    if (i < hostBuffer.size() - 1) {
+      std::cout << ", ";
+    }
+  }
+  std::cout << "]\n";
+}
+
 cudm_state cudm_state::to_density_matrix() const {
   if (!is_initialized()) {
     throw std::runtime_error("State is not initialized.");

runtime/cudaq/dynamics/runge_kutta_integrator.cpp

@@ -30,31 +30,34 @@ void runge_kutta_integrator::integrate(double target_time) {
   while (this->t < target_time) {
     double step_size = std::min(dt, target_time - this->t);
 
-    std::cout << "Runge-Kutta step at time " << this->t
-              << " with step size: " << step_size << std::endl;
+    // std::cout << "Runge-Kutta step at time " << this->t
+    //           << " with step size: " << step_size << std::endl;
 
     if (this->substeps_ == 1) {
       // Euler method (1st order)
       cudm_state k1 = this->stepper->compute(this->state, this->t, step_size);
+      k1 *= step_size;
       this->state += k1;
     } else if (this->substeps_ == 2) {
+      // FIXME: implement it
       // Midpoint method (2nd order)
-      cudm_state k1 =
-          this->stepper->compute(this->state, this->t, step_size / 2.0);
-      cudm_state k2 =
-          this->stepper->compute(k1, this->t + step_size / 2.0, step_size);
-      this->state += (k1 + k2) * 0.5;
+      // cudm_state k1 =
+      //     this->stepper->compute(this->state, this->t, step_size / 2.0);
+      // cudm_state k2 =
+      //     this->stepper->compute(k1, this->t + step_size / 2.0, step_size);
+      // this->state += (k1 + k2) * 0.5;
     } else if (this->substeps_ == 4) {
+      // FIXME: implement it
       // Runge-Kutta method (4th order)
-      cudm_state k1 =
-          this->stepper->compute(this->state, this->t, step_size / 2.0);
-      cudm_state k2 = this->stepper->compute(k1, this->t + step_size / 2.0,
-                                             step_size / 2.0);
-      cudm_state k3 =
-          this->stepper->compute(k2, this->t + step_size / 2.0, step_size);
-      cudm_state k4 =
-          this->stepper->compute(k3, this->t + step_size, step_size);
-      this->state += (k1 + (k2 + k3) * 2.0 + k4) * (1.0 / 6.0);
+      // cudm_state k1 =
+      //     this->stepper->compute(this->state, this->t, step_size / 2.0);
+      // cudm_state k2 = this->stepper->compute(k1, this->t + step_size / 2.0,
+      //                                        step_size / 2.0);
+      // cudm_state k3 =
+      //     this->stepper->compute(k2, this->t + step_size / 2.0, step_size);
+      // cudm_state k4 =
+      //     this->stepper->compute(k3, this->t + step_size, step_size);
+      // this->state += (k1 + (k2 + k3) * 2.0 + k4) * (1.0 / 6.0);
     }
 
     // Update time

unittests/dynamics/test_runge_kutta_integrator.cpp

@@ -147,3 +147,64 @@ TEST_F(RungeKuttaIntegratorTest, InvalidSubsteps) {
                    0.0, time_stepper_, 3),
                std::invalid_argument);
 }
+
+TEST_F(RungeKuttaIntegratorTest, CheckEvolve) {
+  const std::vector<std::complex<double>> initialState = {{1.0, 0.0},
+                                                          {0.0, 0.0}};
+  const std::vector<int64_t> dims = {2};
+  const std::string op_id = "pauli_x";
+  auto func = [](std::vector<int> dimensions,
+                 std::map<std::string, std::complex<double>> _none) {
+    if (dimensions.size() != 1)
+      throw std::invalid_argument("Must have a singe dimension");
+    if (dimensions[0] != 2)
+      throw std::invalid_argument("Must have dimension 2");
+    auto mat = matrix_2(2, 2);
+    mat[{1, 0}] = 1.0;
+    mat[{0, 1}] = 1.0;
+    return mat;
+  };
+  cudaq::matrix_operator::define(op_id, {-1}, func);
+
+  // FIXME: enable all orders
+  // for (int integratorOrder : {1, 2, 4}) {
+  for (int integratorOrder : {1}) {
+    std::cout << "Test RK order " << integratorOrder << "\n";
+    auto op = cudaq::product_operator<matrix_operator>(
+        std::complex<double>{0.0, -1.0} * 2.0 * M_PI * 0.1,
+        cudaq::matrix_operator(op_id, {0}));
+    auto cudmOp =
+        cudaq::convert_to_cudensitymat_operator<cudaq::matrix_operator>(
+            handle_, {}, op, dims);
+    auto time_stepper = std::make_shared<cudm_time_stepper>(handle_, cudmOp);
+
+    auto eulerIntegrator = std::make_unique<runge_kutta_integrator>(
+        cudm_state(handle_, initialState, dims), 0.0, time_stepper,
+        integratorOrder);
+    eulerIntegrator->set_option("dt", 0.001);
+
+    constexpr std::size_t numDataPoints = 10;
+    double t = 0.0;
+    std::vector<std::complex<double>> outputStateVec(2);
+    for (std::size_t i = 1; i < numDataPoints; ++i) {
+      eulerIntegrator->integrate(i);
+      auto [t, state] = eulerIntegrator->get_state();
+      // std::cout << "Time = " << t << "\n";
+      // state.dumpDeviceData();
+      HANDLE_CUDA_ERROR(
+          cudaMemcpy(outputStateVec.data(), state.get_device_pointer(),
+                     outputStateVec.size() * sizeof(std::complex<double>),
+                     cudaMemcpyDefault));
+      // Check state vector norm
+      EXPECT_NEAR(std::norm(outputStateVec[0]) + std::norm(outputStateVec[1]),
+                  1.0, 1e-2);
+      const double expValZ =
+          std::norm(outputStateVec[0]) - std::norm(outputStateVec[1]);
+      // Analytical results
+      EXPECT_NEAR(outputStateVec[0].real(), std::cos(2.0 * M_PI * 0.1 * t),
+                  1e-2);
+    }
+
+    HANDLE_CUDM_ERROR(cudensitymatDestroyOperator(cudmOp));
+  }
+}