[feature/SPO_apply] Adding docstrings to python bindinds for SummedPa…

…uliOp
qognitive · Oct 30, 2024 · 64737a8 · 64737a8
1 parent f83d39d
commit 64737a8
Show file tree

Hide file tree

Showing 3 changed files with 121 additions and 41 deletions.
diff --git a/docs/python_api.rst b/docs/python_api.rst
@@ -27,3 +27,10 @@ PauliOp
 .. autoclass:: fast_pauli.PauliOp
    :members:
    :special-members:
+
+SummedPauliOp
+-------------
+
+.. autoclass:: fast_pauli.SummedPauliOp
+   :members:
+   :special-members:
diff --git a/fast_pauli/cpp/include/__summed_pauli_op.hpp b/fast_pauli/cpp/include/__summed_pauli_op.hpp
@@ -204,8 +204,6 @@ template <std::floating_point T> struct SummedPauliOp
      * @brief \copydoc SummedPauliOp::apply(Tensor<2>, Tensor<2>) const
      *
      * @tparam ExecutionPolicy
-     * @param new_states The output states after applying the SummedPauliOp (n_dim, n_states)
-     * @param states The input states to apply the SummedPauliOp to (n_dim, n_states)
      */
     template <execution_policy ExecutionPolicy>
     void apply(ExecutionPolicy &&, Tensor<2> new_states, Tensor<2> states) const
@@ -281,7 +279,8 @@ template <std::floating_point T> struct SummedPauliOp
      * @tparam data_dtype The floating point type of the weights \f$ x_{kt} \f$ (n_operators, n_states)
      * @param new_states The output states after applying the SummedPauliOp (n_dim, n_states)
      * @param states The input states to apply the SummedPauliOp to (n_dim, n_states)
-     * @param data A 2D std::mdspan of the weights \f$ x_{tk} \f$ in the expression above (n_operators, n_states)
+     * @param data A 2D std::mdspan of the data \f$ x_{tk} \f$ that weights the operators in the expression above
+     * (n_operators, n_states)
      */
     template <std::floating_point data_dtype>
     void apply_weighted(Tensor<2> new_states, Tensor<2> states,

diff --git a/fast_pauli/cpp/src/fast_pauli.cpp b/fast_pauli/cpp/src/fast_pauli.cpp
@@ -945,46 +945,120 @@ SummedPauliOp
     New SummedPauliOp instance
 )%")
 
-        .def_prop_ro("dim", &fp::SummedPauliOp<float_type>::dim)
-        .def_prop_ro("n_operators", &fp::SummedPauliOp<float_type>::n_operators)
-        .def_prop_ro("n_pauli_strings", &fp::SummedPauliOp<float_type>::n_pauli_strings)
+        .def_prop_ro("dim", &fp::SummedPauliOp<float_type>::dim,
+                     R"%(Return the Hilbert space dimension of the SummedPauliOp.
+
+Returns
+-------
+int
+    Hilbert space dimension
+)%")
+        .def_prop_ro("n_operators", &fp::SummedPauliOp<float_type>::n_operators,
+                     R"%(Return the number of Pauli operators in the SummedPauliOp.
+
+Returns
+-------
+int
+    Number of operators
+)%")
+        .def_prop_ro("n_pauli_strings", &fp::SummedPauliOp<float_type>::n_pauli_strings,
+                     R"%(Return the number of PauliStrings in the SummedPauliOp.
+
+Returns
+-------
+int
+    Number of PauliStrings
+)%")
 
         //
