Make assign_states public, faster and remove a bug (#7880)

**Context:**
In
`pennylane.templates.state_preparations.superposition.py::_assign_states`,
some computational basis states are mapped to new computational basis
states.
It is claimed that states that are both in the set of inputs and the set
of outputs are fixed points of the produced map, but this is not true.

`_assign_states` is used in `qualtran_io.py` but is a private function,
so[ it might be better to make it
public.](https://github.com/PennyLaneAI/pennylane/pull/7866/files#r2205072082)

**Description of the Change:**
- Make the mentioned states fixed points of the produced map
- Speed up the function by looping over the set of states fewer times
- Make `assign_states` public-facing.

**Benefits:**
Better import hygiene, faster code, one bug less.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**

---------

Co-authored-by: Yushao Chen (Jerry) <[email protected]>

Author: dwierichs

doc/releases/changelog-dev.md

@@ -64,6 +64,10 @@
 
 <h3>Improvements 🛠</h3>
 
+* Changed how basis states are assigned internally in `qml.Superposition`, improving its
+  decomposition slightly both regarding classical computing time and gate decomposition.
+  [(#7880)](https://github.com/PennyLaneAI/pennylane/pull/7880)
+
 * The printing and drawing of :class:`~.TemporaryAND`, also known as ``qml.Elbow``, and its adjoint
   have been improved to be more legible and consistent with how it's depicted in circuits in the literature.
   [(#8017)](https://github.com/PennyLaneAI/pennylane/pull/8017)

pennylane/io/qualtran_io.py

@@ -32,7 +32,7 @@
 from pennylane.queuing import AnnotatedQueue, QueuingManager
 from pennylane.registers import registers
 from pennylane.tape import make_qscript
-from pennylane.templates.state_preparations.superposition import _assign_states
+from pennylane.templates.state_preparations.superposition import order_states
 from pennylane.wires import WiresLike
 from pennylane.workflow import construct_tape
 from pennylane.workflow.qnode import QNode
@@ -124,7 +124,7 @@ def _(op: qtemps.state_preparations.Superposition):
     size_basis_state = len(bases[0])  # assuming they are all the same size
 
     dic_state = dict(zip(bases, coeffs))
-    perms = _assign_states(bases)
+    perms = order_states(bases)
     new_dic_state = {perms[key]: val for key, val in dic_state.items() if key in perms}
 
     sorted_coefficients = [
@@ -387,7 +387,6 @@ def _(op: qtemps.subroutines.ModExp):
     num_work_wires = len(op.hyperparameters["work_wires"])
     num_x_wires = len(op.hyperparameters["x_wires"])
 
-    mult_resources = {}
     if mod == 2**num_x_wires:
         num_aux_wires = num_x_wires
         num_aux_swap = num_x_wires
@@ -407,12 +406,13 @@ def _(op: qtemps.subroutines.ModExp):
 
     cnot = qt_gates.CNOT()
 
-    mult_resources = {}
-    mult_resources[qft] = 2
-    mult_resources[qft_dag] = 2
-    mult_resources[sequence] = 1
-    mult_resources[sequence_dag] = 1
-    mult_resources[cnot] = min(num_x_wires, num_aux_swap)
+    mult_resources = {
+        qft: 2,
+        qft_dag: 2,
+        sequence: 1,
+        sequence_dag: 1,
+        cnot: min(num_x_wires, num_aux_swap),
+    }
 
     gate_types = defaultdict(int, {})
     ctrl_spec = CtrlSpec(cvs=[1])

pennylane/templates/state_preparations/__init__.py

@@ -20,6 +20,6 @@
 from .basis_qutrit import QutritBasisStatePreparation
 from .cosine_window import CosineWindow
 from .mottonen import MottonenStatePreparation
-from .superposition import Superposition
+from .superposition import Superposition, order_states
 from .qrom_state_prep import QROMStatePreparation
 from .state_prep_mps import MPSPrep, right_canonicalize_mps

pennylane/templates/state_preparations/superposition.py

@@ -14,81 +14,77 @@
 r"""
 Contains the Superposition template.
 """
-
 import pennylane as qml
 from pennylane.operation import Operation
 
 
-def _assign_states(basis_list):
+def order_states(basis_states: list[list[int]]) -> dict[tuple[int], tuple[int]]:
     r"""
-    This function maps a given list of :math:`m` basis states to the first :math:`m` basis states in the
-    computational basis.
-
-    For instance, a given list of :math:`[s_0, s_1, ..., s_m]` where :math:`s` is a basis
-    state of length :math:`4` will be mapped as :math:`{s_0: |0000\rangle, s_1: |0001\rangle, s_2: |0010\rangle, \dots}`.
-
-    Note that if a state in ``basis_list`` is one of the first :math:`m` basis states,
-    this state will be mapped to itself.
+    This function maps a given list of :math:`m` computational basis states to the first
+    :math:`m` computational basis states, except for input states that are among the first
+    :math:`m` computational basis states, which are mapped to themselves.
 
     Args:
-        basis_list (list): list of basis states to be mapped
+        basis_states (list[list[int]]): sequence of :math:`m` basis states to be mapped.
+            Each state is a sequence of 0s and 1s.
 
     Returns:
-        dict: dictionary mapping basis states to the first :math:`m` basis states
+        dict[tuple[int], tuple[int]]: dictionary mapping basis states to the first :math:`m` basis
+        states, except for fixed points (states in the input that already were among the
+        first :math:`m` basis states).
 
+    **Example**
 
-    ** Example **
+    For instance, a given list of :math:`[s_0, s_1, ..., s_m]` where :math:`s` is a basis
+    state of length :math:`4` will be mapped as
+    :math:`\{s_0: |0000\rangle, s_1: |0001\rangle, s_2: |0010\rangle, \dots\}`.
 
     .. code-block:: pycon
 
-        >>> basis_list = [[1, 1, 0, 0], [1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 0, 1]]
-        >>> _assign_states(basis_list)
-        {
-        [1, 1, 0, 0]: [0, 0, 0, 0],
-        [1, 0, 1, 0]: [0, 0, 0, 1],
-        [0, 1, 0, 1]: [0, 0, 1, 0],
-        [1, 0, 0, 1]: [0, 0, 1, 1]
-        }
+        >>> basis_states = [[1, 1, 0, 0], [1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 0, 1]]
+        >>> order_states(basis_states)
+        {(1, 1, 0, 0): (0, 0, 0, 0),
+         (1, 0, 1, 0): (0, 0, 0, 1),
+         (0, 1, 0, 1): (0, 0, 1, 0),
+         (1, 0, 0, 1): (0, 0, 1, 1)}
 
+    If a state in ``basis_states`` is one of the first :math:`m` basis states,
+    this state will be mapped to itself, i.e. it will be a fixed point of the mapping.
 
     .. code-block:: pycon
 
-        >>> basis_list = [[1, 1, 0, 0], [0, 1, 0, 1], [0, 0, 0, 1], [1, 0, 0, 1]]
-        >>> _assign_states(basis_list)
-        {
-        [1, 1, 0, 0]: [0, 0, 0, 0],
-        [0, 1, 0, 1]: [0, 0, 1, 0],
-        [0, 0, 0, 1]: [0, 0, 0, 1],
-        [1, 0, 0, 1]: [0, 0, 1, 1]
-        }
+        >>> basis_states = [[1, 1, 0, 0], [0, 1, 0, 1], [0, 0, 0, 1], [1, 0, 0, 1]]
+        >>> order_states(basis_states)
+        {(0, 0, 0, 1): (0, 0, 0, 1),
+         (1, 1, 0, 0): (0, 0, 0, 0),
+         (0, 1, 0, 1): (0, 0, 1, 0),
+         (1, 0, 0, 1): (0, 0, 1, 1)}
 
     """
 
-    length = len(basis_list[0])
-    smallest_basis_lists = [tuple(map(int, f"{i:0{length}b}")) for i in range(len(basis_list))]
-
-    binary_dict = {}
-    used_smallest = set()
-
-    # Assign keys that can map to themselves
-    for original in basis_list:
-
-        if original in smallest_basis_lists and tuple(original) not in used_smallest:
-
-            binary_dict[tuple(original)] = original
-            used_smallest.add(tuple(original))
+    m = len(basis_states)
+    length = len(basis_states[0])
+    # Create the integers corresponding to the input basis states
+    basis_ints = [int("".join(map(str, state)), 2) for state in basis_states]
 
-    # Assign remaining keys to unused binary lists
-    remaining_keys = [key for key in basis_list if tuple(key) not in binary_dict]
-    remaining_values = [
-        value for value in smallest_basis_lists if tuple(value) not in used_smallest
-    ]
+    basis_states = [tuple(s) for s in basis_states]  # Need hashable objects, so we use tuples
+    state_map = {}  # The map for basis states to be populated
+    unmapped_states = []  # Will collect non-fixed point states
+    unmapped_ints = {i: None for i in range(m)}  # Will remove fixed point states
+    # Map fixed-point states to themselves and collect states and target ints still to be paired
+    for b_int, state in zip(basis_ints, basis_states):
+        if b_int < m:
+            state_map[state] = state
+            unmapped_ints.pop(b_int)
+        else:
+            unmapped_states.append(state)
 
-    for key, value in zip(remaining_keys, remaining_values):
-        binary_dict[tuple(key)] = value
-        used_smallest.add(tuple(value))
+    # Map non-fixed point states
+    for state, new_b_int in zip(unmapped_states, unmapped_ints):
+        # Convert the index of the state to be mapped into a state itself
+        state_map[state] = tuple(map(int, f"{new_b_int:0{length}b}"))
 
-    return binary_dict
+    return state_map
 
 
 def _permutation_operator(basis1, basis2, wires, work_wire):
@@ -292,7 +288,7 @@ def compute_decomposition(coeffs, bases, wires, work_wire):  # pylint: disable=a
         """
 
         dic_state = dict(zip(bases, coeffs))
-        perms = _assign_states(bases)
+        perms = order_states(bases)
         new_dic_state = {perms[key]: dic_state[key] for key in dic_state if key in perms}
 
         sorted_coefficients = [

tests/templates/test_state_preparations/test_superposition.py

@@ -22,6 +22,60 @@
 
 import pennylane as qml
 from pennylane import numpy as pnp
+from pennylane.templates.state_preparations.superposition import order_states
+
+
+def int_to_state(i, length):
+    return tuple(map(int, f"{i:0{length}b}"))
+
+
+@pytest.mark.parametrize(
+    "basis_states, exp_map",
+    (
+        [  # Examples where all basis states are fixed points
+            (
+                [int_to_state(i, L) for i in range(m)],
+                {int_to_state(i, L): int_to_state(i, L) for i in range(m)},
+            )
+            for L, m in [(1, 2), (2, 4), (2, 3), (3, 7), (4, 3), (4, 16)]
+        ]
+        + [  # Examples from docstring
+            (
+                [[1, 1, 0, 0], [1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 0, 1]],
+                {
+                    (1, 1, 0, 0): (0, 0, 0, 0),
+                    (1, 0, 1, 0): (0, 0, 0, 1),
+                    (0, 1, 0, 1): (0, 0, 1, 0),
+                    (1, 0, 0, 1): (0, 0, 1, 1),
+                },
+            ),
+            (
+                [[1, 1, 0, 0], [0, 1, 0, 1], [0, 0, 0, 1], [1, 0, 0, 1]],
+                {
+                    (0, 0, 0, 1): (0, 0, 0, 1),
+                    (1, 1, 0, 0): (0, 0, 0, 0),
+                    (0, 1, 0, 1): (0, 0, 1, 0),
+                    (1, 0, 0, 1): (0, 0, 1, 1),
+                },
+            ),
+        ]
+        + [  # Other examples
+            (
+                [[1, 1, 0, 1], [0, 1, 0, 0], [1, 1, 1, 1], [0, 0, 1, 0], [0, 0, 0, 0]],
+                {
+                    (0, 0, 0, 0): (0, 0, 0, 0),
+                    (0, 0, 1, 0): (0, 0, 1, 0),
+                    (0, 1, 0, 0): (0, 1, 0, 0),
+                    (1, 1, 0, 1): (0, 0, 0, 1),
+                    (1, 1, 1, 1): (0, 0, 1, 1),
+                },
+            ),
+            ([[1, 1, 0], [0, 0, 1]], {(0, 0, 1): (0, 0, 1), (1, 1, 0): (0, 0, 0)}),
+        ]
+    ),
+)
+def test_order_states(basis_states, exp_map):
+    assert order_states(basis_states) == exp_map
 
 
 def test_standard_validity():