Skip to content

Commit

Permalink
[feature/SPO_apply] Adding more python tests for SummedPauliOp
Browse files Browse the repository at this point in the history
  • Loading branch information
@jamesETsmith
jamesETsmith committed Oct 31, 2024
1 parent f117bfd commit 14e366a
Show file tree
Hide file tree
Showing 2 changed files with 115 additions and 23 deletions.
94 changes: 71 additions & 23 deletions fast_pauli/cpp/src/fast_pauli.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -974,11 +974,28 @@ int
.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;
if (states.ndim() == 1)
{
auto states_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 1>(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<cfloat_t, 1>(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<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;
}
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.
Expand All @@ -1000,15 +1017,31 @@ np.ndarray
.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;
if (states.ndim() == 1)
{
auto states_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 1>(states);
auto states_mdspan_2d = std::mdspan(states_mdspan.data_handle(), states_mdspan.size(), 1);
auto data_mdspan = fp::__detail::ndarray_to_mdspan<float_type, 1>(data);
auto data_mdspan_2d = std::mdspan(data_mdspan.data_handle(), data_mdspan.size(), 1);
auto new_states = fp::__detail::owning_ndarray_like_mdspan<cfloat_t, 1>(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<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;
}
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.
Expand All @@ -1031,15 +1064,30 @@ np.ndarray
.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;
if (states.ndim() == 1)
{
auto states_mdspan = fp::__detail::ndarray_to_mdspan<cfloat_t, 1>(states);
auto states_mdspan_2d = std::mdspan(states_mdspan.data_handle(), states_mdspan.size(), 1);
std::array<size_t, 1> out_shape = {1};
auto expected_vals_out = fp::__detail::owning_ndarray_from_shape<cfloat_t, 1>(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<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;
}
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.
Expand Down
44 changes: 44 additions & 0 deletions tests/fast_pauli/test_summed_pauli_op.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand All @@ -43,13 +62,21 @@ 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
new_states = op.apply(psi)

# 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)
Expand Down Expand Up @@ -79,13 +106,22 @@ 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
new_states = op.apply_weighted(psi, data_weights)

# 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]
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down

0 comments on commit 14e366a

Please sign in to comment.