From f77073678bc175c0677de7c0aa06fce933031b53 Mon Sep 17 00:00:00 2001 From: jamesETsmith Date: Thu, 24 Oct 2024 15:12:14 -0400 Subject: [PATCH] [feature/SPO_apply] Adding apply (without x_tk weights) for SummedPauliOp --- fast_pauli/cpp/include/__summed_pauli_op.hpp | 79 +++++++++++++++++++- fast_pauli/cpp/src/fast_pauli.cpp | 10 +++ tests/fast_pauli/test_summed_pauli_op.py | 35 +++++++++ 3 files changed, 122 insertions(+), 2 deletions(-) diff --git a/fast_pauli/cpp/include/__summed_pauli_op.hpp b/fast_pauli/cpp/include/__summed_pauli_op.hpp index a9d63fa..ca64bb6 100644 --- a/fast_pauli/cpp/include/__summed_pauli_op.hpp +++ b/fast_pauli/cpp/include/__summed_pauli_op.hpp @@ -22,6 +22,7 @@ #include #include "__pauli_string.hpp" +#include "__type_traits.hpp" namespace fast_pauli { @@ -34,7 +35,7 @@ template struct SummedPauliOp // std::vector> ops; std::vector pauli_strings; std::vector> coeffs_raw; - Tensor<2> coeffs; + Tensor<2> coeffs; // (n_pauli_strings, n_operators) // TODO dangerous size_t _dim; @@ -173,11 +174,81 @@ template struct SummedPauliOp return pauli_strings.size(); } + void apply(Tensor<2> new_states, Tensor<2> states) const + { + apply(std::execution::seq, new_states, states); + } + + template + void apply(ExecutionPolicy &&, Tensor<2> new_states, Tensor<2> states) const + { + if (states.extent(0) != new_states.extent(0) || states.extent(1) != new_states.extent(1)) + { + throw std::invalid_argument("new_states must have the same dimensions as states"); + } + size_t const n_threads = omp_get_max_threads(); + size_t const n_ps = n_pauli_strings(); + size_t const n_ops = n_operators(); + size_t const n_dim = dim(); + size_t const n_data = states.extent(1); + + std::vector> new_states_thr_raw(n_threads * n_dim * n_data); + std::mdspan, std::dextents> new_states_thr(new_states_thr_raw.data(), n_threads, + n_dim, n_data); + + if constexpr (is_parallel_execution_policy_v) + { + +#pragma omp parallel + { +#pragma omp for schedule(static) + for (size_t j = 0; j < n_ps; ++j) + { + std::complex summed_coeff(0, 0); + + for (size_t k = 0; k < n_ops; ++k) + { + summed_coeff += coeffs(j, k); + } + + std::mdspan new_states_j = + std::submdspan(new_states_thr, omp_get_thread_num(), std::full_extent, std::full_extent); + pauli_strings[j].apply_batch(new_states_j, states, summed_coeff); + } + +#pragma omp for schedule(static) collapse(2) + for (size_t i = 0; i < n_dim; ++i) + { + for (size_t t = 0; t < n_data; ++t) + { + for (size_t j = 0; j < n_threads; ++j) + { + new_states(i, t) += new_states_thr(j, i, t); + } + } + } + } + } + else + { + for (size_t j = 0; j < n_ps; ++j) + { + pauli_strings[j].apply_batch(new_states, states, std::complex(1.)); + std::complex c(0, 0); + for (size_t k = 0; k < n_ops; ++k) + { + c += coeffs(j, k); + } + pauli_strings[j].apply_batch(new_states, states, c); + } + } + } + /** * @brief Apply the SummedPauliOp to a set of weighted states. * * Calculates \f$ - * \bra{\psi_t} \big(\sum_k x_{tk} \sum_i h_{ik} \mathcal{\hat{P}}_i \big) \ket{\psi_t} + * \big(\sum_k x_{tk} \sum_i h_{ik} \mathcal{\hat{P}}_i \big) \ket{\psi_t} * \f$ * * @param new_states @@ -335,6 +406,10 @@ template struct SummedPauliOp * each input states, so the output tensor will have shape (n_operators x * n_states). * + * \f$ + * \bra{\psi_t} \big(\sum_k x_{tk} \sum_i h_{ik} \mathcal{\hat{P}}_i \big) \ket{\psi_t} + * \f$ + * * @param expectation_vals_out Output tensor for the expectation values * (n_operators x n_states) * @param states The states used to calculate the expectation values (n_dim x diff --git a/fast_pauli/cpp/src/fast_pauli.cpp b/fast_pauli/cpp/src/fast_pauli.cpp index 512608a..da01bc1 100644 --- a/fast_pauli/cpp/src/fast_pauli.cpp +++ b/fast_pauli/cpp/src/fast_pauli.cpp @@ -919,7 +919,17 @@ np.ndarray .def_prop_ro("n_operators", &fp::SummedPauliOp::n_operators) .def_prop_ro("n_pauli_strings", &fp::SummedPauliOp::n_pauli_strings) + // .def("apply", + [](fp::SummedPauliOp const &self, nb::ndarray states) { + auto states_mdspan = fp::__detail::ndarray_to_mdspan(states); + auto new_states = fp::__detail::owning_ndarray_like_mdspan(states_mdspan); + auto new_states_mdspan = fp::__detail::ndarray_to_mdspan(new_states); + self.apply(std::execution::par, new_states_mdspan, states_mdspan); + return new_states; + }) + + .def("apply_weighted", [](fp::SummedPauliOp const &self, nb::ndarray states, nb::ndarray data) { auto states_mdspan = fp::__detail::ndarray_to_mdspan(states); auto data_mdspan = fp::__detail::ndarray_to_mdspan(data); diff --git a/tests/fast_pauli/test_summed_pauli_op.py b/tests/fast_pauli/test_summed_pauli_op.py index 19f3cbd..d340164 100644 --- a/tests/fast_pauli/test_summed_pauli_op.py +++ b/tests/fast_pauli/test_summed_pauli_op.py @@ -60,3 +60,38 @@ def test_expectation_values( # Check np.testing.assert_allclose(expectation_vals, expectation_vals_naive, atol=1e-13) + + +@pytest.mark.parametrize( + "summed_pauli_op", [fp.SummedPauliOp], ids=resolve_parameter_repr +) +@pytest.mark.parametrize( + "n_states,n_operators,n_qubits", + [(s, o, q) for s in [1, 10, 1000] for o in [1, 10, 100] for q in [1, 2, 6]], +) +def test_apply( + summed_pauli_op: type[fp.SummedPauliOp], + n_states: int, + n_operators: int, + n_qubits: int, +) -> None: + """Test applying the summed pauli operator method.""" + pauli_strings = fp.helpers.calculate_pauli_strings_max_weight(n_qubits, 2) + n_strings = len(pauli_strings) + + coeffs_2d = np.random.rand(n_strings, n_operators).astype(np.complex128) + psi = np.random.rand(2**n_qubits, n_states).astype(np.complex128) + + op = summed_pauli_op(pauli_strings, coeffs_2d) + + # The new_states we want to check + new_states = op.apply(psi) + + # Trusted new_states + new_states_naive = np.zeros((2**n_qubits, n_states), dtype=np.complex128) + for k in range(n_operators): + A_k = fp.PauliOp(coeffs_2d[:, k].copy(), pauli_strings) + new_states_naive += A_k.apply(psi) + + # Check + np.testing.assert_allclose(new_states, new_states_naive, atol=1e-13)