Qutrit mixed test differentiability (#5186)

**Context:**
Measuring and applying operation functionality have been added to the
`qutrit_mixed` module, these tests ensure that these operations are
differentiable.
Since support for qml.prod for qutrits has been implemented from
[PR](#5400). It will be
possible to change the method of state measurement without issue.
The way expectation value measurement was done may have been overly
complicated so this reducing that complication, while still allowing for
product type observables.

**Description of the Change:**
Tests have been added that ensure that the `apply_operation` and
`measure` functions retain differentiability for the new qutrit
mixed-state device.
Expectation value measurements have been changed to use diagonalizing
gates and eigenvalues instead of trace.

**Benefits:**
Shows that the `apply_operation` and `measure` functions are
differentiable.

**Possible Drawbacks:**
This PR does not remove the abstractions made in apply_operations to
make the trace based measure expectation value work so there is some
code left that may be overly abstract.

**Related GitHub Issues:**
N/A

---------

Co-authored-by: Gabe PC <bottrill@student.ubc.ca>
Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
Co-authored-by: Thomas R. Bromley <49409390+trbromley@users.noreply.github.com>
Co-authored-by: Christina Lee <chrissie.c.l@gmail.com>
Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com>

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)

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:

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)