Merge branch 'main' into uh/sub-braket

.pre-commit-config.yaml

@@ -2,18 +2,18 @@
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.5.0
+  rev: v4.6.0
   hooks:
     - id: check-yaml
       args: ['--unsafe']
     - id: end-of-file-fixer
     - id: trailing-whitespace
 - repo: https://github.com/psf/black
-  rev: 24.1.0
+  rev: 24.4.2
   hooks:
     - id: black
 - repo: https://github.com/charliermarsh/ruff-pre-commit
   # Ruff version.
-  rev: "v0.0.264"
+  rev: "v0.5.5"
   hooks:
     - id: ruff

docs/home/quick_start.md

@@ -10,41 +10,25 @@ As you develop your Bloqade program, you are expected to rely on pop-up "hints"
 
 In [VS Code](https://code.visualstudio.com/) this is automatic, just type the `.` and see what options pop up:
 
-<div align="center">
-<picture>
-  <img src="/assets/quick_start/vscode-hints.gif" alt="VSCode Hints">
-</picture>
-</div>
+![VSCode Hints](../assets/quick_start/vscode-hints.gif)
 
 ### JetBrains PyCharm
 
 The same goes for [JetBrains PyCharm](https://www.jetbrains.com/pycharm/):
 
-<div align="center">
-<picture>
-  <img src="/assets/quick_start/pycharm-hints.gif" alt="PyCharm Hints">
-</picture>
-</div>
+![PyCharm Hints](../assets/quick_start/pycharm-hints.gif)
 
 ### Jupyter Notebook 
 
 In a [Jupyter Notebook](https://jupyter.org/) you'll need to type `.` and then hit tab for the hints to appear:
 
-<div align="center">
-<picture>
-  <img src="/assets/quick_start/jupyter-hints.gif" alt="Jupyter Notebook Hints">
-</picture>
-</div>
+![Jupyter Hints](../assets/quick_start/jupyter-hints.gif)
 
 ### IPython
 
 The same goes for [IPython](https://ipython.readthedocs.io/en/stable/):
 
-<div align="center">
-<picture>
-  <img src="/assets/quick_start/ipython-hints.gif" alt="IPython Hints">
-</picture>
-</div>
+![IPython Hints](../assets/quick_start/ipython-hints.gif)
 
 ## Defining Atom Geometry
 
@@ -62,11 +46,10 @@ You can easily visualize your geometries as well with `.show()`:
 ```python
 more_complex_geometry.show()
 ```
-<div align="center">
-<picture>
-  <img src="/assets/quick_start/geometry-visualization.png" style="width: 50%" alt="Geometry Visualization">
-</picture>
-</div>
+<figure markdown="span">
+![IPython Hints](../assets/quick_start/geometry-visualization.png){ width="50%" }
+</figure>
+
 
 You can also add positions to a pre-defined geometry:
 

src/bloqade/emulate/ir/state_vector.py

@@ -298,6 +298,20 @@ def tocsr(self, time: float) -> csr_matrix:
 
 @dataclass(frozen=True)
 class RydbergHamiltonian:
+    """Hamiltonian for a given task.
+    With the `RydbergHamiltonian` you can convert the Hamiltonian to CSR matrix form
+    as well as obtaining the average energy/variance of a register.
+
+    Attributes:
+        emulator_ir (EmulatorProgram): A copy of the original program
+            used to generate the RydbergHamiltonian
+        space (Space): The Hilbert space of the Hamiltonian, should align with the register the
+            Hamiltonian is being applied on for average energy/variance
+        rydberg (NDArray): Rydberg interaction operator
+        detuning_ops (List[DetuningOperator]): Detuning Operators of the Hamiltonian
+        rabi_ops (List[RabiOperator]): Rabi Operators of the Hamiltonian
+    """
+
     emulator_ir: EmulatorProgram
     space: Space
     rydberg: NDArray

src/bloqade/ir/routine/bloqade.py

@@ -503,7 +503,31 @@ def hamiltonian(
         waveform_runtime: str = "interpret",
         cache_matrices: bool = False,
     ) -> List[BloqadeEmulation]:
+        """
+        Generates a list of BloqadeEmulation objects which contain the Hamiltonian of your program.
+
+        If you have a variable(s) in your program you have assigned multiple values via `batch_assign()`
+        there will be multiple `BloqadeEmulation` objects, one for each value. On the other hand
+        if the program only assumes a singular value per each variable, there will only be
+        one `BloqadeEmulation` object but it will still be encapsulated in a list.
+
+
+        Args:
+            *args (LiteralType): If your program has a variable that was declared as run-time assignable
+                via `.args` you may pass a value to it here. If there are multiple
+                variables declared via `.args` the order in which you assign values to those variables
+                through this argument should follow the order in which the declaration occurred.
+            blockade_radius (float): The radius in which atoms blockade eachother. Default value is 0.0 micrometers.
+            use_hyperfine (bool): Should the Hamiltonian account for hyperfine levels. Default value is False.
+            waveform_runtime (str): Specify which runtime to use for waveforms. If "numba" is specify the waveform
+                is compiled, otherwise it is interpreted via the "interpret" argument. Defaults to "interpret".
+            cache_matrices (bool): Speed up Hamiltonian generation by reusing data (when possible) from previously generated Hamiltonians.
+                Default value is False.
+
+        Returns:
+            List[BloqadeEmulation]
 
+        """
         ir_iter = self._generate_ir(
             args, blockade_radius, waveform_runtime, use_hyperfine
         )

src/bloqade/submission/ir/braket.py

@@ -1,3 +1,6 @@
+"""Helper functions related to IR submission
+co-ordinations between Bloqade and Braket"""
+
 import braket.ir.ahs as braket_ir
 from braket.ahs.pattern import Pattern
 from braket.timings import TimeSeries
@@ -6,31 +9,49 @@
 from braket.ahs.driving_field import DrivingField
 from braket.ahs.shifting_field import ShiftingField
 from braket.ahs.field import Field
-
 from braket.task_result import AnalogHamiltonianSimulationTaskResult
+
+import bloqade.submission.ir.capabilities as cp
 from bloqade.submission.ir.task_results import (
     QuEraTaskResults,
     QuEraTaskStatusCode,
     QuEraShotResult,
     QuEraShotStatusCode,
 )
-
 from bloqade.submission.ir.task_specification import (
     QuEraTaskSpecification,
     GlobalField,
     LocalField,
 )
+
 from typing import Tuple, Union, List
 from pydantic.v1 import BaseModel
 from decimal import Decimal
 
 
 class BraketTaskSpecification(BaseModel):
+    """Class representing geometry of an atom arrangement.
+
+    Attributes:
+        nshots (int): Number of shots
+        program (braket_ir.Program): IR(Intermediate Representation)
+            program suitable for braket
+    """
+
     nshots: int
     program: braket_ir.Program
 
 
 def to_braket_time_series(times: List[Decimal], values: List[Decimal]) -> TimeSeries:
+    """Converts to `TimeSeries` object supported by Braket.
+
+    Args:
+        times (List[Decimal]): Times of the value.
+        values (List[Decimal]): Corresponding values to add to the time series
+
+    Returns:
+        An object of the type `braket.timings.TimeSeries`
+    """
     time_series = TimeSeries()
     for time, value in zip(times, values):
         time_series.put(time, value)
@@ -39,6 +60,18 @@ def to_braket_time_series(times: List[Decimal], values: List[Decimal]) -> TimeSe
 
 
 def to_braket_field(quera_field: Union[GlobalField, LocalField]) -> Field:
+    """Converts to `TimeSeries` object supported by Braket.
+
+    Args:
+        quera_field (Union[GlobalField, LocalField)]:
+            Field supported by Quera
+
+    Returns:
+        An object of the type `braket.ahs.field.Field`
+
+    Raises:
+        TypeError: If field is not of the type `GlobalField` or `LocalField`.
+    """
     if isinstance(quera_field, GlobalField):
         times = quera_field.times
         values = quera_field.values
@@ -56,6 +89,12 @@ def to_braket_field(quera_field: Union[GlobalField, LocalField]) -> Field:
 
 
 def extract_braket_program(quera_task_ir: QuEraTaskSpecification):
+    """Extracts the Braket program.
+
+    Args:
+        quera_task_ir (QuEraTaskSpecification):
+            Quera IR(Intermediate representation) of the task.
+    """
     lattice = quera_task_ir.lattice
 
     rabi_amplitude = (
@@ -90,18 +129,47 @@ def extract_braket_program(quera_task_ir: QuEraTaskSpecification):
 def to_braket_task(
     quera_task_ir: QuEraTaskSpecification,
 ) -> Tuple[int, AnalogHamiltonianSimulation]:
+    """Converts to `Tuple[int, AnalogHamiltonianSimulation]` object supported by Braket.
+
+    Args:
+        quera_task_ir (QuEraTaskSpecification):
+            Quera IR(Intermediate representation) of the task.
+
+    Returns:
+        An tuple  of the type `Tuple[int, AnalogHamiltonianSimulation]`.
+    """
     braket_ahs_program = extract_braket_program(quera_task_ir)
     return quera_task_ir.nshots, braket_ahs_program
 
 
 def to_braket_task_ir(quera_task_ir: QuEraTaskSpecification) -> BraketTaskSpecification:
+    """Converts quera IR(Intermendiate Representation) to
+    to `BraketTaskSpecification` object.
+
+    Args:
+        quera_task_ir (QuEraTaskSpecification):
+            Quera IR(Intermediate representation) of the task.
+
+    Returns:
+        An object of the type `BraketTaskSpecification` in Braket SDK
+
+    """
     nshots, braket_ahs_program = to_braket_task(quera_task_ir)
     return BraketTaskSpecification(nshots=nshots, program=braket_ahs_program.to_ir())
 
 
 def from_braket_task_results(
     braket_task_results: AnalogHamiltonianSimulationTaskResult,
 ) -> QuEraTaskResults:
+    """Get the `QuEraTaskResults` object for working with Bloqade SDK.
+
+    Args:
+        braket_task_results: AnalogHamiltonianSimulationTaskResult
+            Quantum task result of braket system
+
+    Returns:
+        An object of the type `Field` in Braket SDK.
+    """
     shot_outputs = []
     for measurement in braket_task_results.measurements:
         shot_outputs.append(
@@ -117,16 +185,34 @@ def from_braket_task_results(
     )
 
 
-def from_braket_status_codes(braket_message: str) -> QuEraTaskStatusCode:
-    if braket_message == str("QUEUED"):
+def from_braket_status_codes(braket_status: str) -> QuEraTaskStatusCode:
+    """Gets the `QuEraTaskStatusCode` object for working with Bloqade SDK.
+
+    Args:
+        braket_status: str
+            The value of status in metadata() in the Amazon Braket.
+            `GetQuantumTask` operation. If `use_cached_value` is `True`,
+            the value most recently returned from
+            `GetQuantumTask` operation is used
+
+    Returns:
+        An object of the type `Field` in Braket SDK
+    """
+    if braket_status == str("QUEUED"):
         return QuEraTaskStatusCode.Enqueued
     else:
-        return QuEraTaskStatusCode(braket_message.lower().capitalize())
+        return QuEraTaskStatusCode(braket_status.lower().capitalize())
+
 
+def to_quera_capabilities(paradigm) -> cp.QuEraCapabilities:
+    """Converts to `QuEraCapabilities` object supported by Braket.
 
-def to_quera_capabilities(paradigm):
-    import bloqade.submission.ir.capabilities as cp
+    Args:
+        paradigm: The `paradigm` property of the `AwsDevice` object for Aquila
 
+    Returns:
+        An object of the type `QuEraCapabilities` in Bloqade SDK.
+    """
     rydberg_global = paradigm.rydberg.rydbergGlobal
 
     return cp.QuEraCapabilities(