From 6215f5a57994d3590c0406b8aea089274b0b9519 Mon Sep 17 00:00:00 2001 From: Daniel Puzzuoli Date: Tue, 5 Dec 2023 08:05:59 -0800 Subject: [PATCH] Operator collections arraylias integration (#291) Co-authored-by: Kento Ueda <38037695+to24toro@users.noreply.github.com> --- qiskit_dynamics/arraylias/alias.py | 15 +- .../arraylias/register_functions/__init__.py | 3 + .../arraylias/register_functions/conjugate.py | 37 + .../register_functions/linear_combo.py | 51 + .../arraylias/register_functions/transpose.py | 31 + qiskit_dynamics/models/generator_model.py | 19 +- qiskit_dynamics/models/lindblad_model.py | 36 +- .../models/operator_collections.py | 1435 ++++++----------- qiskit_dynamics/type_utils.py | 4 +- test/dynamics/arraylias/test_alias.py | 68 +- test/dynamics/common.py | 20 + .../models/test_operator_collections.py | 422 ++--- test/dynamics/models/test_rotating_frame.py | 7 +- 13 files changed, 921 insertions(+), 1227 deletions(-) create mode 100644 qiskit_dynamics/arraylias/register_functions/conjugate.py create mode 100644 qiskit_dynamics/arraylias/register_functions/linear_combo.py create mode 100644 qiskit_dynamics/arraylias/register_functions/transpose.py diff --git a/qiskit_dynamics/arraylias/alias.py b/qiskit_dynamics/arraylias/alias.py index 8868b98c1..c8912ea10 100644 --- a/qiskit_dynamics/arraylias/alias.py +++ b/qiskit_dynamics/arraylias/alias.py @@ -31,6 +31,9 @@ register_matmul, register_multiply, register_rmatmul, + register_linear_combo, + register_transpose, + register_conjugate, ) # global NumPy and SciPy aliases @@ -60,6 +63,9 @@ register_matmul(alias=DYNAMICS_NUMPY_ALIAS) register_multiply(alias=DYNAMICS_NUMPY_ALIAS) register_rmatmul(alias=DYNAMICS_NUMPY_ALIAS) +register_linear_combo(alias=DYNAMICS_NUMPY_ALIAS) +register_conjugate(alias=DYNAMICS_NUMPY_ALIAS) +register_transpose(alias=DYNAMICS_NUMPY_ALIAS) ArrayLike = Union[Union[DYNAMICS_NUMPY_ALIAS.registered_types()], list] @@ -82,13 +88,16 @@ def _preferred_lib(*args, **kwargs): """ args = list(args) + list(kwargs.values()) if len(args) == 1: - return DYNAMICS_NUMPY_ALIAS.infer_libs(args[0]) + libs = DYNAMICS_NUMPY_ALIAS.infer_libs(args[0]) + return libs[0] if len(libs) > 0 else "numpy" - lib0 = DYNAMICS_NUMPY_ALIAS.infer_libs(args[0])[0] - lib1 = _preferred_lib(args[1:])[0] + lib0 = _preferred_lib(args[0]) + lib1 = _preferred_lib(args[1:]) if lib0 == "numpy" and lib1 == "numpy": return "numpy" + elif lib0 == "jax_sparse" or lib1 == "jax_sparse": + return "jax_sparse" elif lib0 == "jax" or lib1 == "jax": return "jax" diff --git a/qiskit_dynamics/arraylias/register_functions/__init__.py b/qiskit_dynamics/arraylias/register_functions/__init__.py index 4da5fc49c..024c7a633 100644 --- a/qiskit_dynamics/arraylias/register_functions/__init__.py +++ b/qiskit_dynamics/arraylias/register_functions/__init__.py @@ -19,3 +19,6 @@ from .matmul import register_matmul from .rmatmul import register_rmatmul from .multiply import register_multiply +from .linear_combo import register_linear_combo +from .conjugate import register_conjugate +from .transpose import register_transpose diff --git a/qiskit_dynamics/arraylias/register_functions/conjugate.py b/qiskit_dynamics/arraylias/register_functions/conjugate.py new file mode 100644 index 000000000..c2157ba2a --- /dev/null +++ b/qiskit_dynamics/arraylias/register_functions/conjugate.py @@ -0,0 +1,37 @@ +# -*- 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 conjugate. +""" + + +def register_conjugate(alias): + """Register linear functions for each array library.""" + + try: + from jax.dtypes import canonicalize_dtype + import jax.numpy as jnp + from jax.experimental.sparse import sparsify + + # can be changed to sparsify(jnp.conjugate) when implemented + def conj_workaround(x): + if jnp.issubdtype(x.dtype, canonicalize_dtype(jnp.complex128)): + return x.real - 1j * x.imag + return x + + alias.register_function(func=sparsify(conj_workaround), lib="jax_sparse", path="conjugate") + + except ImportError: + pass diff --git a/qiskit_dynamics/arraylias/register_functions/linear_combo.py b/qiskit_dynamics/arraylias/register_functions/linear_combo.py new file mode 100644 index 000000000..9a263381c --- /dev/null +++ b/qiskit_dynamics/arraylias/register_functions/linear_combo.py @@ -0,0 +1,51 @@ +# -*- 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 linear_combo functions to alias. This computes a linear combination of matrices (given +by a 3d array). +""" + +import numpy as np + + +def register_linear_combo(alias): + """Register linear functions for each array library.""" + + @alias.register_default(path="linear_combo") + def _(coeffs, mats): + return np.tensordot(coeffs, mats, axes=1) + + @alias.register_function(lib="numpy", path="linear_combo") + def _(coeffs, mats): + return np.tensordot(coeffs, mats, axes=1) + + try: + import jax.numpy as jnp + + @alias.register_function(lib="jax", path="linear_combo") + def _(coeffs, mats): + return jnp.tensordot(coeffs, mats, axes=1) + + from jax.experimental.sparse import sparsify + + jsparse_sum = sparsify(jnp.sum) + + @alias.register_function(lib="jax_sparse", path="linear_combo") + def _(coeffs, mats): + # pylint: disable=unexpected-keyword-arg + return jsparse_sum(jnp.broadcast_to(coeffs[:, None, None], mats.shape) * mats, axis=0) + + except ImportError: + pass diff --git a/qiskit_dynamics/arraylias/register_functions/transpose.py b/qiskit_dynamics/arraylias/register_functions/transpose.py new file mode 100644 index 000000000..d3361a4b5 --- /dev/null +++ b/qiskit_dynamics/arraylias/register_functions/transpose.py @@ -0,0 +1,31 @@ +# -*- 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 transpose. +""" + + +def register_transpose(alias): + """Register linear functions for each array library.""" + + try: + from jax.experimental.sparse import bcoo_transpose + + @alias.register_function(lib="jax_sparse", path="transpose") + def _(arr, axes=None): + return bcoo_transpose(arr, permutation=axes) + + except ImportError: + pass diff --git a/qiskit_dynamics/models/generator_model.py b/qiskit_dynamics/models/generator_model.py index bd1639ab2..b959113db 100644 --- a/qiskit_dynamics/models/generator_model.py +++ b/qiskit_dynamics/models/generator_model.py @@ -25,10 +25,8 @@ from qiskit import QiskitError from qiskit.quantum_info.operators import Operator from qiskit_dynamics.models.operator_collections import ( - BaseOperatorCollection, - DenseOperatorCollection, - SparseOperatorCollection, - JAXSparseOperatorCollection, + OperatorCollection, + ScipySparseOperatorCollection, ) from qiskit_dynamics.array import Array from qiskit_dynamics.signals import Signal, SignalList @@ -538,7 +536,7 @@ def construct_operator_collection( evaluation_mode: str, static_operator: Union[None, Array, csr_matrix], operators: Union[None, Array, List[csr_matrix]], -) -> BaseOperatorCollection: +) -> Union[OperatorCollection, ScipySparseOperatorCollection]: """Construct an operator collection for :class:`GeneratorModel`. Args: @@ -547,14 +545,15 @@ def construct_operator_collection( operators: Operators for the model. Returns: - BaseOperatorCollection: The relevant operator collection. + Union[OperatorCollection, ScipySparseOperatorCollection]: The relevant operator collection. Raises: NotImplementedError: If the ``evaluation_mode`` is invalid. """ if evaluation_mode == "dense": - return DenseOperatorCollection(static_operator=static_operator, operators=operators) + # return DenseOperatorCollection(static_operator=static_operator, operators=operators) + pass if evaluation_mode == "sparse" and Array.default_backend() == "jax": # warn that sparse mode when using JAX is primarily recommended for use on CPU if jax.default_backend() != "cpu": @@ -563,9 +562,11 @@ def construct_operator_collection( stacklevel=2, ) - return JAXSparseOperatorCollection(static_operator=static_operator, operators=operators) + # return JAXSparseOperatorCollection(static_operator=static_operator, operators=operators) + pass if evaluation_mode == "sparse": - return SparseOperatorCollection(static_operator=static_operator, operators=operators) + # return SparseOperatorCollection(static_operator=static_operator, operators=operators) + pass raise NotImplementedError( f"Evaluation mode '{evaluation_mode}' is not supported. Call " diff --git a/qiskit_dynamics/models/lindblad_model.py b/qiskit_dynamics/models/lindblad_model.py index 9c4d176cd..1c384d87f 100644 --- a/qiskit_dynamics/models/lindblad_model.py +++ b/qiskit_dynamics/models/lindblad_model.py @@ -31,13 +31,10 @@ ) from .hamiltonian_model import HamiltonianModel, is_hermitian from .operator_collections import ( - BaseLindbladOperatorCollection, - DenseLindbladCollection, - DenseVectorizedLindbladCollection, - SparseLindbladCollection, - JAXSparseLindbladCollection, - SparseVectorizedLindbladCollection, - JAXSparseVectorizedLindbladCollection, + LindbladCollection, + ScipySparseLindbladCollection, + VectorizedLindbladCollection, + ScipySparseVectorizedLindbladCollection, ) from .rotating_frame import RotatingFrame @@ -654,7 +651,12 @@ def construct_lindblad_operator_collection( hamiltonian_operators: Union[None, Array, List[csr_matrix]], static_dissipators: Union[None, Array, csr_matrix], dissipator_operators: Union[None, Array, List[csr_matrix]], -) -> BaseLindbladOperatorCollection: +) -> Union[ + LindbladCollection, + ScipySparseLindbladCollection, + VectorizedLindbladCollection, + ScipySparseVectorizedLindbladCollection, +]: """Construct a Lindblad operator collection. Args: @@ -685,19 +687,25 @@ def construct_lindblad_operator_collection( ) if evaluation_mode == "dense": - CollectionClass = DenseLindbladCollection + # CollectionClass = DenseLindbladCollection + pass elif evaluation_mode == "sparse": if Array.default_backend() == "jax": - CollectionClass = JAXSparseLindbladCollection + # CollectionClass = JAXSparseLindbladCollection + pass else: - CollectionClass = SparseLindbladCollection + # CollectionClass = SparseLindbladCollection + pass elif evaluation_mode == "dense_vectorized": - CollectionClass = DenseVectorizedLindbladCollection + # CollectionClass = DenseVectorizedLindbladCollection + pass elif evaluation_mode == "sparse_vectorized": if Array.default_backend() == "jax": - CollectionClass = JAXSparseVectorizedLindbladCollection + # CollectionClass = JAXSparseVectorizedLindbladCollection + pass else: - CollectionClass = SparseVectorizedLindbladCollection + # CollectionClass = SparseVectorizedLindbladCollection + pass else: raise NotImplementedError( f"Evaluation mode '{evaluation_mode}' is not supported. See " diff --git a/qiskit_dynamics/models/operator_collections.py b/qiskit_dynamics/models/operator_collections.py index d85fb539d..7c35b70fb 100644 --- a/qiskit_dynamics/models/operator_collections.py +++ b/qiskit_dynamics/models/operator_collections.py @@ -9,72 +9,76 @@ # 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. +# pylint: disable=invalid-name """Operator collections as math/calculation objects for model classes.""" -from abc import ABC, abstractmethod from typing import Any, Union, List, Optional -from copy import copy import numpy as np from scipy.sparse import csr_matrix, issparse from qiskit import QiskitError -from qiskit.quantum_info.operators.operator import Operator -from qiskit_dynamics.arraylias.alias import _numpy_multi_dispatch -from qiskit_dynamics.array import Array, wrap -from qiskit_dynamics.type_utils import to_array, to_csr, to_BCOO, vec_commutator, vec_dissipator +from qiskit_dynamics import DYNAMICS_NUMPY as unp +from qiskit_dynamics import DYNAMICS_NUMPY_ALIAS as numpy_alias +from qiskit_dynamics.arraylias.alias import ArrayLike, _numpy_multi_dispatch +from qiskit_dynamics.type_utils import to_csr, vec_commutator, vec_dissipator -try: - import jax.numpy as jnp - from jax.experimental import sparse as jsparse - # sparse versions of jax.numpy operations - jsparse_sum = jsparse.sparsify(jnp.sum) - jsparse_matmul = jsparse.sparsify(jnp.matmul) - jsparse_add = jsparse.sparsify(jnp.add) - jsparse_subtract = jsparse.sparsify(jnp.subtract) +def _linear_combo(coeffs, mats): + return _numpy_multi_dispatch(coeffs, mats, path="linear_combo") - BCOO = jsparse.BCOO - def jsparse_linear_combo(coeffs: Array, mats: Array): - """Method for computing a linear combination of sparse arrays. +def _matmul(A, B, **kwargs): + return _numpy_multi_dispatch(A, B, path="matmul", **kwargs) - Args: - coeffs: A vector of coefficients. - mats: A 3-d array. - Returns: - The coefficients multiplied against the first axis of the array, and summed. - """ - # pylint: disable=unexpected-keyword-arg - return jsparse_sum(jnp.broadcast_to(coeffs[:, None, None], mats.shape) * mats, axis=0) - - # sparse version of computing A @ X @ B - jsparse_triple_product = jsparse.sparsify(lambda A, X, B: A @ X @ B) -except ImportError: - BCOO = None +def _to_csr_object_array(ops, decimals): + """Turn a list of matrices into a numpy object array of scipy sparse matrix instances.""" + if ops is None: + return None + return np.array( + [numpy_alias(like="scipy_sparse").asarray(np.round(op, decimals)) for op in ops] + ) -class BaseOperatorCollection(ABC): - r"""Abstract class representing a two-variable matrix function. +class OperatorCollection: + r"""Class for evaluating a linear combination of operators acting on a state. This class represents a function :math:`c,y \mapsto \Lambda(c, y)`, which is assumed to be decomposed as :math:`\Lambda(c, y) = (G_d + \sum_jc_jG_j) y` for matrices :math:`G_d` and :math:`G_j`, with :math:`G_d` referred to as the static operator. - Describes an interface for evaluating the map or its action on ``y``, given the 1d set of values - :math:`c_j`. + This works for ``array_library in ["numpy", "jax", "jax_sparse"]``. """ - def __init__(self, static_operator: Optional[Any] = None, operators: Optional[Any] = None): + def __init__( + self, + static_operator: Optional[ArrayLike] = None, + operators: Optional[ArrayLike] = None, + array_library: Optional[str] = None, + ): """Initialize. Args: operators: ``(k,n,n)`` array specifying the terms :math:`G_j`. static_operator: ``(n,n)`` array specifying the extra static_operator :math:`G_d`. + array_library: Underlying library to use for array operations. + + Raises: + QiskitError: If "scipy_sparse" is passed as array_library. """ - self.operators = operators - self.static_operator = static_operator + if array_library == "scipy_sparse": + raise QiskitError("scipy_sparse is not a valid array_library for OperatorCollection.") + + if static_operator is not None: + self._static_operator = numpy_alias(like=array_library).asarray(static_operator) + else: + self._static_operator = None + + if operators is not None: + self._operators = numpy_alias(like=array_library).asarray(operators) + else: + self._operators = None @property def dim(self) -> int: @@ -85,129 +89,72 @@ def dim(self) -> int: return self.operators[0].shape[-1] @property - def static_operator(self) -> Array: - """The static part of the operator collection.""" - - @static_operator.setter - def static_operator(self, new_static_operator: Optional[Array] = None): - pass - - @property - def operators(self) -> Array: - """The operators of this collection.""" - - @operators.setter - def operators(self, new_operators: Array) -> Array: - pass - - @abstractmethod - def evaluate(self, signal_values: Array) -> Array: - r"""Evaluate the operator :math:`\Lambda(c, \cdot) = (G_d + \sum_jc_jG_j)`. - - Args: - signal_values: The signals values :math:`c` to use on the operators. - - Returns: - An :class:`~Array` that acts on states ``y`` via multiplication. - """ - - @abstractmethod - def evaluate_rhs(self, signal_values: Union[List[Array], Array], y: Array) -> Array: - r"""Evaluate the function and return :math:`\Lambda(c, y) = (G_d + \sum_jc_jG_j) y`. - - Args: - signal_values: The signals :math:`c` to use on the operators. - y: The system state. - - Returns: - The evaluated function. - """ - - def __call__( - self, signal_values: Union[List[Array], Array], y: Optional[Array] = None - ) -> Array: - """Call :meth:`~evaluate` or :meth:`~evaluate_rhs` depending on the presense of ``y``. - - Args: - signal_values: The signals :math:`c` to use on the operators. - y: Optionally, the system state. - - Returns: - The output of :meth:`~evaluate` or :meth:`self.evaluate_rhs`. - """ - return self.evaluate(signal_values) if y is None else self.evaluate_rhs(signal_values, y) - - def copy(self): - """Return a copy of self.""" - return copy(self) - - -class DenseOperatorCollection(BaseOperatorCollection): - r"""An implementation of :class:`~BaseOperatorCollection` in which :math:`G_d` and :math:`G_j` - are dense arrays. - """ - - @property - def static_operator(self) -> Array: + def static_operator(self) -> Union[ArrayLike, None]: """The static part of the operator collection.""" return self._static_operator - @static_operator.setter - def static_operator(self, new_static_operator: Array): - self._static_operator = to_array(new_static_operator) - @property - def operators(self) -> Array: - """Operators in the collection.""" + def operators(self) -> Union[ArrayLike, None]: + """The operators of this collection.""" return self._operators - @operators.setter - def operators(self, new_operators: Array): - self._operators = to_array(new_operators) - - def evaluate(self, signal_values: Union[Array, None]) -> Array: + def evaluate(self, coefficients: Union[ArrayLike, None]) -> ArrayLike: r"""Evaluate the operator :math:`\Lambda(c, \cdot) = (G_d + \sum_jc_jG_j)`. Args: - signal_values: The signals values :math:`c` to use on the operators. + coefficients: The coefficient values :math:`c`. Returns: - An :class:`~Array` that acts on states ``y`` via multiplication. + An ``ArrayLike`` that acts on states ``y`` via multiplication. Raises: QiskitError: If both static_operator and operators are ``None``. """ if self._static_operator is not None and self._operators is not None: - return np.tensordot(signal_values, self._operators, axes=1) + self._static_operator + return _linear_combo(coefficients, self._operators) + self._static_operator elif self._static_operator is None and self._operators is not None: - return np.tensordot(signal_values, self._operators, axes=1) + return _linear_combo(coefficients, self._operators) elif self._static_operator is not None: return self._static_operator raise QiskitError( - f"{type(self).__name__} with None for both static_operator and operators cannot be " + "OperatorCollection with None for both static_operator and operators cannot be " "evaluated." ) - def evaluate_rhs(self, signal_values: Union[Array, None], y: Array) -> Array: + def evaluate_rhs(self, coefficients: Union[ArrayLike, None], y: ArrayLike) -> ArrayLike: r"""Evaluate the function and return :math:`\Lambda(c, y) = (G_d + \sum_jc_jG_j) y`. Args: - signal_values: The signals :math:`c` to use on the operators. + coefficients: The coefficients :math:`c`. y: The system state. Returns: The evaluated function. """ - return np.dot(self.evaluate(signal_values), y) + return _matmul(self.evaluate(coefficients), y) + + def __call__( + self, coefficients: Union[ArrayLike, None], y: Optional[ArrayLike] = None + ) -> ArrayLike: + """Call :meth:`~evaluate` or :meth:`~evaluate_rhs` depending on the presense of ``y``. + + Args: + coefficients: The coefficients :math:`c` to use for evaluation. + y: Optionally, the system state. + + Returns: + The output of :meth:`~evaluate` or :meth:`self.evaluate_rhs`. + """ + return self.evaluate(coefficients) if y is None else self.evaluate_rhs(coefficients, y) -class SparseOperatorCollection(BaseOperatorCollection): - r"""Sparse version of :class:`DenseOperatorCollection`.""" +class ScipySparseOperatorCollection: + r"""Scipy sparse version of :class:`OperatorCollection`.""" def __init__( self, - static_operator: Optional[Union[Array, Operator]] = None, - operators: Optional[Union[Array, List[Operator]]] = None, + static_operator: Optional[ArrayLike] = None, + operators: Optional[ArrayLike] = None, decimals: Optional[int] = 10, ): """Initialize. @@ -217,43 +164,33 @@ def __init__( operators: (k,n,n) Array specifying the terms :math:`G_j`. decimals: Values will be rounded at ``decimals`` places after decimal. """ - self._decimals = decimals - super().__init__(static_operator=static_operator, operators=operators) + if static_operator is not None: + self._static_operator = numpy_alias(like="scipy_sparse").asarray( + np.round(static_operator, decimals) + ) + else: + self._static_operator = None + + self._operators = _to_csr_object_array(operators, decimals) @property - def static_operator(self) -> csr_matrix: + def static_operator(self) -> Union[None, csr_matrix]: """The static part of the operator collection.""" return self._static_operator - @static_operator.setter - def static_operator(self, new_static_operator: csr_matrix): - if new_static_operator is not None: - self._static_operator = np.round(to_csr(new_static_operator), self._decimals) - else: - self._static_operator = None - @property - def operators(self) -> List[csr_matrix]: + def operators(self) -> Union[None, List[csr_matrix]]: + """The operators of this collection.""" if self._operators is None: return None return list(self._operators) - @operators.setter - def operators(self, new_operators: List[csr_matrix]): - if new_operators is not None: - new_operators_to_csr = to_csr(list(new_operators)) - new_operators = np.empty(shape=len(new_operators_to_csr), dtype="O") - for idx, new_op in enumerate(new_operators_to_csr): - new_operators[idx] = csr_matrix(np.round(new_op, self._decimals)) - - self._operators = new_operators - - def evaluate(self, signal_values: Union[Array, None]) -> csr_matrix: + def evaluate(self, coefficients: Union[ArrayLike, None]) -> csr_matrix: r"""Evaluate the operator :math:`\Lambda(c, \cdot) = (G_d + \sum_jc_jG_j)`. Args: - signal_values: The signals values :math:`c` to use on the operators. + coefficients: The coefficient values :math:`c` to use on the operators. Returns: An :class:`~Array` that acts on states ``y`` via multiplication. @@ -263,10 +200,10 @@ def evaluate(self, signal_values: Union[Array, None]) -> csr_matrix: """ if self._static_operator is not None and self._operators is not None: return ( - np.tensordot(signal_values, self._operators, axes=1).item() + self._static_operator + np.tensordot(coefficients, self._operators, axes=1).item() + self._static_operator ) elif self._static_operator is None and self._operators is not None: - return np.tensordot(signal_values, self._operators, axes=1).item() + return np.tensordot(coefficients, self._operators, axes=1).item() elif self.static_operator is not None: return self._static_operator raise QiskitError( @@ -274,11 +211,11 @@ def evaluate(self, signal_values: Union[Array, None]) -> csr_matrix: "evaluated." ) - def evaluate_rhs(self, signal_values: Union[Array, None], y: Array) -> Array: + def evaluate_rhs(self, coefficients: Union[ArrayLike, None], y: ArrayLike) -> ArrayLike: r"""Evaluate the function and return :math:`\Lambda(c, y) = (G_d + \sum_jc_jG_j) y`. Args: - signal_values: The signals :math:`c` to use on the operators. + coefficients: The coefficients :math:`c` to use on the operators. y: The system state. Returns: @@ -288,7 +225,7 @@ def evaluate_rhs(self, signal_values: Union[Array, None], y: Array) -> Array: """ if len(y.shape) == 2: # For 2d array, compute linear combination then multiply - gen = self.evaluate(signal_values) + gen = self.evaluate(coefficients) return gen.dot(y) elif len(y.shape) == 1: # For a 1d array, multiply individual matrices then compute linear combination @@ -296,9 +233,9 @@ def evaluate_rhs(self, signal_values: Union[Array, None], y: Array) -> Array: tmparr[0] = y if self._static_operator is not None and self._operators is not None: - return np.dot(signal_values, self._operators * tmparr) + self.static_operator.dot(y) + return np.dot(coefficients, self._operators * tmparr) + self.static_operator.dot(y) elif self._static_operator is None and self._operators is not None: - return np.dot(signal_values, self._operators * tmparr) + return np.dot(coefficients, self._operators * tmparr) elif self.static_operator is not None: return self.static_operator.dot(y) @@ -310,77 +247,23 @@ def evaluate_rhs(self, signal_values: Union[Array, None], y: Array) -> Array: raise QiskitError(self.__class__.__name__ + """ cannot evaluate RHS for y.ndim > 3.""") - -class JAXSparseOperatorCollection(BaseOperatorCollection): - """Jax version of SparseOperatorCollection built on jax.experimental.sparse.BCOO.""" - - @property - def static_operator(self) -> BCOO: - """The static part of the operator collection.""" - return self._static_operator - - @static_operator.setter - def static_operator(self, new_static_operator: Union[BCOO, None]): - self._static_operator = to_BCOO(new_static_operator) - - @property - def operators(self) -> Union[BCOO, None]: - return self._operators - - @operators.setter - def operators(self, new_operators: Union[BCOO, None]): - self._operators = to_BCOO(new_operators) - - def evaluate(self, signal_values: Union[Array, None]) -> BCOO: - r"""Evaluate the operator :math:`\Lambda(c, \cdot) = (G_d + \sum_jc_jG_j)`. - - Args: - signal_values: The signals values :math:`c` to use on the operators. - - Returns: - An :class:`~Array` that acts on states ``y`` via multiplication. - - Raises: - QiskitError: If the collection cannot be evaluated. - """ - if signal_values is not None and isinstance(signal_values, Array): - signal_values = signal_values.data - - if self._static_operator is not None and self._operators is not None: - return jsparse_linear_combo(signal_values, self._operators) + self._static_operator - elif self._static_operator is None and self._operators is not None: - return jsparse_linear_combo(signal_values, self._operators) - elif self.static_operator is not None: - return self._static_operator - raise QiskitError( - self.__class__.__name__ - + """ with None for both static_operator and - operators cannot be evaluated.""" - ) - - def evaluate_rhs(self, signal_values: Union[Array, None], y: Array) -> Array: - r"""Evaluate the function and return :math:`\Lambda(c, y) = (G_d + \sum_jc_jG_j) y`. + def __call__( + self, coefficients: Union[ArrayLike, None], y: Optional[ArrayLike] = None + ) -> ArrayLike: + """Call :meth:`~evaluate` or :meth:`~evaluate_rhs` depending on the presense of ``y``. Args: - signal_values: The signals :math:`c` to use on the operators. - y: The system state. + coefficients: The coefficients :math:`c` to use for evaluation. + y: Optionally, the system state. Returns: - The evaluated function. - - Raises: - QiskitError: If the function cannot be evaluated. + The output of :meth:`~evaluate` or :meth:`self.evaluate_rhs`. """ - if y.ndim < 3: - if isinstance(y, Array): - y = y.data - return Array(jsparse_matmul(self.evaluate(signal_values), y)) - - raise QiskitError(self.__class__.__name__ + """ cannot evaluate RHS for y.ndim >= 3.""") + return self.evaluate(coefficients) if y is None else self.evaluate_rhs(coefficients, y) -class BaseLindbladOperatorCollection(ABC): - r"""Abstract class representing a two-variable matrix function for evaluating the right hand +class LindbladCollection: + r"""Class representing a two-variable matrix function for evaluating the right hand side of the Lindblad equation. In particular, this object represents the function: @@ -396,14 +279,17 @@ class BaseLindbladOperatorCollection(ABC): Describes an interface for evaluating the map or its action on :math:`\rho`, given a pair of 1d sets of values :math:`c_1, c_2`. + + This class works for ``array_library in ["numpy", "jax", "jax_sparse"]``. """ def __init__( self, - static_hamiltonian: Optional[any] = None, - hamiltonian_operators: Optional[any] = None, - static_dissipators: Optional[any] = None, - dissipator_operators: Optional[any] = None, + static_hamiltonian: Optional[ArrayLike] = None, + hamiltonian_operators: Optional[ArrayLike] = None, + static_dissipators: Optional[ArrayLike] = None, + dissipator_operators: Optional[ArrayLike] = None, + array_library: Optional[str] = None, ): r"""Initialize collection. Argument types depend on concrete subclass. @@ -414,205 +300,120 @@ def __init__( as :math:`H(t) = \sum_j s(t) H_j+H_d` by specifying H_j. (k,n,n) array. static_dissipators: Constant dissipator terms. dissipator_operators: the terms :math:`L_j` in Lindblad equation. (m,n,n) array. - """ - self.static_hamiltonian = static_hamiltonian - self.hamiltonian_operators = hamiltonian_operators - self.static_dissipators = static_dissipators - self.dissipator_operators = dissipator_operators - - @property - @abstractmethod - def static_hamiltonian(self) -> Array: - """The static part of the Hamiltonian.""" - - @static_hamiltonian.setter - @abstractmethod - def static_hamiltonian(self, new_static_operator: Optional[Array] = None): - pass - - @property - @abstractmethod - def hamiltonian_operators(self) -> Array: - """The operators for the non-static part of Hamiltonian.""" - - @hamiltonian_operators.setter - @abstractmethod - def hamiltonian_operators(self, new_hamiltonian_operators: Optional[Array] = None): - pass - - @property - @abstractmethod - def static_dissipators(self) -> Array: - """The operators for the static part of dissipator.""" - - @static_dissipators.setter - @abstractmethod - def static_dissipators(self, new_static_dissipators: Optional[Array] = None): - pass - - @property - @abstractmethod - def dissipator_operators(self) -> Array: - """The operators for the non-static part of dissipator.""" - - @dissipator_operators.setter - @abstractmethod - def dissipator_operators(self, new_dissipator_operators: Optional[Array] = None): - pass - - @abstractmethod - def evaluate_hamiltonian(self, ham_sig_vals: Union[None, Array]) -> Union[csr_matrix, Array]: - r"""Evaluate the Hamiltonian of the model. + array_library: Array library to use for storing arrays in the collection. - Args: - ham_sig_vals: The values of :math:`s_j` in :math:`H = \sum_j s_j(t) H_j + H_d`. - - Returns: - The Hamiltonian. + Raises: + QiskitError: If "scipy_sparse" is passed as the array_library. """ - @abstractmethod - def evaluate( - self, ham_sig_vals: Union[None, Array], dis_sig_vals: Union[None, Array] - ) -> Union[csr_matrix, Array]: - r"""Evaluate the function and return :math:`\Lambda(c_1, c_2, \cdot)`. - - Args: - ham_sig_vals: The signals :math:`c_1` to use on the Hamiltonians. - dis_sig_vals: The signals :math:`c_2` to use on the dissipators. - - Returns: - The evaluated function. - """ + if array_library == "scipy_sparse": + raise QiskitError("scipy_sparse is not a valid array_library for LindbladCollection.") - @abstractmethod - def evaluate_rhs( - self, ham_sig_vals: Union[None, Array], dis_sig_vals: Union[None, Array], y: Array - ) -> Array: - r"""Evaluate the function and return :math:`\Lambda(c_1, c_2, y)`. + if static_hamiltonian is not None: + self._static_hamiltonian = numpy_alias(like=array_library).asarray(static_hamiltonian) + else: + self._static_hamiltonian = None - Args: - ham_sig_vals: The signals :math:`c_1` to use on the Hamiltonians. - dis_sig_vals: The signals :math:`c_2` to use on the dissipators. - y: The system state. + if hamiltonian_operators is not None: + self._hamiltonian_operators = numpy_alias(like=array_library).asarray( + hamiltonian_operators + ) + else: + self._hamiltonian_operators = None - Returns: - The evaluated function. - """ + if static_dissipators is not None: + if array_library == "jax_sparse": + self._static_dissipators = numpy_alias(like="jax").asarray(static_dissipators) + else: + self._static_dissipators = numpy_alias(like=array_library).asarray( + static_dissipators + ) - def __call__( - self, ham_sig_vals: Union[None, Array], dis_sig_vals: Union[None, Array], y: Optional[Array] - ) -> Union[csr_matrix, Array]: - """Call :meth:`~evaluate` or :meth:`~evaluate_rhs` depending on the presense of ``y``. + self._static_dissipators_adj = unp.conjugate( + unp.transpose(self._static_dissipators, [0, 2, 1]) + ) + self._static_dissipators_product_sum = -0.5 * unp.sum( + unp.matmul(self._static_dissipators_adj, self._static_dissipators), axis=0 + ) - Args: - ham_sig_vals: The signals :math:`c_1` to use on the Hamiltonians. - dis_sig_vals: The signals :math:`c_2` to use on the dissipators. - y: Optionally, the system state. + if array_library == "jax_sparse": + from jax.experimental.sparse import BCOO - Returns: - The evaluated function. - """ - if y is None: - return self.evaluate(ham_sig_vals, dis_sig_vals) + self._static_dissipators = BCOO.fromdense(self._static_dissipators, n_batch=1) + self._static_dissipators_adj = BCOO.fromdense( + self._static_dissipators_adj, n_batch=1 + ) + self._static_dissipators_product_sum = BCOO.fromdense( + self._static_dissipators_product_sum + ) + else: + self._static_dissipators = None - return self.evaluate_rhs(ham_sig_vals, dis_sig_vals, y) + if dissipator_operators is not None: + if array_library == "jax_sparse": + self._dissipator_operators = numpy_alias(like="jax").asarray(dissipator_operators) + else: + self._dissipator_operators = numpy_alias(like=array_library).asarray( + dissipator_operators + ) - def copy(self): - """Return a copy of self.""" - return copy(self) + self._dissipator_operators_adj = unp.conjugate( + unp.transpose(self._dissipator_operators, [0, 2, 1]) + ) + self._dissipator_products = -0.5 * unp.matmul( + self._dissipator_operators_adj, self._dissipator_operators + ) + if array_library == "jax_sparse": + from jax.experimental.sparse import BCOO -class DenseLindbladCollection(BaseLindbladOperatorCollection): - r"""Object for computing the right hand side of the Lindblad equation - with dense arrays. - """ + self._dissipator_operators = BCOO.fromdense(self._dissipator_operators, n_batch=1) + self._dissipator_operators_adj = BCOO.fromdense( + self._dissipator_operators_adj, n_batch=1 + ) + self._dissipator_products = BCOO.fromdense(self._dissipator_products, n_batch=1) + else: + self._dissipator_operators = None @property - def static_hamiltonian(self) -> Array: - """The static part of the operator collection.""" + def static_hamiltonian(self) -> ArrayLike: + """The static part of the Hamiltonian.""" return self._static_hamiltonian - @static_hamiltonian.setter - def static_hamiltonian(self, new_static_hamiltonian: Optional[Array] = None): - self._static_hamiltonian = to_array(new_static_hamiltonian) - @property - def hamiltonian_operators(self) -> Array: + def hamiltonian_operators(self) -> ArrayLike: """The operators for the non-static part of Hamiltonian.""" return self._hamiltonian_operators - @hamiltonian_operators.setter - def hamiltonian_operators(self, new_hamiltonian_operators: Optional[Array] = None): - self._hamiltonian_operators = to_array(new_hamiltonian_operators) - @property - def static_dissipators(self) -> Array: + def static_dissipators(self) -> ArrayLike: """The operators for the static part of dissipator.""" return self._static_dissipators - @static_dissipators.setter - def static_dissipators(self, new_static_dissipators: Optional[Array] = None): - self._static_dissipators = to_array(new_static_dissipators) - if self._static_dissipators is not None: - self._static_dissipators_adj = np.conjugate( - np.transpose(self._static_dissipators, [0, 2, 1]) - ).copy() - self._static_dissipators_product_sum = -0.5 * np.sum( - np.matmul(self._static_dissipators_adj, self._static_dissipators), axis=0 - ) - @property - def dissipator_operators(self) -> Array: + def dissipator_operators(self) -> ArrayLike: """The operators for the non-static part of dissipator.""" return self._dissipator_operators - @dissipator_operators.setter - def dissipator_operators(self, new_dissipator_operators: Optional[Array] = None): - self._dissipator_operators = to_array(new_dissipator_operators) - if self._dissipator_operators is not None: - self._dissipator_operators_adj = np.conjugate( - np.transpose(self._dissipator_operators, [0, 2, 1]) - ).copy() - self._dissipator_products = np.matmul( - self._dissipator_operators_adj, self._dissipator_operators - ) - - def evaluate(self, ham_sig_vals: Array, dis_sig_vals: Array) -> Array: - r"""Placeholder to return :math:`\Lambda(c_1, c_2, \cdot)`, which is not possible without - a state. - - Args: - ham_sig_vals: The signals :math:`c_1` to use on the Hamiltonians. - dis_sig_vals: The signals :math:`c_2` to use on the dissipators. - - Returns: - The evaluated function. - - Raises: - ValueError: Always. - """ - raise ValueError("Non-vectorized Lindblad collections cannot be evaluated without a state.") - - def evaluate_hamiltonian(self, ham_sig_vals: Union[None, Array]) -> Array: - r"""Compute the Hamiltonian. + def evaluate_hamiltonian(self, ham_coefficients: Optional[ArrayLike]) -> ArrayLike: + r"""Evaluate the Hamiltonian of the model. Args: - ham_sig_vals: [Real] values of :math:`s_j` in :math:`H = \sum_j s_j(t) H_j + H_d`. + ham_coefficients: The values of :math:`s_j` in :math:`H = \sum_j s_j(t) H_j + H_d`. Returns: - Hamiltonian matrix. + The Hamiltonian. Raises: QiskitError: If collection not sufficiently specified. """ if self._static_hamiltonian is not None and self._hamiltonian_operators is not None: return ( - np.tensordot(ham_sig_vals, self._hamiltonian_operators, axes=1) + _linear_combo(ham_coefficients, self._hamiltonian_operators) + self._static_hamiltonian ) elif self._static_hamiltonian is None and self._hamiltonian_operators is not None: - return np.tensordot(ham_sig_vals, self._hamiltonian_operators, axes=1) + return _linear_combo(ham_coefficients, self._hamiltonian_operators) elif self._static_hamiltonian is not None: return self._static_hamiltonian else: @@ -622,9 +423,29 @@ def evaluate_hamiltonian(self, ham_sig_vals: Union[None, Array]) -> Array: hamiltonian_operators cannot evaluate Hamiltonian.""" ) + def evaluate( + self, ham_coefficients: Optional[ArrayLike], dis_coefficients: Optional[ArrayLike] + ) -> ArrayLike: + r"""Evaluate the function and return :math:`\Lambda(c_1, c_2, \cdot)`. + + Args: + ham_coefficients: The signals :math:`c_1` to use on the Hamiltonians. + dis_coefficients: The signals :math:`c_2` to use on the dissipators. + + Returns: + The evaluated function. + + Raises: + ValueError: Always. + """ + raise ValueError("Non-vectorized Lindblad collections cannot be evaluated without a state.") + def evaluate_rhs( - self, ham_sig_vals: Union[None, Array], dis_sig_vals: Union[None, Array], y: Array - ) -> Array: + self, + ham_coefficients: Optional[ArrayLike], + dis_coefficients: Optional[ArrayLike], + y: ArrayLike, + ) -> ArrayLike: r"""Evaluates Lindblad equation RHS given a pair of signal values for the hamiltonian terms and the dissipator terms. @@ -638,8 +459,8 @@ def evaluate_rhs( C = \sum_j \gamma_j(t) L_j y L_j^\dagger. Args: - ham_sig_vals: Hamiltonian coefficient values, :math:`s_j(t)`. - dis_sig_vals: Dissipator signal values, :math:`\gamma_j(t)`. + ham_coefficients: Hamiltonian coefficient values, :math:`s_j(t)`. + dis_coefficients: Dissipator coefficient values, :math:`\gamma_j(t)`. y: Density matrix as ``(n,n)`` array representing the state at time :math:`t`. Returns: @@ -653,81 +474,98 @@ def evaluate_rhs( hamiltonian_matrix = None if self._static_hamiltonian is not None or self._hamiltonian_operators is not None: - hamiltonian_matrix = -1j * self.evaluate_hamiltonian(ham_sig_vals) # B matrix + hamiltonian_matrix = -1j * self.evaluate_hamiltonian(ham_coefficients) # B matrix # if dissipators present (includes both hamiltonian is None and is not None) if self._dissipator_operators is not None or self._static_dissipators is not None: # A matrix if self._static_dissipators is None: - dissipators_matrix = np.tensordot( - -0.5 * dis_sig_vals, self._dissipator_products, axes=1 - ) + dissipators_matrix = _linear_combo(dis_coefficients, self._dissipator_products) elif self._dissipator_operators is None: dissipators_matrix = self._static_dissipators_product_sum else: - dissipators_matrix = self._static_dissipators_product_sum + np.tensordot( - -0.5 * dis_sig_vals, self._dissipator_products, axes=1 + dissipators_matrix = self._static_dissipators_product_sum + _linear_combo( + dis_coefficients, self._dissipator_products ) if hamiltonian_matrix is not None: - left_mult_contribution = np.matmul(hamiltonian_matrix + dissipators_matrix, y) - right_mult_contribution = np.matmul(y, dissipators_matrix - hamiltonian_matrix) + left_mult_contribution = _matmul(hamiltonian_matrix + dissipators_matrix, y) + right_mult_contribution = _matmul(y, dissipators_matrix - hamiltonian_matrix) else: - left_mult_contribution = np.matmul(dissipators_matrix, y) - right_mult_contribution = np.matmul(y, dissipators_matrix) + left_mult_contribution = _matmul(dissipators_matrix, y) + right_mult_contribution = _matmul(y, dissipators_matrix) if len(y.shape) == 3: # Must do array broadcasting and transposition to ensure vectorization works - y = np.broadcast_to(y, (1, y.shape[0], y.shape[1], y.shape[2])).transpose( + y = unp.broadcast_to(y, (1, y.shape[0], y.shape[1], y.shape[2])).transpose( [1, 0, 2, 3] ) if self._static_dissipators is None: - both_mult_contribution = np.tensordot( - dis_sig_vals, - np.matmul( - self._dissipator_operators, np.matmul(y, self._dissipator_operators_adj) - ), + both_mult_contribution = _numpy_multi_dispatch( + dis_coefficients, + _matmul(self._dissipator_operators, _matmul(y, self._dissipator_operators_adj)), + path="tensordot", axes=(-1, -3), ) elif self._dissipator_operators is None: - both_mult_contribution = np.sum( - np.matmul(self._static_dissipators, np.matmul(y, self._static_dissipators_adj)), + both_mult_contribution = unp.sum( + _matmul(self._static_dissipators, _matmul(y, self._static_dissipators_adj)), axis=-3, ) else: - both_mult_contribution = np.sum( - np.matmul(self._static_dissipators, np.matmul(y, self._static_dissipators_adj)), + both_mult_contribution = unp.sum( + _matmul(self._static_dissipators, _matmul(y, self._static_dissipators_adj)), axis=-3, - ) + np.tensordot( - dis_sig_vals, - np.matmul( - self._dissipator_operators, np.matmul(y, self._dissipator_operators_adj) - ), + ) + _numpy_multi_dispatch( + dis_coefficients, + _matmul(self._dissipator_operators, _matmul(y, self._dissipator_operators_adj)), + path="tensordot", axes=(-1, -3), ) return left_mult_contribution + right_mult_contribution + both_mult_contribution # if just hamiltonian elif hamiltonian_matrix is not None: - return np.dot(hamiltonian_matrix, y) - np.dot(y, hamiltonian_matrix) + return _matmul(hamiltonian_matrix, y) - _matmul(y, hamiltonian_matrix) else: raise QiskitError( - """DenseLindbladCollection with None for static_hamiltonian, + """LindbladCollection with None for static_hamiltonian, hamiltonian_operators, static_dissipators, and dissipator_operators, cannot evaluate rhs.""" ) + def __call__( + self, + ham_coefficients: Optional[ArrayLike], + dis_coefficients: Optional[ArrayLike], + y: Optional[ArrayLike], + ) -> ArrayLike: + """Call :meth:`~evaluate` or :meth:`~evaluate_rhs` depending on the presense of ``y``. + + Args: + ham_coefficients: The coefficients :math:`c_1` to use on the Hamiltonians. + dis_coefficients: The coefficients :math:`c_2` to use on the dissipators. + y: Optionally, the system state. + + Returns: + The evaluated function. + """ + if y is None: + return self.evaluate(ham_coefficients, dis_coefficients) + + return self.evaluate_rhs(ham_coefficients, dis_coefficients, y) -class SparseLindbladCollection(DenseLindbladCollection): - """Sparse version of DenseLindbladCollection.""" + +class ScipySparseLindbladCollection: + """Scipy sparse version of LindbladCollection.""" def __init__( self, - static_hamiltonian: Optional[Union[csr_matrix, Operator]] = None, - hamiltonian_operators: Optional[Union[List[csr_matrix], List[Operator]]] = None, - static_dissipators: Optional[Union[List[csr_matrix], List[Operator]]] = None, - dissipator_operators: Optional[Union[List[csr_matrix], List[Operator]]] = None, + static_hamiltonian: Optional[ArrayLike] = None, + hamiltonian_operators: Optional[ArrayLike] = None, + static_dissipators: Optional[ArrayLike] = None, + dissipator_operators: Optional[ArrayLike] = None, decimals: Optional[int] = 10, ): r"""Initializes sparse lindblad collection. @@ -742,120 +580,73 @@ def __init__( decimal place to avoid excess storage of near-zero values in sparse format. """ - self._decimals = decimals - super().__init__( - static_hamiltonian=static_hamiltonian, - hamiltonian_operators=hamiltonian_operators, - static_dissipators=static_dissipators, - dissipator_operators=dissipator_operators, - ) + if static_hamiltonian is not None: + self._static_hamiltonian = numpy_alias(like="scipy_sparse").asarray( + np.round(static_hamiltonian, decimals) + ) + else: + self._static_hamiltonian = None - @property - def static_hamiltonian(self) -> csr_matrix: - """The static part of the operator collection.""" - return self._static_hamiltonian + self._hamiltonian_operators = _to_csr_object_array(hamiltonian_operators, decimals) + self._static_dissipators = _to_csr_object_array(static_dissipators, decimals) + self._dissipator_operators = _to_csr_object_array(dissipator_operators, decimals) - @static_hamiltonian.setter - def static_hamiltonian(self, new_static_hamiltonian: Optional[csr_matrix] = None): - if new_static_hamiltonian is not None: - new_static_hamiltonian = np.round( - to_csr(new_static_hamiltonian), decimals=self._decimals + if static_dissipators is not None: + # setup adjoints + self._static_dissipators_adj = np.array( + [op.conj().transpose() for op in self._static_dissipators] ) - self._static_hamiltonian = new_static_hamiltonian - @property - def hamiltonian_operators(self) -> np.ndarray: - """The operators for the non-static part of Hamiltonian.""" + # pre-compute product and sum + self._static_dissipators_product_sum = -0.5 * np.sum( + self._static_dissipators_adj * self._static_dissipators, axis=0 + ) + + if self._dissipator_operators is not None: + # setup adjoints + self._dissipator_operators_adj = np.array( + [op.conj().transpose() for op in self._dissipator_operators] + ) + + # pre-compute products + self._dissipator_products = ( + -0.5 * self._dissipator_operators_adj * self._dissipator_operators + ) + + @property + def static_hamiltonian(self) -> Union[None, csr_matrix]: + """The static part of the operator collection.""" + return self._static_hamiltonian + + @property + def hamiltonian_operators(self) -> Union[None, list]: + """The operators for the non-static part of Hamiltonian.""" if self._hamiltonian_operators is None: return None return list(self._hamiltonian_operators) - @hamiltonian_operators.setter - def hamiltonian_operators(self, new_hamiltonian_operators: Optional[List[csr_matrix]] = None): - if new_hamiltonian_operators is not None: - new_hamiltonian_operators = to_csr(new_hamiltonian_operators) - new_hamiltonian_operators = [ - np.round(op, decimals=self._decimals) for op in new_hamiltonian_operators - ] - new_hamiltonian_operators = np.array(new_hamiltonian_operators, dtype="O") - - self._hamiltonian_operators = new_hamiltonian_operators - @property - def static_dissipators(self) -> Union[None, csr_matrix]: + def static_dissipators(self) -> Union[None, list]: """The operators for the static part of dissipator.""" if self._static_dissipators is None: return None return list(self._static_dissipators) - @static_dissipators.setter - def static_dissipators(self, new_static_dissipators: Optional[List[csr_matrix]] = None): - self._static_dissipators = None - if new_static_dissipators is not None: - # setup new dissipators - new_static_dissipators = to_csr(new_static_dissipators) - new_static_dissipators = [ - np.round(op, decimals=self._decimals) for op in new_static_dissipators - ] - - # setup adjoints - static_dissipators_adj = [op.conj().transpose() for op in new_static_dissipators] - - # wrap in object arrays - new_static_dissipators = np.array(new_static_dissipators, dtype="O") - static_dissipators_adj = np.array(static_dissipators_adj, dtype="O") - - # pre-compute products - static_dissipators_product_sum = -0.5 * np.sum( - static_dissipators_adj * new_static_dissipators, axis=0 - ) - - self._static_dissipators = new_static_dissipators - self._static_dissipators_adj = static_dissipators_adj - self._static_dissipators_product_sum = static_dissipators_product_sum - @property - def dissipator_operators(self) -> Union[None, List[csr_matrix]]: + def dissipator_operators(self) -> Union[None, list]: """The operators for the non-static part of dissipator.""" if self._dissipator_operators is None: return None return list(self._dissipator_operators) - @dissipator_operators.setter - def dissipator_operators(self, new_dissipator_operators: Optional[List[csr_matrix]] = None): - """Set up the dissipators themselves, as well as their adjoints, and the product of - adjoint with operator. - """ - self._dissipator_operators = None - if new_dissipator_operators is not None: - # setup new dissipators - new_dissipator_operators = to_csr(new_dissipator_operators) - new_dissipator_operators = [ - np.round(op, decimals=self._decimals) for op in new_dissipator_operators - ] - - # setup adjoints - dissipator_operators_adj = [op.conj().transpose() for op in new_dissipator_operators] - - # wrap in object arrays - new_dissipator_operators = np.array(new_dissipator_operators, dtype="O") - dissipator_operators_adj = np.array(dissipator_operators_adj, dtype="O") - - # pre-compute projducts - dissipator_products = dissipator_operators_adj * new_dissipator_operators - - self._dissipator_operators = new_dissipator_operators - self._dissipator_operators_adj = dissipator_operators_adj - self._dissipator_products = dissipator_products - - def evaluate_hamiltonian(self, ham_sig_vals: Union[None, Array]) -> csr_matrix: + def evaluate_hamiltonian(self, ham_coefficients: Optional[ArrayLike]) -> csr_matrix: r"""Compute the Hamiltonian. Args: - ham_sig_vals: [Real] values of :math:`s_j` in :math:`H = \sum_j s_j(t) H_j + H_d`. + ham_coefficients: The values of :math:`s_j` in :math:`H = \sum_j s_j(t) H_j + H_d`. Returns: Hamiltonian matrix. Raises: @@ -863,11 +654,11 @@ def evaluate_hamiltonian(self, ham_sig_vals: Union[None, Array]) -> csr_matrix: """ if self._static_hamiltonian is not None and self._hamiltonian_operators is not None: return ( - np.sum(ham_sig_vals * self._hamiltonian_operators, axis=-1) + np.sum(ham_coefficients * self._hamiltonian_operators, axis=-1) + self.static_hamiltonian ) elif self._static_hamiltonian is None and self._hamiltonian_operators is not None: - return np.sum(ham_sig_vals * self._hamiltonian_operators, axis=-1) + return np.sum(ham_coefficients * self._hamiltonian_operators, axis=-1) elif self._static_hamiltonian is not None: return self._static_hamiltonian else: @@ -877,15 +668,35 @@ def evaluate_hamiltonian(self, ham_sig_vals: Union[None, Array]) -> csr_matrix: hamiltonian_operators cannot evaluate Hamiltonian.""" ) + def evaluate( + self, ham_coefficients: Optional[ArrayLike], dis_coefficients: Optional[ArrayLike] + ) -> ArrayLike: + r"""Evaluate the function and return :math:`\Lambda(c_1, c_2, \cdot)`. + + Args: + ham_coefficients: The signals :math:`c_1` to use on the Hamiltonians. + dis_coefficients: The signals :math:`c_2` to use on the dissipators. + + Returns: + The evaluated function. + + Raises: + ValueError: Always. + """ + raise ValueError("Non-vectorized Lindblad collections cannot be evaluated without a state.") + def evaluate_rhs( - self, ham_sig_vals: Union[None, Array], dis_sig_vals: Union[None, Array], y: Array - ) -> Array: + self, + ham_coefficients: Optional[ArrayLike], + dis_coefficients: Optional[ArrayLike], + y: ArrayLike, + ) -> ArrayLike: r"""Evaluate the RHS of the Lindblad model for a given list of signal values. Args: - ham_sig_vals: Stores Hamiltonian signal values :math:`s_j(t)`. - dis_sig_vals: Stores dissipator signal values :math:`\gamma_j(t)`. Pass ``None`` if no - dissipator operators are involved. + ham_coefficients: Stores Hamiltonian signal values :math:`s_j(t)`. + dis_coefficients: Stores dissipator signal values :math:`\gamma_j(t)`. Pass ``None`` if + no dissipator operators are involved. y: Density matrix of the system, a ``(k,n,n)`` Array. Returns: @@ -923,7 +734,7 @@ def evaluate_rhs( """ hamiltonian_matrix = None if self._static_hamiltonian is not None or self._hamiltonian_operators is not None: - hamiltonian_matrix = -1j * self.evaluate_hamiltonian(ham_sig_vals) # B matrix + hamiltonian_matrix = -1j * self.evaluate_hamiltonian(ham_coefficients) # B matrix # package (n,n) Arrays as (1) # Arrays of dtype object, or (k,n,n) Arrays as (k,1) Arrays of dtype object @@ -933,14 +744,12 @@ def evaluate_rhs( if self._dissipator_operators is not None or self._static_dissipators is not None: # A matrix if self._static_dissipators is None: - dissipators_matrix = np.sum( - -0.5 * dis_sig_vals * self._dissipator_products, axis=-1 - ) + dissipators_matrix = np.sum(dis_coefficients * self._dissipator_products, axis=-1) elif self._dissipator_operators is None: dissipators_matrix = self._static_dissipators_product_sum else: dissipators_matrix = self._static_dissipators_product_sum + np.sum( - -0.5 * dis_sig_vals * self._dissipator_products, axis=-1 + dis_coefficients * self._dissipator_products, axis=-1 ) if hamiltonian_matrix is not None: @@ -953,7 +762,7 @@ def evaluate_rhs( # both_mult_contribution[i] = \gamma_i L_i\rho L_i^\dagger performed in array language if self._static_dissipators is None: both_mult_contribution = np.sum( - (dis_sig_vals * self._dissipator_operators) + (dis_coefficients * self._dissipator_operators) * y * self._dissipator_operators_adj, axis=-1, @@ -964,7 +773,7 @@ def evaluate_rhs( ) else: both_mult_contribution = np.sum( - (dis_sig_vals * self._dissipator_operators) + (dis_coefficients * self._dissipator_operators) * y * self._dissipator_operators_adj, axis=-1, @@ -976,7 +785,7 @@ def evaluate_rhs( out = (([hamiltonian_matrix] * y) - (y * [hamiltonian_matrix]))[0] else: raise QiskitError( - "SparseLindbladCollection with None for static_hamiltonian, " + "ScipySparseLindbladCollection with None for static_hamiltonian, " "hamiltonian_operators, and dissipator_operators, cannot evaluate rhs." ) if len(y.shape) == 2: @@ -986,233 +795,30 @@ def evaluate_rhs( return out - -class JAXSparseLindbladCollection(BaseLindbladOperatorCollection): - r"""Object for computing the right hand side of the Lindblad equation - with using jax.experimental.sparse.BCOO arrays. - """ - - @property - def static_hamiltonian(self) -> BCOO: - """The static part of the operator collection.""" - return self._static_hamiltonian - - @static_hamiltonian.setter - def static_hamiltonian(self, new_static_hamiltonian: Union[BCOO, None]): - self._static_hamiltonian = to_BCOO(new_static_hamiltonian) - - @property - def hamiltonian_operators(self) -> Union[BCOO, None]: - """The operators for the non-static part of Hamiltonian.""" - return self._hamiltonian_operators - - @hamiltonian_operators.setter - def hamiltonian_operators(self, new_hamiltonian_operators: Union[BCOO, None]): - self._hamiltonian_operators = to_BCOO(new_hamiltonian_operators) - - @property - def static_dissipators(self) -> Union[BCOO, None]: - """The operators for the static part of dissipator.""" - return self._static_dissipators - - @static_dissipators.setter - def static_dissipators(self, new_static_dissipators: Union[BCOO, None]): - self._static_dissipators = to_array(new_static_dissipators) - if self._static_dissipators is not None: - self._static_dissipators_adj = np.conjugate( - np.transpose(self._static_dissipators, [0, 2, 1]) - ).copy() - self._static_dissipators_product_sum = -0.5 * np.sum( - np.matmul(self._static_dissipators_adj, self._static_dissipators), axis=0 - ) - self._static_dissipators = jsparse.BCOO.fromdense( - self._static_dissipators.data, n_batch=1 - ) - self._static_dissipators_adj = jsparse.BCOO.fromdense( - self._static_dissipators_adj.data, n_batch=1 - ) - self._static_dissipators_product_sum = jsparse.BCOO.fromdense( - self._static_dissipators_product_sum.data - ) - - @property - def dissipator_operators(self) -> Union[BCOO, None]: - """The operators for the non-static part of dissipator.""" - return self._dissipator_operators - - @dissipator_operators.setter - def dissipator_operators(self, new_dissipator_operators: Union[BCOO, None]): - """Operators constructed using dense operations.""" - self._dissipator_operators = to_array(new_dissipator_operators) - if self._dissipator_operators is not None: - self._dissipator_operators_adj = np.conjugate( - np.transpose(self._dissipator_operators, [0, 2, 1]) - ).copy() - self._dissipator_products = np.matmul( - self._dissipator_operators_adj, self._dissipator_operators - ) - self._dissipator_operators = jsparse.BCOO.fromdense( - self._dissipator_operators.data, n_batch=1 - ) - self._dissipator_operators_adj = jsparse.BCOO.fromdense( - self._dissipator_operators_adj.data, n_batch=1 - ) - self._dissipator_products = -0.5 * jsparse.BCOO.fromdense( - self._dissipator_products.data, n_batch=1 - ) - - def evaluate(self, ham_sig_vals: Array, dis_sig_vals: Array) -> Array: - r"""Placeholder to return :math:`\Lambda(c_1, c_2, \cdot)`, which is not possible without - a state. + def __call__( + self, + ham_coefficients: Union[None, ArrayLike], + dis_coefficients: Union[None, ArrayLike], + y: Optional[ArrayLike], + ) -> ArrayLike: + """Call :meth:`~evaluate` or :meth:`~evaluate_rhs` depending on the presense of ``y``. Args: - ham_sig_vals: The signals :math:`c_1` to use on the Hamiltonians. - dis_sig_vals: The signals :math:`c_2` to use on the dissipators. + ham_coefficients: The signals :math:`c_1` to use on the Hamiltonians. + dis_coefficients: The signals :math:`c_2` to use on the dissipators. + y: Optionally, the system state. Returns: The evaluated function. - - Raises: - ValueError: Always. - """ - raise ValueError("Non-vectorized Lindblad collections cannot be evaluated without a state.") - - def evaluate_hamiltonian(self, ham_sig_vals: Union[BCOO, None]) -> BCOO: - r"""Compute the Hamiltonian. - - Args: - ham_sig_vals: [Real] values of :math:`s_j` in :math:`H = \sum_j s_j(t) H_j + H_d`. - - Returns: - Hamiltonian matrix. - - Raises: - QiskitError: If collection not sufficiently specified. - """ - if isinstance(ham_sig_vals, Array): - ham_sig_vals = ham_sig_vals.data - - if self._static_hamiltonian is not None and self._hamiltonian_operators is not None: - return ( - jsparse_linear_combo(ham_sig_vals, self._hamiltonian_operators) - + self._static_hamiltonian - ) - elif self._static_hamiltonian is None and self._hamiltonian_operators is not None: - return jsparse_linear_combo(ham_sig_vals, self._hamiltonian_operators) - elif self._static_hamiltonian is not None: - return self._static_hamiltonian - else: - raise QiskitError( - self.__class__.__name__ - + """ with None for both static_hamiltonian and - hamiltonian_operators cannot evaluate Hamiltonian.""" - ) - - @wrap - def evaluate_rhs( - self, ham_sig_vals: Union[None, Array], dis_sig_vals: Union[None, Array], y: Array - ) -> Array: - r"""Evaluates Lindblad equation RHS given a pair of signal values for the Hamiltonian terms - and the dissipator terms. - - Expresses the RHS of the Lindblad equation as :math:`(A+B)y + y(A-B) + C`, where - - .. math:: - A = (-1/2)*\sum_jD_j^\dagger D_j + (-1/2)*\sum_j\gamma_j(t) L_j^\dagger L_j, - - B = -iH, - - C = \sum_j \gamma_j(t) L_j y L_j^\dagger. - - Args: - ham_sig_vals: Hamiltonian coefficient values, :math:`s_j(t)`. - dis_sig_vals: Dissipator signal values, :math:`\gamma_j(t)`. - y: Density matrix as a `(n,n)` Array representing the state at time :math:`t`. - - Returns: - RHS of Lindblad equation - .. math:: - -i[H,y] + \sum_j\gamma_j(t)(L_j y L_j^\dagger - (1/2) * {L_j^\daggerL_j,y}). - - Raises: - QiskitError: If operator collection is underspecified. """ + if y is None: + return self.evaluate(ham_coefficients, dis_coefficients) - hamiltonian_matrix = None - if self._static_hamiltonian is not None or self._hamiltonian_operators is not None: - hamiltonian_matrix = -1j * self.evaluate_hamiltonian(ham_sig_vals) # B matrix - - # if dissipators present (includes both hamiltonian is None and is not None) - if self._dissipator_operators is not None or self._static_dissipators is not None: - # A matrix - if self._static_dissipators is None: - dissipators_matrix = jsparse_linear_combo(dis_sig_vals, self._dissipator_products) - elif self._dissipator_operators is None: - dissipators_matrix = self._static_dissipators_product_sum - else: - dissipators_matrix = self._static_dissipators_product_sum + jsparse_linear_combo( - dis_sig_vals, self._dissipator_products - ) - - if hamiltonian_matrix is not None: - left_mult_contribution = jsparse_matmul(hamiltonian_matrix + dissipators_matrix, y) - right_mult_contribution = jsparse_matmul( - y, dissipators_matrix + (-1 * hamiltonian_matrix) - ) - else: - left_mult_contribution = jsparse_matmul(dissipators_matrix, y) - right_mult_contribution = jsparse_matmul(y, dissipators_matrix) - - if len(y.shape) == 3: - # Must do array broadcasting and transposition to ensure vectorization works - y = jnp.broadcast_to(y, (1, y.shape[0], y.shape[1], y.shape[2])).transpose( - [1, 0, 2, 3] - ) - if self._static_dissipators is None: - both_mult_contribution = jnp.tensordot( - dis_sig_vals, - jsparse_triple_product( - self._dissipator_operators, y, self._dissipator_operators_adj - ), - axes=(-1, -3), - ) - elif self._dissipator_operators is None: - both_mult_contribution = jnp.sum( - jsparse_triple_product( - self._static_dissipators, y, self._static_dissipators_adj - ), - axis=-3, - ) - else: - both_mult_contribution = jnp.sum( - jsparse_triple_product( - self._static_dissipators, y, self._static_dissipators_adj - ), - axis=-3, - ) + jnp.tensordot( - dis_sig_vals, - jsparse_triple_product( - self._dissipator_operators, y, self._dissipator_operators_adj - ), - axes=(-1, -3), - ) - - out = left_mult_contribution + right_mult_contribution + both_mult_contribution - - return out - # if just hamiltonian - elif hamiltonian_matrix is not None: - return jsparse_matmul(hamiltonian_matrix, y) - jsparse_matmul(y, hamiltonian_matrix) - else: - raise QiskitError( - """JAXSparseLindbladCollection with None for static_hamiltonian, - hamiltonian_operators, static_dissipators, and - dissipator_operators, cannot evaluate rhs.""" - ) + return self.evaluate_rhs(ham_coefficients, dis_coefficients, y) -class BaseVectorizedLindbladCollection(BaseLindbladOperatorCollection, BaseOperatorCollection): - """Base class for Vectorized Lindblad collections. +class VectorizedLindbladCollection: + """Vectorized Lindblad collection class. The vectorized Lindblad equation represents the Lindblad master equation in the structure of a linear matrix differential equation in standard form. Hence, this class inherits @@ -1227,14 +833,17 @@ class BaseVectorizedLindbladCollection(BaseLindbladOperatorCollection, BaseOpera - ``evaluation_class``: Class property that returns the subclass of BaseOperatorCollection to be used when evaluating the model, e.g. DenseOperatorCollection or SparseOperatorCollection. + + This class works for ``array_library in ["numpy", "jax", "jax_sparse"]``. """ def __init__( self, - static_hamiltonian: Optional[Array] = None, - hamiltonian_operators: Optional[Array] = None, - static_dissipators: Optional[Array] = None, - dissipator_operators: Optional[Array] = None, + static_hamiltonian: Optional[ArrayLike] = None, + hamiltonian_operators: Optional[ArrayLike] = None, + static_dissipators: Optional[ArrayLike] = None, + dissipator_operators: Optional[ArrayLike] = None, + array_library: Optional[str] = None, ): r"""Initialize collection. @@ -1245,184 +854,221 @@ def __init__( as :math:`H(t) = \sum_j s(t) H_j+H_d` by specifying H_j. (k,n,n) array. static_dissipators: Dissipator terms with coefficient 1. dissipator_operators: the terms :math:`L_j` in Lindblad equation. (m,n,n) array. - """ - self._static_hamiltonian = None - self._hamiltonian_operators = None - self._static_dissipators = None - self._dissipator_operators = None - self._static_operator = None - self._operators = None - super().__init__( - static_hamiltonian=static_hamiltonian, - hamiltonian_operators=hamiltonian_operators, - static_dissipators=static_dissipators, - dissipator_operators=dissipator_operators, - ) - @abstractmethod - def convert_to_internal_type(self, obj: any) -> any: - """Convert either a single operator or a list of operators to an internal representation.""" + Raises: + QiskitError: If "scipy_sparse" is passed as array_library. + """ - @property - @abstractmethod - def evaluation_class(self) -> BaseOperatorCollection: - """The class used for evaluating the vectorized model or RHS.""" + self._array_library = array_library - @property - def static_hamiltonian(self) -> Union[Array, csr_matrix]: - """The static part of the operator collection.""" - return self._static_hamiltonian + if array_library == "scipy_sparse": + raise QiskitError("scipy_sparse is not a valid array_library for OperatorCollection.") - @static_hamiltonian.setter - def static_hamiltonian(self, new_static_hamiltonian: Optional[Union[Array, csr_matrix]] = None): - """Sets static_operator term.""" - self._static_hamiltonian = self.convert_to_internal_type(new_static_hamiltonian) - if self._static_hamiltonian is not None: + if static_hamiltonian is not None: + self._static_hamiltonian = self._convert_to_array_type(static_hamiltonian) self._vec_static_hamiltonian = vec_commutator(self._static_hamiltonian) + else: + self._static_hamiltonian = None - self.concatenate_static_operators() - - @property - def hamiltonian_operators(self) -> Array: - """The operators for the non-static part of Hamiltonian.""" - return self._hamiltonian_operators - - @hamiltonian_operators.setter - def hamiltonian_operators( - self, new_hamiltonian_operators: Optional[Union[Array, csr_matrix]] = None - ): - self._hamiltonian_operators = self.convert_to_internal_type(new_hamiltonian_operators) - if self._hamiltonian_operators is not None: + if hamiltonian_operators is not None: + self._hamiltonian_operators = self._convert_to_array_type(hamiltonian_operators) self._vec_hamiltonian_operators = vec_commutator(self._hamiltonian_operators) + else: + self._hamiltonian_operators = None - self.concatenate_operators() - - @property - def static_dissipators(self) -> Union[Array, List[csr_matrix]]: - """The operators for the static part of dissipator.""" - return self._static_dissipators - - @static_dissipators.setter - def static_dissipators( - self, new_static_dissipators: Optional[Union[Array, List[csr_matrix]]] = None - ): - self._static_dissipators = self.convert_to_internal_type(new_static_dissipators) - if self._static_dissipators is not None: - self._vec_static_dissipators_sum = np.sum( + if static_dissipators is not None: + self._static_dissipators = self._convert_to_array_type(static_dissipators) + self._vec_static_dissipators_sum = unp.sum( vec_dissipator(self._static_dissipators), axis=0 ) + else: + self._static_dissipators = None - self.concatenate_static_operators() - - @property - def dissipator_operators(self) -> Union[Array, List[csr_matrix]]: - """The operators for the non-static part of dissipator.""" - return self._dissipator_operators - - @dissipator_operators.setter - def dissipator_operators( - self, new_dissipator_operators: Optional[Union[Array, List[csr_matrix]]] = None - ): - self._dissipator_operators = self.convert_to_internal_type(new_dissipator_operators) - if self._dissipator_operators is not None: + if dissipator_operators is not None: + self._dissipator_operators = self._convert_to_array_type(dissipator_operators) self._vec_dissipator_operators = vec_dissipator(self._dissipator_operators) + else: + self._dissipator_operators = None - self.concatenate_operators() - - def concatenate_static_operators(self): - """Concatenate static hamiltonian and static dissipators.""" + # concatenate static operators if self._static_hamiltonian is not None and self._static_dissipators is not None: - self._static_operator = self._vec_static_hamiltonian + self._vec_static_dissipators_sum + static_operator = self._vec_static_hamiltonian + self._vec_static_dissipators_sum elif self._static_hamiltonian is None and self._static_dissipators is not None: - self._static_operator = self._vec_static_dissipators_sum + static_operator = self._vec_static_dissipators_sum elif self._static_hamiltonian is not None and self._static_dissipators is None: - self._static_operator = self._vec_static_hamiltonian + static_operator = self._vec_static_hamiltonian else: - self._static_operator = None + static_operator = None - def concatenate_operators(self): - """Concatenate hamiltonian operators and dissipator operators.""" + # concatenate non-static operators if self._hamiltonian_operators is not None and self._dissipator_operators is not None: - self._operators = np.append( + operators = unp.append( self._vec_hamiltonian_operators, self._vec_dissipator_operators, axis=0 ) elif self._hamiltonian_operators is not None and self._dissipator_operators is None: - self._operators = self._vec_hamiltonian_operators + operators = self._vec_hamiltonian_operators elif self._hamiltonian_operators is None and self._dissipator_operators is not None: - self._operators = self._vec_dissipator_operators + operators = self._vec_dissipator_operators else: - self._operators = None + operators = None - def concatenate_signals( - self, ham_sig_vals: Union[None, Array], dis_sig_vals: Union[None, Array] - ) -> Array: - """Concatenate hamiltonian and linblad signals.""" - if self._hamiltonian_operators is not None and self._dissipator_operators is not None: - return _numpy_multi_dispatch(ham_sig_vals, dis_sig_vals, path="append", axis=-1) - if self._hamiltonian_operators is not None and self._dissipator_operators is None: - return ham_sig_vals - if self._hamiltonian_operators is None and self._dissipator_operators is not None: - return dis_sig_vals + # build internally used operator collection + self._operator_collection = self._construct_operator_collection( + static_operator=static_operator, operators=operators + ) - return None + @property + def static_hamiltonian(self) -> Union[ArrayLike, None]: + """The static part of the operator collection.""" + return self._static_hamiltonian + + @property + def hamiltonian_operators(self) -> Union[ArrayLike, None]: + """The operators for the non-static part of Hamiltonian.""" + return self._hamiltonian_operators + + @property + def static_dissipators(self) -> Union[ArrayLike, None]: + """The operators for the static part of dissipator.""" + return self._static_dissipators + + @property + def dissipator_operators(self) -> Union[ArrayLike, None]: + """The operators for the non-static part of dissipator.""" + return self._dissipator_operators + + def evaluate_hamiltonian(self, ham_coefficients: Union[None, ArrayLike]) -> ArrayLike: + r"""Evaluate the Hamiltonian of the model. + + Args: + ham_coefficients: The values of :math:`s_j` in :math:`H = \sum_j s_j(t) H_j + H_d`. + + Returns: + The Hamiltonian. - def evaluate(self, ham_sig_vals: Union[None, Array], dis_sig_vals: Union[None, Array]) -> Array: + Raises: + QiskitError: If collection not sufficiently specified. + """ + if self._static_hamiltonian is not None and self._hamiltonian_operators is not None: + return ( + _linear_combo(ham_coefficients, self._hamiltonian_operators) + + self._static_hamiltonian + ) + elif self._static_hamiltonian is None and self._hamiltonian_operators is not None: + return _linear_combo(ham_coefficients, self._hamiltonian_operators) + elif self._static_hamiltonian is not None: + return self._static_hamiltonian + else: + raise QiskitError( + self.__class__.__name__ + + """ with None for both static_hamiltonian and + hamiltonian_operators cannot evaluate Hamiltonian.""" + ) + + def evaluate( + self, ham_coefficients: Optional[ArrayLike], dis_coefficients: Optional[ArrayLike] + ) -> ArrayLike: r"""Compute and return :math:`\Lambda(c_1, c_2, \cdot)`. Args: - ham_sig_vals: The signals :math:`c_1` to use on the Hamiltonians. - dis_sig_vals: The signals :math:`c_2` to use on the dissipators. + ham_coefficients: The signal coefficients :math:`c_1` to use on the Hamiltonians. + dis_coefficients: The signal coefficients :math:`c_2` to use on the dissipators. Returns: The evaluated function. """ - signal_values = self.concatenate_signals(ham_sig_vals, dis_sig_vals) - return self.evaluation_class.evaluate(self, signal_values) + coeffs = self._concatenate_coefficients(ham_coefficients, dis_coefficients) + return self._operator_collection.evaluate(coeffs) - def evaluate_rhs(self, ham_sig_vals: Array, dis_sig_vals: Array, y: Array) -> Array: + def evaluate_rhs( + self, + ham_coefficients: Optional[ArrayLike], + dis_coefficients: Optional[ArrayLike], + y: ArrayLike, + ) -> ArrayLike: r"""Evaluates the RHS of the Lindblad equation using vectorized maps. Args: - ham_sig_vals: Hamiltonian signal coefficients. - dis_sig_vals: Dissipator signal coefficients. If none involved, pass ``None``. + ham_coefficients: Hamiltonian signal coefficients. + dis_coefficients: Dissipator signal coefficients. If none involved, pass ``None``. y: Density matrix represented as a vector using column-stacking convention. Returns: Vectorized RHS of Lindblad equation :math:`\dot{\rho}` in column-stacking convention. """ - return self.evaluate(ham_sig_vals, dis_sig_vals) @ y + coeffs = self._concatenate_coefficients(ham_coefficients, dis_coefficients) + return self._operator_collection.evaluate_rhs(coeffs, y) + def __call__( + self, + ham_coefficients: Optional[ArrayLike], + dis_coefficients: Optional[ArrayLike], + y: Optional[ArrayLike], + ) -> ArrayLike: + """Call :meth:`~evaluate` or :meth:`~evaluate_rhs` depending on the presense of ``y``. -class DenseVectorizedLindbladCollection( - BaseVectorizedLindbladCollection, DenseLindbladCollection, DenseOperatorCollection -): - r"""Vectorized version of :class:`DenseLindbladCollection`. + Args: + ham_coefficients: The signal coefficients :math:`c_1` to use on the Hamiltonians. + dis_coefficients: The signal coefficients :math:`c_2` to use on the dissipators. + y: Optionally, the system state. - Utilizes :class:`BaseVectorizedLindbladCollection` for property handling, - :class:`DenseLindbladCollection` for ``evaluate_hamiltonian``, and - :class:`DenseOperatorCollection` for operator property handling. - """ + Returns: + The evaluated function. + """ + if y is None: + return self.evaluate(ham_coefficients, dis_coefficients) - def convert_to_internal_type(self, obj: any) -> Array: - return to_array(obj) + return self.evaluate_rhs(ham_coefficients, dis_coefficients, y) - @property - def evaluation_class(self): + def _convert_to_array_type(self, obj: Any) -> ArrayLike: + return numpy_alias(like=self._array_library).asarray(obj) + + def _construct_operator_collection(self, *args, **kwargs): """The class used for evaluating the vectorized model or RHS.""" - return DenseOperatorCollection + return OperatorCollection(*args, **kwargs, array_library=self._array_library) + + def _concatenate_coefficients(self, ham_coefficients, dis_coefficients): + if self._hamiltonian_operators is not None and self._dissipator_operators is not None: + return _numpy_multi_dispatch(ham_coefficients, dis_coefficients, path="append", axis=-1) + if self._hamiltonian_operators is not None and self._dissipator_operators is None: + return ham_coefficients + if self._hamiltonian_operators is None and self._dissipator_operators is not None: + return dis_coefficients + return None -class SparseVectorizedLindbladCollection( - BaseVectorizedLindbladCollection, SparseLindbladCollection, SparseOperatorCollection -): - r"""Vectorized version of :class:`SparseLindbladCollection`. - Utilizes :class:`BaseVectorizedLindbladCollection` for property handling, - :class:`SparseLindbladCollection` for evaluate_hamiltonian, and - :class:`SparseOperatorCollection` for static_operator and operator property handling. - """ +class ScipySparseVectorizedLindbladCollection(VectorizedLindbladCollection): + r"""Scipy sparse version of VectorizedLindbladCollection.""" - def convert_to_internal_type(self, obj: any) -> Union[csr_matrix, List[csr_matrix]]: + def __init__( + self, + static_hamiltonian: Optional[ArrayLike] = None, + hamiltonian_operators: Optional[ArrayLike] = None, + static_dissipators: Optional[ArrayLike] = None, + dissipator_operators: Optional[ArrayLike] = None, + decimals: Optional[int] = 10, + ): + r"""Initialize collection. + + Args: + static_hamiltonian: Constant term :math:`H_d` to be added to the Hamiltonian of the + system. + hamiltonian_operators: Specifies breakdown of Hamiltonian + as :math:`H(t) = \sum_j s(t) H_j+H_d` by specifying H_j. (k,n,n) array. + static_dissipators: Dissipator terms with coefficient 1. + dissipator_operators: the terms :math:`L_j` in Lindblad equation. (m,n,n) array. + decimals: Decimals to round the sparse operators to 0. + """ + self._decimals = decimals + super().__init__( + static_hamiltonian=static_hamiltonian, + hamiltonian_operators=hamiltonian_operators, + static_dissipators=static_dissipators, + dissipator_operators=dissipator_operators, + ) + + def _convert_to_array_type(self, obj: any) -> Union[csr_matrix, List[csr_matrix]]: if obj is None: return None @@ -1432,61 +1078,12 @@ def convert_to_internal_type(self, obj: any) -> Union[csr_matrix, List[csr_matri else: return [np.round(sub_obj, decimals=self._decimals) for sub_obj in obj] - @property - def evaluation_class(self): - """The class used for evaluating the vectorized model or RHS.""" - return SparseOperatorCollection - - -class JAXSparseVectorizedLindbladCollection( - BaseVectorizedLindbladCollection, JAXSparseLindbladCollection, JAXSparseOperatorCollection -): - r"""Vectorized version of :class:`JAXSparseLindbladCollection`. - - Utilizes :class:`BaseVectorizedLindbladCollection` for property handling, - :class:`JAXSparseLindbladCollection` for evaluate_hamiltonian, and - :class:`JAXSparseOperatorCollection` for static_operator and operator property handling. - """ - - def convert_to_internal_type(self, obj: any) -> BCOO: - return to_BCOO(obj) - - @property - def evaluation_class(self): + def _construct_operator_collection(self, *args, **kwargs): """The class used for evaluating the vectorized model or RHS.""" - return JAXSparseOperatorCollection - - def concatenate_static_operators(self): - """Override base class to convert to BCOO again at the end. The vectorization operations are - not implemented for BCOO type, so they automatically get converted to Arrays, and hence need - to be converted back. - """ - super().concatenate_static_operators() - self._static_operator = self.convert_to_internal_type(self._static_operator) - - def concatenate_operators(self): - """Override base class to convert to BCOO again at the end. The vectorization operations are - not implemented for BCOO type, so they automatically get converted to Arrays, and hence need - to be converted back. - """ - super().concatenate_operators() - self._operators = self.convert_to_internal_type(self._operators) - - @wrap - def evaluate_rhs(self, ham_sig_vals: Array, dis_sig_vals: Array, y: Array) -> Array: - r"""Evaluate the function and return :math:`\Lambda(c, y) = (G_d + \sum_jc_jG_j) y`. - - Args: - ham_sig_vals: The signals :math:`c` to use on the Hamiltonians. - dis_sig_vals: The signals :math:`c` to use on the dissipators. - y: The system state. - Returns: - The evaluated function. - """ - return jsparse_matmul(self.evaluate(ham_sig_vals, dis_sig_vals), y) + return ScipySparseOperatorCollection(*args, **kwargs) -def package_density_matrices(y: Array) -> Array: +def package_density_matrices(y: ArrayLike) -> ArrayLike: """Sends an array ``y`` of density matrices to a ``(1,)`` array of dtype object, where entry ``[0]`` is ``y``. Formally avoids for-loops through vectorization. @@ -1505,7 +1102,7 @@ def package_density_matrices(y: Array) -> Array: package_density_matrices = np.vectorize(package_density_matrices, signature="(n,n)->(1)") -def unpackage_density_matrices(y: Array) -> Array: +def unpackage_density_matrices(y: ArrayLike) -> ArrayLike: """Inverse function of :func:`package_density_matrices`. Since this function is much slower than packaging, avoid it unless absolutely needed (as in case diff --git a/qiskit_dynamics/type_utils.py b/qiskit_dynamics/type_utils.py index bd1d9c748..a0f67be86 100644 --- a/qiskit_dynamics/type_utils.py +++ b/qiskit_dynamics/type_utils.py @@ -277,7 +277,7 @@ def vec_commutator( # taken to be 1d array of 2d sparse matrices sp_iden = sparse_identity(A[0].shape[-1], format="csr") out = [-1j * (sparse_kron(sp_iden, mat) - sparse_kron(mat.T, sp_iden)) for mat in A] - return out + return np.array(out) A = to_array(A) iden = Array(np.eye(A.shape[-1])) @@ -317,7 +317,7 @@ def vec_dissipator( * (sparse_kron(sp_iden, mat.conj().T * mat) + sparse_kron(mat.T * mat.conj(), sp_iden)) for mat in L ] - return out + return np.array(out) iden = Array(np.eye(L.shape[-1])) axes = list(range(L.ndim)) diff --git a/test/dynamics/arraylias/test_alias.py b/test/dynamics/arraylias/test_alias.py index c342d0b3a..0ffc0ae4e 100644 --- a/test/dynamics/arraylias/test_alias.py +++ b/test/dynamics/arraylias/test_alias.py @@ -25,8 +25,15 @@ from qiskit_dynamics import DYNAMICS_NUMPY_ALIAS from qiskit_dynamics import DYNAMICS_NUMPY as unp from qiskit_dynamics import DYNAMICS_SCIPY as usp +from qiskit_dynamics.arraylias.alias import _preferred_lib, _numpy_multi_dispatch -from ..common import test_array_backends +from ..common import test_array_backends, JAXTestBase + +try: + import jax.numpy as jnp + from jax.experimental.sparse import BCOO +except ImportError: + pass @partial(test_array_backends, array_libraries=["numpy", "jax", "array_numpy", "array_jax"]) @@ -57,8 +64,8 @@ def test_simple_case(self): self.assertAllClose(output, expected) -class TestDynamicsNumpyAliasType(unittest.TestCase): - """Test cases for which types are registered in dynamics_numpy_alias.""" +class TestDynamicsNumpyScipySparseType(unittest.TestCase): + """Test case verifying scipy sparse types are registered in dynamics_numpy_alias.""" def test_spmatrix_type(self): """Test spmatrix is registered as scipy_sparse.""" @@ -66,13 +73,50 @@ def test_spmatrix_type(self): registered_type_name = "scipy_sparse" self.assertTrue(registered_type_name in DYNAMICS_NUMPY_ALIAS.infer_libs(sp_matrix)) + +class TestDynamicsNumpyJAXBCOOType(JAXTestBase): + """Test case verifying JAX sparse types are registered in dynamics_numpy_alias.""" + def test_bcoo_type(self): - """Test bcoo is registered.""" - try: - from jax.experimental.sparse import BCOO - - bcoo = BCOO.fromdense([[0.0, 1.0], [1.0, 0.0]]) - registered_type_name = "jax_sparse" - self.assertTrue(registered_type_name in DYNAMICS_NUMPY_ALIAS.infer_libs(bcoo)[0]) - except ImportError as err: - raise unittest.SkipTest("Skipping jax tests.") from err + """Test that BCOO type is correctly registered.""" + bcoo = BCOO.fromdense([[0.0, 1.0], [1.0, 0.0]]) + registered_type_name = "jax_sparse" + self.assertTrue(registered_type_name in DYNAMICS_NUMPY_ALIAS.infer_libs(bcoo)[0]) + + +@partial(test_array_backends, array_libraries=["numpy", "jax", "jax_sparse"]) +class Test_linear_combo: + """Test registered linear_combo function.""" + + def test_simple_case(self): + """Simple test case for linear combo.""" + mats = self.asarray([[[0.0, 1.0], [1.0, 0.0]], [[1j, 0.0], [0.0, -1j]]]) + coeffs = np.array([1.0, 2.0]) + out = _numpy_multi_dispatch(coeffs, mats, path="linear_combo") + self.assertArrayType(out) + self.assertAllClose(out, np.array([[2j, 1.0], [1.0, -2j]])) + + +class Test_preferred_lib(JAXTestBase): + """Test class for preferred_lib functions. Inherits from JAXTestBase as this functionality + is primarily to facilitate JAX types. + """ + + def test_defaults_to_numpy(self): + """Test that it defaults to numpy.""" + self.assertEqual(_preferred_lib(None), "numpy") + + def test_prefers_jax_over_numpy(self): + """Test that it chooses jax over numpy.""" + self.assertEqual(_preferred_lib(1.0), "numpy") + self.assertEqual(_preferred_lib(jnp.array(1.0), 1.0), "jax") + + def test_prefers_jax_sparse_over_numpy(self): + """Test that it prefers jax_sparse over numpy.""" + self.assertEqual(_preferred_lib(np.array(1.0)), "numpy") + self.assertEqual(_preferred_lib(np.array(1.0), BCOO.fromdense([1.0, 2.0])), "jax_sparse") + + def test_prefers_jax_sparse_over_jax(self): + """Test that it prefers jax_sparse over jax.""" + self.assertEqual(_preferred_lib(jnp.array(1.0)), "jax") + self.assertEqual(_preferred_lib(jnp.array(1.0), BCOO.fromdense([1.0, 2.0])), "jax_sparse") diff --git a/test/dynamics/common.py b/test/dynamics/common.py index 782148aa6..c504403c2 100644 --- a/test/dynamics/common.py +++ b/test/dynamics/common.py @@ -52,6 +52,8 @@ except ImportError: pass +from qiskit_dynamics import DYNAMICS_NUMPY_ALIAS + from qiskit_dynamics.array import Array, wrap @@ -60,6 +62,11 @@ class QiskitDynamicsTestCase(unittest.TestCase): def assertAllClose(self, A, B, rtol=1e-8, atol=1e-8): """Call np.allclose and assert true.""" + if any("sparse" in x for x in DYNAMICS_NUMPY_ALIAS.infer_libs(A)): + A = A.todense() + if any("sparse" in x for x in DYNAMICS_NUMPY_ALIAS.infer_libs(B)): + B = B.todense() + A = np.array(A) B = np.array(B) @@ -181,6 +188,19 @@ def assertArrayType(self, a): """Assert the correct array type.""" return type(a).__name__ == "BCOO" + def jit_grad(self, func_to_test: Callable) -> Callable: + """Tests whether a function can be graded. Converts + all functions to scalar, real functions if they are not + already. + Args: + func_to_test: The function whose gradient will be graded. + Returns: + JIT-compiled gradient of function. + """ + wf = lambda f: jit(grad(f)) + f = lambda *args: np.sum(func_to_test(*args)).real + return wf(f) + class ArrayNumpyTestBase(QiskitDynamicsTestCase): """Base class for tests working with qiskit_dynamics Arrays with numpy backend.""" diff --git a/test/dynamics/models/test_operator_collections.py b/test/dynamics/models/test_operator_collections.py index 8c99faea9..6d6984b5c 100644 --- a/test/dynamics/models/test_operator_collections.py +++ b/test/dynamics/models/test_operator_collections.py @@ -9,10 +9,12 @@ # 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. -# pylint: disable=invalid-name +# pylint: disable=invalid-name,no-member """Tests for operator_collections.py.""" +from functools import partial + import numpy as np import numpy.random as rand from scipy.sparse import issparse @@ -21,117 +23,137 @@ from qiskit import QiskitError from qiskit.quantum_info.operators import Operator from qiskit_dynamics.models.operator_collections import ( - DenseOperatorCollection, - DenseLindbladCollection, - DenseVectorizedLindbladCollection, - SparseLindbladCollection, - JAXSparseLindbladCollection, - SparseOperatorCollection, - JAXSparseOperatorCollection, - SparseVectorizedLindbladCollection, - JAXSparseVectorizedLindbladCollection, + OperatorCollection, + ScipySparseOperatorCollection, + LindbladCollection, + ScipySparseLindbladCollection, + VectorizedLindbladCollection, + ScipySparseVectorizedLindbladCollection, ) from qiskit_dynamics.array import Array from qiskit_dynamics.type_utils import to_array -from ..common import QiskitDynamicsTestCase, TestJaxBase - -try: - from jax.experimental import sparse as jsparse -except ImportError: - pass +from ..common import test_array_backends, QiskitDynamicsTestCase -class TestDenseOperatorCollection(QiskitDynamicsTestCase): - """Tests for DenseOperatorCollection.""" +@partial(test_array_backends, array_libraries=["numpy", "jax", "jax_sparse"]) +class TestOperatorCollection: + """Test cases for OperatorCollection.""" def setUp(self): - self.X = Array(Operator.from_label("X").data) - self.Y = Array(Operator.from_label("Y").data) - self.Z = Array(Operator.from_label("Z").data) + """Build a simple OperatorCollection.""" + self.X = Operator.from_label("X").data + self.Y = Operator.from_label("Y").data + self.Z = Operator.from_label("Z").data - self.test_operator_list = Array([self.X, self.Y, self.Z]) - self.simple_collection = DenseOperatorCollection( - operators=self.test_operator_list, static_operator=None + self.test_operator_list = [self.X, self.Y, self.Z] + self.simple_collection = OperatorCollection( + operators=self.test_operator_list, + static_operator=None, + array_library=self.array_library(), ) def test_empty_collection_error(self): """Verify that evaluating with no operators or static_operator raises an error.""" - collection = DenseOperatorCollection(operators=None, static_operator=None) + collection = OperatorCollection(operators=None, static_operator=None) with self.assertRaisesRegex(QiskitError, "cannot be evaluated."): collection(None) def test_known_values_basic_functionality(self): - """Test DenseOperatorCollection evaluation against - analytically known values.""" + """Test OperatorCollection evaluation against analytically known values.""" rand.seed(34983) coeffs = rand.uniform(-1, 1, 3) res = self.simple_collection(coeffs) + self.assertArrayType(res) self.assertAllClose(res, coeffs[0] * self.X + coeffs[1] * self.Y + coeffs[2] * self.Z) res = ( - DenseOperatorCollection(operators=self.test_operator_list, static_operator=np.eye(2)) + OperatorCollection( + operators=self.test_operator_list, + static_operator=np.eye(2), + array_library=self.array_library(), + ) )(coeffs) + + self.assertArrayType(res) self.assertAllClose( res, np.eye(2) + coeffs[0] * self.X + coeffs[1] * self.Y + coeffs[2] * self.Z ) def test_basic_functionality_pseudorandom(self): - """Test DenseOperatorCollection evaluation - using pseudorandom arrays.""" + """Test OperatorCollection evaluation using pseudorandom arrays.""" rand.seed(342) vals = rand.uniform(-1, 1, 32) + 1j * rand.uniform(-1, 1, (10, 32)) arr = rand.uniform(-1, 1, (32, 128, 128)) - res = (DenseOperatorCollection(operators=arr))(vals) + collection = OperatorCollection(operators=arr, array_library=self.array_library()) for i in range(10): + res = collection(vals[i]) + self.assertArrayType(res) total = 0 for j in range(32): total = total + vals[i, j] * arr[j] - self.assertAllClose(res[i], total) + self.assertAllClose(res, total) def test_static_collection(self): """Test the case in which only a static operator is present.""" - collection = DenseOperatorCollection(operators=None, static_operator=self.X) + collection = OperatorCollection( + operators=None, static_operator=self.X, array_library=self.array_library() + ) self.assertAllClose(self.X, collection(None)) + def test_collection_with_no_explicit_array_library(self): + """Test when array_library is not explicitly passed.""" + rand.seed(342) + vals = rand.uniform(-1, 1, 32) + 1j * rand.uniform(-1, 1, (10, 32)) + arr = rand.uniform(-1, 1, (32, 128, 128)) + collection = OperatorCollection(operators=self.asarray(arr)) + for i in range(10): + res = collection(vals[i]) + self.assertArrayType(res) + total = 0 + for j in range(32): + total = total + vals[i, j] * arr[j] + self.assertAllClose(res, total) + -class TestDenseOperatorCollectionJax(TestDenseOperatorCollection, TestJaxBase): - """Jax version of TestDenseOperatorCollection tests. +@partial(test_array_backends, array_libraries=["jax", "jax_sparse"]) +class TestOperatorCollectionJAXTransformations: + """Test JAX transformations applied to OperatorCollection evaluation methods.""" - Note: This class has more tests due to inheritance. - """ + def setUp(self): + """Build simple OperatorCollection instance.""" + self.X = Operator.from_label("X").data + self.Y = Operator.from_label("Y").data + self.Z = Operator.from_label("Z").data - def test_functions_jitable(self): - """Tests that all class functions are jittable.""" - doc = DenseOperatorCollection( - operators=Array(self.test_operator_list), - static_operator=Array(self.test_operator_list[0]), + self.test_operator_list = [self.X, self.Y, self.Z] + self.simple_collection = OperatorCollection( + operators=self.test_operator_list, + static_operator=None, + array_library=self.array_library(), ) - rand.seed(3423) - coeffs = rand.uniform(-1, 1, 3) - self.jit_wrap(doc.evaluate)(Array(coeffs)) - self.jit_wrap(doc.evaluate_rhs)(Array(coeffs), self.X) def test_functions_gradable(self): """Tests that all class functions are gradable.""" - doc = DenseOperatorCollection( - operators=Array(self.test_operator_list), - static_operator=Array(self.test_operator_list[0]), + doc = OperatorCollection( + operators=self.test_operator_list, + static_operator=self.test_operator_list[0], + array_library=self.array_library(), ) rand.seed(5433) coeffs = rand.uniform(-1, 1, 3) - self.jit_grad_wrap(doc.evaluate)(Array(coeffs)) - self.jit_grad_wrap(doc.evaluate_rhs)(Array(coeffs), self.X) + self.jit_grad(doc.evaluate)(coeffs) + self.jit_grad(doc.evaluate_rhs)(coeffs, self.X) -class TestSparseOperatorCollection(QiskitDynamicsTestCase): - """Tests for SparseOperatorCollection.""" +class TestScipySparseOperatorCollection(QiskitDynamicsTestCase): + """Tests for ScipySparseOperatorCollection.""" def test_empty_collection_error(self): """Verify that evaluating with no operators or static_operator raises an error.""" - collection = SparseOperatorCollection(operators=None, static_operator=None) + collection = ScipySparseOperatorCollection(operators=None, static_operator=None) with self.assertRaisesRegex(QiskitError, "cannot be evaluated."): collection(None) @@ -141,7 +163,7 @@ def test_empty_collection_error(self): def test_evaluate_simple_case(self): """Simple test case.""" - collection = SparseOperatorCollection(operators=[np.eye(2), [[0.0, 1.0], [1.0, 0.0]]]) + collection = ScipySparseOperatorCollection(operators=[np.eye(2), [[0.0, 1.0], [1.0, 0.0]]]) value = collection(np.array([1.0, 2.0])) self.assertTrue(issparse(value)) @@ -165,8 +187,10 @@ def test_consistency_with_dense_pseudorandom(self): mat = r(4, 16, 16) sigVals = r(4) static_operator = r(16, 16) - dense_collection = DenseOperatorCollection(operators=mat, static_operator=static_operator) - sparse_collection = SparseOperatorCollection(operators=mat, static_operator=static_operator) + dense_collection = OperatorCollection(operators=mat, static_operator=static_operator) + sparse_collection = ScipySparseOperatorCollection( + operators=mat, static_operator=static_operator + ) dense_val = dense_collection(sigVals) sparse_val = sparse_collection(sigVals) self.assertAllClose(dense_val, sparse_val.toarray()) @@ -186,13 +210,13 @@ def test_constructor_takes_operators(self): ham_ops_alt.append(Array(op)) sigVals = r(4) static_operator_numpy_array = r(3, 3) - sparse_collection_operator_list = SparseOperatorCollection( + sparse_collection_operator_list = ScipySparseOperatorCollection( operators=ham_ops, static_operator=Operator(static_operator_numpy_array) ) - sparse_collection_array_list = SparseOperatorCollection( + sparse_collection_array_list = ScipySparseOperatorCollection( operators=ham_ops_alt, static_operator=to_array(static_operator_numpy_array) ) - sparse_collection_pure_array = SparseOperatorCollection( + sparse_collection_pure_array = ScipySparseOperatorCollection( operators=to_array(ham_ops), static_operator=to_array(static_operator_numpy_array) ) a = sparse_collection_operator_list(sigVals) @@ -204,114 +228,32 @@ def test_constructor_takes_operators(self): def test_static_collection(self): """Test the case in which only a static operator is present.""" X = csr_matrix([[0.0, 1.0], [1.0, 0.0]]) - collection = SparseOperatorCollection(static_operator=X) + collection = ScipySparseOperatorCollection(static_operator=X) self.assertAllCloseSparse(X, collection(None)) self.assertAllClose(np.array([0.0, 1.0]), collection(None, np.array([1.0, 0.0]))) -class TestJAXSparseOperatorCollection(QiskitDynamicsTestCase, TestJaxBase): - """Test cases for JAXSparseOperatorCollection.""" - - def setUp(self): - self.X = Array(Operator.from_label("X").data) - self.Y = Array(Operator.from_label("Y").data) - self.Z = Array(Operator.from_label("Z").data) - - self.test_operator_list = Array([self.X, self.Y, self.Z]) - self.simple_collection = JAXSparseOperatorCollection( - operators=self.test_operator_list, static_operator=None - ) - - def test_empty_collection_error(self): - """Verify that evaluating with no operators or static_operator raises an error.""" - - collection = JAXSparseOperatorCollection(operators=None, static_operator=None) - with self.assertRaisesRegex(QiskitError, "cannot be evaluated."): - collection(None) - - def test_known_values_basic_functionality(self): - """Test JAXSparseOperatorCollection evaluation against - analytically known values.""" - rand.seed(34983) - coeffs = rand.uniform(-1, 1, 3) - - res = self.simple_collection(coeffs) - self.assertTrue(isinstance(res, jsparse.BCOO)) - self.assertAllClose( - res.todense(), coeffs[0] * self.X + coeffs[1] * self.Y + coeffs[2] * self.Z - ) - - res = ( - JAXSparseOperatorCollection( - operators=self.test_operator_list, static_operator=np.eye(2) - ) - )(coeffs) - self.assertTrue(isinstance(res, jsparse.BCOO)) - self.assertAllClose( - res.todense(), np.eye(2) + coeffs[0] * self.X + coeffs[1] * self.Y + coeffs[2] * self.Z - ) - - def test_basic_functionality_pseudorandom(self): - """Test JAXSparseOperatorCollection evaluation - using pseudorandom arrays.""" - rand.seed(342) - vals = rand.uniform(-1, 1, 32) + 1j * rand.uniform(-1, 1, (10, 32)) - arr = rand.uniform(-1, 1, (32, 128, 128)) - collection = JAXSparseOperatorCollection(operators=arr) - for i in range(10): - res = collection(vals[i]) - total = 0 - for j in range(32): - total = total + vals[i, j] * arr[j] - self.assertTrue(isinstance(res, jsparse.BCOO)) - self.assertAllClose(res.todense(), total) - - def test_static_collection(self): - """Test the case in which only a static operator is present.""" - collection = JAXSparseOperatorCollection(operators=None, static_operator=self.X) - self.assertTrue(isinstance(collection(None), jsparse.BCOO)) - self.assertAllClose(self.X, collection(None).todense()) - - def test_functions_jitable(self): - """Tests that all class functions are jittable.""" - doc = JAXSparseOperatorCollection( - operators=Array(self.test_operator_list), - static_operator=Array(self.test_operator_list[0]), - ) - rand.seed(3423) - coeffs = rand.uniform(-1, 1, 3) - self.jit_wrap(doc.evaluate)(Array(coeffs)) - self.jit_wrap(doc.evaluate_rhs)(Array(coeffs), self.X) - - def test_functions_gradable(self): - """Tests that all class functions are gradable.""" - doc = JAXSparseOperatorCollection( - operators=Array(self.test_operator_list), - static_operator=Array(self.test_operator_list[0]), - ) - rand.seed(5433) - coeffs = rand.uniform(-1, 1, 3) - self.jit_grad_wrap(doc.evaluate)(Array(coeffs)) - self.jit_grad_wrap(doc.evaluate_rhs)(Array(coeffs), self.X) - - -class TestDenseLindbladCollection(QiskitDynamicsTestCase): - """Tests for DenseLindbladCollection.""" +@partial(test_array_backends, array_libraries=["numpy", "jax", "scipy_sparse", "jax_sparse"]) +class TestLindbladCollection: + """Tests for LindbladCollection and ScipySparseLindbladCollection.""" def setUp(self): - self.X = Array(Operator.from_label("X").data) - self.Y = Array(Operator.from_label("Y").data) - self.Z = Array(Operator.from_label("Z").data) + """Build pseudo-random LindbladCollection instance.""" + self.X = Operator.from_label("X").data + self.Y = Operator.from_label("Y").data + self.Z = Operator.from_label("Z").data rand.seed(2134024) n = 16 k = 8 m = 4 l = 2 - self.hamiltonian_operators = rand.uniform(-1, 1, (k, n, n)) + self.hamiltonian_operators = rand.uniform(-1, 1, (k, n, n)) + 1j * rand.uniform( + -1, 1, (k, n, n) + ) self.dissipator_operators = rand.uniform(-1, 1, (m, n, n)) - self.static_hamiltonian = rand.uniform(-1, 1, (n, n)) - self.rho = rand.uniform(-1, 1, (n, n)) - self.multiple_rho = rand.uniform(-1, 1, (l, n, n)) + self.static_hamiltonian = rand.uniform(-1, 1, (n, n)) + 1j * rand.uniform(-1, 1, (n, n)) + self.rho = rand.uniform(-1, 1, (n, n)) + 1j * rand.uniform(-1, 1, (n, n)) + self.multiple_rho = rand.uniform(-1, 1, (l, n, n)) + 1j * rand.uniform(-1, 1, (l, n, n)) self.ham_sig_vals = rand.uniform(-1, 1, (k)) self.dis_sig_vals = rand.uniform(-1, 1, (m)) self.r = lambda *args: rand.uniform(-1, 1, args) + 1j * rand.uniform(-1, 1, args) @@ -320,7 +262,10 @@ def construct_collection(self, *args, **kwargs): """Construct collection to be tested by this class Used for inheritance. """ - return DenseLindbladCollection(*args, **kwargs) + if self.array_library() == "scipy_sparse": + return ScipySparseLindbladCollection(*args, **kwargs) + + return LindbladCollection(*args, **kwargs, array_library=self.array_library()) def test_empty_collection_error(self): """Test errors get raised for empty collection.""" @@ -444,8 +389,8 @@ def test_static_hamiltonian_only(self): collection = self.construct_collection(static_hamiltonian=self.X) - self.assertAllClose(to_array(collection.evaluate_hamiltonian(None)), self.X) - rho = Array([[1.0, 0.0], [0.0, 0.0]]) + self.assertAllClose(collection.evaluate_hamiltonian(None), self.X) + rho = np.array([[1.0, 0.0], [0.0, 0.0]]) expected = -1j * (self.X @ rho - rho @ self.X) self.assertAllClose(collection.evaluate_rhs(None, None, rho), expected) @@ -475,9 +420,7 @@ def test_dissipators_only(self): def test_static_dissipator_only(self): """Test correct evaluation with just static dissipators.""" - collection = self.construct_collection( - static_dissipators=self.dissipator_operators, - ) + collection = self.construct_collection(static_dissipators=self.dissipator_operators) res = collection(None, None, self.rho) dis_anticommutator = (-1 / 2) * np.tensordot( np.ones_like(self.dis_sig_vals), @@ -539,12 +482,12 @@ def test_operator_type_construction(self): H_i = self.r(3, 3) L_i = self.r(3, 3) ham_op_terms.append(Operator(H_i)) - ham_ar_terms.append(Array(H_i)) + ham_ar_terms.append(H_i) dis_op_terms.append(Operator(L_i)) - dis_ar_terms.append(Array(L_i)) + dis_ar_terms.append(L_i) H_d = self.r(3, 3) op_static_hamiltonian = Operator(H_d) - ar_static_hamiltonian = Array(H_d) + ar_static_hamiltonian = H_d op_collection = self.construct_collection( hamiltonian_operators=ham_op_terms, static_hamiltonian=op_static_hamiltonian, @@ -566,53 +509,47 @@ def test_operator_type_construction(self): ) -class TestDenseLindbladCollectionJax(TestDenseLindbladCollection, TestJaxBase): - """Jax version of TestDenseLindbladCollection tests.""" +@partial(test_array_backends, array_libraries=["jax", "jax_sparse"]) +class TestLindbladCollectionJAXTransformations: + """JAX transformation tests for LindbladCollection.""" - def test_functions_jitable(self): - """Tests that all class functions are jittable""" - dlc = self.construct_collection( - hamiltonian_operators=Array(self.hamiltonian_operators), - static_hamiltonian=Array(self.static_hamiltonian), - dissipator_operators=Array(self.dissipator_operators), - ) - - self.jit_wrap(dlc.evaluate_rhs)( - Array(self.ham_sig_vals), Array(self.dis_sig_vals), self.rho - ) - self.jit_wrap(dlc.evaluate_hamiltonian)(Array(self.ham_sig_vals)) + def setUp(self): + """Build re-usable operators and pseudo-random LindbladCollection instance.""" + self.X = Operator.from_label("X").data + self.Y = Operator.from_label("Y").data + self.Z = Operator.from_label("Z").data + rand.seed(2134024) + n = 16 + k = 8 + m = 4 + l = 2 + self.hamiltonian_operators = rand.uniform(-1, 1, (k, n, n)) + self.dissipator_operators = rand.uniform(-1, 1, (m, n, n)) + self.static_hamiltonian = rand.uniform(-1, 1, (n, n)) + self.rho = rand.uniform(-1, 1, (n, n)) + self.multiple_rho = rand.uniform(-1, 1, (l, n, n)) + self.ham_sig_vals = rand.uniform(-1, 1, (k)) + self.dis_sig_vals = rand.uniform(-1, 1, (m)) + self.r = lambda *args: rand.uniform(-1, 1, args) + 1j * rand.uniform(-1, 1, args) def test_functions_gradable(self): """Tests if all class functions are gradable""" - dlc = self.construct_collection( - hamiltonian_operators=Array(self.hamiltonian_operators), - static_hamiltonian=Array(self.static_hamiltonian), - dissipator_operators=Array(self.dissipator_operators), - ) - self.jit_grad_wrap(dlc.evaluate_rhs)( - Array(self.ham_sig_vals), Array(self.dis_sig_vals), self.rho + dlc = LindbladCollection( + hamiltonian_operators=self.hamiltonian_operators, + static_hamiltonian=self.static_hamiltonian, + dissipator_operators=self.dissipator_operators, + array_library=self.array_library(), ) - self.jit_grad_wrap(dlc.evaluate_hamiltonian)(Array(self.ham_sig_vals)) - - -class TestSparseLindbladCollection(TestDenseLindbladCollection): - """Tests for SparseLindbladCollection.""" - - def construct_collection(self, *args, **kwargs): - return SparseLindbladCollection(*args, **kwargs) + self.jit_grad(dlc.evaluate_rhs)(self.ham_sig_vals, self.dis_sig_vals, self.rho) + self.jit_grad(dlc.evaluate_hamiltonian)(self.ham_sig_vals) -class TestJAXSparseLindbladCollection(TestDenseLindbladCollectionJax): - """Tests for JAXSparseLindbladCollection.""" - - def construct_collection(self, *args, **kwargs): - return JAXSparseLindbladCollection(*args, **kwargs) - - -class TestDenseVectorizedLindbladCollection(QiskitDynamicsTestCase): - """Tests for DenseVectorizedLindbladCollection.""" +@partial(test_array_backends, array_libraries=["numpy", "jax", "jax_sparse", "scipy_sparse"]) +class TestVectorizedLindbladCollection: + """Tests for VectorizedLindbladCollection.""" def setUp(self) -> None: + """Build pseudo random operators.""" rand.seed(123098341) n = 16 k = 4 @@ -628,16 +565,26 @@ def setUp(self) -> None: self.t = r() self.rand_ham_coeffs = r(k) self.rand_dis_coeffs = r(m) - self.vectorized_class = DenseVectorizedLindbladCollection - self.non_vectorized_class = DenseLindbladCollection + + def _build_vectorized_collection(self, *args, **kwargs): + if self.array_library() == "scipy_sparse": + return ScipySparseVectorizedLindbladCollection(*args, **kwargs) + else: + return VectorizedLindbladCollection(*args, **kwargs, array_library=self.array_library()) + + def _build_non_vectorized_collection(self, *args, **kwargs): + if self.array_library() == "scipy_sparse": + return ScipySparseLindbladCollection(*args, **kwargs) + else: + return LindbladCollection(*args, **kwargs, array_library=self.array_library()) def test_empty_collection_error(self): """Test errors get raised for empty collection.""" - collection = self.vectorized_class() - with self.assertRaisesRegex(QiskitError, self.vectorized_class.__name__ + " with None"): + collection = self._build_vectorized_collection() + with self.assertRaisesRegex(QiskitError, "OperatorCollection with None"): collection(None, None, np.array([[1.0, 0.0], [0.0, 0.0]])) - with self.assertRaisesRegex(QiskitError, self.vectorized_class.__name__ + " with None"): + with self.assertRaisesRegex(QiskitError, "LindbladCollection with None"): collection.evaluate_hamiltonian(None) def test_consistency_all_terms(self): @@ -736,13 +683,13 @@ def _consistency_test( ): """Consistency test template for non-vectorized class and vectorized class.""" - collection = self.non_vectorized_class( + collection = self._build_non_vectorized_collection( static_hamiltonian=static_hamiltonian, hamiltonian_operators=hamiltonian_operators, static_dissipators=static_dissipators, dissipator_operators=dissipator_operators, ) - vec_collection = self.vectorized_class( + vec_collection = self._build_vectorized_collection( static_hamiltonian=static_hamiltonian, hamiltonian_operators=hamiltonian_operators, static_dissipators=static_dissipators, @@ -756,58 +703,3 @@ def _consistency_test( self.rand_ham_coeffs, self.rand_dis_coeffs, self.rho.flatten(order="F") ) self.assertAllClose(a, b) - - -class TestDenseVectorizedLindbladCollectionJax(TestDenseVectorizedLindbladCollection, TestJaxBase): - """Jax version of TestDenseVectorizedLindbladCollection tests. - - Note: The evaluation processes for DenseVectorizedLindbladCollection - are not directly jitable or compilable. The compilation of these steps - is taken care of by the tests for LindbladModel. - """ - - -class TestSparseVectorizedLindbladCollection(TestDenseVectorizedLindbladCollection): - """Tests for SparseVectorizedLindbladCollection.""" - - def setUp(self) -> None: - rand.seed(31232) - n = 16 - k = 4 - m = 2 - r = lambda *args: rand.uniform(-1, 1, [*args]) + 1j * rand.uniform(-1, 1, [*args]) - - self.r = r - self.rand_ham = r(k, n, n) - self.rand_static_dis = r(k, n, n) - self.rand_dis = r(m, n, n) - self.rand_dft = r(n, n) - self.rho = r(n, n) - self.t = r() - self.rand_ham_coeffs = r(k) - self.rand_dis_coeffs = r(m) - self.vectorized_class = SparseVectorizedLindbladCollection - self.non_vectorized_class = SparseLindbladCollection - - -class TestJAXSparseVectorizedLindbladCollection(TestDenseVectorizedLindbladCollectionJax): - """Tests for JAXSparseVectorizedLindbladCollection.""" - - def setUp(self) -> None: - rand.seed(31232) - n = 16 - k = 4 - m = 2 - r = lambda *args: rand.uniform(-1, 1, [*args]) + 1j * rand.uniform(-1, 1, [*args]) - - self.r = r - self.rand_ham = r(k, n, n) - self.rand_static_dis = r(k, n, n) - self.rand_dis = r(m, n, n) - self.rand_dft = r(n, n) - self.rho = r(n, n) - self.t = r() - self.rand_ham_coeffs = r(k) - self.rand_dis_coeffs = r(m) - self.vectorized_class = JAXSparseVectorizedLindbladCollection - self.non_vectorized_class = JAXSparseLindbladCollection diff --git a/test/dynamics/models/test_rotating_frame.py b/test/dynamics/models/test_rotating_frame.py index ab1825b74..9cec9b54e 100644 --- a/test/dynamics/models/test_rotating_frame.py +++ b/test/dynamics/models/test_rotating_frame.py @@ -30,7 +30,6 @@ try: from jax import jit import jax.numpy as jnp - from jax.experimental.sparse import BCOO except ImportError: pass @@ -490,7 +489,9 @@ def test_vectorized_frame_basis(self): class TestRotatingFrameTypeHandling(NumpyTestBase): - """Type handling testing for RotatingFrame methods for inputs of type csr_matrix, qutip, Operator, and numpy array.""" + """Type handling testing for RotatingFrame methods for inputs of type csr_matrix, qutip, + Operator, and numpy array. + """ def test_state_transformations_no_frame_csr_matrix_type(self): """Test frame transformations with no frame.""" @@ -608,7 +609,7 @@ def test_operator_into_frame_basis(self): op = np.array([[1.0, -1j], [0.0, 1.0]]) output = rotating_frame.operator_into_frame_basis(self.asarray(op)) expected = rotating_frame.operator_into_frame_basis(op) - + if "jax" in self.array_library(): self.assertTrue(isinstance(output, jnp.ndarray)) else: