Merge branch 'main' into david/229-pyqrack-qubit-reset

Author: david-pl

src/bloqade/analysis/fidelity/analysis.py

@@ -13,22 +13,47 @@
 class FidelityAnalysis(Forward):
     """
     This analysis pass can be used to track the global addresses of qubits and wires.
+
+    ## Usage examples
+
+    ```
+    from bloqade import qasm2
+    from bloqade.noise import native
+    from bloqade.analysis.fidelity import FidelityAnalysis
+    from bloqade.qasm2.passes.noise import NoisePass
+
+    noise_main = qasm2.extended.add(native.dialect)
+
+    @noise_main
+    def main():
+        q = qasm2.qreg(2)
+        qasm2.x(q[0])
+        return q
+
+    NoisePass(main.dialects)(main)
+
+    fid_analysis = FidelityAnalysis(main.dialects)
+    fid_analysis.run_analysis(main, no_raise=False)
+
+    gate_fidelity = fid_analysis.gate_fidelity
+    atom_survival_probs = fid_analysis.atom_survival_probability
+    ```
     """
 
     keys = ["circuit.fidelity"]
     lattice = EmptyLattice
 
+    gate_fidelity: float = 1.0
     """
     The fidelity of the gate set described by the analysed program. It reduces whenever a noise channel is encountered.
     """
-    gate_fidelity: float = 1.0
 
     _current_gate_fidelity: float = field(init=False)
 
+    atom_survival_probability: list[float] = field(init=False)
     """
     The probabilities that each of the atoms in the register survive the duration of the analysed program. The order of the list follows the order they are in the register.
     """
-    atom_survival_probability: list[float] = field(init=False)
 
     _current_atom_survival_probability: list[float] = field(init=False)
 

src/bloqade/pyqrack/device.py

@@ -25,7 +25,11 @@
 
 @dataclass
 class PyQrackSimulatorBase(AbstractSimulatorDevice[PyQrackSimulatorTask]):
+    """PyQrack simulation device base class."""
+
     options: PyQrackOptions = field(default_factory=_default_pyqrack_args)
