fix expand_matrix for qudit matrices

doc/releases/changelog-dev.md

@@ -274,6 +274,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixes the matrices for qutrit operators when expanded to more wires.
+  [(#6398)](https://github.com/PennyLaneAI/pennylane/pull/6398/)
+
 * `MeasurementValue` now raises an error when it is used as a boolean.
   [(#6386)](https://github.com/PennyLaneAI/pennylane/pull/6386)
 

pennylane/math/matrix_manipulation.py

@@ -101,11 +101,14 @@ def expand_matrix(mat, wires, wire_order=None, sparse_format="csr"):
            [0., 0., 1., 0.]])
 
     """
+    wires = Wires(wires)
 
+    qudit_dim = int(qml.math.shape(mat)[-1] ** (1 / (len(wires))))
+    print(qudit_dim)
     if (wire_order is None) or (wire_order == wires):
         return mat
 
-    if not wires and qml.math.shape(mat) == (2, 2):
+    if not wires and qml.math.shape(mat) == (qudit_dim, qudit_dim):
         # global phase
         wires = wire_order[0:1]
 
@@ -118,8 +121,8 @@ def expand_matrix(mat, wires, wire_order=None, sparse_format="csr"):
 
     def eye_interface(dim):
         if interface == "scipy":
-            return eye(2**dim, format="coo")
-        return qml.math.cast_like(qml.math.eye(2**dim, like=interface), mat)
+            return eye(qudit_dim**dim, format="coo")
+        return qml.math.cast_like(qml.math.eye(qudit_dim**dim, like=interface), mat)
 
     def kron_interface(mat1, mat2):
         if interface == "scipy":
@@ -154,7 +157,9 @@ def kron_interface(mat1, mat2):
     if interface == "scipy":
         mat = _permute_sparse_matrix(mat, expanded_wires, subset_wire_order)
     else:
-        mat = _permute_dense_matrix(mat, expanded_wires, subset_wire_order, batch_dim)
+        mat = _permute_dense_matrix(
+            mat, expanded_wires, subset_wire_order, batch_dim, qudit_dim=qudit_dim
+        )
 
     # expand the matrix even further if needed
     if len(expanded_wires) < len(wire_order):
@@ -201,7 +206,7 @@ def _permute_sparse_matrix(matrix, wires, wire_order):
     return matrix
 
 
-def _permute_dense_matrix(matrix, wires, wire_order, batch_dim):
+def _permute_dense_matrix(matrix, wires, wire_order, batch_dim, qudit_dim: int = 2):
     """Permute the matrix to match the wires given in `wire_order`.
 
     Args:
@@ -228,12 +233,14 @@ def _permute_dense_matrix(matrix, wires, wire_order, batch_dim):
 
     # reshape matrix to match wire values e.g. mat[0, 0, 0, 0] = <00|mat|00>
     # with this reshape we can easily swap wires
-    shape = [batch_dim] + [2] * (num_wires * 2) if batch_dim else [2] * (num_wires * 2)
+    shape = (
+        [batch_dim] + [qudit_dim] * (num_wires * 2) if batch_dim else [qudit_dim] * (num_wires * 2)
+    )
     matrix = qml.math.reshape(matrix, shape)
     # transpose matrix
     matrix = qml.math.transpose(matrix, axes=perm)
     # reshape back
-    shape = [batch_dim] + [2**num_wires] * 2 if batch_dim else [2**num_wires] * 2
+    shape = [batch_dim] + [qudit_dim**num_wires] * 2 if batch_dim else [qudit_dim**num_wires] * 2
     return qml.math.reshape(matrix, shape)
 
 

tests/math/test_matrix_manipulation.py

@@ -22,6 +22,7 @@
 
 import pennylane as qml
 from pennylane import numpy as pnp
+from pennylane.math import expand_matrix
 
 # Define a list of dtypes to test
 dtypes = ["complex64", "complex128"]
@@ -545,6 +546,58 @@ def compute_matrix():
         assert np.allclose(op.matrix(wire_order=[0, 1, 2]), expanded_matrix, atol=tol)
 
 
+class TestExpandMatrixQutrit:
+
+    def test_add_wire_at_end(self):
+        """Test that expand_matrix works on qutrit matrices when an additional wire is added at the end."""
+
+        mat = np.reshape(np.arange(9), (3, 3))
+
+        new_mat = expand_matrix(mat, 0, (0, 1))
+        assert qml.math.allclose(new_mat, np.kron(mat, np.eye(3)))
+
+    def test_add_wire_at_start(self):
+        """Test that expand_matrix works on qutrit matrices when an additional wire is added at the start."""
+
+        mat = np.reshape(np.arange(9), (3, 3))
+        new_mat = expand_matrix(mat, 0, (1, 0))
+        assert qml.math.allclose(new_mat, np.kron(np.eye(3), mat))
+
+    def test_wire_permutation(self):
+        """Test that wires can be permuted."""
+        m1 = np.reshape(np.arange(81), (9, 9))
+        m2 = expand_matrix(m1, (0, 1), (1, 0))
+
+        # states across row are 00, 01, 02, 10, 11, 12, 20, 21, 22
+        # extract out right qubit state with mod
+        m1_wire_zero = m1 % 3
+        m2_wire_zero = m2 % 3
+
+        # extract out left qubit state  with floor then mod
+        m1_wire_one = np.floor(m1 / 3) % 3
+        m2_wire_one = np.floor(m2 / 3) % 3
+
+        assert qml.math.allclose(m1_wire_zero, m2_wire_one)
+        assert qml.math.allclose(m1_wire_one, m2_wire_zero)
+
+        # check columns also switched
+        # now matrix numbers indicate row number
+        m1p = np.floor(m1 / 9)
+        m2p = np.floor(m2 / 9)
+
+        # states across row are 00, 01, 02, 10, 11, 12, 20, 21, 22
+        # extract out right qubit state with mod
+        m1_wire_zerop = m1p % 3
+        m2_wire_zerop = m2p % 3
+
+        # extract out left qubit state  with floor then mod
+        m1_wire_onep = np.floor(m1p / 3) % 3
+        m2_wire_onep = np.floor(m2p / 3) % 3
+
+        assert qml.math.allclose(m1_wire_zerop, m2_wire_onep)
+        assert qml.math.allclose(m1_wire_onep, m2_wire_zerop)
+
+
 class TestExpandMatrixSparse:
     """Tests for the _sparse_expand_matrix function."""