Rydberg h factory (#654)

* add rydberg_h + add sample to PythonFn.

* restricing output of `rydberg_h` to always assign variables.

* allow atom arrangements in `rydberg_h`

* fixing type checks.

* renaming decorator to not overshadow `waveform` module internally

* fixing bug in batch_Assign with empty arguments.

* adding unit test.

* adding doc strings. Changing return value of `rydberg_h`

* updating pdm.lock

* updating pydantic dep.

* making Routine's `pydantic` dataclasses.

* updating unit test for routine return value.

* renaming tests for `bloqade.factory`.

* making dataclasses into pydantic.dataclasses.

* implementing mock error directly in backend

* removing mocking through `unittest`

* adding unit tests for errors.

pyproject.toml

@@ -17,7 +17,7 @@ classifiers = [
 dependencies = [
     "juliacall>=0.9.14",
     "numpy>=1.25.2",
-    "pydantic>=1.10.12",
+    "pydantic>=1.10.13",
     "scipy>=1.9.3",
     "pandas>=2.1.0",
     "bokeh>=3.2.2",

src/bloqade/__init__.py

@@ -1,11 +1,13 @@
 from bloqade.ir import var, cast, Variable, Literal, start
+from bloqade.ir import to_waveform as waveform
 from bloqade.serialize import load, save, loads, dumps
 
-from bloqade.builder.factory import (
+from bloqade.factory import (
     piecewise_linear,
     piecewise_constant,
     linear,
     constant,
+    rydberg_h,
 )
 import bloqade.ir as _ir
 from bloqade.constants import RB_C6
@@ -44,4 +46,6 @@ def tree_depth(depth: int = None):
     "save",
     "loads",
     "dumps",
+    "rydberg_h",
+    "waveform",
 ]

src/bloqade/builder/assign.py

@@ -93,6 +93,10 @@ def __init__(
 
         super().__init__(parent)
 
+        if len(assignments) == 0:
+            self._batch_params = []
+            return
+
         circuit = self.parse_circuit()
         variables = ScanVariablesAnalogCircuit().emit(circuit)
 

src/bloqade/builder/backend/quera.py

@@ -100,7 +100,7 @@ def cloud_mock(self):
         """
         return self.parse().quera.cloud_mock()
 
-    def mock(self, state_file: str = ".mock_state.txt"):
+    def mock(self, state_file: str = ".mock_state.txt", submission_error: bool = False):
         """
         Specify mock, testing locally.
 
@@ -123,4 +123,6 @@ def mock(self, state_file: str = ".mock_state.txt"):
 
 
         """
-        return self.parse().quera.mock(state_file)
+        return self.parse().quera.mock(
+            state_file=state_file, submission_error=submission_error
+        )

src/bloqade/factory.py

@@ -0,0 +1,166 @@
+from bloqade.ir.routine.base import Routine
+from bloqade.ir.control.waveform import Waveform, Linear, Constant
+from bloqade.builder.typing import ScalarType
+from beartype import beartype
+from beartype.typing import List, Optional, Union, Dict, Any
+
+
+@beartype
+def linear(duration: ScalarType, start: ScalarType, stop: ScalarType) -> Linear:
+    """Create a Linear waveform.
+
+    Args:
+        duration (ScalarType): duration of linear waveform
+        start (ScalarType): starting value of linear waveform
+        stop (ScalarType): ending value of linear waveform
+
+    Returns:
+        Linear: Linear waveform
+    """
+    return Linear(start, stop, duration)
+
+
+@beartype
+def constant(duration: ScalarType, value: ScalarType) -> Constant:
+    """Create a Constant waveform.
+
+    Args:
+        duration (ScalarType): _description_
+        value (ScalarType): _description_
+
+    Returns:
+        Constant: A Constant waveform.
+    """
+    return Constant(value, duration)
+
+
+@beartype
+def piecewise_linear(durations: List[ScalarType], values: List[ScalarType]) -> Waveform:
+    """Create a piecewise linear waveform.
+
+    Create a piecewise linear waveform from a list of durations and values. The
+    value `duration[i]` is of the linear segment between `values[i]` and `values[i+1]`.
+
+    Args:
+        durations (List[ScalarType]): The duration of each segment
+        values (List[ScalarType]): The values for each segment
+
+    Raises:
+        ValueError: If the length of `values` is not one greater than the length of
+        `durations`.
+
+    Returns:
+        Waveform: The piecewise linear waveform.
+    """
+
+    if len(durations) + 1 != len(values):
+        raise ValueError(
+            "The length of values must be one greater than the length of durations"
+        )
+
+    pwl_wf = None
+    for duration, start, stop in zip(durations, values[:-1], values[1:]):
+        if pwl_wf is None:
+            pwl_wf = Linear(start, stop, duration)
+        else:
+            pwl_wf = pwl_wf.append(Linear(start, stop, duration))
+
+    return pwl_wf
+
+
+@beartype
+def piecewise_constant(
+    durations: List[ScalarType], values: List[ScalarType]
+) -> Waveform:
+    """Create a piecewise linear waveform.
+
+    Create a piecewise constant waveform from a list of durations and values. The
+    value `duration[i]` corresponds to the length of time for the i'th segment
+    with a value of `values[i]`.
+
+    Args:
+        durations (List[ScalarType]): The duration of each segment
+        values (List[ScalarType]): The values for each segment
+
+    Raises:
+        ValueError: If the length of `values` is not the same as the length of
+        `durations`.
+
+    Returns:
+        Waveform: The piecewise linear waveform.
+    """
+    if len(durations) != len(values):
+        raise ValueError(
+            "The length of values must be the same as the length of durations"
+        )
+
+    pwc_wf = None
+    for duration, value in zip(durations, values):
+        if pwc_wf is None:
+            pwc_wf = Constant(value, duration)
+        else:
+            pwc_wf = pwc_wf.append(Constant(value, duration))
+
+    return pwc_wf
+
+
+@beartype
+def rydberg_h(
+    atoms_positions: Any,
+    detuning: Optional[Waveform] = None,
+    amplitude: Optional[Waveform] = None,
+    phase: Optional[Waveform] = None,
+    static_params: Dict[str, Any] = {},
+    batch_params: Union[List[Dict[str, Any]], Dict[str, Any]] = [],
+    args: List[str] = [],
+) -> Routine:
+    """Create a rydberg program with uniform detuning, amplitude, and phase.
+
+    Args:
+        atoms_positions (Any): Description of geometry of atoms in system.
+        detuning (Optional[Waveform], optional): Waveform for detuning.
+            Defaults to None.
+        amplitude (Optional[Waveform], optional): Waveform describing the amplitude of
+            the rabi term. Defaults to None.
+        phase (Optional[Waveform], optional): Waveform describing the phase of rabi
+            term. Defaults to None.
+        static_params (Dict[str, Any], optional): Define static parameters of your
+            program. Defaults to {}.
+        batch_params (Union[List[Dict[str, Any]], Dict[str, Any]], optional):
+            Parmaters for a batch of tasks. Defaults to [].
+        args (List[str], optional): List of arguments to leave till runtime.
+            Defaults to [].
+
+    Returns:
+        Routine: An object that can be used to dispatch a rydberg program to
+            multiple backends.
+    """
+    from bloqade import start
+    from bloqade.atom_arrangement import AtomArrangement
+
+    print(type(atoms_positions))
+
+    if isinstance(atoms_positions, AtomArrangement):
+        prog = atoms_positions
+    else:
+        prog = start.add_position(atoms_positions)
+
+    if detuning is not None:
+        prog = prog.rydberg.detuning.uniform.apply(detuning)
+
+    if amplitude is not None:
+        prog = prog.amplitude.uniform.apply(amplitude)
+
+    if phase is not None:
+        prog = prog.phase.uniform.apply(phase)
+
+    prog = prog.assign(**static_params)
+
+    if isinstance(batch_params, dict):
+        prog = prog.batch_assign(**batch_params)
+    else:
+        prog = prog.batch_assign(batch_params)
+
+    prog = prog.args(args)
+
+    return prog.parse()

src/bloqade/ir/__init__.py

@@ -11,7 +11,7 @@
     Interpolation,
     Sample,
     PythonFn,
-    instruction,
+    to_waveform,
     GaussianKernel,
     LogisticKernel,
     SigmoidKernel,
@@ -68,7 +68,7 @@
     "Sample",
     "Interpolation",
     "PythonFn",
-    "instruction",
+    "to_waveform",
     "GaussianKernel",
     "LogisticKernel",
     "SigmoidKernel",

src/bloqade/ir/control/waveform.py

@@ -1,4 +1,5 @@
 from numbers import Real
+from bloqade.builder.typing import ScalarType
 from bloqade.ir.tree_print import Printer
 from bloqade.ir.scalar import (
     Scalar,
@@ -14,6 +15,7 @@
 from decimal import Decimal
 from pydantic.dataclasses import dataclass
 from beartype.typing import Any, Tuple, Union, List, Callable, Dict
+from beartype import beartype
 from enum import Enum
 
 import numpy as np
@@ -23,7 +25,8 @@
 from bloqade.visualization import display_ir
 
 
-def instruction(duration: Any) -> "PythonFn":
+@beartype
+def to_waveform(duration: ScalarType) -> Callable[[Callable], "PythonFn"]:
     # turn python function into a waveform instruction."""
 
     def waveform_wrapper(fn: Callable) -> "PythonFn":
@@ -486,6 +489,11 @@ def print_node(self):
     def children(self):
         return {"duration": self.duration, **{p.name: p for p in self.parameters}}
 
+    def sample(
+        self, dt: ScalarType, interpolation: Union[str, "Interpolation"]
+    ) -> "Sample":
+        return Sample(self, interpolation, cast(dt))
+
 
 @dataclass
 class SmoothingKernel:

src/bloqade/ir/routine/base.py

@@ -5,8 +5,8 @@
 from bloqade.builder.base import Builder
 from bloqade.ir.routine.params import Params
 
-
-from dataclasses import dataclass
+from pydantic import ConfigDict
+from pydantic.dataclasses import dataclass
 from typing import TYPE_CHECKING, Union
 
 if TYPE_CHECKING:
@@ -29,7 +29,10 @@ def parse(self: "RoutineBase") -> "Routine":
         return self
 
 
-@dataclass(frozen=True)
+__pydantic_dataclass_config__ = ConfigDict(arbitrary_types_allowed=True)
+
+
+@dataclass(frozen=True, config=__pydantic_dataclass_config__)
 class RoutineBase(RoutineParse):
     source: Builder
     circuit: AnalogCircuit
@@ -43,7 +46,7 @@ def __str__(self):
         return out
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, config=__pydantic_dataclass_config__)
 class Routine(RoutineBase):
     """Result of parsing a completed Builder string."""