State vector

@beartype
def run(
    self,
    shots: int = 1,
    solver_name: str = "dop853",
    atol: float = 1e-14,
    rtol: float = 1e-7,
    nsteps: int = 2_147_483_647,
    interaction_picture: bool = False,
    project_hyperfine: bool = True,
) -> NDArray[np.uint8]:
    """Run the emulation with all atoms in the ground state,
    sampling the final state vector."""

    options = dict(
        solver_name=solver_name,
        atol=atol,
        rtol=rtol,
        nsteps=nsteps,
        interaction_picture=interaction_picture,
    )

    state = self.hamiltonian.space.zero_state()
    (result,) = self.apply(state, **options)
    result.normalize()

    return result.sample(shots, project_hyperfine=project_hyperfine)

@beartype
def average(
    self,
    register: StateVector,
    time: Optional[float] = None,
) -> float:
    """Get energy average from RydbergHamiltonian object at time `time` with
    register `register`

    Args:
        register (StateVector): The state vector to take average with
        time (Optional[float], optional): Time value to evaluate average at.
        Defaults to duration of RydbergHamiltonian.

    Returns:
        float: average energy at time `time`
    """
    return np.vdot(register.data, self._apply(register.data, time)).real

@beartype
def average_and_variance(
    self,
    register: StateVector,
    time: Optional[float] = None,
) -> Tuple[float, float]:
    """Get energy average and variance from RydbergHamiltonian object at time `time`
    with register `register`

    Args:
        register (StateVector): The state vector to take average and variance with
        time (Optional[float], optional): Time value to evaluate average at.
        Defaults to duration of RydbergHamiltonian.

    Returns:
        Tuple[float, float]: average and variance of energy at time `time`
        respectively.
    """
    H_register_data = self._apply(register.data, time)

    average = np.vdot(register.data, H_register_data).real
    square_average = np.vdot(H_register_data, H_register_data).real

    return average, square_average - average**2

def tocsr(self, time: float) -> csr_matrix:
    """Return the Hamiltonian as a csr matrix at time `time`.

    Args:
        time (float): time to evaluate the Hamiltonian at.

    Returns:
        csr_matrix: The Hamiltonian as a csr matrix.

    """
    diagonal = sum(
        (detuning.get_diagonal(time) for detuning in self.detuning_ops),
        start=self.rydberg,
    )

    hamiltonian = diags(diagonal).tocsr()
    for rabi_op in self.rabi_ops:
        hamiltonian = hamiltonian + rabi_op.tocsr(time)

    return hamiltonian

@beartype
def variance(
    self,
    register: StateVector,
    time: Optional[float] = None,
) -> float:
    """Get the energy variance from RydbergHamiltonian object at
    time `time` with register `register`

    Args:
        register (StateVector): The state vector to take variance with
        time (Optional[float], optional): Time value to evaluate average at.
        Defaults to duration of RydbergHamiltonian.

    Returns:
        complex: variance of energy at time `time` respectively.
    """

    _, var = self.average_and_variance(register, time)
    return var

@plum.dispatch
def local_trace(  # noqa: F811
    self, matrix: np.ndarray, site_index: Union[int, Tuple[int, int]]
) -> complex:  # noqa: F811
    """return trace of an operator over the StateVector.

    Args:
        matrix (np.ndarray): Square matrix representing operator in the local
            hilbert space.
        site_index (int | Tuple[int, int]): sites to apply one body operator to.

    Returns:
        complex: the trace of the operator over the state-vector.

    Raises:
        ValueError: Error is raised when the dimension of `operator` is not
        consistent with `site` argument. The size of the operator must fit
        the size of the local hilbert space of `site` depending on the number
        of sites and the number of levels inside each atom, e.g. for two site
        expectation value with a three level atom the operator must be a 9 by
        9 array.

        ValueError: Error is raised when the `site` argument is out of bounds.

    """
    ...

def norm(self) -> float:
    """Return the norm of the state vector."""
    return np.linalg.norm(self.data)

def normalize(self) -> None:
    """Normalize the state vector."""
    data = self.data
    data /= np.linalg.norm(data)

def sample(self, shots: int, project_hyperfine: bool = True) -> NDArray:
    """Sample the state vector and return bitstrings."""
    return self.space.sample_state_vector(
        self.data, shots, project_hyperfine=project_hyperfine
    )