Getter and setter for coeffs and split method for SPO

fast_pauli/cpp/include/__nb_helpers.hpp

@@ -193,6 +193,23 @@ nb::ndarray<nb::numpy, T> owning_ndarray_like_mdspan(std::mdspan<T, std::dextent
     return owning_ndarray_from_shape<T, ndim>(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<nb::numpy, T>
+ */
+template <typename T, size_t ndim>
+nb::ndarray<nb::numpy, T> owning_ndarray_from_mdspan(std::mdspan<T, std::dextents<size_t, ndim>> a)
+{
+    auto ndarray = owning_ndarray_like_mdspan<T, ndim>(a);
+    std::memcpy(ndarray.data(), a.data_handle(), a.size() * sizeof(T));
+    return ndarray;
+}
+
 } // namespace fast_pauli::__detail
 
 #endif // __NB_HELPERS_HPP

fast_pauli/cpp/include/__summed_pauli_op.hpp

@@ -613,6 +613,27 @@ template <std::floating_point T> struct SummedPauliOp
         }
     }
 
+    /**
+     * @brief Return the components of the SummedPauliOp as a vector of PauliOps
+     *
+     * @return std::vector<PauliOp<T>>
+     */
+    std::vector<PauliOp<T>> split() const
+    {
+        std::vector<std::complex<T>> op_coeffs(n_pauli_strings());
+        std::vector<PauliOp<T>> 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<T>(op_coeffs, pauli_strings));
+        }
+
+        return ops;
+    }
+
     /**
      * @brief Get the dense representation of the SummedPauliOp as a 3D tensor
      *

fast_pauli/cpp/src/fast_pauli.cpp

@@ -997,8 +997,36 @@ Returns
 int
     Number of PauliStrings
 )%")
+        .def_prop_rw(
+            "coeffs",
+            [](fp::SummedPauliOp<float_type> const &self) {
+                std::vector<cfloat_t> 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<cfloat_t, 2>(coeffs_t);
+            },
+            [](fp::SummedPauliOp<float_type> &self, nb::ndarray<cfloat_t> 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<cfloat_t, 2>(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<float_type> const &self, nb::ndarray<cfloat_t> states) {
@@ -1159,6 +1187,14 @@ Returns
 -------
 SummedPauliOp
     A copy of the SummedPauliOp object
+)%")
+        .def("split", &fp::SummedPauliOp<float_type>::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<float_type>::square, R"%(Square the SummedPauliOp.
 

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)