Add some basic dosctrings and usage examples to pyqrack simulator devices

Author: david-pl

src/bloqade/pyqrack/device.py

@@ -25,6 +25,13 @@
 
 @dataclass
 class PyQrackSimulatorBase(AbstractSimulatorDevice[PyQrackSimulatorTask]):
+    """PyQrack simulation device base class.
+
+    Args:
+       options (PyQrackOptions):
+           Options passed into the PyQrack simulator
+    """
+
     options: PyQrackOptions = field(default_factory=_default_pyqrack_args)
     loss_m_result: Measurement = field(default=Measurement.One, kw_only=True)
     rng_state: np.random.Generator = field(
@@ -96,12 +103,48 @@ def pauli_expectation(pauli: list[Pauli], qubits: list[PyQrackQubit]) -> float:
         if len(qubit_ids) != len(set(qubit_ids)):
             raise ValueError("Qubits must be unique.")
 
-        return sim_reg.pauli_expectation(pauli, qubit_ids)
+        return sim_reg.pauli_expectation(qubit_ids, pauli)
 
 
 @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)
 
@@ -111,6 +154,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 = {}
 
@@ -136,31 +193,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()