diff --git a/fast_pauli/cpp/include/__nb_helpers.hpp b/fast_pauli/cpp/include/__nb_helpers.hpp index cd34f89..0a12b81 100644 --- a/fast_pauli/cpp/include/__nb_helpers.hpp +++ b/fast_pauli/cpp/include/__nb_helpers.hpp @@ -193,6 +193,23 @@ nb::ndarray owning_ndarray_like_mdspan(std::mdspan(shape); } +/** + * @brief Creates a new nb::ndarray that owns the data and has the same content and shape + * as an mdspan. + * + * @tparam T Type of the underlying data in ndarray/mdspan + * @tparam ndim Rank of the array + * @param a The mdspan + * @return nb::ndarray + */ +template +nb::ndarray owning_ndarray_from_mdspan(std::mdspan> a) +{ + auto ndarray = owning_ndarray_like_mdspan(a); + std::memcpy(ndarray.data(), a.data_handle(), a.size() * sizeof(T)); + return ndarray; +} + } // namespace fast_pauli::__detail #endif // __NB_HELPERS_HPP \ No newline at end of file diff --git a/fast_pauli/cpp/include/__summed_pauli_op.hpp b/fast_pauli/cpp/include/__summed_pauli_op.hpp index 05d2930..b00693f 100644 --- a/fast_pauli/cpp/include/__summed_pauli_op.hpp +++ b/fast_pauli/cpp/include/__summed_pauli_op.hpp @@ -613,6 +613,27 @@ template struct SummedPauliOp } } + /** + * @brief Return the components of the SummedPauliOp as a vector of PauliOps + * + * @return std::vector> + */ + std::vector> split() const + { + std::vector> op_coeffs(n_pauli_strings()); + std::vector> ops; + ops.reserve(n_operators()); + + for (size_t k = 0; k < n_operators(); k++) + { + for (size_t i = 0; i < n_pauli_strings(); i++) + op_coeffs[i] = coeffs(i, k); + ops.push_back(PauliOp(op_coeffs, pauli_strings)); + } + + return ops; + } + /** * @brief Get the dense representation of the SummedPauliOp as a 3D tensor * diff --git a/fast_pauli/cpp/src/fast_pauli.cpp b/fast_pauli/cpp/src/fast_pauli.cpp index 66cc532..bf27cc9 100644 --- a/fast_pauli/cpp/src/fast_pauli.cpp +++ b/fast_pauli/cpp/src/fast_pauli.cpp @@ -997,8 +997,36 @@ Returns int Number of PauliStrings )%") + .def_prop_rw( + "coeffs", + [](fp::SummedPauliOp const &self) { + std::vector coeffs_transposed(self.coeffs.size()); + auto coeffs_t = std::mdspan(coeffs_transposed.data(), self.coeffs.extent(1), self.coeffs.extent(0)); + + for (size_t i = 0; i < self.coeffs.extent(0); i++) + for (size_t j = 0; j < self.coeffs.extent(1); j++) + coeffs_t(j, i) = self.coeffs(i, j); + + return fp::__detail::owning_ndarray_from_mdspan(coeffs_t); + }, + [](fp::SummedPauliOp &self, nb::ndarray coeffs_new) { + if (coeffs_new.ndim() != 2 or coeffs_new.shape(0) != self.n_operators() or + coeffs_new.shape(1) != self.n_pauli_strings()) + throw std::invalid_argument( + "The shape of provided coeffs must match the number of operators and PauliStrings"); + + auto coeffs_mdspan = fp::__detail::ndarray_to_mdspan(coeffs_new); + for (size_t i = 0; i < self.coeffs.extent(0); i++) + for (size_t j = 0; j < self.coeffs.extent(1); j++) + self.coeffs(i, j) = coeffs_mdspan(j, i); + }, + nb::rv_policy::automatic, R"%(Getter and setter for coefficients. - // +Returns +------- +np.ndarray + Array of coefficients corresponding with shape (n_operators, n_pauli_strings) +)%") .def( "apply", [](fp::SummedPauliOp const &self, nb::ndarray states) { @@ -1159,6 +1187,14 @@ Returns ------- SummedPauliOp A copy of the SummedPauliOp object +)%") + .def("split", &fp::SummedPauliOp::split, + R"%(Returns all components of the SummedPauliOp expressed as a vector of PauliOps. + +Returns +------- +List[fp.PauliOp] + Components of the SummedPauliOp object )%") .def("square", &fp::SummedPauliOp::square, R"%(Square the SummedPauliOp. diff --git a/tests/fast_pauli/test_summed_pauli_op.py b/tests/fast_pauli/test_summed_pauli_op.py index 27f6e16..a8d606d 100644 --- a/tests/fast_pauli/test_summed_pauli_op.py +++ b/tests/fast_pauli/test_summed_pauli_op.py @@ -262,3 +262,52 @@ def test_clone( op2.to_tensor(), ) assert id(op1) != id(op2) + + +@pytest.mark.parametrize( + "summed_pauli_op", [fp.SummedPauliOp], ids=resolve_parameter_repr +) +@pytest.mark.parametrize( + "n_operators,n_qubits", + [(o, q) for o in [1, 10] for q in [1, 2, 4, 6]], +) +def test_coeffs_prop( + summed_pauli_op: type[fp.SummedPauliOp], + n_operators: int, + n_qubits: int, +) -> None: + """Test getter and setter for coeffs property of the SummedPauliOp.""" + pauli_strings = fp.helpers.calculate_pauli_strings_max_weight(n_qubits, 2) + orig_coeffs = np.random.rand(len(pauli_strings), n_operators).astype(np.complex128) + spo = summed_pauli_op(pauli_strings, orig_coeffs) + + np.testing.assert_allclose(spo.coeffs, orig_coeffs.T) + new_coeffs = np.random.rand(n_operators, len(pauli_strings)).astype(np.complex128) + spo.coeffs = new_coeffs + np.testing.assert_allclose(spo.coeffs, new_coeffs) + np.testing.assert_allclose( + spo.to_tensor(), summed_pauli_op(pauli_strings, new_coeffs.T.copy()).to_tensor() + ) + + +@pytest.mark.parametrize( + "summed_pauli_op", [fp.SummedPauliOp], ids=resolve_parameter_repr +) +@pytest.mark.parametrize( + "n_operators,n_qubits", + [(o, q) for o in [1, 10] for q in [1, 2, 4, 6]], +) +def test_split( + summed_pauli_op: type[fp.SummedPauliOp], + n_operators: int, + n_qubits: int, +) -> None: + """Test split method of the SummedPauliOp.""" + pauli_strings = fp.helpers.calculate_pauli_strings_max_weight(n_qubits, 2) + coeffs_2d = np.random.rand(len(pauli_strings), n_operators).astype(np.complex128) + spo = summed_pauli_op(pauli_strings, coeffs_2d) + + components = spo.split() + for k, (comp, op_dense) in enumerate(zip(components, spo.to_tensor())): + np.testing.assert_allclose(comp.coeffs, spo.coeffs[k]) + np.testing.assert_allclose(comp.to_tensor(), op_dense)