Fix bugs for Waveform visitor and more tests (#287)

* add more tests and enhance cast
1) add tests for scales
2) add tests for bravais
3) to do 2) this commit add cast for tuple type

* tmp

* add more testing cases, bumping to 71%

* fix baravias bugs and more test
1) fix bugs on wrong index*elementary vectors
2) bump testing cases on bravias to 96
3) fix bugs on Rectangular when spacing_y dose not specify

* update to comply ruff

* modify to comply ruff

* add test for double parallel

* modify to comply black

* modify format

* try

* test pre-compile

* pre-commit linter fix

* waveform visitors
1) add more testing
2) fix waveform visitors Error types inconsistent bug
3) fix waveform visitors not waveform ast error bug
4) fix GaussianKernel typo

---------

Co-authored-by: Kai-Hsin Wu <khwu@KHWus-MBP.hsd1.ma.comcast.net>
Co-authored-by: Kai-Hsin Wu <khwu@KHWus-MacBook-Pro.local>

src/bloqade/ir/__init__.py

@@ -12,7 +12,7 @@
     Sample,
     PythonFn,
     instruction,
-    GuassianKernel,
+    GaussianKernel,
     LogisticKernel,
     SigmoidKernel,
     TriangleKernel,
@@ -54,7 +54,7 @@
     "Interpolation",
     "PythonFn",
     "instruction",
-    "GuassianKernel",
+    "GaussianKernel",
     "LogisticKernel",
     "SigmoidKernel",
     "TriangleKernel",

src/bloqade/ir/control/waveform.py

@@ -392,7 +392,7 @@ class InfiniteSmoothingKernel(SmoothingKernel):
     pass
 
 
-class Guassian(InfiniteSmoothingKernel):
+class Gaussian(InfiniteSmoothingKernel):
     def __call__(self, value: float) -> float:
         return np.exp(-(value**2) / 2) / np.sqrt(2 * np.pi)
 
@@ -442,7 +442,7 @@ def __call__(self, value: float) -> float:
         return np.maximum(0, np.pi / 4 * np.cos(np.pi / 2 * value))
 
 
-GuassianKernel = Guassian()
+GaussianKernel = Gaussian()
 LogisticKernel = Logistic()
 SigmoidKernel = Sigmoid()
 TriangleKernel = Triangle()

src/bloqade/ir/visitor/waveform_visitor.py

@@ -19,7 +19,7 @@
 
 class WaveformVisitor:
     def visit_alligned(self, ast: AlignedWaveform) -> Any:
-        raise NotADirectoryError(
+        raise NotImplementedError(
             f"No visitor method implemented in {self.__class__} for AlignedWaveform."
         )
 
@@ -39,7 +39,7 @@ def visit_poly(self, ast: Poly) -> Any:
         )
 
     def visit_smooth(self, ast: Smooth) -> Any:
-        raise NotADirectoryError(
+        raise NotImplementedError(
             f"No visitor method implemented in {self.__class__} for Smooth."
         )
 
@@ -112,4 +112,4 @@ def visit(self, ast: Waveform):
             case Sample():
                 return self.visit_sample(ast)
             case _:
-                raise ValueError()(f"{ast.__class__} is not a Waveform AST type.")
+                raise ValueError(f"{ast.__class__} is not a Waveform AST type.")

tests/test_builder.py

@@ -7,8 +7,8 @@
 # prog.linear(start=1.0, stop=2.0, duration="x")
 # import pytest
 import bloqade.ir as ir
-from bloqade.builder import location
-from bloqade.ir import rydberg, detuning
+from bloqade.builder import location, waveform
+from bloqade.ir import rydberg, detuning, hyperfine, rabi
 from bloqade import start, var, cast
 from bloqade.ir.location import Square, Chain
 import numpy as np
@@ -101,6 +101,15 @@ def test_spatial_var():
 
     assert prog._name == "a"
 
+    prog = prog.piecewise_constant([0.1], [30])
+
+    # test build ast:
+    seq = prog.sequence
+
+    assert seq.value[rydberg].value[detuning].value[
+        ir.RunTimeVector("a")
+    ] == ir.Constant(value=30, duration=0.1)
+
 
 def test_issue_107():
     waveform = (
@@ -140,6 +149,42 @@ def test_issue_150():
     )
 
 
+def test_record():
+    prog = start
+    prog = (
+        prog.rydberg.detuning.location(1)
+        .piecewise_constant([0.1], [30])
+        .record("detuning")
+    )
+
+    assert type(prog) == waveform.Record
+
+    seq = prog.sequence
+    assert seq.value[rydberg].value[detuning].value[
+        ir.ScaledLocations({ir.Location(1): cast(1)})
+    ] == ir.Record(waveform=ir.Constant(value=30, duration=0.1), var=cast("detuning"))
+
+
+def test_hyperfine_phase():
+    prog = start.hyperfine.rabi.phase.location(1).piecewise_constant([0.1], [30])
+
+    seq = prog.sequence
+
+    assert seq.value[hyperfine].value[rabi.phase].value[
+        ir.ScaledLocations({ir.Location(1): cast(1)})
+    ] == ir.Constant(value=30, duration=0.1)
+
+
+def test_hyperfine_amplitude():
+    prog = start.hyperfine.rabi.amplitude.location(1).piecewise_constant([0.1], [30])
+
+    seq = prog.sequence
+
+    assert seq.value[hyperfine].value[rabi.amplitude].value[
+        ir.ScaledLocations({ir.Location(1): cast(1)})
+    ] == ir.Constant(value=30, duration=0.1)
+
+
 prog = start
 prog = (
     prog.rydberg.detuning.location(1)

tests/test_builder_job.py

@@ -0,0 +1,34 @@
+from bloqade.ir.location import Square
+import pytest
+import json
+
+
+def test_braket_unsupport_parallel():
+    prog = Square(3)
+
+    prog = prog.rydberg.detuning.uniform.piecewise_constant([0.1], [32])
+    prog = prog.parallelize(4)
+
+    with pytest.raises(TypeError):
+        prog.braket_local_simulator(10)
+
+
+def test_in_house_simulator_entry():
+    prog = Square(3)
+
+    prog = prog.rydberg.detuning.uniform.piecewise_constant([0.1], [32])
+    prog = prog.parallelize(4)
+
+    with pytest.raises(NotImplementedError):
+        prog.simu(4)
+
+
+def test_braket():
+    prog = Square(3)
+
+    prog = prog.rydberg.detuning.uniform.piecewise_constant([0.1], [32])
+
+    job = prog.braket(4)
+
+    x = json.loads(job.json())
+    assert x["tasks"]["0"]["task_ir"]["nshots"] == 4

tests/test_lattice.py

@@ -1,5 +1,80 @@
-from bloqade.ir.location import ListOfLocations
+import bloqade.ir as ir
+from bloqade.ir.location import ListOfLocations, AtomArrangement
+from bloqade.ir.location import Square
+from bloqade import cast
+import pytest
 
-lp = ListOfLocations(
-    [(0, 0), (0, 1), (1, 0), (1, 1)]
-)  # enumerate method gives a generator that can be taken advantage of
+
+def test_listOfLocatoions():
+    lattice = ListOfLocations(
+        [(0, 0), (0, 1), (1, 0), (1, 1)]
+    )  # enumerate method gives a generator that can be taken advantage of
+
+    positions1 = set(map(lambda info: info.position, lattice.enumerate()))
+    positions_expected = set(
+        cast(
+            [
+                (0, 0),
+                (1, 0),
+                (0, 1),
+                (1, 1),
+            ]
+        )
+    )
+    assert positions1 == positions_expected
+
+    # scale:
+    lattice = lattice.scale(3.2)
+
+    positions1 = set(map(lambda info: info.position, lattice.enumerate()))
+    positions_expected = set(
+        cast(
+            [
+                (0, 0),
+                (3.2, 0),
+                (0, 3.2),
+                (3.2, 3.2),
+            ]
+        )
+    )
+    assert positions1 == positions_expected
+
+
+def test_addlocs_on_square():
+    lattice = Square(2, lattice_spacing=1.0)
+    lattice1 = lattice.add_positions([(9, 6), (4, 4)])
+
+    # type(lattice)
+    assert type(lattice1) == ir.location.ListOfLocations
+
+    positions1 = set(map(lambda info: info.position, lattice1.enumerate()))
+    positions_expected1 = set(cast([(0, 0), (1, 0), (0, 1), (1, 1), (9, 6), (4, 4)]))
+
+    assert positions1 == positions_expected1
+
+
+def test_addlocs_filling_options():
+    lattice = Square(2, lattice_spacing=1.0)
+    lattice = lattice.add_positions([(9, 6), (4, 4)], filling=[False, True])
+
+    # type(lattice)
+    assert type(lattice) == ir.location.ListOfLocations
+
+    for info in lattice.enumerate():
+        if info.position == cast((9, 6)):
+            assert bool(info.filling) is False
+        else:
+            assert bool(info.filling) is True
+
+
+def test_internal_base_listofloc():
+    lattice = AtomArrangement()
+
+    with pytest.raises(NotImplementedError):
+        set(map(lambda info: info.position, lattice.enumerate()))
+
+    with pytest.raises(NotImplementedError):
+        lattice.n_atoms
+
+    with pytest.raises(NotImplementedError):
+        lattice.n_dims

tests/test_parallelize.py

@@ -1,31 +1,70 @@
+import pytest
 from bloqade.ir.location import Square
+from bloqade import start
+
 
 # create lattice
-lattice = Square(3, lattice_spacing=6)
-
-quantum_task = (
-    lattice.rydberg.detuning.uniform.piecewise_linear(
-        durations=["up_time", "anneal_time", "up_time"],
-        values=[
-            "initial_detuning",
-            "initial_detuning",
-            "final_detuning",
-            "final_detuning",
-        ],
-    )
-    .rydberg.rabi.amplitude.uniform.piecewise_linear(
-        durations=["up_time", "anneal_time", "up_time"],
-        values=[0, "rabi_amplitude_max", "rabi_amplitude_max", 0],
-    )
-    .parallelize(10.0)
-    .assign(
-        initial_detuning=-10,
-        up_time=0.1,
-        final_detuning=15,
-        anneal_time=10,
-        rabi_amplitude_max=15,
+def test_parallel_task():
+    lattice = Square(3, lattice_spacing=6)
+
+    quantum_task = (
+        lattice.rydberg.detuning.uniform.piecewise_linear(
+            durations=["up_time", "anneal_time", "up_time"],
+            values=[
+                "initial_detuning",
+                "initial_detuning",
+                "final_detuning",
+                "final_detuning",
+            ],
+        )
+        .rydberg.rabi.amplitude.uniform.piecewise_linear(
+            durations=["up_time", "anneal_time", "up_time"],
+            values=[0, "rabi_amplitude_max", "rabi_amplitude_max", 0],
+        )
+        .parallelize(10.0)
+        .assign(
+            initial_detuning=-10,
+            up_time=0.1,
+            final_detuning=15,
+            anneal_time=10,
+            rabi_amplitude_max=15,
+        )
+        .mock(10)
     )
-    .mock(10)
-)
 
-assert quantum_task.tasks[0].parallel_decoder
+    assert quantum_task.tasks[0].parallel_decoder
+
+
+def test_error_parallel_doublecall():
+    with pytest.raises(TypeError):
+        lattice = Square(3, lattice_spacing=6)
+        (
+            lattice.rydberg.detuning.uniform.piecewise_linear(
+                durations=["up_time", "anneal_time", "up_time"],
+                values=[
+                    "initial_detuning",
+                    "initial_detuning",
+                    "final_detuning",
+                    "final_detuning",
+                ],
+            )
+            .rydberg.rabi.amplitude.uniform.piecewise_linear(
+                durations=["up_time", "anneal_time", "up_time"],
+                values=[0, "rabi_amplitude_max", "rabi_amplitude_max", 0],
+            )
+            .parallelize(10.0)
+            .parallelize(4.0)
+        )
+
+
+def test_error_parallel_noatom():
+    with pytest.raises(ValueError):
+        start.rydberg.detuning.uniform.piecewise_linear(
+            durations=["up_time", "anneal_time", "up_time"],
+            values=[
+                "initial_detuning",
+                "initial_detuning",
+                "final_detuning",
+                "final_detuning",
+            ],
+        ).parallelize(10.0)