Add sparse modules to update dynamics by arraylias

qiskit_dynamics/arraylias/alias.py

@@ -18,12 +18,21 @@
 
 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,
+)
+
 # global NumPy and SciPy aliases
 DYNAMICS_NUMPY_ALIAS = numpy_alias()
 DYNAMICS_SCIPY_ALIAS = scipy_alias()
@@ -35,6 +44,23 @@
 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_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,21 @@
+# -*- 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 .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 np.asarray(arr)
+
+    @alias.register_fallback(path="asarray")
+    def _(arr):
+        return np.asarray(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

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_fallback(path="multiply")
+    def _(x, y):
+        return x * y
+
+    @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

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