Fixing gradient calculation and running tox

bqskit/ir/gates/parameterized/mcry.py

@@ -10,73 +10,85 @@
 from bqskit.qis.unitary.unitary import RealVector
 from bqskit.qis.unitary.unitarymatrix import UnitaryMatrix
 from bqskit.utils.cachedclass import CachedClass
-import logging
 
-def get_indices(index: int, target_qudit, num_qudits):
-    """
-    Get indices for the matrix based on the target qubit.
-    """
+
+def get_indices(
+    index: int,
+    target_qudit: int,
+    num_qudits: int,
+) -> tuple[int, int]:
+    """Get indices for the matrix based on the target qubit."""
     shift_qubit = num_qudits - target_qudit - 1
     shift = 2 ** shift_qubit
     # Split into two parts around target qubit
     # 100 | 111
     left = index // shift
     right = index % shift
 
-    # Now, shift left by one spot to 
+    # Now, shift left by one spot to
     # make room for the target qubit
     left *= (shift * 2)
     # Now add 0 * new_ind and 1 * new_ind to get indices
     return left + right, left + shift + right
 
+
 class MCRYGate(
     QubitGate,
     DifferentiableUnitary,
     CachedClass,
-    LocallyOptimizableUnitary
+    LocallyOptimizableUnitary,
 ):
     """
     A gate representing a multiplexed Y rotation. A multiplexed Y rotation
     uses n - 1 qubits as select qubits and applies a Y rotation to the target.
     If the target qubit is the last qubit, then the unitary is block diagonal.
-    Each block is a 2x2 RY matrix with parameter theta. 
+    Each block is a 2x2 RY matrix with parameter theta.
 
     Since there are n - 1 select qubits, there are 2^(n-1) parameters (thetas).
 
     We allow the target qubit to be specified to any qubit, and the other qubits
-    maintain their order. Qubit 0 is the most significant qubit. 
+    maintain their order. Qubit 0 is the most significant qubit.
 
     See this paper: https://arxiv.org/pdf/quant-ph/0406176
     """
 
     _qasm_name = 'mcry'
 
-    def __init__(self, num_qudits: int, target_qubit: int = -1) -> None:
+    def __init__(
+        self,
+        num_qudits: int,
+        target_qubit: int = -1,
+    ) -> None:
         self._num_qudits = num_qudits
         # 1 param for each configuration of the selec qubits
         self._num_params = 2 ** (num_qudits - 1)
         # By default, the controlled qubit is the last qubit
         if target_qubit == -1:
             target_qubit = num_qudits - 1
-        self.controlled_qubit = target_qubit
+        self.target_qubit = target_qubit
         super().__init__()
 
     def get_unitary(self, params: RealVector = []) -> UnitaryMatrix:
         """Return the unitary for this gate, see :class:`Unitary` for more."""
         self.check_parameters(params)
 
-        matrix = np.zeros((2 ** self.num_qudits, 2 ** self.num_qudits), dtype=np.complex128)
+        matrix = np.zeros(
+            (
+                2 ** self.num_qudits,
+                2 ** self.num_qudits,
+            ), dtype=np.complex128,
+        )
         for i, param in enumerate(params):
             cos = np.cos(param / 2)
             sin = np.sin(param / 2)
 
             # Now, get indices based on target qubit.
-            # i corresponds to the configuration of the 
-            # select qubits (e.g 5 = 101). Now, the 
+            # i corresponds to the configuration of the
+            # select qubits (e.g 5 = 101). Now, the
             # target qubit is 0,1 for both the row and col
             # indices. So, if i = 5 and the target_qubit is 2
             # Then the rows/cols are 1001 and 1101
-            x1, x2 = get_indices(i, self.controlled_qubit, self.num_qudits)
+            x1, x2 = get_indices(i, self.target_qubit, self.num_qudits)
 
             matrix[x1, x1] = cos
             matrix[x2, x2] = cos
@@ -88,36 +100,45 @@ def get_unitary(self, params: RealVector = []) -> UnitaryMatrix:
     def get_grad(self, params: RealVector = []) -> npt.NDArray[np.complex128]:
         """
         Return the gradient for this gate.
+
         See :class:`DifferentiableUnitary` for more info.
         """
         self.check_parameters(params)
 
