Repository: PennyLaneAI/pennylane
Pull request #4861 (2023-11-21 02:57:24)

Qutrit mixed initialization
**Context:**
I am working on implementing qutrit mixed state simulation to PennyLane using the new device API. Adding this helper function is necessary to continue. I based it off the qubit initialize_state helper function.

**Description of the Change:**
Added `qutrit_mixed` module to `qml.devices` to store helper functions for future qutrit mixed state device
Added create_initial_state function which creates initial states for future qutrit mixed device

**Benefits:**
Is a step forward to creating qutrit mixed state device using new device API
The code has been commented to what can be abstracted for qudits in the future (qudit_dimension is not a magic value)

**Possible Drawbacks:**
Structure will need to be changed for abstracting qudits

**Related GitHub Issues:**
N/A


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4861?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `1 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`1f47116`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/1f471160bf6ec57cbed949a6b8c339613869cc3f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.65% compared to head [(`19afcd2`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4861?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.63%.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4861?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/devices/qutrit\_mixed/initialize\_state.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4861?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXV0cml0X21peGVkL2luaXRpYWxpemVfc3RhdGUucHk=) | 96.96% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4861?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #4861      +/-   ##
==========================================
- Coverage   99.65%   99.63%   -0.02%     
==========================================
  Files         382      384       +2     
  Lines       34500    34274     -226     
==========================================
- Hits        34380    34150     -230     
- Misses        120      124       +4     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4861?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

Gabriel-Bottrill: @albi3ro, @glassnotes, @trbromley, @mudit2812, this is the first PR to the new qutrit mixed state device. Would you like me to make an issue or anything to attach this to to keep track of it?
trbromley: Thanks @Gabriel-Bottrill! We'll take a look. What are the follow up PRs you are planning?

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index fc8eec90f..90debbcc3 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -288,6 +288,9 @@
 
   The i-th element in the list refers to the 5 samples generated by the i-th measurement of the algorithm.
 
+* `devices.qutrit_mixed.initialize_state` module added with functionality to initialize qutrit mixed states.
+  [(#4861)](https://github.com/PennyLaneAI/pennylane/pull/4861)
+
 <h3>Improvements 🛠</h3>
 
 * `ClassicalShadow.entropy` now uses the algorithm outlined in 
diff --git a/pennylane/devices/qutrit_mixed/__init__.py b/pennylane/devices/qutrit_mixed/__init__.py
new file mode 100644
index 000000000..e6a696b2f
--- /dev/null
+++ b/pennylane/devices/qutrit_mixed/__init__.py
@@ -0,0 +1,27 @@
+# Copyright 2018-2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Submodule for performing mixed state simulations of qutrit-based quantum circuits.
+
+This submodule is internal and subject to change without a deprecation cycle. Use
+at your own discretion.
+
+.. currentmodule:: pennylane.devices.qutrit_mixed
+.. autosummary::
+    :toctree: api
+
+    create_initial_state
+"""
+
+from .initialize_state import create_initial_state
diff --git a/pennylane/devices/qutrit_mixed/initialize_state.py b/pennylane/devices/qutrit_mixed/initialize_state.py
new file mode 100644
index 000000000..23c21b318
--- /dev/null
+++ b/pennylane/devices/qutrit_mixed/initialize_state.py
@@ -0,0 +1,86 @@
+# Copyright 2018-2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Functions to prepare a qutrit mixed state."""
+
+from typing import Iterable, Union
+import pennylane as qml
+from pennylane.operation import StatePrepBase
+
+qudit_dim = 3  # specifies qudit dimension
+
+
+def create_initial_state(
+    wires: Union[qml.wires.Wires, Iterable],
+    prep_operation: StatePrepBase = None,
+    like: str = None,
+):
+    r"""
+    Returns an initial state, defaulting to :math:`\ket{0}\bra{0}` if no state-prep operator is provided.
+
+    Args:
+        wires (Union[Wires, Iterable]): The wires to be present in the initial state
+        prep_operation (Optional[StatePrepBase]): An operation to prepare the initial state
+        like (Optional[str]): The machine learning interface used to create the initial state.
+            Defaults to None
+
+    Returns:
+        array: The initial state of a circuit
+    """
+    num_wires = len(wires)
+
+    if not prep_operation:
+        rho = _create_basis_state(num_wires, 0)
+
+    else:
+        rho = _apply_state_vector(prep_operation.state_vector(wire_order=wires), num_wires)
+
+    # TODO: add instance for prep_operations as added
+
+    return qml.math.asarray(rho, like=like)
+
+
+def _apply_state_vector(state, num_wires):  # function is easy to abstract for qudit
+    """Initialize the internal state in a specified pure state.
+
+    Args:
+        state (array[complex]): normalized input state of length
+            ``qudit_dim**num_wires``, where ``qudit_dim`` is the dimension of the system.
+        num_wires (int): number of wires that get initialized in the state
+
+    Returns:
+        array[complex]: complex array of shape ``[qudit_dim] * (2 * num_wires)``
+        representing the density matrix of this state, where ``qudit_dim`` is
+        the dimension of the system.
+    """
+
+    # Initialize the entire set of wires with the state
+    rho = qml.math.outer(state, qml.math.conj(state))
+    return qml.math.reshape(rho, [qudit_dim] * 2 * num_wires)
+
+
+def _create_basis_state(num_wires, index):  # function is easy to abstract for qudit
+    """Return the density matrix representing a computational basis state over all wires.
+
+    Args:
+        num_wires (int): number of wires to initialize
+        index (int): integer representing the computational basis state.
+
+    Returns:
+        array[complex]: complex array of shape ``[qudit_dim] * (2 * num_wires)``
+        representing the density matrix of the basis state, where ``qudit_dim`` is
+        the dimension of the system.
+    """
+    rho = qml.math.zeros((qudit_dim**num_wires, qudit_dim**num_wires))
+    rho[index, index] = 1
+    return qml.math.reshape(rho, [qudit_dim] * (2 * num_wires))


Test patch:
diff --git a/tests/devices/qutrit_mixed/test_qutrit_mixed_initialize_state.py b/tests/devices/qutrit_mixed/test_qutrit_mixed_initialize_state.py
new file mode 100644
index 000000000..f297be2e3
--- /dev/null
+++ b/tests/devices/qutrit_mixed/test_qutrit_mixed_initialize_state.py
@@ -0,0 +1,93 @@
+# Copyright 2018-2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for create_initial_state in devices/qutrit_mixed."""
+
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+from pennylane.devices.qutrit_mixed import create_initial_state
+from pennylane import (
+    QutritBasisState,
+)
+from pennylane.operation import StatePrepBase
+
+
+class TestInitializeState:
+    """Test the functions in qutrit_mixed/initialize_state.py"""
+
+    # pylint:disable=unused-argument,too-few-public-methods
+    class DefaultPrep(StatePrepBase):
+        """A dummy class that assumes it was given a state vector."""
+
+        num_wires = qml.operation.AllWires
+
+        def state_vector(self, wire_order=None):
+            return self.parameters[0]
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "jax", "torch", "tensorflow"])
+    def test_create_initial_state_no_state_prep(self, interface):
+        """Tests that create_initial_state works without a state-prep operation."""
+        state = create_initial_state([0, 1], like=interface)
+        expected = np.zeros((3, 3, 3, 3))
+        expected[0, 0, 0, 0] = 1
+        assert qml.math.allequal(state, expected)
+        assert qml.math.get_interface(state) == interface
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "jax", "torch", "tensorflow"])
+    def test_create_initial_state_with_state_prep(self, interface):
+        """Tests that create_initial_state works with a state-prep operation."""
+        prep_op = self.DefaultPrep(
+            qml.math.array([1 / np.sqrt(9)] * 9, like=interface), wires=[0, 1]
+        )
+        state = create_initial_state([0, 1], prep_operation=prep_op)
+        expected = np.reshape([1 / 9] * 81, [3, 3, 3, 3])
+
+        assert qml.math.allequal(state, expected)
+        if interface == "autograd":
+            assert qml.math.get_interface(state) == "numpy"
+        else:
+            assert qml.math.get_interface(state) == interface
+
+    def test_create_initial_state_with_BasisState(self):
+        """Tests that create_initial_state works with a real state-prep operator."""
+        prep_op = QutritBasisState([1, 2, 0], wires=[0, 1, 2])
+        state = create_initial_state([0, 1, 2], prep_operation=prep_op)
+        assert state[1, 2, 0, 1, 2, 0] == 1
+        state[1, 2, 0, 1, 2, 0] = 0  # set to zero to make test below simple
+        assert qml.math.allequal(state, np.zeros(([3] * 6)))
+
+    @pytest.mark.parametrize("wires", [(0, 1), qml.wires.Wires([0, 1])])
+    def test_create_initial_state_wires(self, wires):
+        """Tests that create_initial_state works with qml.Wires object and list."""
+        state = create_initial_state(wires)
+        expected = np.zeros((3, 3, 3, 3))
+        expected[0, 0, 0, 0] = 1
+        assert qml.math.allequal(state, expected)
+
+    # TODO: Add tests for qutrit state prep
+
+    def test_create_initial_state_defaults_to_numpy(self):
+        """Tests that the default interface is vanilla numpy."""
+        state = qml.devices.qubit.create_initial_state((0, 1))
+        assert qml.math.get_interface(state) == "numpy"
+
+    @pytest.mark.torch
+    def test_create_initial_state_casts_to_like_with_prep_op(self):
+        """Tests that the like argument is not ignored when a prep-op is provided."""
+        prep_op = self.DefaultPrep([1 / np.sqrt(9)] * 9, wires=[0, 1])
+        state = create_initial_state([0, 1], prep_operation=prep_op, like="torch")
+        assert qml.math.get_interface(state) == "torch"