Skip to content

Commit df4609e

Browse files
authored
Migrate Python dynamics solver implementation to pybind11 (#2853)
* [WIP] Work on migrating dynamics to C++ binding Signed-off-by: Thien Nguyen <[email protected]> * Fix sporadic segfault due to the Python callbacks Signed-off-by: Thien Nguyen <[email protected]> * Fix a row- and column major mismatches Signed-off-by: Thien Nguyen <[email protected]> * Passing RK step size from Python Signed-off-by: Thien Nguyen <[email protected]> * Init MPI for C++ cudm Signed-off-by: Thien Nguyen <[email protected]> * Init MPI at the context level when the handle is created Signed-off-by: Thien Nguyen <[email protected]> * Trim down CuDensityMatState ctor and remove API that may result in state copy Signed-off-by: Thien Nguyen <[email protected]> * Perf improvement - Make sure we do accumulate + scaling in a single cublas call - Optimize tempory memory allocation with cuda mempool: we allocate quite a lot of temporary memory during the evolution hence make sure to use mempool (as cupy does) Signed-off-by: Thien Nguyen <[email protected]> * Support creating initial state from enum in C++. The created state is MPI-compatible Signed-off-by: Thien Nguyen <[email protected]> * Add performance metric monitor dynamics solver involves many repetitive steps, hence, conventional trace logging is not suitable. Adding a metric map, which tracks the total time across many calls, e.g., to computeAction, to display aggregated data. Refactor the allocator into a helper class for better configurability. Signed-off-by: Thien Nguyen <[email protected]> * Update examples Signed-off-by: Thien Nguyen <[email protected]> * Code format Signed-off-by: Thien Nguyen <[email protected]> * Code format Signed-off-by: Thien Nguyen <[email protected]> * Code clean up Signed-off-by: Thien Nguyen <[email protected]> * Code cleanup Signed-off-by: Thien Nguyen <[email protected]> * add validation Signed-off-by: Thien Nguyen <[email protected]> * Fix skbuild Signed-off-by: Thien Nguyen <[email protected]> * proper translation for nsteps and step size for the native integrator Signed-off-by: Thien Nguyen <[email protected]> * Further slim down dynamics solver Python surface - Get rid of the Pythonic CudmState class. Everything now operates on cudaq::state. - Unify the kron order. - Add a cuBlas-based outter product implementation for to_density_matrix. - Add more tests. Signed-off-by: Thien Nguyen <[email protected]> * Clang-format Signed-off-by: Thien Nguyen <[email protected]> * python yapf format Signed-off-by: Thien Nguyen <[email protected]> * Code review: fix the example preface comment Signed-off-by: Thien Nguyen <[email protected]> * Fix docs build Signed-off-by: Thien Nguyen <[email protected]> * code format Signed-off-by: Thien Nguyen <[email protected]> * Address code review comments/feedback Signed-off-by: Thien Nguyen <[email protected]> * Address code review: rename output file for example Signed-off-by: Thien Nguyen <[email protected]> --------- Signed-off-by: Thien Nguyen <[email protected]>
1 parent 43c635e commit df4609e

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

43 files changed

+1486
-1176
lines changed

docs/sphinx/api/languages/python_api.rst

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ Dynamics
150150
.. autoclass:: cudaq.dynamics.integrator.BaseIntegrator
151151

152152
.. autoclass:: cudaq.dynamics.helpers.InitialState
153-
.. autofunction:: cudaq.dynamics.cudm_state.to_cupy_array
153+
.. autoclass:: cudaq.InitialStateType
154154

155155
Operators
156156
=============================
Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
/*******************************************************************************
2+
* Copyright (c) 2022 - 2025 NVIDIA Corporation & Affiliates. *
3+
* All rights reserved. *
4+
* *
5+
* This source code and the accompanying materials are made available under *
6+
* the terms of the Apache License 2.0 which accompanies this distribution. *
7+
******************************************************************************/
8+
9+
// Compile and run with:
10+
// ```
11+
// nvq++ --target dynamics heisenberg_model_mpi.cpp -o a.out &&
12+
// mpiexec -np <N> ./a.out
13+
// ```
14+
15+
#include "cudaq/algorithms/evolve.h"
16+
#include "cudaq/algorithms/integrator.h"
17+
#include "cudaq/operators.h"
18+
#include "export_csv_helper.h"
19+
#include <cudaq.h>
20+
21+
int main() {
22+
cudaq::mpi::initialize();
23+
std::cout << "Number of ranks = " << cudaq::mpi::num_ranks() << "\n";
24+
// Set up a 15-spin chain, where each spin is a two-level system.
25+
const int num_spins = 15;
26+
cudaq::dimension_map dimensions;
27+
for (int i = 0; i < num_spins; i++) {
28+
dimensions[i] = 2; // Each spin (site) has dimension 2.
29+
}
30+
31+
// Initial state
32+
// Prepare an initial state where the spins are arranged in a staggered
33+
// configuration. Even indices get the value '0' and odd indices get '1'. For
34+
// example, for 15 spins: spins: 0 1 0 1 0 1 0 1 0 1 0 1 0 1 0
35+
std::string spin_state;
36+
for (int i = 0; i < num_spins; i++) {
37+
spin_state.push_back((i % 2 == 0) ? '0' : '1');
38+
}
39+
40+
// Convert the binary string to an integer index
41+
// In the Hilbert space of 15 spins (size 2^15 = 32768), this index
42+
// corresponds to the state |0 1 0 1 0 1 0 1 0 1 0 1 0 1 0>
43+
int initial_state_index = std::stoi(spin_state, nullptr, 2);
44+
45+
// Build the staggered magnetization operator
46+
// The staggered magnetization operator is used to measure antiferromagnetic
47+
// order. It is defined as a sum over all spins of the Z operator, alternating
48+
// in sign. For even sites, we add `sz`; for odd sites, we subtract `sz`.
49+
auto staggered_magnetization_t = cudaq::spin_op::empty();
50+
for (int i = 0; i < num_spins; i++) {
51+
auto sz = cudaq::spin_op::z(i);
52+
if (i % 2 == 0) {
53+
staggered_magnetization_t += sz;
54+
} else {
55+
staggered_magnetization_t -= sz;
56+
}
57+
}
58+
59+
// Normalize the number of spins so that the observable is intensive.
60+
auto staggered_magnetization_op =
61+
(1 / static_cast<double>(num_spins)) * staggered_magnetization_t;
62+
63+
// Each entry will associate a value of g (the `anisotropy` in the Z coupling)
64+
// with its corresponding time-series of expectation values of the staggered
65+
// magnetization.
66+
std::vector<std::pair<double, std::vector<double>>> observe_results;
67+
68+
// Simulate the dynamics over 100 time steps spanning from time 0 to 5.
69+
const int num_steps = 100;
70+
std::vector<double> steps = cudaq::linspace(0.0, 5.0, num_steps);
71+
72+
// For three different values of g, which sets the strength of the Z-Z
73+
// interaction: g = 0.0 (isotropic in the `XY` plane), 0.25, and 4.0 (strongly
74+
// `anisotropy`).
75+
std::vector<double> g_values = {0.0, 0.25, 4.0};
76+
77+
for (auto g : g_values) {
78+
// Set the coupling strengths:
79+
// `Jx` and `Jy` are set to 1.0 (coupling along X and Y axes), while `Jz` is
80+
// set to the current g value (coupling along the Z axis).
81+
double Jx = 1.0, Jy = 1.0, Jz = g;
82+
83+
// The Hamiltonian is built from the nearest-neighbor interactions:
84+
// H = H + `Jx` * `Sx`_i * `Sx`_{i+1}
85+
// H = H + `Jy` * `Sy`_i * `Sy`_{i+1}
86+
// H = H + `Jz` * `Sz`_i * `Sz`_{i+1}
87+
// This is a form of the `anisotropic` Heisenberg (or `XYZ`) model.
88+
auto hamiltonian = cudaq::spin_op::empty();
89+
for (int i = 0; i < num_spins - 1; i++) {
90+
hamiltonian =
91+
hamiltonian + Jx * cudaq::spin_op::x(i) * cudaq::spin_op::x(i + 1);
92+
hamiltonian =
93+
hamiltonian + Jy * cudaq::spin_op::y(i) * cudaq::spin_op::y(i + 1);
94+
hamiltonian =
95+
hamiltonian + Jz * cudaq::spin_op::z(i) * cudaq::spin_op::z(i + 1);
96+
}
97+
98+
// Initial state vector
99+
// For a 9-spin system, the Hilbert space dimension is 2^9 = 512.
100+
// Initialize the state as a vector with all zeros except for a 1 at the
101+
// index corresponding to our staggered state.
102+
const int state_size = 1 << num_spins;
103+
std::vector<std::complex<double>> psi0_data(state_size, {0.0, 0.0});
104+
psi0_data[initial_state_index] = {1.0, 0.0};
105+
auto psi0 = cudaq::state::from_data(psi0_data);
106+
107+
// The schedule is built using the time steps array.
108+
cudaq::schedule schedule(steps);
109+
110+
// Use a Runge-`Kutta` integrator (4`th` order) with a small time step `dt`
111+
// = 0.001.
112+
cudaq::integrators::runge_kutta integrator(4, 0.001);
113+
114+
// Evolve the initial state psi0 under the Hamiltonian, using the specified
115+
// schedule and integrator. No collapse operators are included (closed
116+
// system evolution). Measure the expectation value of the staggered
117+
// magnetization operator at each time step.
118+
auto evolve_result =
119+
cudaq::evolve(hamiltonian, dimensions, schedule, psi0, integrator, {},
120+
{staggered_magnetization_op}, true);
121+
122+
// Lambda to extract expectation values for a given observable index
123+
auto get_expectation = [](int idx, auto &result) -> std::vector<double> {
124+
std::vector<double> expectations;
125+
126+
auto all_exps = result.expectation_values.value();
127+
for (auto exp_vals : all_exps) {
128+
expectations.push_back((double)exp_vals[idx]);
129+
}
130+
return expectations;
131+
};
132+
133+
observe_results.push_back({g, get_expectation(0, evolve_result)});
134+
}
135+
136+
if (observe_results.size() != 3) {
137+
std::cerr << "Unexpected number of g values" << std::endl;
138+
return 1;
139+
}
140+
141+
if (cudaq::mpi::rank() == 0) {
142+
// Only save on the first rank to prevent race condition.
143+
// The `CSV` file "`heisenberg`_model.`csv`" will contain column with:
144+
// - The time steps
145+
// - The expectation values of the staggered magnetization for each g
146+
// value (labeled g_0, g_0.25, g_4).
147+
export_csv("heisenberg_model_mpi_result.csv", "time", steps, "g_0",
148+
observe_results[0].second, "g_0.25", observe_results[1].second,
149+
"g_4", observe_results[2].second);
150+
151+
std::cout << "Simulation complete. The results are saved in "
152+
"heisenberg_model_mpi_result.csv file."
153+
<< std::endl;
154+
}
155+
cudaq::mpi::finalize();
156+
return 0;
157+
}

docs/sphinx/examples/python/dynamics/cavity_qed.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@
3535
# Cavity in a state which has 5 photons initially
3636
cavity_state = cp.zeros((10, 10), dtype=cp.complex128)
3737
cavity_state[5][5] = 1.0
38-
rho0 = cudaq.State.from_data(cp.kron(qubit_state, cavity_state))
38+
rho0 = cudaq.State.from_data(cp.kron(cavity_state, qubit_state))
3939

4040
steps = np.linspace(0, 10, 201)
4141
schedule = Schedule(steps, ["time"])

docs/sphinx/examples/python/dynamics/cross_resonance.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,16 +30,12 @@
3030
# Dimensions of sub-system
3131
dimensions = {0: 2, 1: 2}
3232

33-
# Initial state of the system (ground state).
34-
rho0 = cudaq.State.from_data(
35-
cp.array([[1.0, 0.0], [0.0, 0.0]], dtype=cp.complex128))
36-
3733
# Two initial states: |00> and |10>.
3834
# We show the 'conditional' evolution when controlled qubit is in |1> state.
3935
psi_00 = cudaq.State.from_data(
4036
cp.array([1.0, 0.0, 0.0, 0.0], dtype=cp.complex128))
4137
psi_10 = cudaq.State.from_data(
42-
cp.array([0.0, 0.0, 1.0, 0.0], dtype=cp.complex128))
38+
cp.array([0.0, 1.0, 0.0, 0.0], dtype=cp.complex128))
4339

4440
# Schedule of time steps.
4541
steps = np.linspace(0.0, 1.0, 1001)

python/CMakeLists.txt

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,3 +78,7 @@ if(CUDAQ_BUILD_TESTS)
7878
endif()
7979

8080
add_subdirectory(runtime/interop)
81+
82+
if (CUDENSITYMAT_ROOT AND CUDA_FOUND)
83+
add_subdirectory(runtime/cudaq/dynamics)
84+
endif()

python/cudaq/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,8 @@
9797
from .dynamics.evolution import evolve, evolve_async
9898
from .dynamics.integrators import *
9999

100+
InitialStateType = cudaq_runtime.InitialStateType
101+
100102
# Optimizers + Gradients
101103
optimizers = cudaq_runtime.optimizers
102104
gradients = cudaq_runtime.gradients

0 commit comments

Comments
 (0)