Improved exception handling (#91)

Slight normalization issues no longer raise an error but a warning.
These can occur for circuits with arbitrary rotation gates, where fp32
is used instead of an exact representation.

Parsing errors for invalid Stim circuits now raise useful exceptions.

Author: rafaelha

src/tsim/circuit.py

@@ -16,7 +16,11 @@
 from tsim.noise.dem import get_detector_error_model
 from tsim.utils.clifford import parametric_to_clifford_gates
 from tsim.utils.diagram import render_svg
-from tsim.utils.program_text import shorthand_to_stim, stim_to_shorthand
+from tsim.utils.program_text import (
+    enriched_stim_error,
+    shorthand_to_stim,
+    stim_to_shorthand,
+)
 
 
 class Circuit:
@@ -42,7 +46,11 @@ def __init__(self, stim_program_text: str = ""):
                 empty circuit.
 
         """
-        self._stim_circ = stim.Circuit(shorthand_to_stim(stim_program_text))
+        converted = shorthand_to_stim(stim_program_text)
+        try:
+            self._stim_circ = stim.Circuit(converted)
+        except ValueError as exc:
+            raise enriched_stim_error(exc, converted) from None
 
     @classmethod
     def from_stim_program(cls, stim_circuit: stim.Circuit) -> Circuit:
@@ -64,9 +72,11 @@ def append_from_stim_program_text(self, stim_program_text: str) -> None:
 
         Supports the same shorthand syntax as the constructor.
         """
-        self._stim_circ.append_from_stim_program_text(
-            shorthand_to_stim(stim_program_text)
-        )
+        converted = shorthand_to_stim(stim_program_text)
+        try:
+            self._stim_circ.append_from_stim_program_text(converted)
+        except ValueError as exc:
+            raise enriched_stim_error(exc, converted) from None
 
     @overload
     def append(
@@ -141,18 +151,21 @@ def append(
                 name = "S_DAG"
                 tag = "T"
             elif name in ("R_X", "R_Y", "R_Z"):
-                assert arg is not None, f"For {name} gates, an angle must be provided."
+                if arg is None:
+                    raise ValueError(f"For {name} gates, an angle must be provided.")
                 args = list(arg) if isinstance(arg, Iterable) else [arg]
-                assert (
-                    len(args) == 1
-                ), f"For {name} gates, a single angle must be provided."
+                if len(args) != 1:
+                    raise ValueError(
+                        f"For {name} gates, a single angle must be provided."
+                    )
                 tag = f"{name}(theta={args[0]}*pi)"
                 name = "I"
                 arg = None
             elif name == "U3":
-                assert arg is not None and (
-                    isinstance(arg, Iterable) and len(list(arg)) == 3
-                ), f"For U3 gates, three rotation angles must be provided."
+                if arg is None or not isinstance(arg, Iterable) or len(list(arg)) != 3:
+                    raise ValueError(
+                        "For U3 gates, three rotation angles must be provided."
+                    )
                 theta, phi, lam = list(arg)
                 tag = f"U3(theta={theta}*pi, phi={phi}*pi, lambda={lam}*pi)"
                 name = "I"
@@ -175,7 +188,11 @@ def from_file(cls, filename: str) -> Circuit:
         """
         with open(filename, "r", encoding="utf-8") as f:
             stim_program_text = f.read()
-        stim_circ = stim.Circuit(shorthand_to_stim(stim_program_text))
+        converted = shorthand_to_stim(stim_program_text)
+        try:
+            stim_circ = stim.Circuit(converted)
+        except ValueError as exc:
+            raise enriched_stim_error(exc, converted) from None
         return cls.from_stim_program(stim_circ)
 
     def __repr__(self) -> str:

src/tsim/core/parse.py

@@ -47,8 +47,10 @@ def parse_parametric_tag(tag: str) -> tuple[str, dict[str, Fraction]] | None:
         param = param.strip()
         if not param:
             continue
-        # Match param=value*pi (value can be negative/decimal)
-        param_match = re.match(r"^(\w+)=([-+]?[\d.]+)\*pi$", param)
+        # Match param=value*pi (value can be negative/decimal/scientific)
+        param_match = re.match(
+            r"^(\w+)=([-+]?(?:\d+\.?\d*|\.\d+)(?:[eE][-+]?\d+)?)\*pi$", param
+        )
         if not param_match:
             return None
         param_name = param_match.group(1)

src/tsim/sampler.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import warnings
 from math import ceil
 from typing import TYPE_CHECKING, Literal, overload
 
@@ -133,12 +134,17 @@ def sample_program(
 
     for component in program.components:
         samples, key, max_norm_deviation = sample_component(component, f_params, key)
+        if np.isclose(max_norm_deviation, 1):
+            raise ValueError(
+                "A vanishing marginal probability distributionwas encountered (normalization 0). "
+                "This is likely the result of an underflow error. Please report this "
+                "as a bug at https://github.com/QuEraComputing/tsim/issues/new."
+            )  # pragma: no cover
         if max_norm_deviation > 1e-5:
-            raise AssertionError(
+            warnings.warn(
                 "A marginal probability was not normalized correctly "
                 f"(normalization deviated from 1 by {max_norm_deviation:.1e}). "
-                "This is likely the result of an underflow error. Please report this "
-                "as a bug at https://github.com/QuEraComputing/tsim/issues/new."
+                "This is likely a floating point precision issue."
             )
         results.append(samples)
 

src/tsim/utils/encoder.py

@@ -21,12 +21,11 @@ def _transform_circuit(
     *,
     stride: int,
     offsets: list[int],
-    gate_expansions: dict[str, list[str]] | None = None,
     used_qubits: set[int] | None = None,
     stabilizer_generators: list[list[int]] | None = None,
     observables: list[list[int]] | None = None,
 ) -> stim.Circuit:
-    """Expand and duplicate instructions with broadcast targets for encoding."""
+    """Duplicate instructions with broadcast targets for encoding."""
     stim_circ = tsim.Circuit(program_text)._stim_circ.flattened()
     mod_circ = stim.Circuit()
 
@@ -71,19 +70,12 @@ def _transform_circuit(
             instr.target_groups(), stride=stride, offsets=offsets
         )
 
-        gate_seq = (
-            gate_expansions.get(instr.name, [instr.name])
-            if gate_expansions
-            else [instr.name]
+        mod_circ.append(
+            instr.name,
+            new_ts,
+            instr.gate_args_copy(),
+            tag=instr.tag,
         )
-
-        for g in gate_seq:
-            mod_circ.append(
-                g,
-                new_ts,
-                instr.gate_args_copy(),
-                tag=instr.tag,
-            )
     return mod_circ
 
 
@@ -100,15 +92,13 @@ def __init__(
         encoding_program_text: str | None,
         stabilizer_generators: list[list[int]],
         observables: list[list[int]],
-        logical_gate_expansions: dict[str, list[str]] | None = None,
     ):
         """Initialize the transversal encoder with code parameters."""
         self.n = n
         self.encoding_qubit = encoding_qubit
         self.circuit = tsim.Circuit()
         self.used_qubits: set[int] = set()
         self.encoding_program_text = encoding_program_text
-        self.logical_gate_expansions = logical_gate_expansions or {}
         self.stabilizer_generators = stabilizer_generators
         self.observables = observables
 
@@ -169,7 +159,6 @@ def encode_transversally(self, program_text: str) -> None:
             program_text,
             stride=self.n,
             offsets=list(range(self.n)),
-            gate_expansions=self.logical_gate_expansions,
             stabilizer_generators=self.stabilizer_generators,
             observables=self.observables,
         )
@@ -215,12 +204,6 @@ def __init__(self):
             n=7,
             encoding_qubit=6,
             encoding_program_text=encoding_program,
-            logical_gate_expansions={
-                "SQRT_X": ["SQRT_X", "X"],
-                "SQRT_X_DAG": ["SQRT_X_DAG", "X"],
-                "S": ["S", "Z"],
-                "S_DAG": ["S_DAG", "Z"],
-            },
             stabilizer_generators=[[0, 1, 2, 3], [1, 2, 4, 5], [2, 3, 4, 6]],
             observables=[[0, 1, 5]],
         )

src/tsim/utils/program_text.py

@@ -2,6 +2,30 @@
 
 import re
 
+# Matches valid numeric literals including scientific notation (e.g. 0.5, 4e-4, 1.2e3)
+_FLOAT_RE = r"[-+]?(?:\d+\.?\d*|\.\d+)(?:[eE][-+]?\d+)?"
+
+_TSIM_GATES = {"R_X", "R_Y", "R_Z", "U3"}
+_GATE_NOT_FOUND_RE = re.compile(r"Gate not found: '(\w+)'")
+_GATE_USAGE_RE = re.compile(r"(?<!\[)\b(R_[A-Z]\([^)]*\)|R_[XYZ]\b|U3\([^)]*\)|U3\b)")
+
+
+def enriched_stim_error(exc: ValueError, converted_text: str) -> ValueError:
+    """Improve stim parse errors for tsim-specific gates.
+
+    When stim raises a 'Gate not found' error for a gate that should have been
+    converted by shorthand_to_stim, this searches the converted text for the
+    unconverted usage and returns a more helpful error message.
+    """
+    m = _GATE_NOT_FOUND_RE.search(str(exc))
+    if not m or m.group(1) not in _TSIM_GATES:
+        return exc
+    # Successfully converted gates live inside brackets (e.g. I[R_Z(...)]) and won't match.
+    usage = _GATE_USAGE_RE.search(converted_text)
+    if not usage:
+        return exc
+    return ValueError(f"Could not parse '{usage.group()}' in program text.")
+
 
 def shorthand_to_stim(text: str) -> str:
     """Convert tsim shorthand syntax to valid stim instructions.
@@ -13,27 +37,25 @@ def shorthand_to_stim(text: str) -> str:
         R_X(0.25) 0     → I[R_X(theta=0.25*pi)] 0
         R_Y(-0.5) 0     → I[R_Y(theta=-0.5*pi)] 0
         U3(0.3, 0.24, 0.49) 0 → I[U3(theta=0.3*pi, phi=0.24*pi, lambda=0.49*pi)] 0
+
     """
     # T_DAG must come before T to avoid partial matches
     # (?<!\[) ensures we don't match T inside [T]
     text = re.sub(r"(?<!\[)\bT_DAG\b(?!\[)", "S_DAG[T]", text)
     text = re.sub(r"(?<!\[)\bT\b(?!\[)", "S[T]", text)
 
-    # R_Z(angle), R_X(angle), R_Y(angle)
     def replace_rotation(m: re.Match) -> str:
         axis = m.group(1)
-        angle = m.group(2)
-        return f"I[R_{axis}(theta={angle}*pi)]"
+        return f"I[R_{axis}(theta={float(m.group(2))}*pi)]"
 
-    text = re.sub(r"\bR_([XYZ])\(([-+]?[\d.]+)\)", replace_rotation, text)
+    text = re.sub(rf"\bR_([XYZ])\(({_FLOAT_RE})\)", replace_rotation, text)
 
-    # U3(theta, phi, lambda)
     def replace_u3(m: re.Match) -> str:
-        theta, phi, lam = m.group(1), m.group(2), m.group(3)
+        theta, phi, lam = float(m.group(1)), float(m.group(2)), float(m.group(3))
         return f"I[U3(theta={theta}*pi, phi={phi}*pi, lambda={lam}*pi)]"
 
     text = re.sub(
-        r"\bU3\(([-+]?[\d.]+)\s*,\s*([-+]?[\d.]+)\s*,\s*([-+]?[\d.]+)\)",
+        rf"\bU3\(({_FLOAT_RE})\s*,\s*({_FLOAT_RE})\s*,\s*({_FLOAT_RE})\)",
         replace_u3,
         text,
     )

test/integration/test_sampler.py

@@ -1,4 +1,5 @@
 from test.helpers.gen import gen_stim_circuit
+
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -137,7 +138,7 @@ def fake_sample_component(component, f_params, key):
 
     monkeypatch.setattr(sampler_module, "sample_component", fake_sample_component)
 
-    with pytest.raises(AssertionError, match="underflow error"):
+    with pytest.warns(UserWarning, match="not normalized"):
         sampler_module.sample_program(
             program,
             jnp.zeros((1, 0), dtype=jnp.bool_),

test/unit/utils/test_program_text.py

@@ -1,4 +1,12 @@
-from tsim.utils.program_text import shorthand_to_stim, stim_to_shorthand
+import re
+
+import pytest
+
+from tsim import Circuit
+from tsim.utils.program_text import (
+    shorthand_to_stim,
+    stim_to_shorthand,
+)
 
 
 def test_shorthand_to_stim_t_and_t_dag():
@@ -41,3 +49,37 @@ def test_stim_to_shorthand_rotations_and_u3():
 def test_shorthand_roundtrip():
     text = "T 0\nR_X(0.5) 1\nU3(0.1, 0.2, 0.3) 2"
     assert stim_to_shorthand(shorthand_to_stim(text)) == text
+
+
+def test_shorthand_scientific_notation():
+    result = shorthand_to_stim("R_Z(4e-4) 0")
+    assert "I[R_Z(theta=0.0004*pi)]" in result
+
+
+def test_shorthand_scientific_notation_u3():
+    result = shorthand_to_stim("U3(1e-2, 2.5e1, 3e-3) 0")
+    assert "I[U3(" in result
+
+
+def test_circuit_scientific_notation():
+    c = Circuit("R_Z(4e-4) 0")
+    assert len(c) == 1
+
+
+@pytest.mark.parametrize(
+    "text, snippet",
+    [
+        ("R_Z(a) 0", "R_Z(a)"),
+        ("R_Z(pi) 0", "R_Z(pi)"),
+        ("R_Z(1/3) 0", "R_Z(1/3)"),
+        ("R_Z() 0", "R_Z()"),
+        ("R_Z 0", "R_Z"),
+        ("R_Z(0.5, 0.3) 0", "R_Z(0.5, 0.3)"),
+        ("R_X(abc) 0", "R_X(abc)"),
+        ("U3(0.1, 0.2) 0", "U3(0.1, 0.2)"),
+        ("U3(0.1, 0.2, 0.3, 0.4) 0", "U3(0.1, 0.2, 0.3, 0.4)"),
+    ],
+)
+def test_circuit_parse_error_shows_snippet(text, snippet):
+    with pytest.raises(ValueError, match=re.escape(snippet)):
+        Circuit(text)