diff --git a/unitary/alpha/__init__.py b/unitary/alpha/__init__.py index 24df674a..93bdf1fb 100644 --- a/unitary/alpha/__init__.py +++ b/unitary/alpha/__init__.py @@ -13,6 +13,13 @@ # limitations under the License. # +from unitary.alpha.qudit_state_transform import ( + qubit_to_qudit_state, + qubit_to_qudit_unitary, + qudit_to_qubit_state, + qudit_to_qubit_unitary, +) + from unitary.alpha.quantum_world import ( QuantumWorld, ) diff --git a/unitary/alpha/qudit_state_transform.py b/unitary/alpha/qudit_state_transform.py new file mode 100644 index 00000000..f9a8d522 --- /dev/null +++ b/unitary/alpha/qudit_state_transform.py @@ -0,0 +1,213 @@ +# Copyright 2022 Google +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import itertools + +import numpy as np + + +def _nearest_power_of_two_ceiling(qudit_dim: int) -> int: + """Returns the smallest power of two greater than or equal to qudit_dim.""" + if qudit_dim == 0: + return 0 + result = 1 + while result < qudit_dim: + result = result << 1 + return result + + +def qudit_to_qubit_state( + qudit_dimension: int, + num_qudits: int, + qudit_state_vector: np.ndarray, + _pad_value: np.complex_ = 0, +) -> np.ndarray: + """Converts a qudit-space quantum state vector to m-qubit-per-qudit column vector. + + Each qudit is replaced by a set of qubits. Since the set of qubits can represent a larger state + space than the qudit, the state vector needs to be padded with 0s for those extra elements. + Based on https://drive.google.com/file/d/1n1Ym7JdnM44NnvNQDUXSk5rBLTQZzJ9t and + https://colab.research.google.com/drive/1MC6D4FyOXG0RjvzyV9IFcdRsLbUwwAcx + + Args: + qudit_dimension: The dimension of a single qudit i.e. the number of states it can + represent. + num_qudits: The number of qudits in the given state vector. + qudit_state_vector: A numpy array representing the state vector in qudit form. + Expected shape: `(qudit_dimension ^ num_qudits,)`. + _pad_value: The value to set for the elements in the extra state space created in the + conversion. This field is mostly private, mostly for use by other methods. Do not set + this unless you really need to. + + Returns: + A flat numpy array representing the input state vector using m-qubits-per-qudit. + Expected shape: `((2 ^ m) ^ num_qudits,)`. + """ + # Reshape the state vector to a `num_qudits` rank tensor. + state_tensor = qudit_state_vector.reshape((qudit_dimension,) * num_qudits) + # Number of extra elements needed in each dimension if represented using qubits. + padding_amount = _nearest_power_of_two_ceiling(qudit_dimension) - qudit_dimension + # Expand the number of elements in each dimension by the padding_amount. Fill + # the new elements with the _pad_value. + padded_state_tensor = np.pad( + state_tensor, pad_width=(0, padding_amount), constant_values=_pad_value + ) + # Return a flattened state vector view of the final tensor. + return np.ravel(padded_state_tensor) + + +def qubit_to_qudit_state( + qudit_dimension: int, + num_qudits: int, + qubit_state_vector: np.ndarray, +) -> np.ndarray: + """Converts a m-qubit-per-qudit column vector to a qudit-space quantum state vector. + + Each qudit was replaced by a set of qubits. Since the set of qubits could represent a larger + state space than the qudit, the state vector needs to be sliced up to the qudit length in each + dimension. + Based on https://drive.google.com/file/d/1n1Ym7JdnM44NnvNQDUXSk5rBLTQZzJ9t and + https://colab.research.google.com/drive/1MC6D4FyOXG0RjvzyV9IFcdRsLbUwwAcx + + Args: + qudit_dimension: The dimension of a single qudit i.e. the number of states it can + represent. + num_qudits: The number of qudits in the given/output state vector. + qubit_state_vector: A numpy array representing the state vector in an + m-qubit-per-qudit form. Expected shape: `((2 ^ m) ^ num_qudits,)`. + + Returns: + A flat numpy array representing the input state vector using qudits. + Expected shape: `(qudit_dimension ^ num_qudits,)`. + """ + mbit_dimension = _nearest_power_of_two_ceiling(qudit_dimension) + # Reshape the state vector to a `num_qudits` rank tensor. + state_tensor = qubit_state_vector.reshape((mbit_dimension,) * num_qudits) + # Shrink the number of elements in each dimension up to the qudit_dimension, ignoring the rest. + trimmed_state_tensor = state_tensor[(slice(qudit_dimension),) * num_qudits] + # Return a flattened state vector view of the final tensor. + return np.ravel(trimmed_state_tensor) + + +def qudit_to_qubit_unitary( + qudit_dimension: int, + num_qudits: int, + qudit_unitary: np.ndarray, + memoize: bool = False, +) -> np.ndarray: + """Converts a qudit-space quantum unitary to m-qubit-per-qudit unitary. + + Each qudit is replaced by a set of qubits. Since the set of qubits can represent a larger state + space than the qudit, the unitary needs to be padded with 0s for those extra elements. A + unitary is treated similar to a 2*num_qudits system's state vector and padded using the state + vector protocol. The resulting unitary is updated to have the extra dimensions map to + themselves (identity) to preserve unitarity. + Based on https://drive.google.com/file/d/1n1Ym7JdnM44NnvNQDUXSk5rBLTQZzJ9t and + https://colab.research.google.com/drive/1MC6D4FyOXG0RjvzyV9IFcdRsLbUwwAcx + + Args: + qudit_dimension: The dimension of a single qudit i.e. the number of states it can + represent. + num_qudits: The number of qudits in the given unitary. + qudit_unitary: A 2-D numpy array representing the unitary in qudit form. + Expected shape: `(qudit_dimension ^ num_qudits, qudit_dimension ^ num_qudits)`. + memoize: Currently, this method has two independent implementations. If memoize is True, an + alternate implementation than above is used. A special state vector is passed to the + state vector protocol to get a mapping from qudit state indices to qubit state indices. + This mapping is then iteratively applied to the input unitary's elements. + + Returns: + A numpy array representing the input unitary using m-qubits-per-qudit. + Expected shape: `((2 ^ m) ^ num_qudits, (2 ^ m) ^ num_qudits)`. + """ + dim_qubit_space = _nearest_power_of_two_ceiling(qudit_dimension) ** num_qudits + + if memoize: + # Perform the transform of the below array from qubit to qudit space so that the indices + # represent the position in qudit space and the values represent the position in the qubit + # space. + d_to_b_index_map = qubit_to_qudit_state( + qudit_dimension, + num_qudits, + # An array of ints from 0 to dim_qubit_space. Each element represents the original + # index. + np.arange(dim_qubit_space), + ) + # Initialize the result to the identity unitary in the qubit space. + result = np.identity(dim_qubit_space, dtype=qudit_unitary.dtype) + # Iterate over each element in the qudit space dimension x qudit space dimension. + iter_range = range(qudit_dimension**num_qudits) + for i, j in itertools.product(iter_range, iter_range): + # Use the index map to populate the appropriate element in the qubit representation. + result[d_to_b_index_map[i]][d_to_b_index_map[j]] = qudit_unitary[i][j] + return result + + # Treat the unitary as a num_qudits^2 system's state vector and represent it using qubits (pad + # with 0s). + padded_unitary = qudit_to_qubit_state( + qudit_dimension, num_qudits * 2, np.ravel(qudit_unitary) + ) + # A qubit-based state vector with the extra padding bits having 1s and rest having 0s. This + # vector marks only the bits that are padded. + pad_qubits_vector = qudit_to_qubit_state( + qudit_dimension, + num_qudits, + np.zeros(qudit_dimension**num_qudits), + _pad_value=1, + ) + # Reshape the padded unitary to the final shape and add a diagonal matrix corresponding to the + # pad_qubits_vector. This addition ensures that the invalid states with the "padding" bits map + # to identity, preserving unitarity. + return padded_unitary.reshape(dim_qubit_space, dim_qubit_space) + np.diag( + pad_qubits_vector + ) + + +def qubit_to_qudit_unitary( + qudit_dimension: int, + num_qudits: int, + qubit_unitary: np.ndarray, +): + """Converts a m-qubit-per-qudit unitary to a qudit-space quantum unitary. + + Each qudit was replaced by a set of qubits. Since the set of qubits could represent a larger + state space than the qudit, the unitary needs to be sliced up to the qudit length in each + dimension. A unitary is treated similar to a 2*num_qudits system's state vector. + Based on https://drive.google.com/file/d/1n1Ym7JdnM44NnvNQDUXSk5rBLTQZzJ9t and + https://colab.research.google.com/drive/1MC6D4FyOXG0RjvzyV9IFcdRsLbUwwAcx + + Args: + qudit_dimension: The dimension of a single qudit i.e. the number of states it can + represent. + num_qudits: The number of qudits in the given/output unitary. + qubit_unitary: A 2-D numpy array representing the unitary in m-qubit-per-qudit form. + Expected shape: `((2 ^ m) ^ num_qudits, (2 ^ m) ^ num_qudits)`. + + Returns: + A numpy array representing the input unitary using qudits. + Expected shape: `(qudit_dimension ^ num_qudits, qudit_dimension ^ num_qudits)`. + """ + mbit_dimension = _nearest_power_of_two_ceiling(qudit_dimension) + # Treat unitary as a `num_qudits*2` qudit system state vector. + effective_num_qudits = num_qudits * 2 + # Reshape the state vector to a `num_qudits*2` rank tensor. + unitary_tensor = qubit_unitary.reshape((mbit_dimension,) * effective_num_qudits) + # Shrink the number of elements in each dimension up to the qudit_dimension, ignoring the rest. + trimmed_unitary_tensor = unitary_tensor[ + (slice(qudit_dimension),) * effective_num_qudits + ] + # Return a flat unitary view of the final tensor. + return trimmed_unitary_tensor.reshape( + qudit_dimension**num_qudits, qudit_dimension**num_qudits + ) diff --git a/unitary/alpha/qudit_state_transform_test.py b/unitary/alpha/qudit_state_transform_test.py new file mode 100644 index 00000000..81678d75 --- /dev/null +++ b/unitary/alpha/qudit_state_transform_test.py @@ -0,0 +1,168 @@ +import pytest +import numpy as np +from unitary.alpha import qudit_state_transform + + +@pytest.mark.parametrize("qudit_dim", range(10)) +@pytest.mark.parametrize("num_qudits", range(4)) +def test_qudit_state_and_unitary_transform_equivalence(qudit_dim, num_qudits): + qudit_state_space = qudit_dim**num_qudits + qudit_unitary_shape = (qudit_state_space, qudit_state_space) + # Run each configuration with 3 random states and unitaries. + for i in range(3): + # Random complex state vector in the qudit space. + random_state = np.random.rand(qudit_state_space) + 1j * np.random.rand( + qudit_state_space + ) + # Random complex unitary in the qudit space. + random_unitary = np.random.rand(*qudit_unitary_shape) + 1j * np.random.rand( + *qudit_unitary_shape + ) + # Apply the unitary on the state vector. + expected_product = np.matmul(random_unitary, random_state) + # Qubit space representation of the qudit state vector. + transformed_state = qudit_state_transform.qudit_to_qubit_state( + qudit_dim, num_qudits, random_state + ) + # Qubit space representation of the qudit unitary. Alternate between memoizing or not. + transformed_unitary = qudit_state_transform.qudit_to_qubit_unitary( + qudit_dim, num_qudits, random_unitary, memoize=(i % 2) + ) + # Apply the transformed unitary on the transformed state vector. + transformed_product = np.matmul(transformed_unitary, transformed_state) + # Convert the transformed product back to the qudit space. + product_in_qudit_space = qudit_state_transform.qubit_to_qudit_state( + qudit_dim, num_qudits, transformed_product + ) + # Assert that the transform back from qubit space is the inverse of the transform to qubit + # space. + np.testing.assert_allclose( + qudit_state_transform.qubit_to_qudit_state( + qudit_dim, num_qudits, transformed_state + ), + random_state, + ) + np.testing.assert_allclose( + qudit_state_transform.qubit_to_qudit_unitary( + qudit_dim, num_qudits, transformed_unitary + ), + random_unitary, + ) + # Assert that the operations in the qubit space are equivalent to the operations in the + # qudit space. + np.testing.assert_allclose(product_in_qudit_space, expected_product) + + +@pytest.mark.parametrize( + "qudit_dim, num_qudits, qudit_representation, qubit_representation", + [ + ( + 3, + 2, + np.array([0, 0, 0, 0, 0, 0, 0, 0, 1]), + np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0]), + ), + ( + 4, + 2, + np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0]), + np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0]), + ), + ( + 3, + 2, + np.array([1, 0, 0, 0, 0, 0, 0, 0, 1], dtype=np.complex_), + np.array( + [ + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 1, + 0, + 0, + 0, + 0, + 0, + ] + ), + ), + ], +) +def test_specific_transformations_vectors( + qudit_dim, num_qudits, qudit_representation, qubit_representation +): + transformed_vector = qudit_state_transform.qudit_to_qubit_state( + qudit_dim, num_qudits, qudit_representation + ) + np.testing.assert_allclose(transformed_vector, qubit_representation) + untransformed_vector = qudit_state_transform.qubit_to_qudit_state( + qudit_dim, num_qudits, transformed_vector + ) + np.testing.assert_allclose(untransformed_vector, qudit_representation) + + +a = complex(0.5, 0.5) +b = complex(0.5, -0.5) + + +@pytest.mark.parametrize( + "qudit_dim, num_qudits, qudit_representation, qubit_representation", + [ + # Qutrit square root of iSwap. + ( + 3, + 2, + np.array( + [ + [1, 0, 0, 0, 0, 0, 0, 0, 0], + [0, a, 0, b, 0, 0, 0, 0, 0], + [0, 0, a, 0, 0, 0, b, 0, 0], + [0, b, 0, a, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 1, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 1, 0, 0, 0], + [0, 0, b, 0, 0, 0, a, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 1, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 1], + ] + ), + np.array( + [ + [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, a, 0, 0, b, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, a, 0, 0, 0, 0, 0, b, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, b, 0, 0, a, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, b, 0, 0, 0, 0, 0, a, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0], + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1], + ] + ), + ), + ], +) +def test_specific_transformations_unitaries( + qudit_dim, num_qudits, qudit_representation, qubit_representation +): + transformed_unitary = qudit_state_transform.qudit_to_qubit_unitary( + qudit_dim, num_qudits, qudit_representation + ) + np.testing.assert_allclose(transformed_unitary, qubit_representation) + untransformed_unitary = qudit_state_transform.qubit_to_qudit_unitary( + qudit_dim, num_qudits, transformed_unitary + ) + np.testing.assert_allclose(untransformed_unitary, qudit_representation)