From 4e3fb3fff98112d4bd78d4d60029ea9bf280b7e5 Mon Sep 17 00:00:00 2001
From: Roland Guichard <roland.guichard@pasqal.com>
Date: Wed, 22 Nov 2023 14:34:36 +0000
Subject: [PATCH] Add single qubit projector.

---
 pyqtorch/matrices.py  | 4 ++++
 pyqtorch/primitive.py | 7 +++++++
 2 files changed, 11 insertions(+)

diff --git a/pyqtorch/matrices.py b/pyqtorch/matrices.py
index c4a8c1ae..1e973d8a 100644
--- a/pyqtorch/matrices.py
+++ b/pyqtorch/matrices.py
@@ -15,6 +15,8 @@
     dtype=DEFAULT_MATRIX_DTYPE,
 )
 NMAT = torch.tensor([[0, 0], [0, 1]], dtype=DEFAULT_MATRIX_DTYPE)
+PROJ0MAT = torch.tensor([[1, 0], [0, 0]], dtype=DEFAULT_MATRIX_DTYPE)
+PROJ1MAT = torch.tensor([[0, 0], [0, 1]], dtype=DEFAULT_MATRIX_DTYPE)
 NDIAG = torch.tensor([0, 1], dtype=DEFAULT_MATRIX_DTYPE)
 ZDIAG = torch.tensor([1, -1], dtype=DEFAULT_MATRIX_DTYPE)
 IDIAG = torch.tensor([1, 1], dtype=DEFAULT_MATRIX_DTYPE)
@@ -50,6 +52,8 @@
     "SDAGGER": SDAGGERMAT,
     "T": TMAT,
     "N": NMAT,
+    "PROJ0": PROJ0MAT,
+    "PROJ1": PROJ1MAT,
     "H": HMAT,
     "SWAP": SWAPMAT,
     "CSWAP": CSWAPMAT,
diff --git a/pyqtorch/primitive.py b/pyqtorch/primitive.py
index 0530340d..e1150d6e 100644
--- a/pyqtorch/primitive.py
+++ b/pyqtorch/primitive.py
@@ -76,6 +76,13 @@ def __init__(self, target: int):
         super().__init__(OPERATIONS_DICT["SDAGGER"], target)
 
 
+class Proj(Primitive):
+    def __init__(self, target: int, state: str = "1"):
+        if state == "0":
+            super().__init__(OPERATIONS_DICT["PROJ0"], target)
+        super()._init__(OPERATIONS_DICT["PROJ1"], target)
+
+
 class N(Primitive):
     def __init__(self, target: int):
         super().__init__(OPERATIONS_DICT["N"], target)