More emulator integration tests and fixes (#577)

* adding more integration tests.

* fixing bug.

* adding more integration tests to cover emulator IR codegen.

src/bloqade/emulate/codegen/emulator_ir.py

@@ -143,7 +143,7 @@ def visit_assigned_run_time_vector(
 
     def visit_scaled_locations(self, ast: ScaledLocations) -> Dict[int, Decimal]:
         return {
-            loc.value: coeff(**self.assignments) for loc, coeff in ast.locations.items()
+            loc.value: coeff(**self.assignments) for loc, coeff in ast.value.items()
         }
 
     def visit_detuning(self, ast: Optional[Field]):

tests/test_emulator_codegen.py

@@ -111,7 +111,3 @@ def test_codegen_detuning_and_rabi():
     print(expected_emulator_ir)
 
     assert emulator_ir == expected_emulator_ir
-
-
-test_codegen_detuning()
-test_codegen_detuning_and_rabi()

tests/test_python_emulator.py

@@ -5,14 +5,14 @@
 def test_integration_1():
     (
         start.add_position((0, 0))
-        .add_position((0, 1.0))
+        .add_position((0, 5.0))
         .scale("r")
         .rydberg.detuning.uniform.piecewise_linear(
             [0.1, "ramp_time", 0.1], [-100, -100, 100, 100]
         )
         .amplitude.uniform.piecewise_linear([0.1, "ramp_time", 0.1], [0, 10, 10, 0])
         .assign(ramp_time=3.0)
-        .batch_assign(r=np.linspace(4, 10, 31).tolist())
+        .batch_assign(r=np.linspace(4, 10, 11).tolist())
         .bloqade.python()
         .run(10000, cache_matrices=True, blockade_radius=6.0)
         .report()
@@ -24,7 +24,7 @@ def test_integration_2():
     ramp_time = var("ramp_time")
     (
         start.add_position((0, 0))
-        .add_position((0, 1.0))
+        .add_position((0, 5.0))
         .scale("r")
         .rydberg.detuning.uniform.piecewise_linear(
             [0.1, "ramp_time", 0.1], [-100, -100, 100, 100]
@@ -34,7 +34,100 @@ def test_integration_2():
             [0.1, ramp_time / 2, ramp_time / 2, 0.1], [0, 0, np.pi, np.pi]
         )
         .assign(ramp_time=3.0)
-        .batch_assign(r=np.linspace(4, 10, 31).tolist())
+        .batch_assign(r=np.linspace(4, 10, 11).tolist())
+        .bloqade.python()
+        .run(10000, cache_matrices=True, blockade_radius=6.0)
+        .report()
+        .bitstrings()
+    )
+
+
+def test_integration_3():
+    ramp_time = var("ramp_time")
+    (
+        start.add_position((0, 0))
+        .add_position((0, 5.0))
+        .scale("r")
+        .rydberg.detuning.uniform.piecewise_linear(
+            [0.1, ramp_time, 0.1], [-100, -100, 100, 100]
+        )
+        .amplitude.uniform.piecewise_linear([0.1, ramp_time, 0.1], [0, 10, 10, 0])
+        .phase.uniform.piecewise_constant(
+            [0.1, ramp_time / 2, ramp_time / 2, 0.1], [0, 0, np.pi, np.pi]
+        )
+        .amplitude.var("rabi_mask")
+        .fn(lambda t: 4 * np.sin(3 * t), ramp_time + 0.2)
+        .assign(ramp_time=3.0, rabi_mask=[10.0, 0.1])
+        .batch_assign(r=np.linspace(4, 10, 11).tolist())
+        .bloqade.python()
+        .run(10000, cache_matrices=True, blockade_radius=6.0)
+        .report()
+        .bitstrings()
+    )
+
+
+def test_integration_4():
+    ramp_time = var("ramp_time")
+    (
+        start.add_position((0, 0))
+        .add_position((0, 5.0))
+        .scale("r")
+        .rydberg.detuning.uniform.piecewise_linear(
+            [0.1, ramp_time, 0.1], [-100, -100, 100, 100]
+        )
+        .amplitude.uniform.piecewise_linear([0.1, ramp_time, 0.1], [0, 10, 10, 0])
+        .amplitude.var("rabi_mask")
+        .fn(lambda t: 4 * np.sin(3 * t), ramp_time + 0.2)
+        .amplitude.location(1)
+        .linear(0.0, 1.0, ramp_time + 0.2)
+        .assign(ramp_time=3.0, rabi_mask=[10.0, 0.1])
+        .batch_assign(r=np.linspace(4, 10, 11).tolist())
+        .bloqade.python()
+        .run(10000, cache_matrices=True, blockade_radius=6.0)
+        .report()
+        .bitstrings()
+    )
+
+
+def test_integration_5():
+    ramp_time = var("ramp_time")
+    (
+        start.add_position((0, 0))
+        .add_position((0, 5.0))
+        .scale("r")
+        .rydberg.detuning.uniform.piecewise_linear(
+            [0.1, ramp_time, 0.1], [-100, -100, 100, 100]
+        )
+        .amplitude.uniform.piecewise_linear([0.1, ramp_time, 0.1], [0, 10, 10, 0])
+        .phase.location(1)
+        .linear(0.0, 1.0, ramp_time + 0.2)
+        .assign(ramp_time=3.0)
+        .batch_assign(r=np.linspace(4, 10, 11).tolist())
+        .bloqade.python()
+        .run(10000, cache_matrices=True, blockade_radius=6.0)
+        .report()
+        .bitstrings()
+    )
+
+
+def test_integration_6():
+    ramp_time = var("ramp_time")
+    (
+        start.add_position((0, 0))
+        .add_position((0, 5.0))
+        .scale("r")
+        .rydberg.detuning.uniform.piecewise_linear(
+            [0.1, ramp_time, 0.1], [-100, -100, 100, 100]
+        )
+        .location(1)
+        .constant(3.0, ramp_time + 0.2)
+        .location(0)
+        .linear(2.0, 0, ramp_time + 0.2)
+        .amplitude.uniform.piecewise_linear([0.1, ramp_time, 0.1], [0, 10, 10, 0])
+        .phase.location(1)
+        .linear(0.0, 1.0, ramp_time + 0.2)
+        .assign(ramp_time=3.0)
+        .batch_assign(r=np.linspace(4, 10, 11).tolist())
         .bloqade.python()
         .run(10000, cache_matrices=True, blockade_radius=6.0)
         .report()