-        .def("apply",
-             [](fp::SummedPauliOp<float_type> const &self, nb::ndarray<cfloat_t> states) {
-                 auto states_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 2>(states);
-                 auto new_states = fp::__detail::owning_ndarray_like_mdspan<cfloat_t, 2>(states_mdspan);
-                 auto new_states_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 2>(new_states);
-                 self.apply(std::execution::par, new_states_mdspan, states_mdspan);
-                 return new_states;
-             })
-
-        .def("apply_weighted",
-             [](fp::SummedPauliOp<float_type> const &self, nb::ndarray<cfloat_t> states, nb::ndarray<float_type> data) {
-                 auto states_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 2>(states);
-                 auto data_mdspan = fp::__detail::ndarray_to_mdspan<float_type, 2>(data);
-
-                 auto new_states = fp::__detail::owning_ndarray_like_mdspan<cfloat_t, 2>(states_mdspan);
-                 auto new_states_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 2>(new_states);
-
-                 self.apply_weighted(std::execution::par, new_states_mdspan, states_mdspan, data_mdspan);
-
-                 return new_states;
-             })
-        .def("expectation_value",
-             [](fp::SummedPauliOp<float_type> const &self, nb::ndarray<cfloat_t> states) {
-                 auto states_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 2>(states);
-
-                 std::array<size_t, 2> out_shape = {self.n_operators(), states_mdspan.extent(1)};
-                 auto expected_vals_out = fp::__detail::owning_ndarray_from_shape<cfloat_t, 2>(out_shape);
-                 auto expected_vals_out_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 2>(expected_vals_out);
-
-                 self.expectation_value(std::execution::par, expected_vals_out_mdspan, states_mdspan);
-
-                 return expected_vals_out;
-             })
-        //
-        ;
+        .def(
+            "apply",
+            [](fp::SummedPauliOp<float_type> const &self, nb::ndarray<cfloat_t> states) {
+                auto states_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 2>(states);
+                auto new_states = fp::__detail::owning_ndarray_like_mdspan<cfloat_t, 2>(states_mdspan);
+                auto new_states_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 2>(new_states);
+                self.apply(std::execution::par, new_states_mdspan, states_mdspan);
+                return new_states;
+            },
+            "states"_a,
+            R"%(Apply the SummedPauliOp to a batch of states.
+
+.. math::
+    \big(\sum_k \sum_i h_{ik} \mathcal{\hat{P}}_i \big) \ket{\psi_t}
+
+Parameters
+----------
+states : np.ndarray
+    The original state(s) represented as 2D numpy array (n_operators, n_states) for batched calculation.
+
+Returns
+-------
+np.ndarray
+    New state(s) in a form of 2D numpy array (n_operators, n_states) according to the shape of input states
+)%")
+
+        .def(
+            "apply_weighted",
+            [](fp::SummedPauliOp<float_type> const &self, nb::ndarray<cfloat_t> states, nb::ndarray<float_type> data) {
+                auto states_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 2>(states);
+                auto data_mdspan = fp::__detail::ndarray_to_mdspan<float_type, 2>(data);
+
+                auto new_states = fp::__detail::owning_ndarray_like_mdspan<cfloat_t, 2>(states_mdspan);
+                auto new_states_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 2>(new_states);
+
+                self.apply_weighted(std::execution::par, new_states_mdspan, states_mdspan, data_mdspan);
+
+                return new_states;
+            },
+            "states"_a, "data"_a,
+            R"%(Apply the SummedPauliOp to a batch of states with corresponding weights.
+
+.. math::
+    \big(\sum_k x_{tk} \sum_i h_{ik} \mathcal{\hat{P}}_i \big) \ket{\psi_t}
+
+Parameters
+----------
+states : np.ndarray
+    The original state(s) represented as 2D numpy array (n_operators, n_states) for batched calculation.
+data : np.ndarray
+    The data to weight the operators corresponding to the states (n_operators, n_states)
+
+Returns
+-------
+np.ndarray
+    New state(s) in a form of 2D numpy array (n_operators, n_states) according to the shape of input states
+)%")
+        .def(
+            "expectation_value",
+            [](fp::SummedPauliOp<float_type> const &self, nb::ndarray<cfloat_t> states) {
+                auto states_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 2>(states);
+
+                std::array<size_t, 2> out_shape = {self.n_operators(), states_mdspan.extent(1)};
+                auto expected_vals_out = fp::__detail::owning_ndarray_from_shape<cfloat_t, 2>(out_shape);
+                auto expected_vals_out_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 2>(expected_vals_out);
+
+                self.expectation_value(std::execution::par, expected_vals_out_mdspan, states_mdspan);
+
+                return expected_vals_out;
+            },
+            "states"_a,
+            R"%(Calculate expectation value(s) for a given batch of states.
+
+.. math::
+    \bra{\psi_t} \big(\sum_k \sum_i h_{ik} \mathcal{\hat{P}}_i \big) \ket{\psi_t}
+
+Parameters
+----------
+states : np.ndarray
+    The state(s) represented as 2D numpy array (n_operators, n_states) for batched calculation.
+
+Returns
+-------
+np.ndarray
+    Expectation value(s) in a form of 2D numpy array (n_operators, n_states) according to the shape of input states
+)%");
+    //
+    ;
 
     //
     // Helpers