Add sparse modules to update dynamics by arraylias

qiskit_dynamics/arraylias/alias.py

@@ -18,12 +18,24 @@
 
 from typing import Union
 
+from scipy.sparse import spmatrix
+
 from arraylias import numpy_alias, scipy_alias
 
 from qiskit import QiskitError
 
 from qiskit_dynamics.array import Array
 
+from .register_functions import (
+    register_asarray,
+    register_matmul,
+    register_multiply,
+    register_rmatmul,
+    register_todense,
+    register_to_numeric_matrix_type,
+    register_tosparse,
+)
+
 # global NumPy and SciPy aliases
 DYNAMICS_NUMPY_ALIAS = numpy_alias()
 DYNAMICS_SCIPY_ALIAS = scipy_alias()
@@ -35,6 +47,26 @@
 DYNAMICS_NUMPY = DYNAMICS_NUMPY_ALIAS()
 DYNAMICS_SCIPY = DYNAMICS_SCIPY_ALIAS()
 
+# register required custom versions of functions for sparse type here
+DYNAMICS_NUMPY_ALIAS.register_type(spmatrix, lib="scipy_sparse")
+
+try:
+    from jax.experimental.sparse import BCOO
+
+    # register required custom versions of functions for BCOO type here
+    DYNAMICS_NUMPY_ALIAS.register_type(BCOO, lib="jax_sparse")
+except ImportError:
+    pass
+
+# register custom functions for numpy_alias
+register_asarray(alias=DYNAMICS_NUMPY_ALIAS)
+register_todense(alias=DYNAMICS_NUMPY_ALIAS)
+register_to_numeric_matrix_type(alias=DYNAMICS_NUMPY_ALIAS)
+register_tosparse(alias=DYNAMICS_NUMPY_ALIAS)
+register_matmul(alias=DYNAMICS_NUMPY_ALIAS)
+register_multiply(alias=DYNAMICS_NUMPY_ALIAS)
+register_rmatmul(alias=DYNAMICS_NUMPY_ALIAS)
+
 
 ArrayLike = Union[Union[DYNAMICS_NUMPY_ALIAS.registered_types()], list]
 

qiskit_dynamics/arraylias/register_functions/__init__.py

@@ -0,0 +1,24 @@
+# -*- coding: utf-8 -*-
+
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+"""
+Register custom functions using alias
+"""
+
+from .asarray import register_asarray
+from .to_dense import register_todense
+from .to_numeric_matrix_type import register_to_numeric_matrix_type
+from .to_sparse import register_tosparse
+from .matmul import register_matmul
+from .rmatmul import register_rmatmul
+from .multiply import register_multiply

qiskit_dynamics/arraylias/register_functions/asarray.py

@@ -0,0 +1,50 @@
+# -*- coding: utf-8 -*-
+
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""
+Registering asarray functions to alias
+"""
+
+import numpy as np
+from scipy.sparse import csr_matrix, issparse
+
+
+def register_asarray(alias):
+    """register asarray functions to each array libraries"""
+
+    @alias.register_default(path="asarray")
+    def _(arr):
+        return arr
+
+    @alias.register_function(lib="scipy_sparse", path="asarray")
+    def _(arr):
+        if issparse(arr):
+            return arr
+        return csr_matrix(arr)
+
+    try:
+        from jax.experimental.sparse import BCOO
+
+        @alias.register_function(lib="jax_sparse", path="asarray")
+        def _(arr):
+            if type(arr).__name__ == "BCOO":
+                return arr
+            return BCOO.fromdense(arr)
+
+    except ImportError:
+        pass
+
+    @alias.register_fallback(path="asarray")
+    def _(arr):
+        return np.asarray(arr)

qiskit_dynamics/arraylias/register_functions/matmul.py

@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""
+Register matmul functions to alias.
+"""
+
+
+def register_matmul(alias):
+    """Register matmul functions to required array libraries."""
+
+    @alias.register_function(lib="scipy_sparse", path="matmul")
+    def _(x, y):
+        return x * y
+
+    try:
+        from jax.experimental import sparse as jsparse
+        import jax.numpy as jnp
+
+        jsparse_matmul = jsparse.sparsify(jnp.matmul)
+
+        @alias.register_function(lib="jax_sparse", path="matmul")
+        def _(x, y):
+            return jsparse_matmul(x, y)
+
+    except ImportError:
+        pass

qiskit_dynamics/arraylias/register_functions/multiply.py