+    """options (PyQrackOptions): options passed into the pyqrack simulator."""
+
     loss_m_result: Measurement = field(default=Measurement.One, kw_only=True)
     rng_state: np.random.Generator = field(
         default_factory=np.random.default_rng, kw_only=True
@@ -99,7 +103,43 @@ def pauli_expectation(pauli: list[Pauli], qubits: list[PyQrackQubit]) -> float:
 
 @dataclass
 class StackMemorySimulator(PyQrackSimulatorBase):
-    """PyQrack simulator device with precalculated stack of qubits."""
+    """
+    PyQrack simulator device with preallocated stack of qubits.
+
+    This can be used to simulate kernels where the number of qubits is known
+    ahead of time.
+
+    ## Usage examples
+
+    ```
+    # Define a kernel
+    @qasm2.main
+    def main():
+        q = qasm2.qreg(2)
+        c = qasm2.creg(2)
+
+        qasm2.h(q[0])
+        qasm2.cx(q[0], q[1])
+
+        qasm2.measure(q, c)
+        return q
+
+    # Create the simulator object
+    sim = StackMemorySimulator(min_qubits=2)
+
+    # Execute the kernel
+    qubits = sim.run(main)
+    ```
+
+    You can also obtain other information from it, such as the state vector:
+
+    ```
+    ket = sim.state_vector(main)
+
+    from pyqrack.pauli import Pauli
+    expectation_vals = sim.pauli_expectation([Pauli.PauliX, Pauli.PauliI], qubits)
+    ```
+    """
 
     min_qubits: int = field(default=0, kw_only=True)
 
@@ -109,6 +149,20 @@ def task(
         args: tuple[Any, ...] = (),
         kwargs: dict[str, Any] | None = None,
     ):
+        """
+        Args:
+            kernel (ir.Method):
+                The kernel method to run.
+            args (tuple[Any, ...]):
+                Positional arguments to pass to the kernel method.
+            kwargs (dict[str, Any] | None):
+                Keyword arguments to pass to the kernel method.
+
+        Returns:
+            PyQrackSimulatorTask:
+                The task object used to track execution.
+
+        """
         if kwargs is None:
             kwargs = {}
 
@@ -134,31 +188,67 @@ def task(
 
 @dataclass
 class DynamicMemorySimulator(PyQrackSimulatorBase):
-    """PyQrack simulator device with dynamic qubit allocation."""
+    """
+
+    PyQrack simulator device with dynamic qubit allocation.
+
+    This can be used to simulate kernels where the number of qubits is not known
+    ahead of time.
+
+    ## Usage examples
+
+    ```
+    # Define a kernel
+    @qasm2.main
+    def main():
+        q = qasm2.qreg(2)
+        c = qasm2.creg(2)
+
+        qasm2.h(q[0])
+        qasm2.cx(q[0], q[1])
+
+        qasm2.measure(q, c)
+        return q
+
+    # Create the simulator object
+    sim = DynamicMemorySimulator()
+
+    # Execute the kernel
+    qubits = sim.run(main)
+    ```
+
+    You can also obtain other information from it, such as the state vector:
+
+    ```
+    ket = sim.state_vector(main)
+
+    from pyqrack.pauli import Pauli
+    expectation_vals = sim.pauli_expectation([Pauli.PauliX, Pauli.PauliI], qubits)
+
+    """
 
     def task(
         self,
         kernel: ir.Method[Params, RetType],
         args: tuple[Any, ...] = (),
         kwargs: dict[str, Any] | None = None,
     ):
+        """
+        Args:
+            kernel (ir.Method):
+                The kernel method to run.
+            args (tuple[Any, ...]):
+                Positional arguments to pass to the kernel method.
+            kwargs (dict[str, Any] | None):
+                Keyword arguments to pass to the kernel method.
+
+        Returns:
+            PyQrackSimulatorTask:
+                The task object used to track execution.
+
+        """
         if kwargs is None:
             kwargs = {}
 
         memory = DynamicMemory(self.options.copy())
         return self.new_task(kernel, args, kwargs, memory)
-
-
-def test():
-    from bloqade.qasm2 import extended
-
-    @extended
-    def main():
-        return 1
-
-    @extended
-    def obs(result: int) -> int:
-        return result
-
-    res = DynamicMemorySimulator().task(main)
-    return res.run()

src/bloqade/qasm2/passes/noise.py

@@ -21,6 +21,34 @@ class NoisePass(Pass):
     NOTE: This pass is not guaranteed to be supported long-term in bloqade. We will be
     moving towards a more general approach to noise modeling in the future.
 
+    ## Usage examples
+
+    ```
+    from bloqade import qasm2
+    from bloqade.noise import native
+    from bloqade.qasm2.passes.noise import NoisePass
+
+    noise_main = qasm2.extended.add(native.dialect)
+
+    @noise_main
+    def main():
+        q = qasm2.qreg(2)
+        qasm2.h(q[0])
+        qasm2.cx(q[0], q[1])
+        return q
+
+    # simple IR without any nosie
+    main.print()
+
+    noise_pass = NoisePass(noise_main)
+
+    # rewrite stuff in-place
+    noise_pass.unsafe_run(main)
+
+    # now, we do have noise channels in the IR
+    main.print()
+    ```
+
     """
 
     noise_model: native.MoveNoiseModelABC = field(

src/bloqade/stim/__init__.py

@@ -1,6 +1,6 @@
 from .groups import main as main
 from ._wrappers import *  # noqa: F403
-from .dialects.aux import *  # noqa F403
 from .dialects.gate import *  # noqa F403
 from .dialects.noise import *  # noqa F403
 from .dialects.collapse import *  # noqa F403
+from .dialects.auxiliary import *  # noqa F403

src/bloqade/stim/_wrappers.py

@@ -2,7 +2,7 @@
 
 from kirin.lowering import wraps
 
-from .dialects import aux, gate, noise, collapse
+from .dialects import gate, noise, collapse, auxiliary
 
 
 # dialect:: gate
@@ -69,32 +69,34 @@ def cz(
 
 ## pp
 @wraps(gate.SPP)
-def spp(targets: tuple[aux.PauliString, ...], dagger=False) -> None: ...
+def spp(targets: tuple[auxiliary.PauliString, ...], dagger=False) -> None: ...
 
 
 # dialect:: aux
-@wraps(aux.GetRecord)
-def rec(id: int) -> aux.RecordResult: ...
+@wraps(auxiliary.GetRecord)
+def rec(id: int) -> auxiliary.RecordResult: ...
 
 
-@wraps(aux.Detector)
+@wraps(auxiliary.Detector)
 def detector(
-    coord: tuple[Union[int, float], ...], targets: tuple[aux.RecordResult, ...]
+    coord: tuple[Union[int, float], ...], targets: tuple[auxiliary.RecordResult, ...]
 ) -> None: ...
 
 
-@wraps(aux.ObservableInclude)
-def observable_include(idx: int, targets: tuple[aux.RecordResult, ...]) -> None: ...
+@wraps(auxiliary.ObservableInclude)
+def observable_include(
+    idx: int, targets: tuple[auxiliary.RecordResult, ...]
+) -> None: ...
 
 
-@wraps(aux.Tick)
+@wraps(auxiliary.Tick)
 def tick() -> None: ...
 
 
-@wraps(aux.NewPauliString)
+@wraps(auxiliary.NewPauliString)
 def pauli_string(
     string: tuple[str, ...], flipped: tuple[bool, ...], targets: tuple[int, ...]
-) -> aux.PauliString: ...
+) -> auxiliary.PauliString: ...
 
 
 # dialect:: collapse
@@ -123,7 +125,7 @@ def mxx(p: float, targets: tuple[int, ...]) -> None: ...
 
 
 @wraps(collapse.PPMeasurement)
-def mpp(p: float, targets: tuple[aux.PauliString, ...]) -> None: ...
+def mpp(p: float, targets: tuple[auxiliary.PauliString, ...]) -> None: ...
 
 
 @wraps(collapse.RZ)

src/bloqade/stim/dialects/__init__.py

@@ -1,5 +1,5 @@
-from . import aux as aux, gate as gate, noise as noise, collapse as collapse
-from .aux.stmts import *  # noqa F403
+from . import gate as gate, noise as noise, collapse as collapse, auxiliary as auxiliary
 from .gate.stmts import *  # noqa F403
 from .noise.stmts import *  # noqa F403
 from .collapse.stmts import *  # noqa F403
+from .auxiliary.stmts import *  # noqa F403