-        matrix = np.zeros((2 ** self.num_qudits, 2 ** self.num_qudits), dtype=np.complex128)
+        orig_utry = self.get_unitary(params).numpy
+        grad = []
+
+        # For each parameter, calculate the derivative
+        # with respect to that parameter
         for i, param in enumerate(params):
             dcos = -np.sin(param / 2) / 2
             dsin = -1j * np.cos(param / 2) / 2
 
             # Again, get indices based on target qubit.
-            x1, x2 = get_indices(i, self.controlled_qubit, self.num_qudits)
+            x1, x2 = get_indices(i, self.target_qubit, self.num_qudits)
+
+            matrix = orig_utry.copy()
 
             matrix[x1, x1] = dcos
             matrix[x2, x2] = dcos
             matrix[x2, x1] = dsin
             matrix[x1, x2] = -1 * dsin
 
-        return UnitaryMatrix(matrix)
+            grad.append(matrix)
 
+        return np.array(grad)
 
     def optimize(self, env_matrix: npt.NDArray[np.complex128]) -> list[float]:
         """
         Return the optimal parameters with respect to an environment matrix.
+
         See :class:`LocallyOptimizableUnitary` for more info.
         """
         self.check_env_matrix(env_matrix)
-        thetas = [0] * self.num_params
+        thetas: list[float] = [0] * self.num_params
 
         for i in range(self.num_params):
-            x1, x2 = get_indices(i, self.controlled_qubit, self.num_qudits)
+            x1, x2 = get_indices(i, self.target_qubit, self.num_qudits)
             a = np.real(env_matrix[x1, x1] + env_matrix[x2, x2])
             b = np.real(env_matrix[x2, x1] - env_matrix[x1, x2])
             theta = 2 * np.arccos(a / np.sqrt(a ** 2 + b ** 2))
@@ -127,12 +148,17 @@ def optimize(self, env_matrix: npt.NDArray[np.complex128]) -> list[float]:
         return thetas
 
     @staticmethod
-    def get_decomposition(params: RealVector = []) -> tuple[RealVector, RealVector]:
-        '''
-        Get the corresponding parameters for one level of decomposition
-        of a multiplexed gate. This is used in the decomposition of both
+    def get_decomposition(params: RealVector = []) -> tuple[
+        RealVector,
+        RealVector,
+    ]:
+        """
+        Get the corresponding parameters for one level of decomposition of a
+        multiplexed gate.
+
+        This is used in the decomposition of both
         the MCRY and MCRZ gates. See :class:`MGDPass` for more info.
-        '''
+        """
         new_num_params = len(params) // 2
         left_params = np.zeros(new_num_params)
         right_params = np.zeros(new_num_params)
@@ -148,4 +174,4 @@ def get_decomposition(params: RealVector = []) -> tuple[RealVector, RealVector]:
     def name(self) -> str:
         """The name of this gate, with the number of qudits appended."""
         base_name = getattr(self, '_name', self.__class__.__name__)
-        return f"{base_name}_{self.num_qudits}"
+        return f'{base_name}_{self.num_qudits}'

bqskit/ir/gates/parameterized/mcrz.py

@@ -4,31 +4,31 @@
 import numpy as np
 import numpy.typing as npt
 
+from bqskit.ir.gates.parameterized.mcry import get_indices
 from bqskit.ir.gates.qubitgate import QubitGate
 from bqskit.qis.unitary.differentiable import DifferentiableUnitary
 from bqskit.qis.unitary.optimizable import LocallyOptimizableUnitary
 from bqskit.qis.unitary.unitary import RealVector
 from bqskit.qis.unitary.unitarymatrix import UnitaryMatrix
 from bqskit.utils.cachedclass import CachedClass
-from bqskit.ir.gates.parameterized.mcry import get_indices
-from typing import Any
+
 
 class MCRZGate(
     QubitGate,
     DifferentiableUnitary,
     CachedClass,
-    LocallyOptimizableUnitary
+    LocallyOptimizableUnitary,
 ):
     """
     A gate representing a multiplexed Z rotation. A multiplexed Z rotation
     uses n - 1 qubits as select qubits and applies a Z rotation to the target.
     If the target qubit is the last qubit, then the unitary is block diagonal.
