Repository: unitaryfoundation/mitiq

Issue:
Add docstrings to Observable class
## Issue Description

The [`Observable class`](https://github.com/unitaryfund/mitiq/blob/main/mitiq/observable/observable.py#L19) lacks docstrings for most of its methods. While some might be considered self-explanatory (depending on a person's background), many are not, such as [`Observable.partition()`](https://github.com/unitaryfund/mitiq/blob/main/mitiq/observable/observable.py#L99). Even short `@property` functions could still benefit from having fully written out explanations of what the property is. 

## Proposed Solution

Add descriptive docstrings for each method.

Hints from issue comments:
@natestemen: To close this issue we should document the following methods.

1. `measure_in`
2. The arguments of `matrix` should be flushed out using the `Args: ...` pattern.
3. `expectation`
4. `partition`
5. `from_pauli_string_collections`

Documenting these will also ensure they show up in the API-doc!@Shivansh20128: Hi! I have made a PR for this issue. You can assign it to me.
Thank you

Solution patch:
diff --git a/mitiq/observable/observable.py b/mitiq/observable/observable.py
index ce7d978b..f5a6fcfa 100644
--- a/mitiq/observable/observable.py
+++ b/mitiq/observable/observable.py
@@ -34,6 +34,16 @@ class Observable:
     def from_pauli_string_collections(
         *pauli_string_collections: PauliStringCollection,
     ) -> "Observable":
+        """Creates an ``Observable`` from one or more ``PauliStringCollection``
+            instances.
+
+        Args:
+            pauli_string_collections: One or more collections of Pauli strings
+                used to define the observable.
+
+        Returns:
+            An ``Observable`` containing the given Pauli string collections.
+        """
         obs = Observable()
         obs._groups = list(pauli_string_collections)
         obs._ngroups = len(pauli_string_collections)
@@ -93,6 +103,20 @@ class Observable:
         return self._ngroups
 
     def partition(self, seed: Optional[int] = None) -> None:
+        """Partitions the observable's Pauli strings into commuting groups.
+
+        This method groups the ``PauliStringCollection`` instances such that
+        each group consists of mutually commuting operators, which can be
+        measured together in a quantum circuit.
+
+        Note:
+            This method randomizes the way in which the list of
+            paulis is partitioned.
+
+        Args:
+            seed: An optional seed for shuffling to ensure deterministic
+                behavior when partitioning.
+        """
         rng = np.random.RandomState(seed)
 
         psets: List[PauliStringCollection] = []
@@ -115,13 +139,34 @@ class Observable:
         self._ngroups = len(self._groups)
 
     def measure_in(self, circuit: QPROGRAM) -> List[QPROGRAM]:
+        """Given a quantum circuit, this method returns a list of circuits
+        where each circuit corresponds to a different group of commuting Pauli
+        strings, which allows measurement in the appropriate basis.
+
+        Args:
+            circuit: The quantum circuit in which the observable should be
+            measured.
+
+        Returns:
+            A list of quantum circuits with the appropriate measurement
+            settings for each group of Pauli strings.
+        """
         return [pset.measure_in(circuit) for pset in self._groups]
 
     def matrix(
         self,
         qubit_indices: Optional[List[int]] = None,
     ) -> npt.NDArray[np.complex64]:
-        """Returns the (potentially very large) matrix of the Observable."""
+        """Returns the (potentially very large) matrix of the ``Observable``.
+
+        Args:
+            qubit_indices: Optional list of qubit indices specifying the order
+            of qubits in the matrix representation. If None, the default
+            ordering from `self.qubit_indices` is used.
+
+        Returns:
+            A ``NumPy`` array representing the matrix form of the observable.
+        """
         if qubit_indices is None:
             qubit_indices = self.qubit_indices
         n = len(qubit_indices)
@@ -135,6 +180,19 @@ class Observable:
     def expectation(
         self, circuit: QPROGRAM, execute: Callable[[QPROGRAM], QuantumResult]
     ) -> complex:
+        """Computes the expectation value of the observable.
+
+        This method executes the given quantum circuit and estimates the
+        expectation value of the observable based on the measurement results.
+
+        Args:
+            circuit: The quantum circuit to be executed.
+            execute: A function that takes a quantum circuit as input and
+            returns a QuantumResult containing measurement outcomes.
+
+        Returns:
+            The expectation value of the observable.
+        """
         from mitiq.executor import Executor  # Avoid circular import.
 
         return Executor(execute).evaluate(circuit, observable=self)[0]
diff --git a/mitiq/vd/vd_utils.py b/mitiq/vd/vd_utils.py
index 68de2103..4a7cf5c5 100644
--- a/mitiq/vd/vd_utils.py
+++ b/mitiq/vd/vd_utils.py
@@ -1,23 +1,19 @@
 from typing import List, cast
 
 import cirq
+import numpy as np
 
 
 def _copy_circuit_parallel(
     circuit: cirq.Circuit, num_copies: int = 2
 ) -> cirq.Circuit:
-    """Copies a circuit num_copies times in parallel.
-
-    Given a circuit that acts on N qubits,
-    this function returns a circuit
-    that copies the circuit num_copies times in parallel.
-    This means the resulting circuit has N * num_copies qubits.
+    """Given a circuit that acts on N qubits, this function returns a circuit
+    that copies the circuit num_copies times in parallel. This means the
+    resulting circuit has N * num_copies qubits.
 
     Args:
-        circuit:
-            The circuit to be copied.
-        num_copies:
-            The number of copies of circuit to be made.
+        circuit: The circuit to be copied.
+        num_copies: The number of copies of circuit to be made.
 
     Returns:
         A cirq circuit that is the parallel composition of
@@ -51,3 +47,64 @@ def _copy_circuit_parallel(
             new_circuit += circuit.transform_qubits(map_for_grid_qubits)
 
     return new_circuit
+
+
+def _apply_diagonalizing_gate(
+    circuit: cirq.Circuit, num_copies: int
+) -> cirq.Circuit:
+    """Apply the diagonalizing gate to a circuit, assuming the circuit has been
+    copied in parallel ``num_copies`` times.
+
+    The diagonalizing gate is applied to the first qubit of all copies, to the
+    second qubit of all copies, etc.
+
+    Args:
+        circuit: The circuit to apply the diagonalizing gate to.
+        num_copies: The number of copies of the original circuit that the input
+            circuit consists of.
+
+    Returns:
+        The circuit with the diagonalizing gate applied.
+    """
+
+    new_circuit = circuit.copy()
+
+    # this is the number of qubits in the original circuit
+    N = len(circuit.all_qubits()) // num_copies
+
+    diag_gate = _generate_diagonalizing_gate(num_copies)
+
+    for i in range(N):
+        qubits = [
+            cirq.LineQubit(i + N * j) for j in range(num_copies)
+        ]  # select qubit i of each copy
+
+        new_circuit.append(diag_gate(*qubits))
+
+    return new_circuit
+
+
+def _generate_diagonalizing_gate(num_copies: int = 2) -> cirq.Gate:
+    """Generate the diagonalizing gate for the VD algorithm. Currently only
+    ``num_copies=2`` is supported.
+
+    Args:
+        num_qubits: The number of qubits that the gate acts on.
+
+    Returns: The diagonalizing gate.
+    """
+    if num_copies == 2:
+        diagonalizing_matrix = np.array(
+            [
+                [1, 0, 0, 0],
+                [0, np.sqrt(2) / 2, np.sqrt(2) / 2, 0],
+                [0, np.sqrt(2) / 2, -np.sqrt(2) / 2, 0],
+                [0, 0, 0, 1],
+            ]
+        )
+    else:
+        raise NotImplementedError(
+            "Only num_copies = 2 is currently supported."
+        )
+
+    return cirq.MatrixGate(diagonalizing_matrix)
diff --git a/requirements/requirements-braket.txt b/requirements/requirements-braket.txt
index d89b4827..c97ed667 100644
--- a/requirements/requirements-braket.txt
+++ b/requirements/requirements-braket.txt
@@ -1,2 +1,2 @@
-amazon-braket-sdk~=1.90.2
+amazon-braket-sdk~=1.91.0
 cirq-ionq>=1.4.0,<1.5.0
diff --git a/requirements/requirements-qiskit.txt b/requirements/requirements-qiskit.txt
index 7f043474..3f3ff41b 100644
--- a/requirements/requirements-qiskit.txt
+++ b/requirements/requirements-qiskit.txt
@@ -1,4 +1,4 @@
 qiskit~=1.4.2
-qiskit-aer~=0.15.1
-qiskit-ibm-runtime~=0.36.1
+qiskit-aer~=0.17.0
+qiskit-ibm-runtime~=0.37.0
 ply==3.11
\ No newline at end of file


Test patch:
diff --git a/mitiq/vd/tests/test_vd_utils.py b/mitiq/vd/tests/test_vd_utils.py
index 29141bf8..544900c9 100644
--- a/mitiq/vd/tests/test_vd_utils.py
+++ b/mitiq/vd/tests/test_vd_utils.py
@@ -1,7 +1,12 @@
 import cirq
 import numpy as np
+import pytest
 
-from mitiq.vd.vd_utils import _copy_circuit_parallel
+from mitiq.vd.vd_utils import (
+    _apply_diagonalizing_gate,
+    _copy_circuit_parallel,
+    _generate_diagonalizing_gate,
+)
 
 
 def test_copy_circuit_parallel_lengths():
@@ -88,3 +93,59 @@ def test_copy_circuit_parallel_gridqubits():
     new_unitary = cirq.unitary(new_circuit)
     expected_unitary = cirq.unitary(expected_circuit)
     assert np.allclose(new_unitary, expected_unitary)
+
+
+def test_apply_diagonalizing_gate_keeps_circuit_structure():
+    num_copies = 2
+    qubits = cirq.LineQubit.range(3)
+    circuit = cirq.Circuit(
+        cirq.H(qubits[0]),
+        cirq.CNOT(qubits[0], qubits[1]),
+        cirq.X(qubits[2]),
+    )
+    num_qubits = len(circuit.all_qubits())
+
+    copied_circuit = _copy_circuit_parallel(circuit, num_copies)
+    new_circuit = _apply_diagonalizing_gate(copied_circuit, num_copies)
+
+    num_ops = len(list(copied_circuit.all_operations()))
+    copied_circuit_ops = list(new_circuit.all_operations())
+
+    assert len(copied_circuit_ops) == num_ops + num_qubits
+    assert set(new_circuit.all_qubits()) == set(copied_circuit.all_qubits())
+
+
+@pytest.mark.xfail(reason="VD does not yet support grid qubits")
+def test_apply_diagonalizing_gate_fails_on_grid_qubits():
+    num_copies = 2
+    qubits = cirq.GridQubit.rect(2, 2)
+    circuit = cirq.Circuit(
+        cirq.H(qubits[0]),
+        cirq.CNOT(qubits[0], qubits[1]),
+        cirq.X(qubits[2]),
+    )
+    num_qubits = len(circuit.all_qubits())
+
+    copied_circuit = _copy_circuit_parallel(circuit, num_copies)
+    new_circuit = _apply_diagonalizing_gate(copied_circuit, num_copies)
+
+    num_ops = len(list(copied_circuit.all_operations()))
+    copied_circuit_ops = list(new_circuit.all_operations())
+
+    assert len(copied_circuit_ops) == num_ops + num_qubits  # this passes
+    assert set(new_circuit.all_qubits()) == set(copied_circuit.all_qubits())
+
+
+def test_generate_diagonalizing_gate():
+    expected_matrix = np.array(
+        [
+            [1, 0, 0, 0],
+            [0, np.sqrt(2) / 2, np.sqrt(2) / 2, 0],
+            [0, np.sqrt(2) / 2, -np.sqrt(2) / 2, 0],
+            [0, 0, 0, 1],
+        ]
+    )
+
+    assert np.allclose(
+        _generate_diagonalizing_gate(2)._matrix, expected_matrix
+    )