Merge pull request #247 from QunaSys/operator_to_sparse

packages/core/quri_parts/core/operator/__init__.py

@@ -19,6 +19,7 @@
     pauli_product,
 )
 from .representation import transition_amp_comp_basis, transition_amp_representation
+from .sparse import get_sparse_matrix
 from .trotter_suzuki import trotter_suzuki_decomposition
 
 PAULI_IDENTITY = PAULI_IDENTITY
@@ -44,4 +45,5 @@
     "transition_amp_representation",
     "zero",
     "trotter_suzuki_decomposition",
+    "get_sparse_matrix",
 ]

packages/core/quri_parts/core/operator/sparse.py

@@ -0,0 +1,144 @@
+# 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
+#      http://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.
+
+from functools import reduce
+from typing import Literal, Optional, Union
+
+import numpy as np
+import scipy.sparse as sparse
+from typing_extensions import TypeAlias
+
+from . import PAULI_IDENTITY, Operator, PauliLabel, SinglePauli, zero
+
+_sparse_pauli_x = sparse.csc_matrix([[0.0, 1.0], [1.0, 0.0]], dtype=np.complex128)
+_sparse_pauli_y = sparse.csc_matrix([[0.0, -1.0j], [1.0j, 0.0]], dtype=np.complex128)
+_sparse_pauli_z = sparse.csc_matrix([[1.0, 0.0], [0.0, -1.0]], dtype=np.complex128)
+
+_pauli_map = {
+    SinglePauli.X: _sparse_pauli_x,
+    SinglePauli.Y: _sparse_pauli_y,
+    SinglePauli.Z: _sparse_pauli_z,
+}
+
+
+_sparse_matrix_map: dict[str, sparse.spmatrix] = {
+    "csc": sparse.csc_matrix,
+    "csr": sparse.csr_matrix,
+    "bsr": sparse.bsr_matrix,
+    "coo": sparse.coo_matrix,
+    "dok": sparse.dok_matrix,
+    "dia": sparse.dia_matrix,
+    "lil": sparse.lil_matrix,
+}
+
+SparseMatrixName: TypeAlias = Literal["csc", "csr", "bsr", "coo", "dok", "dia", "lil"]
+
+
+def _convert_pauli_map_to_other_format(format: SparseMatrixName) -> None:
+    if (
+        isinstance(_sparse_pauli_x, _sparse_matrix_map[format])
+        and isinstance(_sparse_pauli_y, _sparse_matrix_map[format])
+        and isinstance(_sparse_pauli_z, _sparse_matrix_map[format])
+    ):
+        return
+
+    if format == "csc":
+        for i in _pauli_map:
+            _pauli_map[i] = _pauli_map[i].tocsc()
+    elif format == "csr":
+        for i in _pauli_map:
+            _pauli_map[i] = _pauli_map[i].tocsr()
+    elif format == "bsr":
+        for i in _pauli_map:
+            _pauli_map[i] = _pauli_map[i].tobsr()
+    elif format == "coo":
+        for i in _pauli_map:
+            _pauli_map[i] = _pauli_map[i].tocoo()
+    elif format == "dok":
+        for i in _pauli_map:
+            _pauli_map[i] = _pauli_map[i].todok()
+    elif format == "dia":
+        for i in _pauli_map:
+            _pauli_map[i] = _pauli_map[i].todia()
+    elif format == "lil":
+        for i in _pauli_map:
+            _pauli_map[i] = _pauli_map[i].tolil()
+    else:
+        assert False, f"format {format} is not supported."
+
+
+def _convert_pauli_label_to_sparse(
+    single_pauli_label: PauliLabel,
+    n_qubits: Optional[int] = None,
+    format: SparseMatrixName = "csc",
+) -> sparse.spmatrix:
+    """Convert :class:`~PauliLabel` into scipy sparse matrix."""
+    _convert_pauli_map_to_other_format(format)
+    if n_qubits is None:
+        assert single_pauli_label != PAULI_IDENTITY, (
+            "n_qubits needs to be specified for PAULI_IDENTITY"
+            " to be converted to matrix."
+        )
+        n_qubits = max(single_pauli_label.qubit_indices()) + 1
+
+    single_pauli_list = [
+        sparse.identity(2, np.complex128, format=format) for _ in range(n_qubits)
+    ]
+
+    if single_pauli_label != PAULI_IDENTITY:
+        assert n_qubits >= max(single_pauli_label.qubit_indices()) + 1, (
+            "The specified number of qubits should not be less then the length"
+            " of the pauli operator."
+        )
+        for bit, pauli in zip(*single_pauli_label.index_and_pauli_id_list):
+            single_pauli_list[n_qubits - bit - 1] = _pauli_map[SinglePauli(pauli)]
+
+    return reduce(lambda o1, o2: sparse.kron(o1, o2, format), single_pauli_list)
+
+
+def _convert_operator_to_sparse(
+    operator: Operator, n_qubits: Optional[int] = None, format: SparseMatrixName = "csc"
+) -> sparse.spmatrix:
+    """Convert :class:`~Operator` into scipy sparse matrix."""
+    if operator == zero():
+        return (
+            _sparse_matrix_map[format](np.zeros((1, 1), dtype=np.complex128))
+            if n_qubits is None
+            else _sparse_matrix_map[format](
+                np.zeros((2**n_qubits, 2**n_qubits), dtype=np.complex128)
+            )
+        )
+
+    if n_qubits is None:
+        n_qubits = max(
+            [max(op.qubit_indices()) + 1 for op in operator if op != PAULI_IDENTITY]
+        )
+
+    return sum(
+        [
+            coeff * _convert_pauli_label_to_sparse(op, n_qubits, format)
+            for op, coeff in operator.items()
+        ]
+    )
+
+
+def get_sparse_matrix(
+    operator: Union[PauliLabel, Operator],
+    n_qubits: Optional[int] = None,
+    format: SparseMatrixName = "csc",
+) -> sparse.spmatrix:
+    """Convert :class:`~PauliLabel` and :class:`~Operator` into scipy sparse
+    matrix."""
+    if isinstance(operator, PauliLabel):
+        return _convert_pauli_label_to_sparse(operator, n_qubits, format)
+    elif isinstance(operator, Operator):
+        return _convert_operator_to_sparse(operator, n_qubits, format)
+    else:
+        assert False, "operator should be either a PauliLabel or an Operator object."