-    Each block is a 2x2 RZ matrix with parameter theta. 
+    Each block is a 2x2 RZ matrix with parameter theta.
 
     Since there are n - 1 select qubits, there are 2^(n-1) parameters (thetas).
 
     We allow the target qubit to be specified to any qubit, and the other qubits
-    maintain their order. Qubit 0 is the most significant qubit. 
+    maintain their order. Qubit 0 is the most significant qubit.
 
 
     Why is 0 the MSB? Typically, in the QSD diagram, we see the block drawn
@@ -40,11 +40,14 @@ class MCRZGate(
 
     _qasm_name = 'mcrz'
 
-    def __init__(self, num_qudits: int, target_qubit: int = -1) -> None:
-        '''
-        Create a new MCRZGate with `num_qudits` qubits and 
-        `target_qubit` as the target qubit. We then have 2^(n-1) parameters
-        for this gate.
+    def __init__(
+        self,
+        num_qudits: int,
+        target_qubit: int = -1,
+    ) -> None:
+        """
+        Create a new MCRZGate with `num_qudits` qubits and `target_qubit` as the
+        target qubit. We then have 2^(n-1) parameters for this gate.
 
         For Example:
         `num_qudits` = 3, `target_qubit` = 1
@@ -56,7 +59,7 @@ def __init__(self, num_qudits: int, target_qubit: int = -1) -> None:
 
         If the input vector is |1x0> then the selection is 01, and
         RZ(theta_1) is applied to the target qubit.
-        '''
+        """
         self._num_qudits = num_qudits
         # 1 param for each configuration of the selec qubits
         self._num_params = 2 ** (num_qudits - 1)
@@ -69,7 +72,12 @@ def __init__(self, num_qudits: int, target_qubit: int = -1) -> None:
     def get_unitary(self, params: RealVector = []) -> UnitaryMatrix:
         """Return the unitary for this gate, see :class:`Unitary` for more."""
         self.check_parameters(params)
-        matrix = np.zeros((2 ** self.num_qudits, 2 ** self.num_qudits), dtype=np.complex128)
+        matrix = np.zeros(
+            (
+                2 ** self.num_qudits,
+                2 ** self.num_qudits,
+            ), dtype=np.complex128,
+        )
         for i, param in enumerate(params):
             pos = np.exp(1j * param / 2)
             neg = np.exp(-1j * param / 2)
@@ -90,27 +98,37 @@ def get_grad(self, params: RealVector = []) -> npt.NDArray[np.complex128]:
         See :class:`DifferentiableUnitary` for more info.
         """
         self.check_parameters(params)
-        matrix = np.zeros((2 ** self.num_qudits, 2 ** self.num_qudits), dtype=np.complex128)
+        orig_utry = self.get_unitary(params).numpy
+        grad = []
+
+        # For each parameter, calculate the derivative
+        # with respect to that parameter
         for i, param in enumerate(params):
-            dpos = 1j / 2 * np.exp(1j * param / 2)
-            dneg = -1j / 2 * np.exp(-1j * param / 2)
+            dcos = -np.sin(param / 2) / 2
+            dsin = -1j * np.cos(param / 2) / 2
 
             # Again, get indices based on target qubit.
             x1, x2 = get_indices(i, self.target_qubit, self.num_qudits)
 
-            matrix[x1, x1] = dpos
-            matrix[x2, x2] = dneg
+            matrix = orig_utry.copy()
 
-        return UnitaryMatrix(matrix)
+            matrix[x1, x1] = dcos
+            matrix[x2, x2] = dcos
+            matrix[x2, x1] = dsin
+            matrix[x1, x2] = -1 * dsin
 
+            grad.append(matrix)
+
+        return np.array(grad)
 
     def optimize(self, env_matrix: npt.NDArray[np.complex128]) -> list[float]:
         """
         Return the optimal parameters with respect to an environment matrix.
+
         See :class:`LocallyOptimizableUnitary` for more info.
         """
         self.check_env_matrix(env_matrix)
-        thetas = [0] * self.num_params
+        thetas: list[float] = [0] * self.num_params
 
         for i in range(self.num_params):
             x1, x2 = get_indices(i, self.target_qubit, self.num_qudits)
@@ -127,4 +145,4 @@ def optimize(self, env_matrix: npt.NDArray[np.complex128]) -> list[float]:
     def name(self) -> str:
         """The name of this gate, with the number of qudits appended."""
         base_name = getattr(self, '_name', self.__class__.__name__)
-        return f"{base_name}_{self.num_qudits}"
+        return f'{base_name}_{self.num_qudits}'

testing_mcr_gates.py

@@ -0,0 +1,80 @@
+import numpy as np
+from scipy.linalg import block_diag
+
+from bqskit.ir.gates.parameterized import MCRYGate, RYGate, MCRZGate, RZGate
+from bqskit.ir.gates.constant import PermutationGate
+
+def test_get_unitary_mcry(thetas: list[float]) -> None:
+    '''
+    Test the get_unitary method of the MCRYGate class.
+    Use the default target qubit.
+    '''
+    # Ensure that len(thetas) is a power of 2
+    # There are 2 ** (n - 1) parameters
+    num_qudits = int(np.log2(len(thetas))) + 1
+    thetas = thetas[:2 ** (num_qudits - 1)]
+
+    mcry = MCRYGate(num_qudits=num_qudits)
+    block_unitaries = [RYGate().get_unitary([theta]) for theta in thetas]
+    blocked_unitary = block_diag(*block_unitaries)
+    dist = mcry.get_unitary(thetas).get_distance_from(blocked_unitary)
+    assert dist < 1e-7
+
+def test_get_unitary_mcrz(thetas: list[float]) -> None:
+    '''
+    Test the get_unitary method of the MCRYGate class.
+    Use the default target qubit.
+    '''
+    # Ensure that len(thetas) is a power of 2
+    # There are 2 ** (n - 1) parameters
+    num_qudits = int(np.log2(len(thetas))) + 1
+    thetas = thetas[:2 ** (num_qudits - 1)]
+
+    mcry = MCRZGate(num_qudits=num_qudits)
+    block_unitaries = [RZGate().get_unitary([theta]) for theta in thetas]
+    blocked_unitary = block_diag(*block_unitaries)
+    dist = mcry.get_unitary(thetas).get_distance_from(blocked_unitary)
+    assert dist < 1e-7
+
+def test_get_unitary_target_select_mcry(target_qubit: int) -> None:
+    '''
+    Test the get_unitary method of the MCRYGate class when
+    the target qubit is set.
+    '''
+    # Create an MCRY gate with 6 qubits and random parameters
+    num_qudits = 6
+    mcry = MCRYGate(num_qudits=num_qudits, target_qubit=target_qubit)
+    thetas = list(np.random.rand(2 ** (num_qudits - 1)) * 2 * np.pi)
+
+    # Create the block diagonal matrix
+    block_unitaries = [RYGate().get_unitary([theta]) for theta in thetas]
+    blocked_unitary = block_diag(*block_unitaries)
+
+    # Apply a permutation transformation 
+    # to the block diagonal matrix
+    # Swap the target qubit with the last qubit
+    # perm = np.arange(num_qudits)
+    perm = list(range(num_qudits))
+    for i in range(target_qubit, num_qudits):
+        perm[i] = i + 1
+    perm[-1] = target_qubit
+
+    perm_gate = PermutationGate(num_qudits, perm) 
+
+    full_utry = perm_gate.get_unitary().conj().T @ blocked_unitary @ perm_gate.get_unitary()
+
+    dist = mcry.get_unitary(thetas).get_distance_from(full_utry)
+    assert dist < 1e-7
+
+
+
+
+for num_params in [2,4,8, 20]:
+    params = np.random.rand(num_params) * 2 * np.pi
+    test_get_unitary_mcry(params)
+    test_get_unitary_mcrz(params)
+
+np.printoptions(precision=3, threshold=np.inf, linewidth=np.inf)
+
+for target_qubit in [0,1,2,3,4,5]:
+    test_get_unitary_target_select_mcry(target_qubit)