Bloqade

@beartype
def run(
    self,
    shots: int,
    args: Tuple[LiteralType, ...] = (),
    name: Optional[str] = None,
    blockade_radius: float = 0.0,
    waveform_runtime: str = "interpret",
    interaction_picture: bool = False,
    cache_matrices: bool = False,
    multiprocessing: bool = False,
    num_workers: Optional[int] = None,
    solver_name: str = "dop853",
    atol: float = 1e-7,
    rtol: float = 1e-14,
    nsteps: int = 2_147_483_647,
) -> LocalBatch:
    """Run the current program using bloqade python backend

    Args:
        shots (int): number of shots after running state vector simulation
        args (Tuple[LiteralType, ...], optional): The values for parameters defined
        in `args`. Defaults to ().
        name (Optional[str], optional): Name to give this run. Defaults to None.
        blockade_radius (float, optional): Use the Blockade subspace given a
        particular radius. Defaults to 0.0.
        waveform_runtime: (bool, optional): Use Numba to compile the waveforms,
        Defaults to False.
        interaction_picture (bool, optional): Use the interaction picture when
        solving schrodinger equation. Defaults to False.
        cache_matrices (bool, optional): Reuse previously evaluated matrcies when
        possible. Defaults to False.
        multiprocessing (bool, optional): Use multiple processes to process the
        batches. Defaults to False.
        num_workers (Optional[int], optional): Number of processes to run with
        multiprocessing. Defaults to None.
        solver_name (str, optional): Which SciPy Solver to use. Defaults to
        "dop853".
        atol (float, optional): Absolute tolerance for ODE solver. Defaults to
        1e-14.
        rtol (float, optional): Relative tolerance for adaptive step in ODE solver.
        Defaults to 1e-7.
        nsteps (int, optional): Maximum number of steps allowed per integration
        step. Defaults to 2_147_483_647, the maximum value.

    Raises:
        ValueError: Cannot use multiprocessing and cache_matrices at the same time.

    Returns:
        LocalBatch: Batch of local tasks that have been executed.
    """
    if multiprocessing and cache_matrices:
        raise ValueError(
            "Cannot use multiprocessing and cache_matrices at the same time."
        )

    compile_options = dict(
        shots=shots,
        args=args,
        name=name,
        blockade_radius=blockade_radius,
        cache_matrices=cache_matrices,
        waveform_runtime=waveform_runtime,
    )

    solver_options = dict(
        multiprocessing=multiprocessing,
        num_workers=num_workers,
        solver_name=solver_name,
        atol=atol,
        rtol=rtol,
        nsteps=nsteps,
        interaction_picture=interaction_picture,
    )

    batch = self._compile(**compile_options)
    batch._run(**solver_options)

    return batch

@beartype
def run_callback(
    self,
    callback: Callable[[StateVector, NamedTuple, RydbergHamiltonian, Any], Any],
    program_args: Tuple[LiteralType, ...] = (),
    callback_args: Tuple = (),
    ignore_exceptions: bool = False,
    blockade_radius: float = 0.0,
    waveform_runtime: str = "interpret",
    interaction_picture: bool = False,
    cache_matrices: bool = False,
    multiprocessing: bool = False,
    num_workers: Optional[int] = None,
    solver_name: str = "dop853",
    atol: float = 1e-7,
    rtol: float = 1e-14,
    nsteps: int = 2_147_483_647,
) -> List:
    """Run state-vector simulation with a callback to access full state-vector from
    emulator

    Args:
        callback (Callable[[StateVector, Metadata, RydbergHamiltonian, Any], Any]):
        The callback function to run for each task in batch. See note below for more
        details about the signature of the function.
        program_args (Tuple[LiteralType, ...], optional): The values for parameters
        defined in `args`. Defaults to ().
        callback_args (Tuple[Any,...], optional): Extra arguments to pass into
        ignore_exceptions: (bool, optional) If `True` any exception raised during
        a task will be saved instead of the resulting output of the callback,
        otherwise the first exception by task number will be raised after *all*
        tasks have executed. Defaults to False.
        blockade_radius (float, optional): Use the Blockade subspace given a
        particular radius. Defaults to 0.0.
        waveform_runtime: (str, optional): Specify which runtime to use for
        waveforms. Defaults to "interpret".
        interaction_picture (bool, optional): Use the interaction picture when
        solving schrodinger equation. Defaults to False.
        cache_matrices (bool, optional): Reuse previously evaluated matrcies when
        possible. Defaults to False.
        multiprocessing (bool, optional): Use multiple processes to process the
        batches. Defaults to False.
        num_workers (Optional[int], optional): Number of processes to run with
        multiprocessing. Defaults to None.
        solver_name (str, optional): Which SciPy Solver to use. Defaults to
        "dop853".
        atol (float, optional): Absolute tolerance for ODE solver. Defaults to
        1e-14.
        rtol (float, optional): Relative tolerance for adaptive step in ODE solver.
        Defaults to 1e-7.
        nsteps (int, optional): Maximum number of steps allowed per integration
        step. Defaults to 2_147_483_647, the maximum value.

    Returns:
        List: List of resulting outputs from the callbacks

    Raises:
        RuntimeError: Raises the first error that occurs, only if
        `ignore_exceptions=False`.

    Note:
        For the `callback` function, first argument is the many-body wavefunction
        as a 1D complex numpy array, the second argument is of type `Metadata` which
        is a Named Tuple where the fields correspond to the parameters of that given
        task, RydbergHamiltonian is the object that contains the Hamiltonian used to
        generate the evolution for that task, Finally any optional positional
        arguments are allowed after that. The return value can be anything, the
        results will be collected in a list for each task in the batch.


    """
    if multiprocessing:
        from multiprocessing import Process, Queue, cpu_count
    else:
        from queue import Queue

    if cache_matrices:
        compile_cache = CompileCache()
    else:
        compile_cache = None

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

    runner = self.EmuRunner(
        compile_cache=compile_cache,
        solver_args=solver_args,
        callback=callback,
        callback_args=callback_args,
    )

    tasks = Queue()
    results = Queue()

    total_tasks = 0
    it_iter = self._generate_ir(program_args, blockade_radius, waveform_runtime)
    for task_number, emulator_ir, metadata in it_iter:
        total_tasks += 1
        tasks.put((task_number, (emulator_ir, metadata)))

    workers = []
    if multiprocessing:
        num_workers = max(int(num_workers or cpu_count()), 1)
        num_workers = min(total_tasks, num_workers)

        for _ in range(num_workers):
            worker = Process(
                target=BloqadePythonRoutine.process_tasks,
                args=(runner, tasks, results),
            )
            worker.start()

            workers.append(worker)
    else:
        self.process_tasks(runner, tasks, results)

    # blocks until all
    # results have been fetched
    # from the id_results Queue
    id_results = []
    for i in range(total_tasks):
        id_results.append(results.get())

    if workers:
        for worker in workers:
            worker.join()

        tasks.close()
        results.close()

    id_results.sort(key=lambda x: x[0])
    results = []

    for task_id, result in id_results:
        if not ignore_exceptions and isinstance(result, BaseException):
            try:
                raise result
            except BaseException:
                raise RuntimeError(
                    f"{result.__class__.__name__} occured during child process "
                    f"running for task number {task_id}:\n{traceback.format_exc()}"
                )

        results.append(result)

    return results