Get observable idx from measurement id analysis (#712)

Adresses #671 

Removes the global observable index from squin. This index is still
required to emit stim -- here it is extracted from the
MeasurementIDAnalysis pass.

---------

Co-authored-by: Phillip Weinberg <pweinberg@quera.com>

Author: rafaelha

pyproject.toml

@@ -11,7 +11,7 @@ authors = [
 ]
 requires-python = ">=3.10"
 dependencies = [
-    "bloqade-decoders~=0.3.0",
+    "bloqade-decoders~=0.4.0",
     "numpy>=1.22.0",
     "scipy>=1.13.1",
     "kirin-toolchain~=0.22.2",

src/bloqade/stim/rewrite/set_observable_to_stim.py

@@ -1,11 +1,12 @@
 from dataclasses import dataclass
 
 from kirin import ir
+from kirin.dialects import py
 from kirin.rewrite.abc import RewriteRule, RewriteResult
 
 from bloqade.analysis.measure_id import MeasureIDFrame
 from bloqade.stim.dialects.auxiliary import ObservableInclude
-from bloqade.analysis.measure_id.lattice import MeasureIdTuple
+from bloqade.analysis.measure_id.lattice import ObservableId, MeasureIdTuple
 from bloqade.decoders.dialects.annotate.stmts import SetObservable
 
 from ..rewrite.get_record_util import insert_get_records
@@ -29,13 +30,18 @@ def rewrite_Statement(self, node: ir.Statement) -> RewriteResult:
     def rewrite_SetObservable(self, node: SetObservable) -> RewriteResult:
         measure_ids = self.measure_id_frame.entries[node.measurements]
         assert isinstance(measure_ids, MeasureIdTuple)
+        observable_id = self.measure_id_frame.entries.get(node.result)
+        if not isinstance(observable_id, ObservableId):
+            return RewriteResult()
 
         get_record_list = insert_get_records(
             node, measure_ids, self.measure_id_frame.num_measures_at_stmt[node]
         )
 
+        idx_stmt = py.Constant(observable_id.idx)
+        idx_stmt.insert_before(node)
         observable_include_stmt = ObservableInclude(
-            idx=node.idx, targets=tuple(get_record_list)
+            idx=idx_stmt.result, targets=tuple(get_record_list)
         )
 
         node.replace_by(observable_include_stmt)

test/stim/passes/test_annotation_to_stim.py

@@ -50,7 +50,7 @@ def main():
         squin.set_detector([ms[0], ms[1]], coordinates=[0.0, 0.0])
         squin.set_detector([ms[1], ms[2]], coordinates=[1.0, 0.0])
 
-        squin.set_observable(measurements=[ms[2]], idx=0)
+        squin.set_observable(measurements=[ms[2]])
 
         return
 
@@ -81,7 +81,7 @@ def main():
 
         ms1 = squin.broadcast.measure(q)
         squin.set_detector([ms1[0], ms1[1]], coordinates=[0.0, 0.0])
-        squin.set_observable(measurements=[ms1[2]], idx=0)
+        squin.set_observable(measurements=[ms1[2]])
 
         return
 
@@ -342,7 +342,7 @@ def rep_code():
         )
 
         # Now we want to dictate a measurement as the observable
-        squin.set_observable(measurements=[data_meas_res[-1]], idx=0)
+        squin.set_observable(measurements=[data_meas_res[-1]])
 
     SquinToStimPass(rep_code.dialects)(rep_code)
 

test/test_annotate.py

@@ -23,7 +23,7 @@ def test():
         qs = squin.qalloc(4)
         ms = squin.broadcast.measure(qs)
         squin.set_detector([ms[0], ms[1], ms[2]], coordinates=(0, 0))
-        squin.set_observable([ms[3]], 0)
+        squin.set_observable([ms[3]])
 
     SquinToStimPass(dialects=test.dialects)(test)
     codegen_output = codegen(test)
@@ -35,22 +35,22 @@ def test():
     assert codegen_output == expected_output
 
 
-def test_set_observable_with_multiple_idx():
+def test_set_observable_with_multiple_observables():
 
     @squin.kernel
     def test():
         qs = squin.qalloc(4)
         ms = squin.broadcast.measure(qs)
-        squin.set_observable([ms[0], ms[1]], 0)
-        squin.set_observable([ms[2]], 42)
-        squin.set_observable([ms[3]], 0)
+        squin.set_observable([ms[0], ms[1]])
+        squin.set_observable([ms[2]])
+        squin.set_observable([ms[3]])
 
     SquinToStimPass(dialects=test.dialects)(test)
     codegen_output = codegen(test)
     expected_output = (
         "MZ(0.00000000) 0 1 2 3\n"
         "OBSERVABLE_INCLUDE(0) rec[-4] rec[-3]\n"
-        "OBSERVABLE_INCLUDE(42) rec[-2]\n"
-        "OBSERVABLE_INCLUDE(0) rec[-1]"
+        "OBSERVABLE_INCLUDE(1) rec[-2]\n"
+        "OBSERVABLE_INCLUDE(2) rec[-1]"
     )
     assert codegen_output == expected_output