Fix stimflow start="auto" failing for observables

glue/stimflow/src/stimflow/_chunk/_chunk_builder_test.py

@@ -311,3 +311,31 @@ def test_partial_observable_include_memory_experiment():
     assert np.count_nonzero(dets) == 0
     assert 200 <= np.count_nonzero(obs[:, 0]) <= 300  # Some measurements included in partial X obs
     assert np.count_nonzero(obs[:, 1]) == 0  # No measurements included in partial Z obs
+
+
+def test_auto_obs():
+    builder = stimflow.ChunkBuilder(allowed_qubits=[0, 1])
+    builder.append("R", [0])
+    builder.add_flow(
+        start="auto",
+        end=stimflow.PauliMap({0: 'Z', 1: 'Z'}, obs_name='test'),
+    )
+    builder.add_flow(
+        start=stimflow.PauliMap({1: 'Z'}, obs_name='test2'),
+        end="auto",
+    )
+    chunk = builder.finish_chunk()
+    chunk.verify()
+
+    assert chunk.flows == (
+        stimflow.Flow(
+            start=stimflow.PauliMap({1: 'Z'}, obs_name='test'),
+            end=stimflow.PauliMap({0: 'Z', 1: 'Z'}, obs_name='test'),
+            center=0.5,
+        ),
+        stimflow.Flow(
+            start=stimflow.PauliMap({1: 'Z'}, obs_name='test2'),
+            end=stimflow.PauliMap({1: 'Z'}, obs_name='test2'),
+            center=1,
+        ),
+    )

glue/stimflow/src/stimflow/_chunk/_flow_util.py

