diff --git a/fast_pauli/cpp/src/fast_pauli.cpp b/fast_pauli/cpp/src/fast_pauli.cpp index a53cd36..dd27bbf 100644 --- a/fast_pauli/cpp/src/fast_pauli.cpp +++ b/fast_pauli/cpp/src/fast_pauli.cpp @@ -974,11 +974,28 @@ int .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; + if (states.ndim() == 1) + { + auto states_mdspan = fp::__detail::ndarray_to_mdspan(states); + auto states_mdspan_2d = std::mdspan(states_mdspan.data_handle(), states_mdspan.extent(0), 1); + auto new_states = fp::__detail::owning_ndarray_like_mdspan(states_mdspan); + auto new_states_mdspan = std::mdspan(new_states.data(), new_states.size(), 1); + self.apply(std::execution::par, new_states_mdspan, states_mdspan_2d); + return new_states; + } + else if (states.ndim() == 2) + { + 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; + } + else + { + throw std::invalid_argument( + fmt::format("apply: expected 1 or 2 dimensions, got {}", states.ndim())); + } }, "states"_a, R"%(Apply the SummedPauliOp to a batch of states. @@ -1000,15 +1017,31 @@ np.ndarray .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); - - auto new_states = fp::__detail::owning_ndarray_like_mdspan(states_mdspan); - auto new_states_mdspan = fp::__detail::ndarray_to_mdspan(new_states); - - self.apply_weighted(std::execution::par, new_states_mdspan, states_mdspan, data_mdspan); - - return new_states; + if (states.ndim() == 1) + { + auto states_mdspan = fp::__detail::ndarray_to_mdspan(states); + auto states_mdspan_2d = std::mdspan(states_mdspan.data_handle(), states_mdspan.size(), 1); + auto data_mdspan = fp::__detail::ndarray_to_mdspan(data); + auto data_mdspan_2d = std::mdspan(data_mdspan.data_handle(), data_mdspan.size(), 1); + auto new_states = fp::__detail::owning_ndarray_like_mdspan(states_mdspan); + auto new_states_mdspan = std::mdspan(new_states.data(), new_states.size(), 1); + self.apply_weighted(std::execution::par, new_states_mdspan, states_mdspan_2d, data_mdspan_2d); + return new_states; + } + else if (states.ndim() == 2) + { + auto states_mdspan = fp::__detail::ndarray_to_mdspan(states); + auto data_mdspan = fp::__detail::ndarray_to_mdspan(data); + auto new_states = fp::__detail::owning_ndarray_like_mdspan(states_mdspan); + auto new_states_mdspan = fp::__detail::ndarray_to_mdspan(new_states); + self.apply_weighted(std::execution::par, new_states_mdspan, states_mdspan, data_mdspan); + return new_states; + } + else + { + throw std::invalid_argument( + fmt::format("apply_weighted: expected 1 or 2 dimensions for states, got {}", states.ndim())); + } }, "states"_a, "data"_a, R"%(Apply the SummedPauliOp to a batch of states with corresponding weights. @@ -1031,15 +1064,30 @@ np.ndarray .def( "expectation_value", [](fp::SummedPauliOp const &self, nb::ndarray states) { - auto states_mdspan = fp::__detail::ndarray_to_mdspan(states); - - std::array out_shape = {self.n_operators(), states_mdspan.extent(1)}; - auto expected_vals_out = fp::__detail::owning_ndarray_from_shape(out_shape); - auto expected_vals_out_mdspan = fp::__detail::ndarray_to_mdspan(expected_vals_out); - - self.expectation_value(std::execution::par, expected_vals_out_mdspan, states_mdspan); - - return expected_vals_out; + if (states.ndim() == 1) + { + auto states_mdspan = fp::__detail::ndarray_to_mdspan(states); + auto states_mdspan_2d = std::mdspan(states_mdspan.data_handle(), states_mdspan.size(), 1); + std::array out_shape = {1}; + auto expected_vals_out = fp::__detail::owning_ndarray_from_shape(out_shape); + auto expected_vals_out_mdspan = std::mdspan(expected_vals_out.data(), expected_vals_out.size(), 1); + self.expectation_value(std::execution::par, expected_vals_out_mdspan, states_mdspan_2d); + return expected_vals_out; + } + else if (states.ndim() == 2) + { + auto states_mdspan = fp::__detail::ndarray_to_mdspan(states); + std::array out_shape = {self.n_operators(), states_mdspan.extent(1)}; + auto expected_vals_out = fp::__detail::owning_ndarray_from_shape(out_shape); + auto expected_vals_out_mdspan = fp::__detail::ndarray_to_mdspan(expected_vals_out); + self.expectation_value(std::execution::par, expected_vals_out_mdspan, states_mdspan); + return expected_vals_out; + } + else + { + throw std::invalid_argument( + fmt::format("expectation_value: expected 1 or 2 dimensions for states, got {}", states.ndim())); + } }, "states"_a, R"%(Calculate expectation value(s) for a given batch of states. diff --git a/tests/fast_pauli/test_summed_pauli_op.py b/tests/fast_pauli/test_summed_pauli_op.py index 9e2016a..8faed20 100644 --- a/tests/fast_pauli/test_summed_pauli_op.py +++ b/tests/fast_pauli/test_summed_pauli_op.py @@ -23,6 +23,25 @@ from tests.conftest import resolve_parameter_repr +@pytest.mark.parametrize("n_qubits,n_operators", [(2, 2), (3, 2), (4, 3)]) +def test_ctors(n_qubits: int, n_operators: int) -> None: + """Test the constructor of SummedPauliOp.""" + # Test with PauliStrings + 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) + op = fp.SummedPauliOp(pauli_strings, coeffs_2d) + assert op.dim == 2**n_qubits + assert op.n_operators == n_operators + assert op.n_pauli_strings == len(pauli_strings) + + # Test with list of strings + pauli_strings_str = [str(s) for s in pauli_strings] + op = fp.SummedPauliOp(pauli_strings_str, coeffs_2d) + assert op.dim == 2**n_qubits + assert op.n_operators == n_operators + assert op.n_pauli_strings == len(pauli_strings) + + @pytest.mark.parametrize( "summed_pauli_op", [fp.SummedPauliOp], ids=resolve_parameter_repr ) @@ -43,6 +62,10 @@ def test_apply( coeffs_2d = np.random.rand(n_strings, n_operators).astype(np.complex128) psi = np.random.rand(2**n_qubits, n_states).astype(np.complex128) + # For a single state, test that a 1D array works + if n_states == 1: + psi = psi[:, 0] + op = summed_pauli_op(pauli_strings, coeffs_2d) # The new_states we want to check @@ -50,6 +73,10 @@ def test_apply( # Trusted new_states new_states_naive = np.zeros((2**n_qubits, n_states), dtype=np.complex128) + # For a single state, test that a 1D array works + if n_states == 1: + new_states_naive = new_states_naive[:, 0] + for k in range(n_operators): A_k = fp.PauliOp(coeffs_2d[:, k].copy(), pauli_strings) new_states_naive += A_k.apply(psi) @@ -79,6 +106,11 @@ def test_apply_weighted( psi = np.random.rand(2**n_qubits, n_states).astype(np.complex128) data_weights = np.random.rand(n_operators, n_states).astype(np.float64) + # For a single state, test that a 1D array works + if n_states == 1: + psi = psi[:, 0] + data_weights = data_weights[:, 0] + op = summed_pauli_op(pauli_strings, coeffs_2d) # The new_states we want to check @@ -86,6 +118,10 @@ def test_apply_weighted( # Trusted new_states new_states_naive = np.zeros((2**n_qubits, n_states), dtype=np.complex128) + # For a single state, test that a 1D array works + if n_states == 1: + new_states_naive = new_states_naive[:, 0] + for k in range(n_operators): A_k = fp.PauliOp(coeffs_2d[:, k].copy(), pauli_strings) new_states_naive += A_k.apply(psi) * data_weights[k] @@ -116,6 +152,10 @@ def test_expectation_values( coeffs_2d = np.random.rand(n_strings, n_operators).astype(np.complex128) psi = np.random.rand(2**n_qubits, n_states).astype(np.complex128) + # For a single state, test that a 1D array works + if n_states == 1: + psi = psi[:, 0] + op = summed_pauli_op(pauli_strings, coeffs_2d) # The expectation values we want to test @@ -124,6 +164,10 @@ def test_expectation_values( # The "trusted" expectation_values expectation_vals_naive = np.zeros((n_operators, n_states), dtype=np.complex128) + # For a single state, test that a 1D array works + if n_states == 1: + expectation_vals_naive = expectation_vals_naive[:, 0] + # Calculate using brute force for k in range(n_operators): A_k = pp.helpers.naive_pauli_operator(coeffs_2d[:, k], pauli_strings_str)