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,
):
    """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 /= np.linalg.norm(result)

    return self.hamiltonian.space.sample_state_vector(
        result, 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

@plum.dispatch
def expectation_value(  # noqa: F811
    self, register: np.ndarray, operator: np.ndarray, site_indices: int
) -> complex:
    """Calculate expectation values of one and two body operators.

    Args:
        register (np.ndarray): Register to evaluate expectation value with
        operator (np.ndarray): Operator to take expectation value of.
        site_indices (int, Tuple[int, int]): site/sites to evaluate `operator` at.
            It can either a single integer or a tuple of two integers for one and
            two body operator respectively.

    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 v
        alue with a three level atom the operator must be a 9 by 9 array.

    Returns:
        complex: The expectation value.
    """
    self._check_register(register)

@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

def local_trace(
    self, matrix: np.ndarray, site_indices: Union[int, Tuple[int, int]]
) -> complex:
    """return trace of an operator over the StateVector.

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

    Returns:
        complex: the trace of the operator over the state-vector.
    """
    return self._trace(matrix, site_indices)