"""State vector representation module.
This module provides:
- `StateVector`: State vector representation class.
"""
from __future__ import annotations
import math
from typing import TYPE_CHECKING
import numpy as np
import typing_extensions
from graphqomb.matrix import is_hermitian
from graphqomb.simulator_backend import BaseFullStateSimulator, QubitIndexManager
if TYPE_CHECKING:
from collections.abc import Sequence
from numpy.typing import ArrayLike, DTypeLike, NDArray
from graphqomb.common import MeasBasis
CZ_TENSOR = np.asarray(
[
[1, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, -1],
]
)
[docs]
class StateVector(BaseFullStateSimulator):
r"""State vector representation."""
__state: NDArray[np.complex128]
__qindex_mng: QubitIndexManager
[docs]
def __init__(self, state: ArrayLike | None = None, *, copy: bool | None = None) -> None:
if state is not None:
state = np.asarray(state, dtype=np.complex128, copy=copy)
size = state.size
if size & (size - 1) or size == 0:
msg = "State vector must have a size that is a power of 2."
raise ValueError(msg)
num_qubits = (size - 1).bit_length()
self.__state = state.reshape((2,) * num_qubits)
else:
num_qubits = 0
self.__state = np.zeros((2,) * 0, dtype=np.complex128)
# Internal qubit ordering: maps external qubit index to internal index
self.__qindex_mng = QubitIndexManager(num_qubits)
def __array__(self, dtype: DTypeLike | None = None, copy: bool | None = None) -> NDArray[np.complex128]:
return np.asarray(self.state(), dtype=dtype, copy=copy)
@property
@typing_extensions.override
def num_qubits(self) -> int:
"""Get the number of qubits in the state vector.
Returns
-------
`int`
The number of qubits in the state vector.
"""
return self.__state.ndim
[docs]
@typing_extensions.override
def state(self) -> NDArray[np.complex128]:
r"""Get the state vector in external qubit order.
Returns
-------
`numpy.typing.NDArray`\[`numpy.complex128`\]
The state vector as a numpy array in external qubit order.
"""
# If internal order matches external order, return directly
if self.__qindex_mng.match(range(self.num_qubits)):
return self.__state
# Otherwise, reorder to external qubit order
axes = self.__qindex_mng.inverse_permutation()
return self.__state.transpose(axes)
[docs]
def copy(self) -> StateVector:
"""Create a copy of the state vector.
Returns
-------
`StateVector`
A new `StateVector` instance with the same state.
"""
return StateVector(self.state(), copy=True)
[docs]
@staticmethod
def from_num_qubits(num_qubits: int) -> StateVector:
"""Create a state vector in the plus state with given number of qubits.
Parameters
----------
num_qubits : `int`
Number of qubits in the state vector.
Returns
-------
`StateVector`
The resulting state vector in the plus state.
Raises
------
ValueError
If num_qubits is negative.
"""
if num_qubits < 0:
msg = "Number of qubits must be non-negative."
raise ValueError(msg)
return StateVector(np.full((2,) * num_qubits, 1 / math.sqrt(2**num_qubits), dtype=np.complex128))
[docs]
@staticmethod
def tensor_product(a: StateVector, b: StateVector) -> StateVector:
"""Tensor product with other state vector, self ⊗ other.
Parameters
----------
a : `StateVector`
first state vector
b : `StateVector`
second state vector
Returns
-------
`StateVector`
The resulting state vector after tensor product.
"""
return StateVector(np.kron(a.state(), b.state()))
[docs]
@typing_extensions.override
def evolve(self, operator: NDArray[np.complex128], qubits: int | Sequence[int]) -> None:
r"""Evolve the state by applying an operator to a subset of qubits.
Parameters
----------
operator : `numpy.typing.NDArray`\[`numpy.complex128`\]
The operator to apply.
qubits : `int` | `collections.abc.Sequence`\[`int`\]
The qubits to apply the operator to.
"""
# Convert external qubit indices to internal indices
internal_qubits = self.__qindex_mng.external_to_internal(qubits)
internal_qubits = (internal_qubits,) if isinstance(internal_qubits, int) else internal_qubits
k = len(internal_qubits)
rest = tuple(i for i in range(self.num_qubits) if i not in set(internal_qubits))
perm = internal_qubits + rest
inv_perm = np.argsort(perm)
op_tensor = operator.reshape((2,) * (2 * k))
contracted = np.tensordot(op_tensor, self.__state, axes=(range(k, 2 * k), internal_qubits))
contracted = contracted.transpose(inv_perm)
self.__state = np.asarray(contracted, dtype=np.complex128, copy=False) # for type checker
[docs]
@typing_extensions.override
def measure(self, qubit: int, meas_basis: MeasBasis, result: int) -> None:
"""Measure a qubit in a given measurement basis.
Parameters
----------
qubit : `int`
The qubit to measure.
meas_basis : `MeasBasis`
The measurement basis to use.
result : `int`
The measurement result.
"""
# Convert external qubit index to internal index
internal_qubit = self.__qindex_mng.external_to_internal(qubit)
meas_basis = meas_basis.flip() if result else meas_basis
basis_vector = meas_basis.vector()
new_state = np.tensordot(basis_vector.conjugate(), self.__state, axes=(0, internal_qubit))
self.__state = np.asarray(new_state, dtype=np.complex128, copy=False) # for type checker
# Update qubit order: remove the measured qubit
self.__qindex_mng.remove_qubit(internal_qubit)
self.normalize()
[docs]
@typing_extensions.override
def add_node(self, num_qubits: int) -> None:
"""Add plus state to the end of state vector.
Parameters
----------
num_qubits : `int`
number of qubits to add
"""
flat_state: NDArray[np.complex128] = self.__state.ravel()
repeated_flat = np.repeat(flat_state, 1 << num_qubits) / math.sqrt(2**num_qubits)
repeated_state = np.asarray(repeated_flat, dtype=np.complex128)
self.__state = repeated_state.reshape((2,) * (self.num_qubits + num_qubits))
# Append new qubits to the end of the qubit order
self.__qindex_mng.add_qubits(num_qubits)
[docs]
@typing_extensions.override
def entangle(self, qubit1: int, qubit2: int) -> None:
r"""Entangle two qubits.
Parameters
----------
qubit1 : `int`
first qubit index
qubit2 : `int`
second qubit index
"""
self.evolve(CZ_TENSOR, (qubit1, qubit2))
[docs]
def normalize(self) -> None:
"""Normalize the state."""
self.__state /= np.linalg.norm(self.__state)
[docs]
@typing_extensions.override
def reorder(self, permutation: Sequence[int]) -> None:
r"""Permute qubits.
if permutation is [2, 0, 1], then
# [q0, q1, q2] -> [q1, q2, q0]
Parameters
----------
permutation : `collections.abc.Sequence`\[`int`\]
permutation list
"""
# Update the internal qubit order only (no state reordering)
self.__qindex_mng.reorder(permutation)
[docs]
@typing_extensions.override
def norm(self) -> float:
"""Get norm of state vector.
Returns
-------
`float`
norm of state vector
"""
return float(np.linalg.norm(self.__state))
[docs]
@typing_extensions.override
def expectation(self, operator: NDArray[np.complex128], qubits: int | Sequence[int]) -> float:
r"""Calculate expectation value of operator.
Parameters
----------
operator : `numpy.typing.NDArray`\[`numpy.complex128`\]
Hermitian operator matrix
qubits : `int` | `collections.abc.Sequence`\[`int`\]
target qubits
Returns
-------
`float`
expectation value
Raises
------
ValueError
if operator is not Hermitian
"""
if not is_hermitian(operator):
msg = "Operator must be Hermitian"
raise ValueError(msg)
# Convert external qubit indices to internal indices
internal_qubits = self.__qindex_mng.external_to_internal(qubits)
internal_qubits = (internal_qubits,) if isinstance(internal_qubits, int) else internal_qubits
k = len(internal_qubits)
rest = tuple(i for i in range(self.num_qubits) if i not in internal_qubits)
perm = internal_qubits + rest
state_perm = self.__state.transpose(perm)
op_tensor = operator.reshape((2,) * (2 * k))
# Apply operator: O|ψ⟩
transformed_state = np.tensordot(op_tensor, state_perm, axes=(range(k, 2 * k), range(k)))
# Calculate expectation value: ⟨ψ|O|ψ⟩
norm_squared = np.real(np.vdot(state_perm, state_perm))
expectation = np.real(np.vdot(state_perm, transformed_state))
return float(expectation / norm_squared)