@@ -0,0 +1,42 @@
+# -*- coding: utf-8 -*-
+
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""
+Register multiply functions to alias.
+"""
+
+
+def register_multiply(alias):
+    """Register multiply functions to each array library."""
+
+    @alias.register_function(lib="scipy_sparse", path="multiply")
+    def _(x, y):
+        return x.multiply(y)
+
+    try:
+        from jax.experimental import sparse as jsparse
+        import jax.numpy as jnp
+
+        jsparse_multiply = jsparse.sparsify(jnp.multiply)
+
+        @alias.register_function(lib="jax_sparse", path="multiply")
+        def _(x, y):
+            return jsparse_multiply(x, y)
+
+    except ImportError:
+        pass
+
+    @alias.register_fallback(path="multiply")
+    def _(x, y):
+        return x * y

qiskit_dynamics/arraylias/register_functions/rmatmul.py

@@ -0,0 +1,49 @@
+# -*- coding: utf-8 -*-
+
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+
+"""
+Register rmatmul functions to alias.
+"""
+
+import numpy as np
+
+
+def register_rmatmul(alias):
+    """Register rmatmul functions to each array library."""
+
+    @alias.register_function(lib="numpy", path="rmatmul")
+    def _(x, y):
+        return np.matmul(y, x)
+
+    @alias.register_function(lib="scipy_sparse", path="rmatmul")
+    def _(x, y):
+        return y * x
+
+    try:
+        from jax.experimental import sparse as jsparse
+        import jax.numpy as jnp
+
+        jsparse_matmul = jsparse.sparsify(jnp.matmul)
+
+        @alias.register_function(lib="jax", path="rmatmul")
+        def _(x, y):
+            return jnp.matmul(y, x)
+
+        @alias.register_function(lib="jax_sparse", path="rmatmul")
+        def _(x, y):
+            return jsparse_matmul(y, x)
+
+    except ImportError:
+        pass

qiskit_dynamics/arraylias/register_functions/to_dense.py

@@ -0,0 +1,55 @@
+# -*- coding: utf-8 -*-
+
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""
+Register to_dense functions to alias.
+"""
+
+import numpy as np
+from arraylias.exceptions import LibraryError
+
+
+def register_todense(alias):
+    """Register to_dense functions to each array library."""
+
+    @alias.register_default(path="to_dense")
+    def _(arr):
+        if arr is None:
+            return None
+        return arr
+
+    @alias.register_function(lib="numpy", path="to_dense")
+    def _(arr):
+        return arr
+
+    try:
+
+        @alias.register_function(lib="jax", path="to_dense")
+        def _(arr):
+            return arr
+
+        @alias.register_function(lib="jax_sparse", path="to_dense")
+        def _(arr):
+            return arr.todense()
+
+    except LibraryError:
+        pass
+
+    @alias.register_function(lib="scipy_sparse", path="to_dense")
+    def _(arr):
+        return arr.toarray()
+
+    @alias.register_fallback(path="to_dense")
+    def _(arr):
+        return np.asarray(arr)

qiskit_dynamics/arraylias/register_functions/to_numeric_matrix_type.py

@@ -0,0 +1,57 @@
+# -*- coding: utf-8 -*-
+
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2023.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+"""
+Registering to_numeric_matrix_type functions to alias
+"""
+
+from arraylias.exceptions import LibraryError
+from qiskit_dynamics.type_utils import isinstance_qutip_qobj
+
+
+def register_to_numeric_matrix_type(alias):
+    """register to_numeric_matrix_type functions to each array libraries"""
+
+    @alias.register_default(path="to_numeric_matrix_type")
+    def _(arr):
+        if arr is None:
+            return None
+        if isinstance_qutip_qobj(arr):
+            return alias().to_sparse(arr.data)
+        return arr
+
+    @alias.register_function(lib="numpy", path="to_numeric_matrix_type")
+    def _(arr):
+        return arr
+
+    try:
+
+        @alias.register_function(lib="jax", path="to_numeric_matrix_type")
+        def _(arr):
+            return arr
+
+        @alias.register_function(lib="jax_sparse", path="to_numeric_matrix_type")
+        def _(arr):
+            return arr
+
+    except LibraryError:
+        pass
+
+    @alias.register_function(lib="scipy_sparse", path="to_numeric_matrix_type")
+    def _(arr):
+        return arr
+
+    @alias.register_fallback(path="to_numeric_matrix_type")
+    def _(arr):
+        return arr