From 8eea9ef551228ca84cb13e024fa90340975a2b64 Mon Sep 17 00:00:00 2001
From: Astral Cai <astralcai@gmail.com>
Date: Tue, 16 Apr 2024 10:31:51 -0400
Subject: [PATCH] Add Prod observable support (#64)

* Add Prod observable support

* update workflow file to include pytest-benchmark
---
 .github/workflows/tests.yml       | 1 +
 pennylane_qulacs/qulacs_device.py | 3 +++
 tests/conftest.py                 | 2 +-
 tests/test_apply.py               | 4 +++-
 4 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 38fa070..400fadb 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -66,6 +66,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install -r requirements.txt
+          pip install pytest-benchmark
           pip install wheel pytest pytest-cov pytest-mock flaky --upgrade
 
       - name: Install Plugin
diff --git a/pennylane_qulacs/qulacs_device.py b/pennylane_qulacs/qulacs_device.py
index 88628ca..e8b9031 100644
--- a/pennylane_qulacs/qulacs_device.py
+++ b/pennylane_qulacs/qulacs_device.py
@@ -122,6 +122,7 @@ class QulacsDevice(QubitDevice):
         "Identity": "I",
         "Hadamard": None,
         "Hermitian": None,
+        "Prod": None,
     }
 
     operations = _operation_map.keys()
@@ -348,6 +349,8 @@ def expval(self, observable, **kwargs):
             qulacs_observable = Observable(self.num_wires)
             if isinstance(observable.name, list):
                 observables = [self._observable_map[obs] for obs in observable.name]
+            elif observable.name == "Prod":
+                observables = [self._observable_map[obs.name] for obs in observable.operands]
             else:
                 observables = [self._observable_map[observable.name]]
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 31b49c7..3c62c00 100755
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -59,7 +59,7 @@ def init_state():
     """Fixture to create an n-qubit initial state"""
 
     def _init_state(n):
-        state = np.random.random([2 ** n]) + np.random.random([2 ** n]) * 1j
+        state = np.random.random([2**n]) + np.random.random([2**n]) * 1j
         state /= np.linalg.norm(state)
         return state
 
diff --git a/tests/test_apply.py b/tests/test_apply.py
index 3b24e3a..1b9b02a 100755
--- a/tests/test_apply.py
+++ b/tests/test_apply.py
@@ -180,7 +180,9 @@ def test_qubit_state_vector(self, init_state, state_prep_op, tol):
     @pytest.mark.parametrize("state_prep_op", (qml.QubitStateVector, qml.StatePrep))
     @pytest.mark.parametrize("device_wires", [3, 4, 5])
     @pytest.mark.parametrize("op_wires", [[0], [2], [0, 1], [1, 0], [2, 0]])
-    def test_qubit_state_vector_on_wires_subset(self, init_state, device_wires, op_wires, state_prep_op, tol):
+    def test_qubit_state_vector_on_wires_subset(
+        self, init_state, device_wires, op_wires, state_prep_op, tol
+    ):
         """Test QubitStateVector and StatePrep application on a subset of device wires"""
         dev = QulacsDevice(device_wires)
         state = init_state(len(op_wires))