[Feature] Adding digital noise

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
 

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
 
 

docs/noise.md

@@ -0,0 +1,101 @@
+## Digital Noise
+
+In the description of closed quantum systems, a pure state vector is used to represent the complete quantum state. Thus, pure quantum states are represented by state vectors $|\psi \rangle $.
+
+However, this description is not sufficient to study open quantum systems. When the system interacts with its environment, quantum systems can be in 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 environment (noise) 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 the Kraus operators, and satisfy the 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 `NoiseInstance` providing the `NoiseType` and the `error_probability` (either float or tuple of float):
+
+```python exec="on" source="material-block" html="1"
+from horqrux.noise import NoiseInstance, NoiseType
+
+noise_prob = 0.3
+AmpD = NoiseInstance(NoiseType.AMPLITUDE_DAMPING, error_probability=noise_prob)
+
+```
+
+Then a gate can be instantiated by providing a tuple of `NoiseInstance` instances. Let’s show this through the simulation of a realistic $X$ gate.
+
+We know that an $X$ gate flips the state of the qubit, for instance $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 a `BitFlip` `NoiseInstance`, 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 NoiseInstance, NoiseType
+from horqrux.utils import density_mat, product_state
+from horqrux.primitive import X
+
+noise = (NoiseInstance(NoiseType.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("Noiseless samples", noiseless_samples) # markdown-exec: hide
+print("Noiseless samples", noisy_samples) # markdown-exec: hide
+```

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 (

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

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,128 @@
 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, observable_to_matrix
+from horqrux.utils import (
+    DensityMatrix,
+    DiffMode,
+    ForwardMode,
+    OperationType,
+    State,
+    get_probas,
+    inner,
+    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.")
+    output_circuit = apply_gate(state, gates, values)
 
-    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
-        }
-    )
+    if isinstance(output_circuit, DensityMatrix):
+        n_qubits = len(output_circuit.array.shape) // 2
+        d = 2**n_qubits
+        output_circuit.array = output_circuit.array.reshape((d, d))
+    else:
+        n_qubits = len(output_circuit.shape)
+
+    probs = get_probas(output_circuit)
+    return sample_from_probs(probs, n_qubits, n_shots)
 
 
+@singledispatch
 def __ad_expectation_single_observable(
-    state: Array, gates: GateSequence, observable: Primitive, values: dict[str, float]
+    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:
-    """
-    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
+    projected_state = apply_gate(
+        state,
+        observable,
+        values,
+        OperationType.UNITARY,
+    )
+    return inner(state, projected_state).real
+
+
+@__ad_expectation_single_observable.register
+def _(
+    state: DensityMatrix,
+    observable: Primitive,
+    values: dict[str, float],
+) -> Array:
+    n_qubits = len(state.array.shape) // 2
+    mat_obs = observable_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 +142,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 ValueError("Adjoint does not support density matrices.")
         return adjoint_expectation(state, gates, observables, values)
     elif diff_mode == DiffMode.GPSR:
         checkify.check(
@@ -105,4 +174,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,
+            gates,
+            observables,
+            values,
+            n_shots=n_shots,
+            key=key,
+        )