packages/core/tests/core/operator/test_sparse.py

@@ -0,0 +1,115 @@
+# 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
+#      http://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.
+
+from functools import reduce
+
+import numpy as np
+import pytest
+import scipy.sparse as sparse
+
+from quri_parts.core.operator import (
+    PAULI_IDENTITY,
+    Operator,
+    get_sparse_matrix,
+    pauli_label,
+    zero,
+)
+
+I = np.eye(2, dtype=np.complex128)  # noqa: E741
+X = np.array([[0, 1], [1, 0]], dtype=np.complex128)
+Y = np.array([[0, -1j], [1j, 0]], dtype=np.complex128)
+Z = np.array([[1, 0], [0, -1]], dtype=np.complex128)
+
+
+class TestGetSparseMatrix:
+    def test_get_sparse_matrix_pauli_label(self) -> None:
+        test_pauli_label = pauli_label("X0 X1 Y3 Z4")
+        expected_matrix = reduce(np.kron, [Z, Y, I, X, X])
+        assert np.allclose(
+            get_sparse_matrix(test_pauli_label).toarray(), expected_matrix
+        )
+        assert isinstance(get_sparse_matrix(test_pauli_label), sparse.csc_matrix)
+
+    def test_get_sparse_matrix_sparse_matrix_type(self) -> None:
+        test_pauli_label = pauli_label("X0 X1 Y3 Z4")
+        assert isinstance(get_sparse_matrix(test_pauli_label), sparse.csc_matrix)
+
+        assert isinstance(
+            get_sparse_matrix(test_pauli_label, format="csr"), sparse.csr_matrix
+        )
+
+        assert isinstance(
+            get_sparse_matrix(test_pauli_label, format="bsr"), sparse.bsr_matrix
+        )
+
+        assert isinstance(
+            get_sparse_matrix(test_pauli_label, format="coo"), sparse.coo_matrix
+        )
+
+        assert isinstance(
+            get_sparse_matrix(test_pauli_label, format="dia"), sparse.dia_matrix
+        )
+
+        assert isinstance(
+            get_sparse_matrix(test_pauli_label, format="dok"), sparse.dok_matrix
+        )
+
+        assert isinstance(
+            get_sparse_matrix(test_pauli_label, format="lil"), sparse.lil_matrix
+        )
+
+    def test_get_sparse_matrix_nqubit_specified(self) -> None:
+        test_pauli_label = pauli_label("X0 X1 Y3 Z4")
+        expected_matrix = reduce(np.kron, [I, Z, Y, I, X, X])
+        assert np.allclose(
+            get_sparse_matrix(test_pauli_label, 6).toarray(), expected_matrix
+        )
+
+        with pytest.raises(
+            AssertionError,
+            match=(
+                "The specified number of qubits should not be less then the length"
+                " of the pauli operator."
+            ),
+        ):
+            get_sparse_matrix(test_pauli_label, 3)
+
+    def test_get_sparse_matrix_for_operator(self) -> None:
+        operator = Operator({pauli_label("X0 Y3"): -3, PAULI_IDENTITY: 8})
+        expected_matrix = -3 * reduce(np.kron, [Y, I, I, X]) + 8 * np.eye(16)
+        assert np.allclose(get_sparse_matrix(operator).toarray(), expected_matrix)
+
+    def test_get_sparse_matrix_for_zero(self) -> None:
+        operator = zero()
+        converted = get_sparse_matrix(operator).toarray()
+        assert np.allclose(converted, np.zeros((1, 1)))
+        assert converted.shape == (1, 1)
+
+        operator = zero()
+        converted = get_sparse_matrix(operator, 3).toarray()
+        assert np.allclose(converted, np.zeros((8, 8)))
+        assert converted.shape == (8, 8)
+
+    def test_get_sparse_matrix_for_identity(self) -> None:
+        test_pauli_label = PAULI_IDENTITY
+        assert np.allclose(get_sparse_matrix(test_pauli_label, 3).toarray(), np.eye(8))
+        with pytest.raises(
+            AssertionError,
+            match=(
+                "n_qubits needs to be specified for PAULI_IDENTITY"
+                " to be converted to matrix."
+            ),
+        ):
+            get_sparse_matrix(test_pauli_label)
+
+        operator = Operator({PAULI_IDENTITY: 8})
+        assert np.allclose(get_sparse_matrix(operator, 3).toarray(), 8 * np.eye(8))
+        with pytest.raises(ValueError):
+            get_sparse_matrix(operator)