Add some basic dosctrings and usage examples to pyqrack simulator (#254)

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(