diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 2713c21fc81..b117f4b3f82 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -230,6 +230,9 @@ * `qml.transforms.split_non_commuting` will now work with single-term operator arithmetic. [(#5314)](https://github.com/PennyLaneAI/pennylane/pull/5314) +* Fixed differentiability for Hamiltonian measurements in new `qutrit_mixed` module. + [(#5186)](https://github.com/PennyLaneAI/pennylane/pull/5186) + * Added `simulate` function to the new `qutrit_mixed` module in `qml.devices`. This allows for simulation of a noisy qutrit circuit with measurement and sampling. [(#5213)](https://github.com/PennyLaneAI/pennylane/pull/5213) diff --git a/pennylane/devices/qutrit_mixed/measure.py b/pennylane/devices/qutrit_mixed/measure.py index d05bd85da40..236071bfce0 100644 --- a/pennylane/devices/qutrit_mixed/measure.py +++ b/pennylane/devices/qutrit_mixed/measure.py @@ -27,66 +27,16 @@ ProbabilityMP, VarianceMP, ) -from pennylane.operation import Observable from pennylane.typing import TensorLike -from .utils import ( - get_einsum_mapping, - reshape_state_as_matrix, - get_num_wires, - get_new_state_einsum_indices, - QUDIT_DIM, -) +from .utils import reshape_state_as_matrix, get_num_wires, QUDIT_DIM from .apply_operation import apply_operation -def _map_indices_apply_operation(**kwargs): - """Map indices to wires. - - Args: - **kwargs (dict): Stores indices calculated in `get_einsum_mapping`: - state_indices (str): Indices that are summed. - row_indices (str): Indices that must be replaced with sums. - new_row_indices (str): Tensor indices of the state. - - Returns: - String of einsum indices to complete einsum calculations. - """ - op_1_indices = f"{kwargs['new_row_indices']}{kwargs['row_indices']}" - - new_state_indices = get_new_state_einsum_indices( - old_indices=kwargs["row_indices"], - new_indices=kwargs["new_row_indices"], - state_indices=kwargs["state_indices"], - ) - - return f"{op_1_indices},...{kwargs['state_indices']}->...{new_state_indices}" - - -def apply_observable_einsum(obs: Observable, state, is_state_batched: bool = False): - r"""Applies an observable to a density matrix rho, giving obs@state. - - Args: - obs (Operator): Operator to apply to the quantum state. - state (array[complex]): Input quantum state. - is_state_batched (bool): Boolean representing whether the state is batched or not. - - Returns: - TensorLike: the result of obs@state. - """ - - num_ch_wires = len(obs.wires) - einsum_indices = get_einsum_mapping(obs, state, _map_indices_apply_operation, is_state_batched) - obs_mat = obs.matrix() - obs_shape = [QUDIT_DIM] * num_ch_wires * 2 - obs_mat = math.cast(math.reshape(obs_mat, obs_shape), complex) - return math.einsum(einsum_indices, obs_mat, state) - - def calculate_expval( measurementprocess: ExpectationMP, state: TensorLike, is_state_batched: bool = False ) -> TensorLike: - """Measure the expectation value of an observable by finding the trace of obs@rho. + """Measure the expectation value of an observable. Args: measurementprocess (ExpectationMP): measurement process to apply to the state. @@ -96,21 +46,13 @@ def calculate_expval( Returns: TensorLike: expectation value of observable wrt the state. """ - obs = measurementprocess.obs - rho_mult_obs = apply_observable_einsum(obs, state, is_state_batched) - - # using einsum since trace function axis selection parameter names - # are not consistent across interfaces, they don't exist for torch - - num_wires = get_num_wires(state, is_state_batched) - rho_mult_obs_reshaped = reshape_state_as_matrix(rho_mult_obs, num_wires) - if is_state_batched: - return math.real(math.stack([math.sum(math.diagonal(dm)) for dm in rho_mult_obs_reshaped])) - - return math.real(math.sum(math.diagonal(rho_mult_obs_reshaped))) + probs = calculate_probability(measurementprocess, state, is_state_batched) + eigvals = math.asarray(measurementprocess.eigvals(), dtype="float64") + # In case of broadcasting, `probs` has two axes and these are a matrix-vector products + return math.dot(probs, eigvals) -def calculate_reduced_density_matrix( # TODO: ask if I should have state diagonalization gates? +def calculate_reduced_density_matrix( measurementprocess: StateMeasurement, state: TensorLike, is_state_batched: bool = False ) -> TensorLike: """Get the state or reduced density matrix. @@ -260,8 +202,6 @@ def get_measurement_function( """ if isinstance(measurementprocess, StateMeasurement): if isinstance(measurementprocess, ExpectationMP): - # TODO add faster methods - # TODO add support for sparce Hamiltonians if isinstance(measurementprocess.obs, (Hamiltonian, Sum)): return calculate_expval_sum_of_terms if measurementprocess.obs.has_matrix: diff --git a/tests/devices/qutrit_mixed/test_qutrit_mixed_apply_operation.py b/tests/devices/qutrit_mixed/test_qutrit_mixed_apply_operation.py index c5bb0255ab2..30e89df56cd 100644 --- a/tests/devices/qutrit_mixed/test_qutrit_mixed_apply_operation.py +++ b/tests/devices/qutrit_mixed/test_qutrit_mixed_apply_operation.py @@ -20,7 +20,7 @@ import pennylane as qml from pennylane import math from pennylane.operation import Channel -from pennylane.devices.qutrit_mixed import apply_operation +from pennylane.devices.qutrit_mixed import apply_operation, measure ml_frameworks_list = [ "numpy", @@ -31,6 +31,24 @@ ] subspaces = [(0, 1), (0, 2), (1, 2)] +kraus_matrix = np.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]], dtype=complex) + + +class CustomChannel(Channel): # pylint: disable=too-few-public-methods + num_params = 1 + num_wires = 1 + + def __init__(self, p, wires, id=None): + super().__init__(p, wires=wires, id=id) + + @staticmethod + def compute_kraus_matrices(p): + if math.get_interface(p) == "tensorflow": + p = math.cast_like(p, 1j) + + K0 = math.sqrt(1 - p + math.eps) * math.convert_like(math.eye(3, dtype=complex), p) + K1 = math.sqrt(p + math.eps) * math.convert_like(kraus_matrix, p) + return [K0, K1] def test_custom_operator_with_matrix(one_qutrit_state): @@ -320,3 +338,223 @@ def test_broadcasted_state(self, ml_framework, wire, two_qutrit_batched_state): assert qml.math.get_interface(res) == ml_framework assert qml.math.allclose(res, expected) + + +@pytest.mark.parametrize("subspace", [(0, 1), (0, 2), (1, 2)]) +class TestTRXCalcGrad: + """Tests the application and differentiation of a TRX gate in the different interfaces.""" + + phi = 0.325 + + @staticmethod + def compare_expected_result(phi, state, probs, subspace, jacobian): + """Compare the expected result for this circuit and gradient with observed values""" + trx = qml.TRX.compute_matrix(phi, subspace) + trx_adj = qml.TRX.compute_matrix(-phi, subspace) + state = math.reshape(state, (9, 9)) + + expected_probs = math.diagonal( + np.kron(trx, np.eye(3)) @ state @ np.kron(trx_adj, np.eye(3)) + ) + assert qml.math.allclose(probs, expected_probs) + + if subspace[0] == 0: + gell_mann_index = 1 if subspace[1] == 1 else 4 + else: + gell_mann_index = 6 + gell_mann_matrix = qml.GellMann.compute_matrix(gell_mann_index) + trx_derivative = -0.5j * gell_mann_matrix @ trx + trx_adj_derivative = 0.5j * gell_mann_matrix @ trx_adj + + expected_derivative_state = ( + np.kron(trx_derivative, np.eye(3)) @ state @ np.kron(trx_adj, np.eye(3)) + ) + (np.kron(trx, np.eye(3)) @ state @ np.kron(trx_adj_derivative, np.eye(3))) + expected_derivative = np.diagonal(expected_derivative_state) + assert qml.math.allclose(jacobian, expected_derivative) + + @pytest.mark.autograd + def test_trx_grad_autograd(self, two_qutrit_state, subspace): + """Test that the application of a trx gate is differentiable with autograd.""" + + state = qml.numpy.array(two_qutrit_state) + + def f(phi): + op = qml.TRX(phi, wires=0, subspace=subspace) + new_state = apply_operation(op, state) + return measure(qml.probs(), new_state) + + phi = qml.numpy.array(self.phi, requires_grad=True) + + probs = f(phi) + jacobian = qml.jacobian(lambda x: qml.math.real(f(x)))(phi) + self.compare_expected_result(phi, state, probs, subspace, jacobian) + + @pytest.mark.jax + @pytest.mark.parametrize("use_jit", (True, False)) + def test_trx_grad_jax(self, use_jit, two_qutrit_state, subspace): + """Test that the application of a trx gate is differentiable with jax.""" + + import jax + + state = jax.numpy.array(two_qutrit_state) + + def f(phi): + op = qml.TRX(phi, wires=0, subspace=subspace) + new_state = apply_operation(op, state) + return measure(qml.probs(), new_state) + + if use_jit: + f = jax.jit(f) + + probs = f(self.phi) + jacobian = jax.jacobian(f)(self.phi) + self.compare_expected_result(self.phi, state, probs, subspace, jacobian) + + @pytest.mark.torch + def test_trx_grad_torch(self, two_qutrit_state, subspace): + """Tests the application and differentiation of a trx gate with torch.""" + + import torch + + state = torch.tensor(two_qutrit_state) + + def f(phi): + op = qml.TRX(phi, wires=0, subspace=subspace) + new_state = apply_operation(op, state) + return measure(qml.probs(), new_state) + + phi = torch.tensor(self.phi, requires_grad=True) + + probs = f(phi) + jacobian = torch.autograd.functional.jacobian(f, phi) + + self.compare_expected_result( + phi.detach().numpy(), + state.detach().numpy(), + probs.detach().numpy(), + subspace, + jacobian.detach().numpy(), + ) + + @pytest.mark.tf + def test_trx_grad_tf(self, two_qutrit_state, subspace): + """Tests the application and differentiation of a trx gate with tensorflow""" + import tensorflow as tf + + state = tf.Variable(two_qutrit_state) + phi = tf.Variable(0.8589, trainable=True) + + with tf.GradientTape() as grad_tape: + op = qml.TRX(phi, wires=0, subspace=subspace) + new_state = apply_operation(op, state) + probs = measure(qml.probs(), new_state) + + jacobians = grad_tape.jacobian(probs, [phi]) + phi_jacobian = jacobians[0] + + self.compare_expected_result(phi, state, probs, subspace, phi_jacobian) + + +class TestChannelCalcGrad: + """Tests the application and differentiation of a Channel in the different interfaces.""" + + p = 0.325 + + @staticmethod + def compare_expected_result(p, state, new_state, jacobian): + """Compare the expected result for this channel and gradient with observed values""" + kraus_matrix_two_qutrits = np.kron(np.eye(3), kraus_matrix) + kraus_matrix_two_qutrits_adj = kraus_matrix_two_qutrits.transpose() + state = math.reshape(state, (9, 9)) + + state_kraus_applied = kraus_matrix_two_qutrits @ state @ kraus_matrix_two_qutrits_adj + + expected_state = (1 - p) * state + (p * state_kraus_applied) + expected_probs = np.diagonal(expected_state) + assert qml.math.allclose(new_state, expected_probs) + + expected_derivative_state = state_kraus_applied - state + expected_derivative = np.diagonal(expected_derivative_state) + assert qml.math.allclose(jacobian, expected_derivative) + + @pytest.mark.autograd + def test_channel_grad_autograd(self, two_qutrit_state): + """Test that the application of a channel is differentiable with autograd.""" + + state = qml.numpy.array(two_qutrit_state) + + def f(p): + channel = CustomChannel(p, wires=1) + new_state = apply_operation(channel, state) + return measure(qml.probs(), new_state) + + p = qml.numpy.array(self.p, requires_grad=True) + + probs = f(p) + jacobian = qml.jacobian(lambda x: qml.math.real(f(x)))(p) + self.compare_expected_result(p, state, probs, jacobian) + + @pytest.mark.jax + @pytest.mark.parametrize("use_jit", (True, False)) + def test_channel_grad_jax(self, use_jit, two_qutrit_state): + """Test that the application of a channel is differentiable with jax.""" + + import jax + + state = jax.numpy.array(two_qutrit_state) + + def f(p): + op = CustomChannel(p, wires=1) + new_state = apply_operation(op, state) + return measure(qml.probs(), new_state) + + if use_jit: + f = jax.jit(f) + + probs = f(self.p) + jacobian = jax.jacobian(f)(self.p) + self.compare_expected_result(self.p, state, probs, jacobian) + + @pytest.mark.torch + def test_channel_grad_torch(self, two_qutrit_state): + """Tests the application and differentiation of a channel with torch.""" + + import torch + + state = torch.tensor(two_qutrit_state) + + def f(p): + op = CustomChannel(p, wires=1) + new_state = apply_operation(op, state) + return measure(qml.probs(), new_state) + + p = torch.tensor(self.p, requires_grad=True) + + probs = f(p) + jacobian = torch.autograd.functional.jacobian(f, p) + + self.compare_expected_result( + p.detach().numpy(), + state.detach().numpy(), + probs.detach().numpy(), + jacobian.detach().numpy(), + ) + + @pytest.mark.tf + def test_channel_grad_tf(self, two_qutrit_state): + """Tests the application and differentiation of a channel with tensorflow""" + import tensorflow as tf + + state = tf.Variable(two_qutrit_state) + p = tf.Variable(0.8589 + 0j, trainable=True) + + with tf.GradientTape() as grad_tape: + op = CustomChannel(p, wires=1) + new_state = apply_operation(op, state) + probs = measure(qml.probs(), new_state) + + jacobians = grad_tape.jacobian(probs, [p]) + # tf takes gradient with respect to conj(z), so we need to conj the gradient + phi_jacobian = tf.math.conj(jacobians[0]) + + self.compare_expected_result(p, state, probs, phi_jacobian) diff --git a/tests/devices/qutrit_mixed/test_qutrit_mixed_measure.py b/tests/devices/qutrit_mixed/test_qutrit_mixed_measure.py index d2801ffd213..58c13661aa4 100644 --- a/tests/devices/qutrit_mixed/test_qutrit_mixed_measure.py +++ b/tests/devices/qutrit_mixed/test_qutrit_mixed_measure.py @@ -20,9 +20,8 @@ from pennylane import numpy as np from pennylane import math -from pennylane.devices.qutrit_mixed import measure +from pennylane.devices.qutrit_mixed import measure, create_initial_state, apply_operation from pennylane.devices.qutrit_mixed.measure import ( - apply_observable_einsum, get_measurement_function, calculate_expval, calculate_expval_sum_of_terms, @@ -118,49 +117,6 @@ def test_var_compute_variance(self): assert get_measurement_function(qml.var(obs)) is calculate_variance -@pytest.mark.parametrize( - "obs", - [ - qml.GellMann(2, 2), - qml.GellMann(1, 8), - qml.GellMann(0, 5), - (qml.GellMann(0, 5) @ qml.GellMann(1, 8)), - ], -) -@pytest.mark.parametrize("ml_framework", ml_frameworks_list) -class TestApplyObservableEinsum: - """Tests that observables are applied correctly for calculate expval method.""" - - num_qutrits = 3 - dims = (3**num_qutrits, 3**num_qutrits) - - def test_apply_observable_einsum(self, obs, three_qutrit_state, ml_framework): - """Tests that unbatched observables are applied correctly to a unbatched state.""" - res = apply_observable_einsum( - obs, qml.math.asarray(three_qutrit_state, like=ml_framework), is_state_batched=False - ) - expected = get_expanded_op_mult_state(obs, three_qutrit_state) - expected = expected.reshape([3] * (3 * 2)) - - assert qml.math.get_interface(res) == ml_framework - assert qml.math.allclose(res, expected) - - def test_apply_observable_einsum_batched(self, obs, three_qutrit_batched_state, ml_framework): - """Tests that unbatched observables are applied correctly to a batched state.""" - res = apply_observable_einsum( - obs, - qml.math.asarray(three_qutrit_batched_state, like=ml_framework), - is_state_batched=True, - ) - expected = [ - get_expanded_op_mult_state(obs, state).reshape([3] * (3 * 2)) - for state in three_qutrit_batched_state - ] - - assert qml.math.get_interface(res) == ml_framework - assert qml.math.allclose(res, expected) - - class TestMeasurements: """Test that measurements on unbatched states work as expected.""" @@ -471,4 +427,177 @@ def test_variance_measurement(self, observable, ml_framework, two_qutrit_batched assert np.allclose(res, expected) -# TODO TestSumOfTermsDifferentiability in future PR (with other differentiabilty tests) +@pytest.mark.usefixtures("use_legacy_and_new_opmath") +class TestSumOfTermsDifferentiability: + x = 0.52 + + @staticmethod + def f(scale, coeffs, n_wires=5, offset=0.1): + """Function to differentiate that implements a circuit with a SumOfTerms operator""" + ops = [qml.TRX(offset + scale * i, wires=i, subspace=(0, 2)) for i in range(n_wires)] + H = qml.Hamiltonian( + coeffs, + [ + reduce(lambda x, y: x @ y, (qml.GellMann(i, 3) for i in range(n_wires))), + reduce(lambda x, y: x @ y, (qml.GellMann(i, 5) for i in range(n_wires))), + ], + ) + state = create_initial_state(range(n_wires), like=math.get_interface(scale)) + for op in ops: + state = apply_operation(op, state) + return measure(qml.expval(H), state) + + @staticmethod + def expected(scale, coeffs, n_wires=5, offset=0.1, like="numpy"): + """Get the expected expval of the class' circuit.""" + phase = offset + scale * qml.math.asarray(range(n_wires), like=like) + cosines = math.cos(phase / 2) ** 2 + sines = -math.sin(phase) + return coeffs[0] * qml.math.prod(cosines) + coeffs[1] * qml.math.prod(sines) + + @pytest.mark.autograd + @pytest.mark.parametrize( + "coeffs", + [ + (qml.numpy.array(2.5), qml.numpy.array(6.2)), + (qml.numpy.array(2.5, requires_grad=False), qml.numpy.array(6.2, requires_grad=False)), + ], + ) + def test_autograd_backprop(self, coeffs): + """Test that backpropagation derivatives work in autograd with + Hamiltonians using new and old math.""" + + x = qml.numpy.array(self.x) + out = self.f(x, coeffs) + expected_out = self.expected(x, coeffs) + assert qml.math.allclose(out, expected_out) + + gradient = qml.grad(self.f)(x, coeffs) + expected_gradient = qml.grad(self.expected)(x, coeffs) + assert qml.math.allclose(expected_gradient, gradient) + + @pytest.mark.autograd + def test_autograd_backprop_coeffs(self): + """Test that backpropagation derivatives work in autograd with + the coefficients of Hamiltonians using new and old math.""" + + coeffs = qml.numpy.array((2.5, 6.2), requires_grad=True) + gradient = qml.grad(self.f, argnum=1)(self.x, coeffs) + expected_gradient = qml.grad(self.expected)(self.x, coeffs) + + assert len(gradient) == 2 + assert qml.math.allclose(expected_gradient, gradient) + + @pytest.mark.jax + @pytest.mark.parametrize("use_jit", (True, False)) + def test_jax_backprop(self, use_jit): + """Test that backpropagation derivatives work with jax with + Hamiltonians using new and old math.""" + import jax + + jax.config.update("jax_enable_x64", True) + + x = jax.numpy.array(self.x, dtype=jax.numpy.float64) + coeffs = (5.2, 6.7) + f = jax.jit(self.f, static_argnums=(1, 2, 3, 4)) if use_jit else self.f + + out = f(x, coeffs) + expected_out = self.expected(x, coeffs) + assert qml.math.allclose(out, expected_out) + + gradient = jax.grad(f)(x, coeffs) + expected_gradient = jax.grad(self.expected)(x, coeffs) + assert qml.math.allclose(expected_gradient, gradient) + + @pytest.mark.jax + def test_jax_backprop_coeffs(self): + """Test that backpropagation derivatives work with jax with + the coefficients of Hamiltonians using new and old math.""" + import jax + + jax.config.update("jax_enable_x64", True) + coeffs = jax.numpy.array((5.2, 6.7), dtype=jax.numpy.float64) + + gradient = jax.grad(self.f, argnums=1)(self.x, coeffs) + expected_gradient = jax.grad(self.expected, argnums=1)(self.x, coeffs) + assert len(gradient) == 2 + assert qml.math.allclose(expected_gradient, gradient) + + @pytest.mark.torch + def test_torch_backprop(self): + """Test that backpropagation derivatives work with torch with + Hamiltonians using new and old math.""" + import torch + + coeffs = [ + torch.tensor(9.2, requires_grad=False, dtype=torch.float64), + torch.tensor(6.2, requires_grad=False, dtype=torch.float64), + ] + + x = torch.tensor(-0.289, requires_grad=True, dtype=torch.float64) + x2 = torch.tensor(-0.289, requires_grad=True, dtype=torch.float64) + out = self.f(x, coeffs) + expected_out = self.expected(x2, coeffs, like="torch") + assert qml.math.allclose(out, expected_out) + + out.backward() + expected_out.backward() + assert qml.math.allclose(x.grad, x2.grad) + + @pytest.mark.torch + def test_torch_backprop_coeffs(self): + """Test that backpropagation derivatives work with torch with + the coefficients of Hamiltonians using new and old math.""" + import torch + + coeffs = torch.tensor((9.2, 6.2), requires_grad=True, dtype=torch.float64) + coeffs_expected = torch.tensor((9.2, 6.2), requires_grad=True, dtype=torch.float64) + + x = torch.tensor(-0.289, requires_grad=False, dtype=torch.float64) + out = self.f(x, coeffs) + expected_out = self.expected(x, coeffs_expected, like="torch") + assert qml.math.allclose(out, expected_out) + + out.backward() + expected_out.backward() + assert len(coeffs.grad) == 2 + assert qml.math.allclose(coeffs.grad, coeffs_expected.grad) + + @pytest.mark.tf + def test_tf_backprop(self): + """Test that backpropagation derivatives work with tensorflow with + Hamiltonians using new and old math.""" + import tensorflow as tf + + x = tf.Variable(self.x) + coeffs = [8.3, 5.7] + + with tf.GradientTape() as tape1: + out = self.f(x, coeffs) + + with tf.GradientTape() as tape2: + expected_out = self.expected(x, coeffs) + + assert qml.math.allclose(out, expected_out) + gradient = tape1.gradient(out, x) + expected_gradient = tape2.gradient(expected_out, x) + assert qml.math.allclose(expected_gradient, gradient) + + @pytest.mark.tf + def test_tf_backprop_coeffs(self): + """Test that backpropagation derivatives work with tensorflow with + the coefficients of Hamiltonians using new and old math.""" + import tensorflow as tf + + coeffs = tf.Variable([8.3, 5.7]) + + with tf.GradientTape() as tape1: + out = self.f(self.x, coeffs) + + with tf.GradientTape() as tape2: + expected_out = self.expected(self.x, coeffs) + + gradient = tape1.gradient(out, coeffs) + expected_gradient = tape2.gradient(expected_out, coeffs) + assert len(gradient) == 2 + assert qml.math.allclose(expected_gradient, gradient)