[Feature] Adding digital noise (#34)

* [Feature] Use single precision by default * refac expectation like pyq * Add sample, increase atol * rework circ * Remove spurious import. * Lint. * adding noisy operators * change interface * test tuple length noise in gate * fix noise * separate tests between noisy and non noisy * lint * fix order noise and param by inheritance * fix union * rm densitymatrix object - to replace by more functional way * add raises notimplementederrors * add boolean in fwd fcts * add current implementation of channel apply * change dagger * apply_gate correct on density matrices before permutation * fix permutation with density matrices * fix permute_basis * adding tests with dm * adding test shots * add permute_basis at beginning of apply * using jax lax transpose * adding test dm parameteric * fix controlled ops with dm * checking tests expectation work for dm * separate apply_operator with a density matrix version * add sampling from density mat * fix expectation with density matrices * test also shots with noise * adding docs * fix lint * rm circuit methods in favor of api * rm comment * change for isdensity * Tuple to tuple * change from default noise tuple to None * more docstr in apply * change tuple call new state dims * fix union * mention density matrix simulator * using single dispatch - breaking * shots not working * add fixmes * fix single dispatch methods * adding values to observable to matrix * fix shots dm * fix doc strings * update State typing * update typing primitive noiseprotocol * fix docs is_density * Union instead of pipe * change union pipe is_controlled * union of primitive * Update docs/noise.md Co-authored-by: RolandMacDoland <[email protected]> * Update docs/noise.md Co-authored-by: RolandMacDoland <[email protected]> * Update docs/noise.md Co-authored-by: RolandMacDoland <[email protected]> * Update docs/noise.md Co-authored-by: RolandMacDoland <[email protected]> * Update docs/noise.md Co-authored-by: RolandMacDoland <[email protected]> * Update pyproject.toml Co-authored-by: RolandMacDoland <[email protected]> * test errors noise * Update docs/noise.md Co-authored-by: RolandMacDoland <[email protected]> * Update horqrux/api.py Co-authored-by: RolandMacDoland <[email protected]> * rm dunder * value error to type * Update horqrux/utils_noise.py Co-authored-by: RolandMacDoland <[email protected]> * Update horqrux/primitive.py Co-authored-by: RolandMacDoland <[email protected]> * Update horqrux/shots.py Co-authored-by: RolandMacDoland <[email protected]> * change name of noise to include digital * finiteshots fwd named args * more renaming * change error message noise * fix raise * Update horqrux/utils.py Co-authored-by: RolandMacDoland <[email protected]> * add num_qubits * bump horqrux --------- Co-authored-by: seitzdom <[email protected]> Co-authored-by: Roland Guichard <[email protected]> Co-authored-by: Charles MOUSSA <[email protected]> Co-authored-by: RolandMacDoland <[email protected]>
pasqal-io · Jan 8, 2025 · 5b1de24 · 5b1de24
1 parent b20e094
commit 5b1de24
Show file tree

Hide file tree

Showing 20 changed files with 1,476 additions and 213 deletions.
diff --git a/README.md b/README.md
@@ -2,7 +2,7 @@
 [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 [![Pypi](https://badge.fury.io/py/horqrux.svg)](https://pypi.org/project/horqrux/)
 
-`horqrux` is a [JAX](https://jax.readthedocs.io/en/latest/)-based state vector simulator designed for quantum machine learning and acts as a backend for [`Qadence`](https://github.com/pasqal-io/qadence), a digital-analog quantum programming interface.
+`horqrux` is a [JAX](https://jax.readthedocs.io/en/latest/)-based state vector and density matrix simulator designed for quantum machine learning and acts as a backend for [`Qadence`](https://github.com/pasqal-io/qadence), a digital-analog quantum programming interface.
 
 ## Installation
 

diff --git a/docs/index.md b/docs/index.md
@@ -1,6 +1,6 @@
 # Welcome to horqrux
 
-**horqrux** is a state vector simulator designed for quantum machine learning written in [JAX](https://jax.readthedocs.io/).
+**horqrux** is a state vector and density matrix simulator designed for quantum machine learning written in [JAX](https://jax.readthedocs.io/).
 
 ## Setup
 
@@ -110,10 +110,11 @@ from operator import add
 from typing import Any, Callable
 from uuid import uuid4
 
-from horqrux.circuit import QuantumCircuit, hea, expectation
+from horqrux import expectation
+from horqrux import Z, RX, RY, NOT, zero_state, apply_gate
+from horqrux.circuit import QuantumCircuit, hea
 from horqrux.primitive import Primitive
 from horqrux.parametric import Parametric
-from horqrux import Z, RX, RY, NOT, zero_state, apply_gate
 from horqrux.utils import DiffMode
 
 

diff --git a/docs/noise.md b/docs/noise.md
@@ -0,0 +1,101 @@
+## Digital Noise
+
+In the description of closed quantum systems, the complete quantum state is a pure quantum state represented by a state vector $|\psi \rangle $.
+
+However, this description is not sufficient to describe open quantum systems. When the system interacts with its surrounding environment, it transitions into a mixed state where quantum information is no longer entirely contained in a single state vector but is distributed probabilistically.
+
+To address these more general cases, we consider a probabilistic combination $p_i$ of possible pure states $|\psi_i \rangle$. Thus, the system is described by a density matrix $\rho$ defined as follows:
+
+$$
+\rho = \sum_i p_i |\psi_i\rangle \langle \psi_i|
+$$
+
+The transformations of the density operator of an open quantum system interacting with its noisy environment are represented by the super-operator $S: \rho \rightarrow S(\rho)$, often referred to as a quantum channel.
+Quantum channels, due to the conservation of the probability distribution, must be CPTP (Completely Positive and Trace Preserving). Any CPTP super-operator can be written in the following form:
+
+$$
+S(\rho) = \sum_i K_i \rho K^{\dagger}_i
+$$
+
+Where $K_i$ are Kraus operators satisfying the closure property $\sum_i K_i K^{\dagger}_i = \mathbb{I}$. As noise is the result of system interactions with its environment, it is therefore possible to simulate noisy quantum circuit with noise type gates.
+
+Thus, `horqrux` implements a large selection of single qubit noise gates such as:
+
+- The bit flip channel defined as: $\textbf{BitFlip}(\rho) =(1-p) \rho + p X \rho X^{\dagger}$
+- The phase flip channel defined as: $\textbf{PhaseFlip}(\rho) = (1-p) \rho + p Z \rho Z^{\dagger}$
+- The depolarizing channel defined as: $\textbf{Depolarizing}(\rho) = (1-p) \rho + \frac{p}{3} (X \rho X^{\dagger} + Y \rho Y^{\dagger} + Z \rho Z^{\dagger})$
+- The pauli channel defined as: $\textbf{PauliChannel}(\rho) = (1-p_x-p_y-p_z) \rho
+            + p_x X \rho X^{\dagger}
+            + p_y Y \rho Y^{\dagger}
+            + p_z Z \rho Z^{\dagger}$
+- The amplitude damping channel defined as: $\textbf{AmplitudeDamping}(\rho) =  K_0 \rho K_0^{\dagger} + K_1 \rho K_1^{\dagger}$
+    with:
+    $\begin{equation*}
+    K_{0} \ =\begin{pmatrix}
+    1 & 0\\
+    0 & \sqrt{1-\ \gamma }
+    \end{pmatrix} ,\ K_{1} \ =\begin{pmatrix}
+    0 & \sqrt{\ \gamma }\\
+    0 & 0
+    \end{pmatrix}
+    \end{equation*}$
+- The phase damping channel defined as: $\textbf{PhaseDamping}(\rho) = K_0 \rho K_0^{\dagger} + K_1 \rho K_1^{\dagger}$
+    with:
+    $\begin{equation*}
+    K_{0} \ =\begin{pmatrix}
+    1 & 0\\
+    0 & \sqrt{1-\ \gamma }
+    \end{pmatrix}, \ K_{1} \ =\begin{pmatrix}
+    0 & 0\\
+    0 & \sqrt{\ \gamma }
+    \end{pmatrix}
+    \end{equation*}$
+* The generalize amplitude damping channel is defined as: $\textbf{GeneralizedAmplitudeDamping}(\rho) = K_0 \rho K_0^{\dagger} + K_1 \rho K_1^{\dagger} + K_2 \rho K_2^{\dagger} + K_3 \rho K_3^{\dagger}$
+    with:
+$\begin{cases}
+K_{0} \ =\sqrt{p} \ \begin{pmatrix}
+1 & 0\\
+0 & \sqrt{1-\ \gamma }
+\end{pmatrix} ,\ K_{1} \ =\sqrt{p} \ \begin{pmatrix}
+0 & 0\\
+0 & \sqrt{\ \gamma }
+\end{pmatrix} \\
+K_{2} \ =\sqrt{1\ -p} \ \begin{pmatrix}
+\sqrt{1-\ \gamma } & 0\\
+0 & 1
+\end{pmatrix} ,\ K_{3} \ =\sqrt{1-p} \ \begin{pmatrix}
+0 & 0\\
+\sqrt{\ \gamma } & 0
+\end{pmatrix}
+\end{cases}$
+
+Noise protocols can be added to gates by instantiating `DigitalNoiseInstance` providing the `DigitalNoiseType` and the `error_probability` (either float or tuple of float):
+
+```python exec="on" source="material-block" html="1"
+from horqrux.noise import DigitalNoiseInstance, DigitalNoiseType
+
+noise_prob = 0.3
+AmpD = DigitalNoiseInstance(DigitalNoiseType.AMPLITUDE_DAMPING, error_probability=noise_prob)
+
+```
+
+Then a gate can be instantiated by providing a tuple of `DigitalNoiseInstance` instances. Let’s show this through the simulation of a realistic $X$ gate.
+
+For instance, an $X$ gate flips the state of the qubit: $X|0\rangle = |1\rangle$. In practice, it's common for the target qubit to stay in its original state after applying $X$ due to the interactions between it and its environment. The possibility of failure can be represented by the `BitFlip` subclass of `DigitalNoiseInstance`, which flips the state again after the application of the $X$ gate, returning it to its original state with a probability `1 - gate_fidelity`.
+
+```python exec="on" source="material-block"
+from horqrux.api import sample
+from horqrux.noise import DigitalNoiseInstance, DigitalNoiseType
+from horqrux.utils import density_mat, product_state
+from horqrux.primitive import X
+
+noise = (DigitalNoiseInstance(DigitalNoiseType.BITFLIP, 0.1),)
+ops = [X(0)]
+noisy_ops = [X(0, noise=noise)]
+state = product_state("0")
+
+noiseless_samples = sample(state, ops)
+noisy_samples = sample(density_mat(state), noisy_ops)
+print(f"Noiseless samples: {noiseless_samples}") # markdown-exec: hide
+print(f"Noiseless samples: {noisy_samples}") # markdown-exec: hide
+```
diff --git a/horqrux/__init__.py b/horqrux/__init__.py
@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from .api import expectation
+from .api import expectation, run, sample
 from .apply import apply_gate, apply_operator
-from .circuit import QuantumCircuit, sample
+from .circuit import QuantumCircuit
 from .parametric import PHASE, RX, RY, RZ
 from .primitive import NOT, SWAP, H, I, S, T, X, Y, Z
 from .utils import (

diff --git a/horqrux/adjoint.py b/horqrux/adjoint.py
@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from typing import Tuple
-
 from jax import Array, custom_vjp
 
 from horqrux.apply import apply_gate
@@ -37,14 +35,14 @@ def adjoint_expectation(
 
 def adjoint_expectation_single_observable_fwd(
     state: Array, gates: list[Primitive], observable: list[Primitive], values: dict[str, float]
-) -> Tuple[Array, Tuple[Array, Array, list[Primitive], dict[str, float]]]:
+) -> tuple[Array, tuple[Array, Array, list[Primitive], dict[str, float]]]:
     out_state = apply_gate(state, gates, values, OperationType.UNITARY)
     projected_state = apply_gate(out_state, observable, values, OperationType.UNITARY)
     return inner(out_state, projected_state).real, (out_state, projected_state, gates, values)
 
 
 def adjoint_expectation_single_observable_bwd(
-    res: Tuple[Array, Array, list[Primitive], dict[str, float]], tangent: Array
+    res: tuple[Array, Array, list[Primitive], dict[str, float]], tangent: Array
 ) -> tuple:
     """Implementation of Algorithm 1 of https://arxiv.org/abs/2009.02823
     which computes the vector-jacobian product in O(P) time using O(1) state vectors

diff --git a/horqrux/api.py b/horqrux/api.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections import Counter
+from functools import singledispatch
 from typing import Any, Optional
 
 import jax
@@ -11,74 +12,126 @@
 from horqrux.adjoint import adjoint_expectation
 from horqrux.apply import apply_gate
 from horqrux.primitive import GateSequence, Primitive
-from horqrux.shots import finite_shots_fwd
-from horqrux.utils import DiffMode, ForwardMode, OperationType, inner
+from horqrux.shots import finite_shots_fwd, to_matrix
+from horqrux.utils import (
+    DensityMatrix,
+    DiffMode,
+    ForwardMode,
+    OperationType,
+    State,
+    inner,
+    num_qubits,
+    probabilities,
+    sample_from_probs,
+)
 
 
 def run(
     circuit: GateSequence,
-    state: Array,
+    state: State,
     values: dict[str, float] = dict(),
-) -> Array:
+) -> State:
     return apply_gate(state, circuit, values)
 
 
 def sample(
-    state: Array,
+    state: State,
     gates: GateSequence,
     values: dict[str, float] = dict(),
     n_shots: int = 1000,
 ) -> Counter:
+    """Sample from a quantum program.
+
+    Args:
+        state (State): Input state vector or density matrix.
+        gates (GateSequence): Sequence of gates.
+        values (dict[str, float], optional): _description_. Defaults to dict().
+        n_shots (int, optional): Parameter values.. Defaults to 1000.
+
+    Raises:
+        ValueError: If n_shots < 1.
+
+    Returns:
+        Counter: Bitstrings and frequencies.
+    """
     if n_shots < 1:
-        raise ValueError("You can only call sample with n_shots>0.")
-
-    wf = apply_gate(state, gates, values)
-    probs = jnp.abs(jnp.float_power(wf, 2.0)).ravel()
-    key = jax.random.PRNGKey(0)
-    n_qubits = len(state.shape)
-    # JAX handles pseudo random number generation by tracking an explicit state via a random key
-    # For more details, see https://jax.readthedocs.io/en/latest/random-numbers.html
-    samples = jax.vmap(
-        lambda subkey: jax.random.choice(key=subkey, a=jnp.arange(0, 2**n_qubits), p=probs)
-    )(jax.random.split(key, n_shots))
-
-    return Counter(
-        {
-            format(k, "0{}b".format(n_qubits)): count.item()
-            for k, count in enumerate(jnp.bincount(samples))
-            if count > 0
-        }
+        raise ValueError("You can only sample with non-negative 'n_shots'.")
+    output_circuit = apply_gate(state, gates, values)
+    n_qubits = num_qubits(output_circuit)
+    if isinstance(output_circuit, DensityMatrix):
+        d = 2**n_qubits
+        output_circuit.array = output_circuit.array.reshape((d, d))
+
+    probs = probabilities(output_circuit)
+    return sample_from_probs(probs, n_qubits, n_shots)
+
+
+@singledispatch
+def _ad_expectation_single_observable(
+    state: Any,
+    observable: Primitive,
+    values: dict[str, float],
+) -> Any:
+    raise NotImplementedError("_ad_expectation_single_observable is not implemented")
+
+
+@_ad_expectation_single_observable.register
+def _(
+    state: Array,
+    observable: Primitive,
+    values: dict[str, float],
+) -> Array:
+    projected_state = apply_gate(
+        state,
+        observable,
+        values,
+        OperationType.UNITARY,
     )
+    return inner(state, projected_state).real
 
 
-def __ad_expectation_single_observable(
-    state: Array, gates: GateSequence, observable: Primitive, values: dict[str, float]
+@_ad_expectation_single_observable.register
+def _(
+    state: DensityMatrix,
+    observable: Primitive,
+    values: dict[str, float],
 ) -> Array:
-    """
-    Run 'state' through a sequence of 'gates' given parameters 'values'
-    and compute the expectation given an observable.
-    """
-    out_state = apply_gate(state, gates, values, OperationType.UNITARY)
-    projected_state = apply_gate(out_state, observable, values, OperationType.UNITARY)
-    return inner(out_state, projected_state).real
+    n_qubits = num_qubits(state)
+    mat_obs = to_matrix(observable, n_qubits, values)
+    d = 2**n_qubits
+    prod = jnp.matmul(mat_obs, state.array.reshape((d, d)))
+    return jnp.trace(prod, axis1=-2, axis2=-1).real
 
 
 def ad_expectation(
-    state: Array, gates: GateSequence, observables: list[Primitive], values: dict[str, float]
+    state: State,
+    gates: GateSequence,
+    observables: list[Primitive],
+    values: dict[str, float],
 ) -> Array:
-    """
-    Run 'state' through a sequence of 'gates' given parameters 'values'
-    and compute the expectation given an observable.
+    """Run 'state' through a sequence of 'gates' given parameters 'values'
+       and compute the expectation given an observable.
+
+    Args:
+        state (State): Input state vector or density matrix.
+        gates (GateSequence): Sequence of gates.
+        observables (list[Primitive]): List of observables.
+        values (dict[str, float]): Parameter values.
+
+    Returns:
+        Array: Expectation values.
     """
     outputs = [
-        __ad_expectation_single_observable(state, gates, observable, values)
+        _ad_expectation_single_observable(
+            apply_gate(state, gates, values, OperationType.UNITARY), observable, values
+        )
         for observable in observables
     ]
     return jnp.stack(outputs)
 
 
 def expectation(
-    state: Array,
+    state: State,
     gates: GateSequence,
     observables: list[Primitive],
     values: dict[str, float],
@@ -87,13 +140,27 @@ def expectation(
     n_shots: Optional[int] = None,
     key: Any = jax.random.PRNGKey(0),
 ) -> Array:
-    """
-    Run 'state' through a sequence of 'gates' given parameters 'values'
+    """Run 'state' through a sequence of 'gates' given parameters 'values'
     and compute the expectation given an observable.
+
+    Args:
+        state (State): Input state vector or density matrix.
+        gates (GateSequence): Sequence of gates.
+        observables (list[Primitive]): List of observables.
+        values (dict[str, float]): Parameter values.
+        diff_mode (DiffMode, optional): Differentiation mode. Defaults to DiffMode.AD.
+        forward_mode (ForwardMode, optional): Type of forward method. Defaults to ForwardMode.EXACT.
+        n_shots (Optional[int], optional): Number of shots. Defaults to None.
+        key (Any, optional): Random key. Defaults to jax.random.PRNGKey(0).
+
+    Returns:
+        Array: Expectation values.
     """
     if diff_mode == DiffMode.AD:
         return ad_expectation(state, gates, observables, values)
     elif diff_mode == DiffMode.ADJOINT:
+        if isinstance(state, DensityMatrix):
+            raise TypeError("Adjoint does not support density matrices.")
         return adjoint_expectation(state, gates, observables, values)
     elif diff_mode == DiffMode.GPSR:
         checkify.check(
@@ -105,4 +172,11 @@ def expectation(
         )
         # Type checking is disabled because mypy doesn't parse checkify.check.
         # type: ignore
-        return finite_shots_fwd(state, gates, observables, values, n_shots=n_shots, key=key)
+        return finite_shots_fwd(
+            state=state,
+            gates=gates,
+            observables=observables,
+            values=values,
+            n_shots=n_shots,
+            key=key,
+        )