Projecting parts of `SeleneQuestState`s and representing substates

[GBoydQ]

In guppyalgos, we have found it very convenient to have methods that return new SeleneQuestStates but acting on a smaller number of qubits, either through projection, or simply discarding qubits that are not entangled. This is because it makes it much easier to analyse and compare smaller statevectors, and often a lot of the total statevector is irrelevant anyway (ancilla that are returned to 0 etc).
The current way that projection works is that you supply some bitstring that is the desired state of the specified_qubits, and then the relevant projected state on the unspecified qubits is returned.
The problem with this is if you want to do a series of these operations, then your simulation level specified_qubits may have indices that are larger than the total number of qubits in this subsystem state. You have to then keep track of which qubit indices are no longer present and recalculate indices to the 'local' subsystem state.

Some vague ideas for solutions I have thrown together examples of are:

• a wrapper for SeleneQuestState that silently (probably not desirable) does this reindexing
• a ProjectedState that is essentially a TracedState but also has a set_specified_qubits() method to handle this logic

(happy to contribute this if it would be helpful, I would need access as I currently cannot add a branch to remote)

[GBoydQ]

To give an idea of what we already have in guppyalgos

here self._reshape_sv_to_matrix() is just the first part of SeleneQuestState.get_density_matrix() that reshapes the statevector

    def get_projected_state(
        self,
        bitstring: list[bool],
        zero_threshold: float = 1e-12,
        new_specified_qubits: list[int] | None = None,
        renormalize: bool = True,
    ) -> TracedState[SubQuestState]:
        """Projects a state onto the desired bitstring on the specified qubits."""
        if len(bitstring) != len(self.specified_qubits):
            raise ValueError(f"""Projection bitstring {bitstring} must have the same length
            as specified qubits {self.specified_qubits}""")

        n_specified = len(self.specified_qubits)
        n_unspecified = self.total_qubits - n_specified

        reshaped = self._reshape_sv_to_matrix()
        # get the projected state
        index = sum(2**i if bit else 0 for i, bit in enumerate(bitstring))
        projected_state = switch_endianness(reshaped[:, index])

        # renormalise
        if renormalize:
            projected_norm = float(np.linalg.norm(projected_state))
        else:
            projected_norm = 1.0

        
        new_indices_after_projection: list[int] = []
        if new_specified_qubits is not None:
            # work out indices of new specified qubits now that we have projected out old ones
            # TODO this logic needs changing from global to local indices

        if projected_norm > zero_threshold:
            projected_state = projected_state / projected_norm

        # build output SubQuestState and TracedState
        sqs = SubQuestState(
            projected_state,
            total_qubits=n_unspecified,
            specified_qubits=new_indices_after_projection,
            projected_out_qubits=self.specified_qubits,
        )
        return TracedState(projected_norm**2, sqs)

[GBoydQ]

and a draft of a wrapper for SeleneQuestState that does the reindexing

class SubQuestState(SeleneQuestState):
    projected_out_qubits: list[int]
    _specified_qubits: list[int]

    def __init__(
        self,
        state: NDArray[np.complex128],
        total_qubits: int,
        specified_qubits: list[int],
        projected_out_qubits: list[int],
    ):
        # need to control init order because projected_out_qubits needs to be init before specified_qubits
        self.state = state
        self.projected_out_qubits = projected_out_qubits
        self.total_qubits = total_qubits
        self.specified_qubits = specified_qubits

    @property
    def specified_qubits(self) -> list[int]:
        return self._specified_qubits

    @specified_qubits.setter
    def specified_qubits(self, spec_qubits_outer: list[int]) -> None:
        """Set specified qubits of inner SeleneQuestState.

        Reindexes using the fact that some qubits have already been removed.
        """
        if not set(spec_qubits_outer).isdisjoint(self.projected_out_qubits):
            raise ValueError(
                "Global specified qubits contain some that have already been discarded."
            )

        # Reindex spec_qubits_outer to account for projected out qubits
        spec_qubits_reindexed = self._reindex_global_to_local(spec_qubits_outer)

        self._specified_qubits = spec_qubits_reindexed

    def _reindex_global_to_local(self, global_idxs: list[int]):
        local_idxs = []
        for qubitidx in global_idxs:
            decrement = sum(1 for pq in self.projected_out_qubits if pq < qubitidx)
            local_idxs.append(qubitidx - decrement)
        return local_idxs

    def _reindex_local_to_global(self, local_idxs: list[int]):
        global_idxs = []
        for local_idx in local_idxs:
            global_idx = local_idx
            for pq in sorted(self.projected_out_qubits):
                if pq <= global_idx:
                    global_idx += 1
            global_idxs.append(global_idx)
        return global_idxs

        def get_projected_state(
        self,
        bitstring: list[bool],
        zero_threshold: float = 1e-12,
        new_specified_qubits: list[int] | None = None,
        renormalize: bool = True,
    ) -> TracedState[SubQuestState]:
        if not set(self.projected_out_qubits).isdisjoint(
            self._reindex_local_to_global(self.specified_qubits)
        ):
            raise ValueError(
                "Trying to project out qubits that have already been discarded"
            )

        new_projected = super().get_projected_state(
            bitstring, zero_threshold, new_specified_qubits, renormalize
        )
        new_projected.state.projected_out_qubits = (
            self._reindex_local_to_global(self.specified_qubits)
            + self.projected_out_qubits
        )
        return new_projected