@@ -29,7 +29,7 @@ def _solve_auto_flow_starts(
         except ValueError:
             failure_out.append(flow)
             continue
-        start = PauliMap({i2q[q]: "_XYZ"[stim_start[q]] for q in stim_start.pauli_indices()})
+        start = PauliMap({i2q[q]: "_XYZ"[stim_start[q]] for q in stim_start.pauli_indices()}, obs_name=flow.end.obs_name)
         new_flows.append(flow.with_edits(start=start))
 
     return new_flows
@@ -43,7 +43,7 @@ def _solve_auto_flow_ends(
     failure_out: list[Flow],
 ) -> list[Flow]:
 
-    num_qubits = circuit.num_qubits
+    num_qubits = max(q2i.values(), default=-1) + 1
     i2q = {i: q for q, i in q2i.items()}
 
     new_flows = []
@@ -54,7 +54,7 @@ def _solve_auto_flow_ends(
         except ValueError:
             failure_out.append(flow)
             continue
-        end = PauliMap({i2q[q]: "_XYZ"[stim_end[q]] for q in stim_end.pauli_indices()})
+        end = PauliMap({i2q[q]: "_XYZ"[stim_end[q]] for q in stim_end.pauli_indices()}, obs_name=flow.start.obs_name)
         new_flows.append(flow.with_edits(end=end))
 
     return new_flows

glue/stimflow/src/stimflow/_chunk/_stabilizer_code.py

@@ -269,7 +269,7 @@ def x_basis_subset(self) -> StabilizerCode:
 
     def z_basis_subset(self) -> StabilizerCode:
         return StabilizerCode(
-            stabilizers=self.stabilizers.with_only_x_tiles(),
+            stabilizers=self.stabilizers.with_only_z_tiles(),
             logicals=self.list_pure_basis_observables("Z"),
         )
 

glue/stimflow/src/stimflow/_core/_flow.py

@@ -209,14 +209,21 @@ def __str__(self) -> str:
         return result
 
     def __repr__(self):
-        return (
-            f"stimflow.Flow(start={self.start!r}, "
-            f"end={self.end!r}, "
-            f"measurement_indices={self.measurement_indices!r}, "
-            f"flags={self.flags!r}, "
-            f"center={self.center!r}, "
-            f"sign={self.sign!r}"
-        )
+        lines = ["stimflow.Flow("]
+        if self.start:
+            lines.append(f"start={self.start!r},")
+        if self.end:
+            lines.append(f"end={self.end!r},")
+        if self.measurement_indices:
+            lines.append(f"measurement_indices={self.measurement_indices!r},")
+        if self.flags:
+            lines.append(f"flags={self.flags!r},")
+        if self.center is not None:
+            lines.append(f"center={self.center!r},")
+        if self.sign is not None:
+            lines.append(f"sign={self.sign!r},")
+        lines.append(")")
+        return '\n'.join(lines)
 
     def with_xz_flipped(self) -> Flow:
         return self.with_edits(start=self.start.with_xz_flipped(), end=self.end.with_xz_flipped())

mock.py

@@ -0,0 +1,26 @@
+import sys
+
+import stim
+import numpy as np
+
+
+def test_x():
+    print("TEST PROGRESSION: START", file=sys.stderr)
+    circuit = stim.Circuit(
+        """
+        H 0 1
+        CX 0 2 1 3
+        MPP X0*X1 Y0*Y1 Z0*Z1
+        """
+    )
+    print("TEST PROGRESSION: A", file=sys.stderr)
+    aa = circuit.reference_sample()
+    print("TEST PROGRESSION: B", file=sys.stderr)
+    bb = circuit.reference_sample()
+    print("TEST PROGRESSION: C", file=sys.stderr)
+    np.testing.assert_array_equal(aa, bb)
+    print("TEST PROGRESSION: D", file=sys.stderr)
+
+
+if __name__ == '__main__':
+    test_x()

src/stim/circuit/circuit_pybind_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import io
 import pathlib
+import sys
 import tempfile
 from typing import cast
 
@@ -1497,31 +1498,46 @@ def test_num_ticks():
 
 
 def test_reference_sample():
+    print("TEST PROGRESSION: A", file=sys.stderr)
     circuit = stim.Circuit(
         """
         H 0
         CNOT 0 1
     """
     )
+    print("TEST PROGRESSION: B", file=sys.stderr)
     ref = circuit.reference_sample()
+    print("TEST PROGRESSION: C", file=sys.stderr)
     assert len(ref) == 0
+    print("TEST PROGRESSION: D", file=sys.stderr)
     circuit = stim.Circuit(
         """
         H 0 1
         CX 0 2 1 3
         MPP X0*X1 Y0*Y1 Z0*Z1
         """
     )
+    print("TEST PROGRESSION: E", file=sys.stderr)
     np.testing.assert_array_equal(circuit.reference_sample(), circuit.reference_sample())
+    print("TEST PROGRESSION: F", file=sys.stderr)
     assert np.sum(circuit.reference_sample()) % 2 == 1
+    print("TEST PROGRESSION: G", file=sys.stderr)
     circuit.append("X", (i for i in range(0, 100, 2)))
+    print("TEST PROGRESSION: H", file=sys.stderr)
     circuit.append("M", (i for i in range(100)))
+    print("TEST PROGRESSION: I", file=sys.stderr)
     ref_sample = circuit.reference_sample(bit_packed=True)
+    print("TEST PROGRESSION: J", file=sys.stderr)
     unpacked = np.unpackbits(ref_sample, bitorder="little")
+    print("TEST PROGRESSION: K", file=sys.stderr)
     expected = circuit.reference_sample(bit_packed=False)
+    print("TEST PROGRESSION: L", file=sys.stderr)
     expected_padded = np.zeros_like(unpacked)
+    print("TEST PROGRESSION: M", file=sys.stderr)
     expected_padded[:len(expected)] = expected
+    print("TEST PROGRESSION: N", file=sys.stderr)
     np.testing.assert_array_equal(unpacked, expected_padded)
+    print("TEST PROGRESSION: O", file=sys.stderr)
 
 
 def test_max_mix_depolarization_is_allowed_in_dem_conversion_without_args():

src/stim/simulators/frame_simulator_pybind_test.py

@@ -587,8 +587,6 @@ def test_generate_bernoulli_samples():
     assert np.all(v == 0)
 
     sim.generate_bernoulli_samples(256 - 101, p=1, bit_packed=True, out=v[1:-11])
-    for k in v:
-        print(k)
     assert np.all(v[1:-12] == 0xFF)
     assert v[-12] == 7
     assert np.all(v[-11:] == 0)