From ba29e02910ee14ea184b5ac945962551becbb63e Mon Sep 17 00:00:00 2001
From: chMoussa <charles.moussa@pasqal.com>
Date: Tue, 26 Nov 2024 16:02:17 +0100
Subject: [PATCH] [Feature] Add quantum dropout  (#623)

---
 qadence/backends/pyqtorch/backend.py   | 19 ++++++++++++++-----
 qadence/backends/pyqtorch/config.py    |  9 ++++++---
 qadence/types.py                       |  3 ++-
 tests/backends/pyq/test_quantum_pyq.py | 13 +++++++++++++
 4 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/qadence/backends/pyqtorch/backend.py b/qadence/backends/pyqtorch/backend.py
index 529c2c6a..6924e2a2 100644
--- a/qadence/backends/pyqtorch/backend.py
+++ b/qadence/backends/pyqtorch/backend.py
@@ -124,11 +124,20 @@ def circuit(self, circuit: QuantumCircuit) -> ConvertedCircuit:
             if self.config.noise
             else None
         )
-        native = pyq.QuantumCircuit(
-            circuit.n_qubits,
-            ops,
-            readout_noise,
-        )
+        if self.config.dropout_probability == 0:
+            native = pyq.QuantumCircuit(
+                circuit.n_qubits,
+                ops,
+                readout_noise,
+            )
+        else:
+            native = pyq.DropoutQuantumCircuit(
+                circuit.n_qubits,
+                ops,
+                readout_noise,
+                dropout_prob=self.config.dropout_probability,
+                dropout_mode=self.config.dropout_mode,
+            )
         return ConvertedCircuit(native=native, abstract=circuit, original=original_circ)
 
     def observable(self, observable: AbstractBlock, n_qubits: int) -> ConvertedObservable:
diff --git a/qadence/backends/pyqtorch/config.py b/qadence/backends/pyqtorch/config.py
index c7d18160..08958ccd 100644
--- a/qadence/backends/pyqtorch/config.py
+++ b/qadence/backends/pyqtorch/config.py
@@ -4,8 +4,6 @@
 from logging import getLogger
 from typing import Callable
 
-from pyqtorch.utils import SolverType
-
 from qadence.analog import add_background_hamiltonian
 from qadence.backend import BackendConfiguration
 from qadence.noise import NoiseHandler
@@ -15,7 +13,7 @@
     flatten,
     scale_primitive_blocks_only,
 )
-from qadence.types import AlgoHEvo
+from qadence.types import AlgoHEvo, DropoutMode, SolverType
 
 logger = getLogger(__name__)
 
@@ -67,3 +65,8 @@ class Configuration(BackendConfiguration):
 
     noise: NoiseHandler | None = None
     """NoiseHandler containing readout noise applied in backend."""
+
+    dropout_probability: float = 0.0
+    """Quantum dropout probability (0 means no dropout)."""
+    dropout_mode: DropoutMode = DropoutMode.ROTATIONAL
+    """Type of quantum dropout to perform."""
diff --git a/qadence/types.py b/qadence/types.py
index 80d08d2d..9bc21a9e 100644
--- a/qadence/types.py
+++ b/qadence/types.py
@@ -11,7 +11,7 @@
 from numpy.typing import ArrayLike
 from pyqtorch.noise import NoiseType as DigitalNoise
 from pyqtorch.noise.readout import WhiteNoise
-from pyqtorch.utils import SolverType
+from pyqtorch.utils import DropoutMode, SolverType
 from torch import Tensor, pi
 from torch.nn import Module
 
@@ -55,6 +55,7 @@
     "SerializationFormat",
     "PI",
     "SolverType",
+    "DropoutMode",
     "NoiseProtocol",
     "WhiteNoise",
 ]  # type: ignore
diff --git a/tests/backends/pyq/test_quantum_pyq.py b/tests/backends/pyq/test_quantum_pyq.py
index aa4332af..48e42aa0 100644
--- a/tests/backends/pyq/test_quantum_pyq.py
+++ b/tests/backends/pyq/test_quantum_pyq.py
@@ -12,6 +12,7 @@
 from hypothesis import given, settings
 from pyqtorch import U as pyqU
 from pyqtorch import zero_state as pyq_zero_state
+from pyqtorch.circuit import DropoutQuantumCircuit as PyQDropoutQuantumCircuit
 from pyqtorch.circuit import QuantumCircuit as PyQQuantumCircuit
 from sympy import acos
 from torch import Tensor
@@ -842,3 +843,15 @@ def test_move_to_dtype(
     assert wf.dtype == dtype
     expval = qm.expectation(inputs, state=state)
     assert expval.dtype == torch.float64 if dtype == torch.cdouble else torch.float32
+
+
+@given(st.batched_digital_circuits())
+@settings(deadline=None)
+def test_dropout_qc(circuit_and_inputs: tuple[QuantumCircuit, dict[str, torch.Tensor]]) -> None:
+    circuit, inputs = circuit_and_inputs
+    config = PyqConfig(dropout_probability=0.1)
+    bknd_pyqtorch = backend_factory(
+        backend=BackendName.PYQTORCH, diff_mode=DiffMode.AD, configuration=config
+    )
+    (circ_pyqtorch, _, _, _) = bknd_pyqtorch.convert(circuit)
+    assert isinstance(circ_pyqtorch.native, PyQDropoutQuantumCircuit)