Add simulate of analytic case for DefaultMixedNewAPI

doc/releases/changelog-dev.md

@@ -90,6 +90,10 @@ added `binary_mapping()` function to map `BoseWord` and `BoseSentence` to qubit
 * Added submodule `devices.qubit_mixed.measure` as a necessary step for the new API, featuring a `measure` function for measuring qubits in mixed-state devices.
   [(#6637)](https://github.com/PennyLaneAI/pennylane/pull/6637)
 
+* Added submodule `devices.qubit_mixed.simulate` as a necessary step for the new API,
+featuring a `simulate` function for simulating mixed states in analytic mode.
+  [(#6618)](https://github.com/PennyLaneAI/pennylane/pull/6618)
+
 * Support is added for `if`/`else` statements and `for` and `while` loops in circuits executed with `qml.capture.enabled`, via `autograph`
   [(#6406)](https://github.com/PennyLaneAI/pennylane/pull/6406)
   [(#6413)](https://github.com/PennyLaneAI/pennylane/pull/6413)

pennylane/devices/default_mixed.py

@@ -182,8 +182,7 @@ def stopping_condition(op: qml.operation.Operator) -> bool:
 def warn_readout_error_state(
     tape: qml.tape.QuantumTape,
 ) -> tuple[Sequence[qml.tape.QuantumTape], Callable]:
-    """If a measurement in the QNode is an analytic state or density_matrix, and a readout error
-    parameter is defined, warn that readout error will not be applied.
+    """If a measurement in the QNode is an analytic state or density_matrix, warn that readout error will not be applied.
 
     Args:
         tape (QuantumTape, .QNode, Callable): a quantum circuit.

pennylane/devices/qubit_mixed/__init__.py

@@ -24,7 +24,9 @@
     apply_operation
     create_initial_state
     measure
+    simulate
 """
 from .apply_operation import apply_operation
 from .initialize_state import create_initial_state
 from .measure import measure
+from .simulate import get_final_state, measure_final_state, simulate

pennylane/devices/qubit_mixed/initialize_state.py

@@ -14,6 +14,7 @@
 """Functions to prepare a state."""
 
 from collections.abc import Iterable
+from typing import Union
 
 import pennylane as qml
 import pennylane.numpy as np
@@ -22,8 +23,8 @@
 
 def create_initial_state(
     # pylint: disable=unsupported-binary-operation
-    wires: qml.wires.Wires | Iterable,
-    prep_operation: qml.operation.StatePrepBase | qml.QubitDensityMatrix = None,
+    wires: Union[qml.wires.Wires, Iterable],
+    prep_operation: Union[qml.operation.StatePrepBase, qml.QubitDensityMatrix] = None,
     like: str = None,
 ):
     r"""
@@ -60,7 +61,7 @@ def _post_process(density_matrix, num_axes, like):
     r"""
     This post processor is necessary to ensure that the density matrix is in the correct format, i.e. the original tensor form, instead of the pure matrix form, as requested by all the other more fundamental chore functions in the module (again from some legacy code).
     """
-    density_matrix = np.reshape(density_matrix, (-1,) + (2,) * num_axes)
+    density_matrix = np.reshape(density_matrix, (2,) * num_axes)
     dtype = str(density_matrix.dtype)
     floating_single = "float32" in dtype or "complex64" in dtype
     dtype = "complex64" if floating_single else "complex128"

pennylane/devices/qubit_mixed/measure.py

@@ -87,6 +87,7 @@ def calculate_reduced_density_matrix(
         state (TensorLike): state to apply the measurement to.
         is_state_batched (bool): whether the state is batched or not.
         readout_errors (List[Callable]): List of channels to apply to each wire being measured
+            to simulate readout errors. These are not applied on this type of measurement.
 
     Returns:
         TensorLike: state or reduced density matrix.

pennylane/devices/qubit_mixed/simulate.py

@@ -0,0 +1,188 @@
+# Copyright 2018-2024 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.
+"""Simulate a quantum script for a qubit mixed state device."""
+import pennylane as qml
+from pennylane.typing import Result
+
+from .apply_operation import apply_operation
+from .initialize_state import create_initial_state
+from .measure import measure
+
+INTERFACE_TO_LIKE = {
+    # map interfaces known by autoray to themselves
+    "numpy": "numpy",
+    "autograd": "autograd",
+    "jax": "jax",
+    "torch": "torch",
+    "tensorflow": "tensorflow",
+    # map non-standard interfaces to those known by autoray
+    "auto": None,
+    "scipy": "numpy",
+    "jax-jit": "jax",
+    "jax-python": "jax",
+    "JAX": "jax",
+    "pytorch": "torch",
+    "tf": "tensorflow",
+    "tensorflow-autograph": "tensorflow",
+    "tf-autograph": "tensorflow",
+}
+
+
+def get_final_state(circuit, debugger=None, interface=None, **kwargs):
+    """
+    Get the final state that results from executing the given quantum script.
+
+    This is an internal function that will be called by ``default.mixed``.
+
+    Args:
+        circuit (.QuantumScript): The single circuit to simulate
+        debugger (._Debugger): The debugger to use
+        interface (str): The machine learning interface to create the initial state with
+
+    Returns:
+        Tuple[TensorLike, bool]: A tuple containing the final state of the quantum script and
+            whether the state has a batch dimension.
+
+    """
+    circuit = circuit.map_to_standard_wires()
+
+    prep = None
+    if len(circuit) > 0 and isinstance(circuit[0], qml.operation.StatePrepBase):
+        prep = circuit[0]
+
+    state = create_initial_state(
+        sorted(circuit.op_wires), prep, like=INTERFACE_TO_LIKE[interface] if interface else None
+    )
+
+    # initial state is batched only if the state preparation (if it exists) is batched
+    is_state_batched = bool(prep and prep.batch_size is not None)
+    for op in circuit.operations[bool(prep) :]:
+        state = apply_operation(
+            op,
+            state,
+            is_state_batched=is_state_batched,
+            debugger=debugger,
+            tape_shots=circuit.shots,
+            **kwargs,
+        )
+
+        # new state is batched if i) the old state is batched, or ii) the new op adds a batch dim
+        is_state_batched = is_state_batched or op.batch_size is not None
+
+    num_operated_wires = len(circuit.op_wires)
+    for i in range(len(circuit.wires) - num_operated_wires):
+        # If any measured wires are not operated on, we pad the density matrix with zeros.
+        # We know they belong at the end because the circuit is in standard wire-order
+        # Since it is a dm, we must pad it with 0s on the last row and last column
+        current_axis = num_operated_wires + i + is_state_batched
+        state = qml.math.stack(([state] + [qml.math.zeros_like(state)]), axis=current_axis)
+        state = qml.math.stack(([state] + [qml.math.zeros_like(state)]), axis=-1)
+
+    return state, is_state_batched
+
+
+# pylint: disable=too-many-arguments, too-many-positional-arguments, unused-argument
+def measure_final_state(
+    circuit, state, is_state_batched, rng=None, prng_key=None, readout_errors=None
+) -> Result:
+    """
+    Perform the measurements required by the circuit on the provided state.
+
+    This is an internal function that will be called by ``default.mixed``.
+
+    Args:
+        circuit (.QuantumScript): The single circuit to simulate
+        state (TensorLike): The state to perform measurement on
+        is_state_batched (bool): Whether the state has a batch dimension or not.
+        rng (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
+            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
+            If no value is provided, a default RNG will be used.
+        prng_key (Optional[jax.random.PRNGKey]): An optional ``jax.random.PRNGKey``. This is
+            the key to the JAX pseudo random number generator. Only for simulation using JAX.
+            If None, the default ``sample_state`` function and a ``numpy.random.default_rng``
+            will be for sampling.
+        readout_errors (List[Callable]): List of channels to apply to each wire being measured
+        to simulate readout errors.
+
+    Returns:
+        Tuple[TensorLike]: The measurement results
+    """
+
+    circuit = circuit.map_to_standard_wires()
+
+    if not circuit.shots:
+        # analytic case
+        if len(circuit.measurements) == 1:
+            return measure(circuit.measurements[0], state, is_state_batched, readout_errors)
+
+        return tuple(
+            measure(mp, state, is_state_batched, readout_errors) for mp in circuit.measurements
+        )
+    # !TODO: add finite-shot branch afterwards. After implementation of finite-shot scenario, del this line.
+    raise NotImplementedError
+
+
+# pylint: disable=too-many-arguments, too-many-positional-arguments
+def simulate(
+    circuit: qml.tape.QuantumScript,
+    rng=None,
+    prng_key=None,
+    debugger=None,
+    interface=None,
+    readout_errors=None,
+) -> Result:
+    """Simulate a single quantum script.
+
+        This is an internal function that will be called by ``default.mixed``.
+
+        Args:
+            circuit (QuantumScript): The single circuit to simulate
+            rng (Optional[Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]]): A
+                seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
+                If no value is provided, a default RNG will be used.
+            prng_key (Optional[jax.random.PRNGKey]): An optional ``jax.random.PRNGKey``. This is
+                the key to the JAX pseudo random number generator. If None, a random key will be
+                generated. Only for simulation using JAX.
+            debugger (_Debugger): The debugger to use
+            interface (str): The machine learning interface to create the initial state with
+            readout_errors (List[Callable]): List of channels to apply to each wire being measured
+            to simulate readout errors.
+
+        Returns:
+            tuple(TensorLike): The results of the simulation
+
+        Note that this function can return measurements for non-commuting observables simultaneously.
+
+        This function assumes that all operations provide matrices.
+
+    >>> qs = qml.tape.QuantumScript(
+    ...     [qml.RX(1.2, wires=0)],
+    ...     [qml.expval(qml.PauliX(0)), qml.probs(wires=(0, 1))]
+    ... )
+    >>> simulate(qs)
+    (0.0, array([0.68117888, 0.        , 0.31882112, 0.        ]))
+
+
+    """
+    state, is_state_batched = get_final_state(
+        circuit, debugger=debugger, interface=interface, rng=rng, prng_key=prng_key
+    )
+    return measure_final_state(
+        circuit,
+        state,
+        is_state_batched,
+        rng=rng,
+        prng_key=prng_key,
+        readout_errors=readout_errors,
+    )