diff --git a/qiskit_dynamics/models/generator_model.py b/qiskit_dynamics/models/generator_model.py index b959113db..b47fde94d 100644 --- a/qiskit_dynamics/models/generator_model.py +++ b/qiskit_dynamics/models/generator_model.py @@ -16,21 +16,20 @@ """ from abc import ABC, abstractmethod -from typing import Tuple, Union, List, Optional +from typing import Union, List, Optional from warnings import warn -from copy import copy -import numpy as np -from scipy.sparse import csr_matrix, issparse, diags +from scipy.sparse import diags from qiskit import QiskitError -from qiskit.quantum_info.operators import Operator + +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 from qiskit_dynamics.models.operator_collections import ( OperatorCollection, ScipySparseOperatorCollection, ) -from qiskit_dynamics.array import Array from qiskit_dynamics.signals import Signal, SignalList -from qiskit_dynamics.type_utils import to_numeric_matrix_type from .rotating_frame import RotatingFrame try: @@ -44,15 +43,6 @@ class BaseGeneratorModel(ABC): :math:`\dot{y}(t) = \Lambda(t, y)`, where :math:`\Lambda` is linear in :math:`y`. The core functionality is evaluation of :math:`\Lambda(t, y)`, as well as, if possible, evaluation of the map :math:`\Lambda(t, \cdot)`. - - Additionally, the class defines interfaces for: - - * Setting a "rotating frame", specified either directly as a :class:`RotatingFrame` - instance, or an operator from which a :class:`RotatingFrame` instance can be constructed. - The exact meaning of this transformation is determined by the structure of - :math:`\Lambda(t, y)`, and is therefore by handled by concrete subclasses. - * Setting an internal "evaluation mode", to set the specific numerical methods to use - when evaluating :math:`\Lambda(t, y)`. """ @property @@ -65,33 +55,18 @@ def dim(self) -> int: def rotating_frame(self) -> RotatingFrame: """The rotating frame.""" - @rotating_frame.setter - @abstractmethod - def rotating_frame(self, rotating_frame: RotatingFrame): - pass - @property @abstractmethod def in_frame_basis(self) -> bool: """Whether or not the model is evaluated in the basis in which the frame is diagonalized.""" - @in_frame_basis.setter - @abstractmethod - def in_frame_basis(self, in_frame_basis: bool): - pass - @property @abstractmethod - def evaluation_mode(self) -> str: - """Numerical evaluation mode of the model.""" - - @evaluation_mode.setter - @abstractmethod - def evaluation_mode(self, new_mode: str): - pass + def array_library(self) -> Union[None, str]: + """Array library used to store the operators in the model.""" @abstractmethod - def evaluate(self, time: float) -> Array: + def evaluate(self, time: float) -> ArrayLike: r"""If possible, evaluate the map :math:`\Lambda(t, \cdot)`. Args: @@ -99,7 +74,7 @@ def evaluate(self, time: float) -> Array: """ @abstractmethod - def evaluate_rhs(self, time: float, y: Array) -> Array: + def evaluate_rhs(self, time: float, y: ArrayLike) -> ArrayLike: r"""Evaluate the right hand side :math:`\dot{y}(t) = \Lambda(t, y)`. Args: @@ -107,11 +82,7 @@ def evaluate_rhs(self, time: float, y: Array) -> Array: y: State of the differential equation. """ - def copy(self): - """Return a copy of self.""" - return copy(self) - - def __call__(self, time: float, y: Optional[Array] = None) -> Array: + def __call__(self, time: float, y: Optional[ArrayLike] = None) -> ArrayLike: r"""Evaluate generator RHS functions. If ``y is None``, attemps to evaluate :math:`\Lambda(t, \cdot)`, otherwise, calculates :math:`\Lambda(t, y)`. @@ -144,12 +115,12 @@ class GeneratorModel(BaseGeneratorModel): def __init__( self, - static_operator: Optional[Array] = None, - operators: Optional[Array] = None, + static_operator: Optional[ArrayLike] = None, + operators: Optional[ArrayLike] = None, signals: Optional[Union[SignalList, List[Signal]]] = None, - rotating_frame: Optional[Union[Operator, Array, RotatingFrame]] = None, + rotating_frame: Optional[Union[ArrayLike, RotatingFrame]] = None, in_frame_basis: bool = False, - evaluation_mode: str = "dense", + array_library: Optional[str] = None, ): """Initialize. @@ -161,8 +132,9 @@ def __init__( rotating_frame: Rotating frame operator. in_frame_basis: Whether to represent the model in the basis in which the rotating frame operator is diagonalized. - evaluation_mode: Evaluation mode to use. Supported options are ``'dense'`` and - ``'sparse'``. Call ``help(GeneratorModel.evaluation_mode)`` for more details. + array_library: Array library for storing the operators in the model. Supported options + are ``'numpy'``, ``'jax'``, ``'jax_sparse'``, and ``'scipy_sparse'``. If ``None``, + the arrays will be handled by general dispatching rules. Raises: QiskitError: If model not sufficiently specified. """ @@ -172,21 +144,26 @@ def __init__( "be specified at construction." ) - # initialize internal operator representation - self._operator_collection = construct_operator_collection( - evaluation_mode=evaluation_mode, static_operator=static_operator, operators=operators + self._array_library = array_library + self._rotating_frame = RotatingFrame(rotating_frame) + self._in_frame_basis = in_frame_basis + + # set up internal operators + static_operator = _static_operator_into_frame_basis( + static_operator=static_operator, + rotating_frame=self._rotating_frame, + array_library=array_library, ) - self._evaluation_mode = evaluation_mode - # set frame - self._rotating_frame = None - self.rotating_frame = rotating_frame + operators = _operators_into_frame_basis( + operators=operators, rotating_frame=self._rotating_frame, array_library=array_library + ) - # set operation in frame basis - self._in_frame_basis = None - self.in_frame_basis = in_frame_basis + self._operator_collection = _get_operator_collection( + static_operator=static_operator, operators=operators, array_library=array_library + ) - # initialize signal-related attributes + self._signals = None self.signals = signals @property @@ -195,64 +172,45 @@ def dim(self) -> int: return self._operator_collection.dim @property - def evaluation_mode(self) -> str: - """Numerical evaluation mode of the model. - - Possible values: - - * ``"dense"``: Stores/evaluates operators using dense Arrays. - * ``"sparse"``: Stores/evaluates operators using sparse matrices. If the default Array - backend is JAX, implemented with JAX BCOO arrays, otherwise uses scipy - :class:`csr_matrix` sparse type. Note that JAX sparse mode is only recommended for use on - CPU. - - Raises: - NotImplementedError: If, when being set, ``new_mode`` is not one of the above supported - evaluation modes. - """ - return self._evaluation_mode - - @evaluation_mode.setter - def evaluation_mode(self, new_mode: str): - if new_mode != self._evaluation_mode: - self._operator_collection = construct_operator_collection( - new_mode, - self._operator_collection.static_operator, - self._operator_collection.operators, - ) - self._evaluation_mode = new_mode + def rotating_frame(self) -> RotatingFrame: + """The rotating frame.""" + return self._rotating_frame @property - def rotating_frame(self) -> RotatingFrame: - """The rotating frame. + def in_frame_basis(self) -> bool: + """Whether or not the model is evaluated in the basis in which the frame is diagonalized.""" + return self._in_frame_basis - This property can be set with a :class:`RotatingFrame` instance, or any valid constructor - argument to this class. + @in_frame_basis.setter + def in_frame_basis(self, in_frame_basis: bool): + self._in_frame_basis = in_frame_basis - Setting this property updates the internal operator to use the new frame. - """ - return self._rotating_frame + @property + def array_library(self) -> Union[None, str]: + """Array library used to store the operators in the model.""" + return self._array_library - @rotating_frame.setter - def rotating_frame(self, rotating_frame: Union[Operator, Array, RotatingFrame]): - new_frame = RotatingFrame(rotating_frame) + @property + def static_operator(self) -> Union[ArrayLike, None]: + """The static operator.""" + if self._operator_collection.static_operator is None: + return None - new_static_operator = transfer_static_operator_between_frames( - self._get_static_operator(in_frame_basis=True), - new_frame=new_frame, - old_frame=self.rotating_frame, - ) - new_operators = transfer_operators_between_frames( - self._get_operators(in_frame_basis=True), - new_frame=new_frame, - old_frame=self.rotating_frame, + if self.in_frame_basis: + return self._operator_collection.static_operator + return self.rotating_frame.operator_out_of_frame_basis( + self._operator_collection.static_operator ) - self._rotating_frame = new_frame + @property + def operators(self) -> Union[ArrayLike, None]: + """The operators in the model.""" + if self._operator_collection.operators is None: + return None - self._operator_collection = construct_operator_collection( - self.evaluation_mode, new_static_operator, new_operators - ) + if self.in_frame_basis: + return self._operator_collection.operators + return self.rotating_frame.operator_out_of_frame_basis(self._operator_collection.operators) @property def signals(self) -> SignalList: @@ -289,83 +247,7 @@ def signals(self, signals: Union[SignalList, List[Signal]]): self._signals = signals - @property - def in_frame_basis(self) -> bool: - """Whether the model is represented in the basis that diagonalizes the frame operator.""" - return self._in_frame_basis - - @in_frame_basis.setter - def in_frame_basis(self, in_frame_basis: bool): - self._in_frame_basis = in_frame_basis - - @property - def operators(self) -> Array: - """The operators in the model.""" - return self._get_operators(in_frame_basis=self._in_frame_basis) - - @property - def static_operator(self) -> Array: - """The static operator.""" - return self._get_static_operator(in_frame_basis=self._in_frame_basis) - - @static_operator.setter - def static_operator(self, static_operator: Array): - self._set_static_operator( - new_static_operator=static_operator, operator_in_frame_basis=self._in_frame_basis - ) - - def _get_operators(self, in_frame_basis: Optional[bool] = False) -> Array: - """Get the operators used in the model construction. - - Args: - in_frame_basis: Flag indicating whether to return the operators - in the basis in which the rotating frame operator is diagonal. - Returns: - The operators in the basis specified by ``in_frame_basis``. - """ - ops = self._operator_collection.operators - if ops is not None and not in_frame_basis and self.rotating_frame is not None: - return self.rotating_frame.operator_out_of_frame_basis(ops) - return ops - - def _get_static_operator(self, in_frame_basis: Optional[bool] = False) -> Array: - """Get the constant term. - - Args: - in_frame_basis: Whether the returned static_operator should be in the basis in which the - frame is diagonal. - Returns: - The static operator term. - """ - op = self._operator_collection.static_operator - if not in_frame_basis and self.rotating_frame is not None: - return self.rotating_frame.operator_out_of_frame_basis(op) - return op - - def _set_static_operator( - self, - new_static_operator: Array, - operator_in_frame_basis: Optional[bool] = False, - ): - """Sets static term. Note that if the model has a rotating frame this will override any - contributions to the static term due to the frame transformation. - - Args: - new_static_operator: The static operator operator. - operator_in_frame_basis: Whether `new_static_operator` is already in the rotating frame - basis. - """ - if new_static_operator is None: - self._operator_collection.static_operator = None - else: - if not operator_in_frame_basis and self.rotating_frame is not None: - new_static_operator = self.rotating_frame.operator_into_frame_basis( - new_static_operator - ) - - self._operator_collection.static_operator = new_static_operator - - def evaluate(self, time: float) -> Array: + def evaluate(self, time: float) -> ArrayLike: """Evaluate the model in array format as a matrix, independent of state. Args: @@ -390,7 +272,7 @@ def evaluate(self, time: float) -> Array: time, op_combo, operator_in_frame_basis=True, return_in_frame_basis=self._in_frame_basis ) - def evaluate_rhs(self, time: float, y: Array) -> Array: + def evaluate_rhs(self, time: float, y: ArrayLike) -> ArrayLike: r"""Evaluate ``G(t) @ y``. Args: @@ -428,133 +310,75 @@ def evaluate_rhs(self, time: float, y: Array) -> Array: return self._operator_collection(sig_vals, y) -def transfer_static_operator_between_frames( - static_operator: Union[None, Array, csr_matrix], - new_frame: Optional[Union[Array, RotatingFrame]] = None, - old_frame: Optional[Union[Array, RotatingFrame]] = None, -) -> Tuple[Union[None, Array]]: - """Helper function for transforming the static operator of a model from one frame basis into - another. - - ``static_operator`` is additionally transformed as a generator: the old frame operator is added, - and the new one is subtracted. - - Args: - static_operator: Static operator of the model. None is treated as 0. - new_frame: New rotating frame. - old_frame: Old rotating frame. - - Returns: - static_operator: Transformed as described above. +def _static_operator_into_frame_basis( + static_operator: Union[None, ArrayLike], + rotating_frame: RotatingFrame, + array_library: Optional[ArrayLike] = None, +) -> Union[None, ArrayLike]: + """Converts the static_operator into the frame basis, including a subtraction of the frame + operator. This function also enforces typing via array_library. """ - new_frame = RotatingFrame(new_frame) - old_frame = RotatingFrame(old_frame) - - static_operator = to_numeric_matrix_type(static_operator) - - # transform out of old frame basis, and add the old frame operator - if static_operator is not None: - static_operator = old_frame.generator_out_of_frame( - t=0.0, - operator=static_operator, - operator_in_frame_basis=True, - return_in_frame_basis=False, - ) - else: - # "add" the frame operator to 0 - if old_frame.frame_operator is not None: - if issparse(static_operator): - if old_frame.frame_operator.ndim == 1: - static_operator = diags(old_frame.frame_operator, format="csr") - else: - static_operator = csr_matrix(old_frame.frame_operator) - else: - if old_frame.frame_operator.ndim == 1: - static_operator = np.diag(old_frame.frame_operator) - else: - static_operator = old_frame.frame_operator - # transform into new frame basis, and add the new frame operator - if static_operator is not None: - static_operator = new_frame.generator_into_frame( - t=0.0, - operator=static_operator, - operator_in_frame_basis=False, - return_in_frame_basis=True, - ) - else: - # "subtract" the frame operator from 0 - if new_frame.frame_operator is not None: - if issparse(static_operator): - static_operator = -diags(new_frame.frame_diag, format="csr") - else: - static_operator = -np.diag(new_frame.frame_diag) - return static_operator + # handle static_operator is None case + if static_operator is None: + if rotating_frame.frame_operator is None: + return None + if array_library == "scipy_sparse": + return -diags(rotating_frame.frame_diag, format="csr") + return unp.diag(-rotating_frame.frame_diag) + static_operator = numpy_alias(like=array_library).asarray(static_operator) -def transfer_operators_between_frames( - operators: Union[None, Array, List[csr_matrix]], - new_frame: Optional[Union[Array, RotatingFrame]] = None, - old_frame: Optional[Union[Array, RotatingFrame]] = None, -) -> Tuple[Union[None, Array]]: - """Helper function for transforming a list of operators for a model from one frame basis into - another. + return rotating_frame.generator_into_frame( + t=0.0, operator=static_operator, return_in_frame_basis=True + ) - Args: - operators: Operators of the model. If ``None``, remains as ``None``. - new_frame: New rotating frame. - old_frame: Old rotating frame. - Returns: - operators: Transformed as described above. +def _operators_into_frame_basis( + operators: Union[None, list, ArrayLike], + rotating_frame: RotatingFrame, + array_library: Optional[str] = None, +) -> Union[None, ArrayLike]: + """Converts operators into the frame basis. This function also enforces typing via + array_library. """ - new_frame = RotatingFrame(new_frame) - old_frame = RotatingFrame(old_frame) - - operators = to_numeric_matrix_type(operators) + if operators is None: + return None - # transform out of old frame basis - if operators is not None: - # For sparse case, if list, loop - if isinstance(operators, list): - operators = [old_frame.operator_out_of_frame_basis(op) for op in operators] - else: - operators = old_frame.operator_out_of_frame_basis(operators) - - # transform into new frame basis - if operators is not None: - # For sparse case, if list, loop - if isinstance(operators, list): - operators = [new_frame.operator_into_frame_basis(op) for op in operators] - else: - operators = new_frame.operator_into_frame_basis(operators) + if array_library == "scipy_sparse" or ( + array_library is None and "scipy_sparse" in numpy_alias.infer_libs(operators) + ): + ops = [] + for op in operators: + op = numpy_alias(like="scipy_sparse").asarray(op) + ops.append(rotating_frame.operator_into_frame_basis(op)) + return ops - return operators + return rotating_frame.operator_into_frame_basis( + numpy_alias(like=array_library).asarray(operators) + ) -def construct_operator_collection( - evaluation_mode: str, - static_operator: Union[None, Array, csr_matrix], - operators: Union[None, Array, List[csr_matrix]], +def _get_operator_collection( + static_operator: Union[None, ArrayLike], + operators: Union[None, ArrayLike], + array_library: Optional[str] = None, ) -> Union[OperatorCollection, ScipySparseOperatorCollection]: """Construct an operator collection for :class:`GeneratorModel`. Args: - evaluation_mode: Evaluation mode. static_operator: Static operator of the model. operators: Operators for the model. + array_library: Array library to use. Returns: 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) - pass - if evaluation_mode == "sparse" and Array.default_backend() == "jax": + if array_library == "scipy_sparse": + return ScipySparseOperatorCollection(static_operator=static_operator, operators=operators) + + if array_library == "jax_sparse": # warn that sparse mode when using JAX is primarily recommended for use on CPU if jax.default_backend() != "cpu": warn( @@ -562,13 +386,6 @@ def construct_operator_collection( stacklevel=2, ) - # return JAXSparseOperatorCollection(static_operator=static_operator, operators=operators) - pass - if evaluation_mode == "sparse": - # return SparseOperatorCollection(static_operator=static_operator, operators=operators) - pass - - raise NotImplementedError( - f"Evaluation mode '{evaluation_mode}' is not supported. Call " - "help(GeneratorModel.evaluation_mode) for available options." + return OperatorCollection( + static_operator=static_operator, operators=operators, array_library=array_library ) diff --git a/qiskit_dynamics/models/hamiltonian_model.py b/qiskit_dynamics/models/hamiltonian_model.py index 632b4b98b..e47a0fac8 100644 --- a/qiskit_dynamics/models/hamiltonian_model.py +++ b/qiskit_dynamics/models/hamiltonian_model.py @@ -17,20 +17,16 @@ from typing import Union, List, Optional import numpy as np -from scipy.sparse import csr_matrix, issparse +from scipy.sparse import issparse from scipy.sparse.linalg import norm as spnorm from qiskit import QiskitError -from qiskit.quantum_info.operators import Operator -from qiskit_dynamics.array import Array + +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 from qiskit_dynamics.signals import Signal, SignalList -from qiskit_dynamics.type_utils import to_numeric_matrix_type, to_array -from .generator_model import ( - GeneratorModel, - transfer_static_operator_between_frames, - transfer_operators_between_frames, - construct_operator_collection, -) +from .generator_model import GeneratorModel from .rotating_frame import RotatingFrame @@ -49,9 +45,11 @@ class HamiltonianModel(GeneratorModel): .. code-block:: python - hamiltonian = HamiltonianModel(static_operator=static_operator, - operators=operators, - signals=signals) + hamiltonian = HamiltonianModel( + static_operator=static_operator, + operators=operators, + signals=signals + ) This class inherits most functionality from :class:`GeneratorModel`, with the following modifications: @@ -64,12 +62,12 @@ class HamiltonianModel(GeneratorModel): def __init__( self, - static_operator: Optional[Array] = None, - operators: Optional[List[Operator]] = None, + static_operator: Optional[ArrayLike] = None, + operators: Optional[ArrayLike] = None, signals: Optional[Union[SignalList, List[Signal]]] = None, - rotating_frame: Optional[Union[Operator, Array, RotatingFrame]] = None, + rotating_frame: Optional[Union[ArrayLike, RotatingFrame]] = None, in_frame_basis: bool = False, - evaluation_mode: str = "dense", + array_library: Optional[str] = None, validate: bool = True, ): """Initialize, ensuring that the operators are Hermitian. @@ -78,31 +76,39 @@ def __init__( static_operator: Time-independent term in the Hamiltonian. operators: List of Operator objects. signals: List of coefficients :math:`s_i(t)`. Not required at instantiation, but - necessary for evaluation of the model. - rotating_frame: Rotating frame operator. - If specified with a 1d array, it is interpreted as the - diagonal of a diagonal matrix. Assumed to store - the antihermitian matrix F = -iH. + necessary for evaluation of the model. + rotating_frame: Rotating frame operator. If specified with a 1d array, it is interpreted + as the diagonal of a diagonal matrix. Assumed to store the antihermitian matrix + F = -iH. in_frame_basis: Whether to represent the model in the basis in which the rotating - frame operator is diagonalized. - evaluation_mode: Evaluation mode to use. Supported options are ``'dense'`` and - ``'sparse'``. Call ``help(HamiltonianModel.evaluation_mode)`` for more - details. - validate: If True check input operators are Hermitian. + frame operator is diagonalized. + array_library: Array library for storing the operators in the model. Supported options + are ``'numpy'``, ``'jax'``, ``'jax_sparse'``, and ``'scipy_sparse'``. If ``None``, + the arrays will be handled by general dispatching rules. Call + ``help(GeneratorModel.array_library)`` for more details. + validate: If ``True`` check input operators are Hermitian. Note that this is + incompatible with JAX transformations. Raises: QiskitError: if operators are not Hermitian """ - # verify operators are Hermitian, and if so instantiate - operators = to_numeric_matrix_type(operators) - static_operator = to_numeric_matrix_type(static_operator) + # prepare and validate operators + if static_operator is not None: + if validate and not is_hermitian(static_operator): + raise QiskitError("""HamiltonianModel static_operator must be Hermitian.""") + static_operator = -1j * numpy_alias(like=array_library).asarray(static_operator) - if validate: - if (operators is not None) and (not is_hermitian(operators)): + if operators is not None: + if validate and any(not is_hermitian(op) for op in operators): raise QiskitError("""HamiltonianModel operators must be Hermitian.""") - if (static_operator is not None) and (not is_hermitian(static_operator)): - raise QiskitError("""HamiltonianModel static_operator must be Hermitian.""") + + if array_library == "scipy_sparse" or ( + array_library is None and "scipy_sparse" in numpy_alias.infer_libs(operators) + ): + operators = [-1j * numpy_alias(like=array_library).asarray(op) for op in operators] + else: + operators = -1j * numpy_alias(like=array_library).asarray(operators) super().__init__( static_operator=static_operator, @@ -110,111 +116,63 @@ def __init__( signals=signals, rotating_frame=rotating_frame, in_frame_basis=in_frame_basis, - evaluation_mode=evaluation_mode, + array_library=array_library, ) @property - def rotating_frame(self) -> RotatingFrame: - """The rotating frame. - - This property can be set with a :class:`RotatingFrame` instance, or any valid constructor - argument to this class. It can either be Hermitian or anti-Hermitian, and in either case - only the anti-Hermitian version :math:`F=-iH` will be stored. - - Setting this property updates the internal operator list to use the new frame. - """ - return super().rotating_frame - - @rotating_frame.setter - def rotating_frame(self, rotating_frame: Union[Operator, Array, RotatingFrame]) -> Array: - new_frame = RotatingFrame(rotating_frame) - - # convert static operator to new frame setup - static_op = self._get_static_operator(in_frame_basis=True) - if static_op is not None: - static_op = -1j * static_op - - new_static_operator = transfer_static_operator_between_frames( - static_op, - new_frame=new_frame, - old_frame=self.rotating_frame, + def static_operator(self) -> Union[ArrayLike, None]: + """The static operator.""" + if self._operator_collection.static_operator is None: + return None + + if self.in_frame_basis: + return self._operator_collection.static_operator + return 1j * self.rotating_frame.operator_out_of_frame_basis( + self._operator_collection.static_operator ) - if new_static_operator is not None: - new_static_operator = 1j * new_static_operator - - # convert operators to new frame set up - new_operators = transfer_operators_between_frames( - self._get_operators(in_frame_basis=True), - new_frame=new_frame, - old_frame=self.rotating_frame, - ) - - self._rotating_frame = new_frame - - self._operator_collection = construct_operator_collection( - self.evaluation_mode, new_static_operator, new_operators - ) - - def evaluate(self, time: float) -> Array: - """Evaluate the model in array format as a matrix, independent of state. - - Args: - time: The time to evaluate the model at. - - Returns: - Array: The evaluated model as an anti-Hermitian matrix. - - Raises: - QiskitError: If model cannot be evaluated. - """ - return -1j * super().evaluate(time) - - def evaluate_rhs(self, time: float, y: Array) -> Array: - r"""Evaluate ``-1j * H(t) @ y``. + @property + def operators(self) -> Union[ArrayLike, None]: + """The operators in the model.""" + if self._operator_collection.operators is None: + return None - Args: - time: The time to evaluate the model at . - y: Array specifying system state. + if self.in_frame_basis: + ops = self._operator_collection.operators + else: + ops = self.rotating_frame.operator_out_of_frame_basis( + self._operator_collection.operators + ) - Returns: - Array defined by :math:`G(t) \times y`. + if isinstance(ops, list): + return [1j * op for op in ops] - Raises: - QiskitError: If model cannot be evaluated. - """ - return -1j * super().evaluate_rhs(time, y) + return 1j * ops -def is_hermitian( - operators: Union[Array, csr_matrix, List[csr_matrix], "BCOO"], tol: Optional[float] = 1e-10 -) -> bool: - """Validate that operators are Hermitian. +def is_hermitian(operator: ArrayLike, tol: Optional[float] = 1e-10) -> bool: + """Validate that an operator is Hermitian. Args: - operators: Either a 2d array representing a single operator, a 3d array representing a list - of operators, a ``csr_matrix``, or a list of ``csr_matrix``. + operator: A 2d array representing a single operator. tol: Tolerance for checking zeros. Returns: - bool: Whether or not the operators are Hermitian to within tolerance. + bool: Whether or not the operator is Hermitian to within tolerance. Raises: QiskitError: If an unexpeted type is received. """ - if isinstance(operators, (np.ndarray, Array)): - adj = None - if operators.ndim == 2: - adj = np.transpose(np.conjugate(operators)) - elif operators.ndim == 3: - adj = np.transpose(np.conjugate(operators), (0, 2, 1)) - return np.linalg.norm(adj - operators) < tol - elif issparse(operators): - return spnorm(operators - operators.conj().transpose()) < tol - elif isinstance(operators, list) and issparse(operators[0]): - return all(spnorm(op - op.conj().transpose()) < tol for op in operators) - elif type(operators).__name__ == "BCOO": + operator = unp.asarray(operator) + + if issparse(operator): + return spnorm(operator - operator.conj().transpose()) < tol + elif type(operator).__name__ == "BCOO": # fall back on array case for BCOO - return is_hermitian(to_array(operators)) + return is_hermitian(operator.todense()) + elif isinstance(operator, ArrayLike): + adj = None + adj = np.transpose(np.conjugate(operator)) + return np.linalg.norm(adj - operator) < tol raise QiskitError("is_hermitian got an unexpected type.") diff --git a/qiskit_dynamics/models/lindblad_model.py b/qiskit_dynamics/models/lindblad_model.py index 1c384d87f..5d604ee70 100644 --- a/qiskit_dynamics/models/lindblad_model.py +++ b/qiskit_dynamics/models/lindblad_model.py @@ -17,17 +17,14 @@ from typing import Tuple, Union, List, Optional from warnings import warn -from scipy.sparse import csr_matrix from qiskit import QiskitError -from qiskit.quantum_info.operators import Operator -from qiskit_dynamics.array import Array -from qiskit_dynamics.type_utils import to_numeric_matrix_type +from qiskit_dynamics.arraylias.alias import ArrayLike from qiskit_dynamics.signals import Signal, SignalList from .generator_model import ( BaseGeneratorModel, - transfer_static_operator_between_frames, - transfer_operators_between_frames, + _static_operator_into_frame_basis, + _operators_into_frame_basis, ) from .hamiltonian_model import HamiltonianModel, is_hermitian from .operator_collections import ( @@ -102,15 +99,16 @@ class LindbladModel(BaseGeneratorModel): def __init__( self, - static_hamiltonian: Optional[Union[Array, csr_matrix]] = None, - hamiltonian_operators: Optional[Union[Array, List[csr_matrix]]] = None, + static_hamiltonian: Optional[ArrayLike] = None, + hamiltonian_operators: Optional[ArrayLike] = None, hamiltonian_signals: Optional[Union[List[Signal], SignalList]] = None, - static_dissipators: Optional[Union[Array, csr_matrix]] = None, - dissipator_operators: Optional[Union[Array, List[csr_matrix]]] = None, + static_dissipators: Optional[ArrayLike] = None, + dissipator_operators: Optional[ArrayLike] = None, dissipator_signals: Optional[Union[List[Signal], SignalList]] = None, - rotating_frame: Optional[Union[Operator, Array, RotatingFrame]] = None, + rotating_frame: Optional[Union[ArrayLike, RotatingFrame]] = None, in_frame_basis: bool = False, - evaluation_mode: Optional[str] = "dense", + array_library: Optional[str] = None, + vectorized: bool = False, validate: bool = True, ): """Initialize. @@ -127,8 +125,13 @@ def __init__( assumed that all operators were already in the frame basis. in_frame_basis: Whether to represent the model in the basis in which the rotating frame operator is diagonalized. - evaluation_mode: Evaluation mode to use. Call ``help(LindbladModel.evaluation_mode)`` - for more details. + array_library: Array library for storing the operators in the model. Supported options + are ``'numpy'``, ``'jax'``, ``'jax_sparse'``, and ``'scipy_sparse'``. If ``None``, + the arrays will be handled by general dispatching rules. + vectorized: Whether or not to setup the Lindblad equation in vectorized mode. + If ``True``, the operators in the model are stored as :math:`(dim^2,dim^2)` matrices + that act on vectorized density matrices by left-multiplication. Setting this to + ``True`` is necessary for ``SuperOp`` simulation. validate: If True check input hamiltonian_operators and static_hamiltonian are Hermitian. @@ -148,28 +151,61 @@ def __init__( "to be specified at construction." ) - static_hamiltonian = to_numeric_matrix_type(static_hamiltonian) - hamiltonian_operators = to_numeric_matrix_type(hamiltonian_operators) if validate: - if (hamiltonian_operators is not None) and (not is_hermitian(hamiltonian_operators)): - raise QiskitError("""LinbladModel hamiltonian_operators must be Hermitian.""") if (static_hamiltonian is not None) and (not is_hermitian(static_hamiltonian)): raise QiskitError("""LinbladModel static_hamiltonian must be Hermitian.""") + if (hamiltonian_operators is not None) and any( + not is_hermitian(op) for op in hamiltonian_operators + ): + raise QiskitError("""LindbladModel hamiltonian_operators must be Hermitian.""") + + self._array_library = array_library + self._vectorized = vectorized + self._rotating_frame = RotatingFrame(rotating_frame) + self._in_frame_basis = in_frame_basis + + # jax sparse arrays cannot be used directly at this stage + setup_library = array_library + if array_library == "jax_sparse": + setup_library = "jax" + + # set up internal operators + if static_hamiltonian is not None: + static_hamiltonian = -1j * static_hamiltonian + static_hamiltonian = _static_operator_into_frame_basis( + static_operator=static_hamiltonian, + rotating_frame=self._rotating_frame, + array_library=setup_library, + ) + if static_hamiltonian is not None: + static_hamiltonian = 1j * static_hamiltonian - self._operator_collection = construct_lindblad_operator_collection( - evaluation_mode=evaluation_mode, + hamiltonian_operators = _operators_into_frame_basis( + operators=hamiltonian_operators, + rotating_frame=self._rotating_frame, + array_library=setup_library, + ) + + static_dissipators = _operators_into_frame_basis( + operators=static_dissipators, + rotating_frame=self._rotating_frame, + array_library=setup_library, + ) + + dissipator_operators = _operators_into_frame_basis( + operators=dissipator_operators, + rotating_frame=self._rotating_frame, + array_library=setup_library, + ) + + self._operator_collection = _get_lindblad_operator_collection( + array_library=array_library, + vectorized=vectorized, static_hamiltonian=static_hamiltonian, hamiltonian_operators=hamiltonian_operators, static_dissipators=static_dissipators, dissipator_operators=dissipator_operators, ) - self._evaluation_mode = evaluation_mode - self.vectorized_operators = "vectorized" in evaluation_mode - - self._rotating_frame = None - self.rotating_frame = rotating_frame - self._in_frame_basis = None - self.in_frame_basis = in_frame_basis self.signals = (hamiltonian_signals, dissipator_signals) @@ -177,10 +213,11 @@ def __init__( def from_hamiltonian( cls, hamiltonian: HamiltonianModel, - static_dissipators: Optional[Union[Array, csr_matrix]] = None, - dissipator_operators: Optional[Union[Array, List[csr_matrix]]] = None, - dissipator_signals: Optional[Union[List[Signal], SignalList]] = None, - evaluation_mode: Optional[str] = None, + static_dissipators: Optional[ArrayLike] = None, + dissipator_operators: Optional[ArrayLike] = None, + dissipator_signals: Optional[ArrayLike] = None, + array_library: Optional[str] = None, + vectorized: bool = False, ): """Construct from a :class:`HamiltonianModel`. @@ -189,29 +226,35 @@ def from_hamiltonian( static_dissipators: List of dissipators with coefficient 1. dissipator_operators: List of dissipators with time-dependent coefficients. dissipator_signals: List time-dependent coefficients for dissipator_operators. - evaluation_mode: Evaluation mode. Call ``help(LindbladModel.evaluation_mode)`` for more - information. + array_library: Array library to use. + vectorized: Whether or not to vectorize the Lindblad equation. Returns: LindbladModel: Linblad model from parameters. """ - if evaluation_mode is None: - evaluation_mode = hamiltonian.evaluation_mode + # store whether hamiltonian is in_frame_basis and set to False + in_frame_basis = hamiltonian.in_frame_basis + hamiltonian.in_frame_basis = False - ham_copy = hamiltonian.copy() - ham_copy.rotating_frame = None + # get the operators + static_hamiltonian = hamiltonian.static_operator + hamiltonian_operators = hamiltonian.operators + + # return to previous value + hamiltonian.in_frame_basis = in_frame_basis return cls( - static_hamiltonian=ham_copy._get_static_operator(in_frame_basis=False), - hamiltonian_operators=ham_copy._get_operators(in_frame_basis=False), - hamiltonian_signals=ham_copy.signals, + static_hamiltonian=static_hamiltonian, + hamiltonian_operators=hamiltonian_operators, + hamiltonian_signals=hamiltonian.signals, static_dissipators=static_dissipators, dissipator_operators=dissipator_operators, dissipator_signals=dissipator_signals, rotating_frame=hamiltonian.rotating_frame, in_frame_basis=hamiltonian.in_frame_basis, - evaluation_mode=evaluation_mode, + array_library=array_library, + vectorized=vectorized, ) @property @@ -305,216 +348,69 @@ def in_frame_basis(self, in_frame_basis: bool): self._in_frame_basis = in_frame_basis @property - def static_hamiltonian(self) -> Array: + def static_hamiltonian(self) -> ArrayLike: """The static Hamiltonian term.""" - return self._get_static_hamiltonian(in_frame_basis=self._in_frame_basis) + if self._operator_collection.static_hamiltonian is None: + return None - @static_hamiltonian.setter - def static_hamiltonian(self, static_hamiltonian: Array): - self._set_static_hamiltonian( - new_static_hamiltonian=static_hamiltonian, operator_in_frame_basis=self._in_frame_basis + if self.in_frame_basis: + return self._operator_collection.static_hamiltonian + return self.rotating_frame.operator_out_of_frame_basis( + self._operator_collection.static_hamiltonian ) @property - def hamiltonian_operators(self) -> Array: + def hamiltonian_operators(self) -> ArrayLike: """The Hamiltonian operators.""" - return self._get_hamiltonian_operators(in_frame_basis=self._in_frame_basis) + if self._operator_collection.hamiltonian_operators is None: + return None + + if self.in_frame_basis: + return self._operator_collection.hamiltonian_operators + return self.rotating_frame.operator_out_of_frame_basis( + self._operator_collection.hamiltonian_operators + ) @property - def static_dissipators(self) -> Array: + def static_dissipators(self) -> ArrayLike: """The static dissipator operators.""" - return self._get_static_dissipators(in_frame_basis=self._in_frame_basis) + if self._operator_collection.static_dissipators is None: + return None + + if self.in_frame_basis: + return self._operator_collection.static_dissipators + return self.rotating_frame.operator_out_of_frame_basis( + self._operator_collection.static_dissipators + ) @property - def dissipator_operators(self) -> Array: + def dissipator_operators(self) -> ArrayLike: """The dissipator operators.""" - return self._get_dissipator_operators(in_frame_basis=self._in_frame_basis) - - def _get_static_hamiltonian(self, in_frame_basis: Optional[bool] = False) -> Array: - """Get the constant Hamiltonian term. - - Args: - in_frame_basis: Flag for whether the returned ``static_operator`` should be in the basis - in which the frame is diagonal. - - Returns: - The static operator term. - """ - op = self._operator_collection.static_hamiltonian - if not in_frame_basis and self.rotating_frame is not None: - return self.rotating_frame.operator_out_of_frame_basis(op) - else: - return op - - def _set_static_hamiltonian( - self, - new_static_hamiltonian: Array, - operator_in_frame_basis: Optional[bool] = False, - ): - """Set the constant Hamiltonian term. - - Note that if the model has a rotating frame this will override any contributions to the - static term due to the frame transformation. - - Args: - new_static_hamiltonian: The static operator operator. - operator_in_frame_basis: Whether ``new_static_operator`` is already in the rotating - frame basis. - """ - if new_static_hamiltonian is None: - self._operator_collection.static_hamiltonian = None - else: - if not operator_in_frame_basis and self.rotating_frame is not None: - new_static_hamiltonian = self.rotating_frame.operator_into_frame_basis( - new_static_hamiltonian - ) - - self._operator_collection.static_hamiltonian = new_static_hamiltonian - - def _get_hamiltonian_operators(self, in_frame_basis: Optional[bool] = False) -> Tuple[Array]: - """Get the Hamiltonian operators, either in the frame basis or not. - - Args: - in_frame_basis: Whether to return in frame basis or not. - - Returns: - Hamiltonian operators. - """ - ham_ops = self._operator_collection.hamiltonian_operators - if not in_frame_basis and self.rotating_frame is not None: - ham_ops = self.rotating_frame.operator_out_of_frame_basis(ham_ops) - - return ham_ops + if self._operator_collection.dissipator_operators is None: + return None - def _get_static_dissipators(self, in_frame_basis: Optional[bool] = False) -> Tuple[Array]: - """Get the static dissipators, either in the frame basis or not. - - Args: - in_frame_basis: Whether to return in frame basis or not. - - Returns: - Dissipator operators. - """ - diss_ops = self._operator_collection.static_dissipators - if not in_frame_basis and self.rotating_frame is not None: - diss_ops = self.rotating_frame.operator_out_of_frame_basis(diss_ops) - - return diss_ops - - def _get_dissipator_operators(self, in_frame_basis: Optional[bool] = False) -> Tuple[Array]: - """Get the dissipator operators, either in the frame basis or not. - - Args: - in_frame_basis: Whether to return in frame basis or not. - - Returns: - Dissipator operators. - """ - diss_ops = self._operator_collection.dissipator_operators - if not in_frame_basis and self.rotating_frame is not None: - diss_ops = self.rotating_frame.operator_out_of_frame_basis(diss_ops) - - return diss_ops + if self.in_frame_basis: + return self._operator_collection.dissipator_operators + return self.rotating_frame.operator_out_of_frame_basis( + self._operator_collection.dissipator_operators + ) @property - def evaluation_mode(self) -> str: - """Numerical evaluation mode of the model. - - Available options: - - * ``'dense'``: Stores Hamiltonian and dissipator terms as dense Array types. - * ``'dense_vectorized'``: Stores the Hamiltonian and dissipator terms as - :math:`(dim^2,dim^2)` matrices that acts on a vectorized density matrix by - left-multiplication. Allows for direct evaluate generator. - * ``'sparse'``: Like dense, but matrices stored in sparse format. If the default Array - backend is JAX, uses JAX BCOO sparse arrays, otherwise uses scipy :class:`csr_matrix` - sparse type. - * ```sparse_vectorized```: Like dense_vectorized, but matrices stored in sparse format. If - the default Array backend is JAX, uses JAX BCOO sparse arrays, otherwise uses scipy - :class:`csr_matrix` sparse type. Note that JAX sparse mode is only recommended for use - on CPU. - - Raises: - NotImplementedError: If this property is set to something other than one of the above - modes. - """ - return self._evaluation_mode - - @evaluation_mode.setter - def evaluation_mode(self, new_mode: str): - if new_mode != self._evaluation_mode: - self._operator_collection = construct_lindblad_operator_collection( - evaluation_mode=new_mode, - static_hamiltonian=self._operator_collection.static_hamiltonian, - hamiltonian_operators=self._operator_collection.hamiltonian_operators, - static_dissipators=self._operator_collection.static_dissipators, - dissipator_operators=self._operator_collection.dissipator_operators, - ) + def array_library(self) -> Union[None, str]: + """Array library used to store the operators in the model.""" + return self._array_library - self.vectorized_operators = "vectorized" in new_mode - self._evaluation_mode = new_mode + @property + def vectorized(self) -> bool: + """Whether or not the Lindblad equation is vectorized.""" + return self._vectorized @property def rotating_frame(self): - """The rotating frame. - - This property can be set with a :class:`RotatingFrame` instance, or any valid constructor - argument to this class. - - Setting this property updates the internal operator to use the new frame. - """ + """The rotating frame.""" return self._rotating_frame - @rotating_frame.setter - def rotating_frame(self, rotating_frame: Union[Operator, Array, RotatingFrame]): - new_frame = RotatingFrame(rotating_frame) - - # convert static hamiltonian to new frame setup - static_ham = self._get_static_hamiltonian(in_frame_basis=True) - if static_ham is not None: - static_ham = -1j * static_ham - - new_static_hamiltonian = transfer_static_operator_between_frames( - static_ham, - new_frame=new_frame, - old_frame=self.rotating_frame, - ) - - if new_static_hamiltonian is not None: - new_static_hamiltonian = 1j * new_static_hamiltonian - - # convert hamiltonian operators and dissipator operators - ham_ops = self._get_hamiltonian_operators(in_frame_basis=True) - static_diss_ops = self._get_static_dissipators(in_frame_basis=True) - diss_ops = self._get_dissipator_operators(in_frame_basis=True) - - new_hamiltonian_operators = transfer_operators_between_frames( - ham_ops, - new_frame=new_frame, - old_frame=self.rotating_frame, - ) - new_static_dissipators = transfer_operators_between_frames( - static_diss_ops, - new_frame=new_frame, - old_frame=self.rotating_frame, - ) - new_dissipator_operators = transfer_operators_between_frames( - diss_ops, - new_frame=new_frame, - old_frame=self.rotating_frame, - ) - - self._rotating_frame = new_frame - - self._operator_collection = construct_lindblad_operator_collection( - evaluation_mode=self.evaluation_mode, - static_hamiltonian=new_static_hamiltonian, - hamiltonian_operators=new_hamiltonian_operators, - static_dissipators=new_static_dissipators, - dissipator_operators=new_dissipator_operators, - ) - - def evaluate_hamiltonian(self, time: float) -> Array: + def evaluate_hamiltonian(self, time: float) -> ArrayLike: """Evaluates Hamiltonian matrix at a given time. Args: @@ -535,12 +431,12 @@ def evaluate_hamiltonian(self, time: float) -> Array: ham, operator_in_frame_basis=True, return_in_frame_basis=self._in_frame_basis, - vectorized_operators=self.vectorized_operators, + vectorized_operators=self.vectorized, ) return ham - def evaluate(self, time: float) -> Array: + def evaluate(self, time: float) -> ArrayLike: """Evaluate the model in array format as a matrix, independent of state. Args: @@ -571,7 +467,7 @@ def evaluate(self, time: float) -> Array: "without dissipator signals." ) - if self.vectorized_operators: + if self.vectorized: out = self._operator_collection.evaluate(hamiltonian_sig_vals, dissipator_sig_vals) return self.rotating_frame.vectorized_map_into_frame( time, out, operator_in_frame_basis=True, return_in_frame_basis=self._in_frame_basis @@ -581,7 +477,7 @@ def evaluate(self, time: float) -> Array: "Non-vectorized Lindblad models cannot be represented without a given state." ) - def evaluate_rhs(self, time: Union[float, int], y: Array) -> Array: + def evaluate_rhs(self, time: float, y: ArrayLike) -> ArrayLike: """Evaluates the Lindblad model at a given time. Args: @@ -621,7 +517,7 @@ def evaluate_rhs(self, time: Union[float, int], y: Array) -> Array: y, operator_in_frame_basis=self._in_frame_basis, return_in_frame_basis=True, - vectorized_operators=self.vectorized_operators, + vectorized_operators=self.vectorized, ) rhs = self._operator_collection.evaluate_rhs( @@ -634,7 +530,7 @@ def evaluate_rhs(self, time: Union[float, int], y: Array) -> Array: rhs, operator_in_frame_basis=True, return_in_frame_basis=self._in_frame_basis, - vectorized_operators=self.vectorized_operators, + vectorized_operators=self.vectorized, ) else: @@ -645,12 +541,13 @@ def evaluate_rhs(self, time: Union[float, int], y: Array) -> Array: return rhs -def construct_lindblad_operator_collection( - evaluation_mode: str, - static_hamiltonian: Union[None, Array, csr_matrix], - hamiltonian_operators: Union[None, Array, List[csr_matrix]], - static_dissipators: Union[None, Array, csr_matrix], - dissipator_operators: Union[None, Array, List[csr_matrix]], +def _get_lindblad_operator_collection( + array_library: Optional[str], + vectorized: bool, + static_hamiltonian: Optional[ArrayLike], + hamiltonian_operators: Optional[ArrayLike], + static_dissipators: Optional[ArrayLike], + dissipator_operators: Optional[ArrayLike], ) -> Union[ LindbladCollection, ScipySparseLindbladCollection, @@ -660,61 +557,44 @@ def construct_lindblad_operator_collection( """Construct a Lindblad operator collection. Args: - evaluation_mode: String specifying new mode. Available options are ``'dense'``, - ``'sparse'``, ``'dense_vectorized'``, and ``'sparse_vectorized'``. + array_library: Array library to use. + vectorized: Whether or not to vectorize the Lindblad equation. static_hamiltonian: Constant part of the Hamiltonian. hamiltonian_operators: Operators in Hamiltonian with time-dependent coefficients. static_dissipators: Dissipation operators with coefficient 1. dissipator_operators: Dissipation operators with variable coefficients. Returns: - BaseLindbladOperatorCollection: Right-hand side evaluation object. - - Raises: - NotImplementedError: If ``evaluation_mode`` is not one of the above supported evaluation - modes. + Union[ + LindbladCollection, + ScipySparseLindbladCollection, + VectorizedLindbladCollection, + ScipySparseVectorizedLindbladCollection, + ]: Right-hand side evaluation object. """ - # raise warning if sparse mode set with JAX not on cpu - if ( - Array.default_backend() == "jax" - and "sparse" in evaluation_mode - and jax.default_backend() != "cpu" - ): - warn( - """Using sparse mode with JAX is primarily recommended for use on CPU.""", - stacklevel=2, - ) + operator_kwargs = { + "static_hamiltonian": static_hamiltonian, + "hamiltonian_operators": hamiltonian_operators, + "static_dissipators": static_dissipators, + "dissipator_operators": dissipator_operators, + } + + if array_library == "scipy_sparse": + if vectorized: + return ScipySparseVectorizedLindbladCollection(**operator_kwargs) + + return ScipySparseLindbladCollection(**operator_kwargs) + + if array_library == "jax_sparse": + # warn that sparse mode when using JAX is primarily recommended for use on CPU + if jax.default_backend() != "cpu": + warn( + """Using sparse mode with JAX is primarily recommended for use on CPU.""", + stacklevel=2, + ) - if evaluation_mode == "dense": - # CollectionClass = DenseLindbladCollection - pass - elif evaluation_mode == "sparse": - if Array.default_backend() == "jax": - # CollectionClass = JAXSparseLindbladCollection - pass - else: - # CollectionClass = SparseLindbladCollection - pass - elif evaluation_mode == "dense_vectorized": - # CollectionClass = DenseVectorizedLindbladCollection - pass - elif evaluation_mode == "sparse_vectorized": - if Array.default_backend() == "jax": - # CollectionClass = JAXSparseVectorizedLindbladCollection - pass - else: - # CollectionClass = SparseVectorizedLindbladCollection - pass - else: - raise NotImplementedError( - f"Evaluation mode '{evaluation_mode}' is not supported. See " - "help(LindbladModel.evaluation_mode) for available options." - ) + if vectorized: + return VectorizedLindbladCollection(**operator_kwargs, array_library=array_library) - return CollectionClass( - static_hamiltonian=static_hamiltonian, - hamiltonian_operators=hamiltonian_operators, - static_dissipators=static_dissipators, - dissipator_operators=dissipator_operators, - ) + return LindbladCollection(**operator_kwargs, array_library=array_library) diff --git a/releasenotes/notes/model-changes-221fb43b3e9cc0dd.yaml b/releasenotes/notes/model-changes-221fb43b3e9cc0dd.yaml new file mode 100644 index 000000000..05ff10a34 --- /dev/null +++ b/releasenotes/notes/model-changes-221fb43b3e9cc0dd.yaml @@ -0,0 +1,12 @@ +--- +upgrade: + - | + The interface for :class:`.GeneratorModel`, :class:`.HamiltonianModel`, and + :class:`.LindbladModel` has been modified. The ``copy`` method has been removed, and all setter + methods other than ``in_frame_basis`` and ``signals`` have been removed. The ``evaluation_mode`` + construction argument has been replaced by ``array_library``, which controls which array library + is used internally to store and evaluate operations, and the additional ``vectorized`` boolean + argument has been added to :class:`.LindbladModel` to control whether the equation is evaluated + in vectorized mode. Note that, regardless of array library used, dense arrays must be supplied + to the constructors of these classes, due to peculiarities of the internal setup for sparse + libraries. \ No newline at end of file diff --git a/test/dynamics/common.py b/test/dynamics/common.py index c504403c2..a08883428 100644 --- a/test/dynamics/common.py +++ b/test/dynamics/common.py @@ -63,9 +63,15 @@ 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 isinstance(A, list): + A = [x.todense() for x in A] + else: + A = A.todense() if any("sparse" in x for x in DYNAMICS_NUMPY_ALIAS.infer_libs(B)): - B = B.todense() + if isinstance(B, list): + B = [x.todense() for x in B] + else: + B = B.todense() A = np.array(A) B = np.array(B) diff --git a/test/dynamics/models/test_generator_model.py b/test/dynamics/models/test_generator_model.py index 5d1b6bb93..34ba70e87 100644 --- a/test/dynamics/models/test_generator_model.py +++ b/test/dynamics/models/test_generator_model.py @@ -8,25 +8,26 @@ # # 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 generator_models.py. """ -from scipy.sparse import issparse, csr_matrix +from functools import partial + from scipy.linalg import expm import numpy as np import numpy.random as rand from qiskit import QiskitError from qiskit.quantum_info.operators import Operator +from qiskit_dynamics import DYNAMICS_NUMPY as unp +from qiskit_dynamics.arraylias.alias import ArrayLike from qiskit_dynamics.models import GeneratorModel, RotatingFrame from qiskit_dynamics.models.generator_model import ( - transfer_static_operator_between_frames, - transfer_operators_between_frames, + _static_operator_into_frame_basis, + _operators_into_frame_basis, ) from qiskit_dynamics.signals import Signal, SignalList -from qiskit_dynamics.array import Array, wrap -from qiskit_dynamics.type_utils import to_array -from ..common import QiskitDynamicsTestCase, TestJaxBase +from ..common import QiskitDynamicsTestCase, test_array_backends class TestGeneratorModelErrors(QiskitDynamicsTestCase): @@ -72,31 +73,108 @@ def test_signals_bad_format(self): model.signals = lambda t: t -class TestGeneratorModel(QiskitDynamicsTestCase): +@partial(test_array_backends, array_libraries=["numpy", "jax", "jax_sparse", "scipy_sparse"]) +class Test_into_frame_basis_functions: + """Tests for _static_operator_into_frame_basis and _operators_into_frame_basis.""" + + def test_all_None(self): + """Test all arguments being None.""" + + static_operator = _static_operator_into_frame_basis( + None, RotatingFrame(None), array_library=self.array_library() + ) + operators = _operators_into_frame_basis( + None, RotatingFrame(None), array_library=self.array_library() + ) + + self.assertTrue(static_operator is None) + self.assertTrue(operators is None) + + def test_array_inputs_diagonal_frame(self): + """Test correct handling when operators are arrays for diagonal frames.""" + + static_operator = -1j * np.array([[1.0, 0.0], [0.0, -1.0]]) + operators = -1j * np.array([[[0.0, 1.0], [1.0, 0.0]], [[0.0, -1j], [1j, 0.0]]]) + rotating_frame = RotatingFrame(-1j * np.array([1.0, -1.0])) + + out_static = _static_operator_into_frame_basis( + static_operator, rotating_frame=rotating_frame, array_library=self.array_library() + ) + out_operators = _operators_into_frame_basis( + operators, rotating_frame=rotating_frame, array_library=self.array_library() + ) + + self.assertArrayType(out_static) + self.assertArrayType(out_operators) + + self.assertAllClose(out_operators, operators) + self.assertAllClose(out_static, np.zeros((2, 2))) + + def test_array_inputs_pseudo_random(self): + """Test correct handling when operators are pseudo random arrays.""" + + rng = np.random.default_rng(34233) + num_terms = 3 + dim = 5 + b = 1.0 # bound on size of random terms + static_operator = rng.uniform(low=-b, high=b, size=(dim, dim)) + 1j * rng.uniform( + low=-b, high=b, size=(dim, dim) + ) + operators = rng.uniform(low=-b, high=b, size=(num_terms, dim, dim)) + 1j * rng.uniform( + low=-b, high=b, size=(dim, dim) + ) + + rotating_frame = rng.uniform(low=-b, high=b, size=(dim, dim)) + 1j * rng.uniform( + low=-b, high=b, size=(dim, dim) + ) + rotating_frame = RotatingFrame(rotating_frame - rotating_frame.conj().transpose()) + + out_static = _static_operator_into_frame_basis( + static_operator, rotating_frame=rotating_frame, array_library=self.array_library() + ) + out_operators = _operators_into_frame_basis( + operators, rotating_frame=rotating_frame, array_library=self.array_library() + ) + + self.assertArrayType(out_static) + self.assertArrayType(out_operators) + + expected_static = rotating_frame.operator_into_frame_basis( + static_operator - rotating_frame.frame_operator + ) + expected_operators = [rotating_frame.operator_into_frame_basis(op) for op in operators] + + self.assertAllClose(out_static, expected_static) + self.assertAllClose(out_operators, expected_operators) + + +@partial(test_array_backends, array_libraries=["numpy", "jax", "jax_sparse", "scipy_sparse"]) +class TestGeneratorModel: """Tests for GeneratorModel.""" 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 simple model elements.""" + self.X = Operator.from_label("X").data + self.Y = Operator.from_label("Y").data + self.Z = Operator.from_label("Z").data - # define a basic model - w = 2.0 - r = 0.5 - operators = [-1j * 2 * np.pi * self.Z / 2, -1j * 2 * np.pi * r * self.X / 2] - signals = [w, Signal(1.0, w)] + # define basic model elements + self.w = 2.0 + self.r = 0.5 + self.operators = [-1j * 2 * np.pi * self.Z / 2, -1j * 2 * np.pi * self.r * self.X / 2] + self.signals = [self.w, Signal(1.0, self.w)] - self.w = 2 - self.r = r - self.basic_model = GeneratorModel(operators=operators, signals=signals) + self.basic_model = GeneratorModel( + operators=self.operators, signals=self.signals, array_library=self.array_library() + ) def test_diag_frame_operator_basic_model(self): """Test setting a diagonal frame operator for the internally set up basic model. """ - self._basic_frame_evaluate_test(Array([1j, -1j]), 1.123) - self._basic_frame_evaluate_test(Array([1j, -1j]), np.pi) + self._basic_frame_evaluate_test(np.array([1j, -1j]), 1.123) + self._basic_frame_evaluate_test(np.array([1j, -1j]), np.pi) def test_non_diag_frame_operator_basic_model(self): """Test setting a non-diagonal frame operator for the internally @@ -109,15 +187,21 @@ def _basic_frame_evaluate_test(self, frame_operator, t): """Routine for testing setting of valid frame operators using the basic_model. """ - self.basic_model.rotating_frame = frame_operator + basic_model = GeneratorModel( + operators=self.operators, + signals=self.signals, + rotating_frame=frame_operator, + array_library=self.array_library(), + ) # convert to 2d array if isinstance(frame_operator, Operator): - frame_operator = Array(frame_operator.data) - if isinstance(frame_operator, Array) and frame_operator.ndim == 1: + frame_operator = frame_operator.data + + if isinstance(frame_operator, ArrayLike) and frame_operator.ndim == 1: frame_operator = np.diag(frame_operator) - value = self.basic_model(t) + value = basic_model(t) i2pi = -1j * 2 * np.pi @@ -133,6 +217,7 @@ def _basic_frame_evaluate_test(self, frame_operator, t): - frame_operator ) + self.assertArrayType(value) self.assertAllClose(value, expected) def test_evaluate_no_frame_basic_model(self): @@ -142,26 +227,31 @@ def test_evaluate_no_frame_basic_model(self): value = self.basic_model(t) i2pi = -1j * 2 * np.pi d_coeff = self.r * np.cos(2 * np.pi * self.w * t) - expected = i2pi * self.w * self.Z.data / 2 + i2pi * d_coeff * self.X.data / 2 + expected = i2pi * self.w * self.Z / 2 + i2pi * d_coeff * self.X / 2 + self.assertArrayType(value) self.assertAllClose(value, expected) def test_evaluate_generator_in_frame_basis_basic_model(self): """Test generator evaluation in frame basis in the basic_model.""" - frame_op = -1j * (self.X + 0.2 * self.Y + 0.1 * self.Z).data + frame_op = -1j * (self.X + 0.2 * self.Y + 0.1 * self.Z) - # enter the frame given by the -1j * X - self.basic_model.rotating_frame = frame_op + basic_model = GeneratorModel( + operators=self.operators, + signals=self.signals, + rotating_frame=frame_op, + array_library=self.array_library(), + ) # set to operate in frame basis - self.basic_model.in_frame_basis = True + basic_model.in_frame_basis = True # get the frame basis that is used in model - _, U = wrap(np.linalg.eigh)(1j * frame_op) + _, U = unp.linalg.eigh(1j * frame_op) t = 3.21412 - value = self.basic_model(t) + value = basic_model(t) # compose the frame basis transformation with the exponential # frame rotation (this will be multiplied on the right) @@ -170,11 +260,7 @@ def test_evaluate_generator_in_frame_basis_basic_model(self): i2pi = -1j * 2 * np.pi d_coeff = self.r * np.cos(2 * np.pi * self.w * t) - expected = ( - Uadj - @ (i2pi * self.w * self.Z.data / 2 + i2pi * d_coeff * self.X.data / 2 - frame_op) - @ U - ) + expected = Uadj @ (i2pi * self.w * self.Z / 2 + i2pi * d_coeff * self.X / 2 - frame_op) @ U self.assertAllClose(value, expected) @@ -188,17 +274,16 @@ def test_evaluate_pseudorandom(self): rand_op = rng.uniform(low=-b, high=b, size=(dim, dim)) + 1j * rng.uniform( low=-b, high=b, size=(dim, dim) ) - frame_op = Array(rand_op - rand_op.conj().transpose()) + frame_op = rand_op - rand_op.conj().transpose() randoperators = rng.uniform(low=-b, high=b, size=(num_terms, dim, dim)) + 1j * rng.uniform( low=-b, high=b, size=(num_terms, dim, dim) ) - rand_coeffs = Array( - rng.uniform(low=-b, high=b, size=(num_terms)) - + 1j * rng.uniform(low=-b, high=b, size=(num_terms)) + rand_coeffs = rng.uniform(low=-b, high=b, size=(num_terms)) + 1j * rng.uniform( + low=-b, high=b, size=(num_terms) ) - rand_carriers = Array(rng.uniform(low=-b, high=b, size=(num_terms))) - rand_phases = Array(rng.uniform(low=-b, high=b, size=(num_terms))) + rand_carriers = rng.uniform(low=-b, high=b, size=(num_terms)) + rand_phases = rng.uniform(low=-b, high=b, size=(num_terms)) self._test_evaluate(frame_op, randoperators, rand_coeffs, rand_carriers, rand_phases) @@ -209,18 +294,16 @@ def test_evaluate_pseudorandom(self): rand_op = rng.uniform(low=-b, high=b, size=(dim, dim)) + 1j * rng.uniform( low=-b, high=b, size=(dim, dim) ) - frame_op = Array(rand_op - rand_op.conj().transpose()) - randoperators = Array( - rng.uniform(low=-b, high=b, size=(num_terms, dim, dim)) - + 1j * rng.uniform(low=-b, high=b, size=(num_terms, dim, dim)) + frame_op = rand_op - rand_op.conj().transpose() + randoperators = rng.uniform(low=-b, high=b, size=(num_terms, dim, dim)) + 1j * rng.uniform( + low=-b, high=b, size=(num_terms, dim, dim) ) - rand_coeffs = Array( - rng.uniform(low=-b, high=b, size=(num_terms)) - + 1j * rng.uniform(low=-b, high=b, size=(num_terms)) + rand_coeffs = rng.uniform(low=-b, high=b, size=(num_terms)) + 1j * rng.uniform( + low=-b, high=b, size=(num_terms) ) - rand_carriers = Array(rng.uniform(low=-b, high=b, size=(num_terms))) - rand_phases = Array(rng.uniform(low=-b, high=b, size=(num_terms))) + rand_carriers = rng.uniform(low=-b, high=b, size=(num_terms)) + rand_phases = rng.uniform(low=-b, high=b, size=(num_terms)) self._test_evaluate(frame_op, randoperators, rand_coeffs, rand_carriers, rand_phases) @@ -237,8 +320,12 @@ def env(t): return env sig_list.append(Signal(get_env_func(), freq, phase)) - model = GeneratorModel(operators=operators, signals=sig_list) - model.rotating_frame = frame_op + model = GeneratorModel( + operators=operators, + signals=sig_list, + rotating_frame=frame_op, + array_library=self.array_library(), + ) value = model(1.0) coeffs = np.real(coefficients * np.exp(1j * 2 * np.pi * carriers * 1.0 + 1j * phases)) @@ -258,16 +345,6 @@ def env(t): ) self.assertAllClose(value, expected) - if value.backend != "jax": - model.evaluation_mode = "sparse" - state = np.arange(operators.shape[-1] * 4).reshape(operators.shape[-1], 4) - value = model(1.0) - if issparse(value): - value = value.toarray() - self.assertAllClose(value, expected) - - val_with_state = model(1.0, state) - self.assertAllClose(val_with_state, value @ state) def test_signal_setting(self): """Test updating the signals.""" @@ -280,7 +357,8 @@ def test_signal_setting(self): i2pi = -1j * 2 * np.pi Z_coeff = (2 * t) * np.cos(2 * np.pi * 1 * t) X_coeff = self.r * (t**2) * np.cos(2 * np.pi * 2 * t) - expected = i2pi * Z_coeff * self.Z.data / 2 + i2pi * X_coeff * self.X.data / 2 + expected = i2pi * Z_coeff * self.Z / 2 + i2pi * X_coeff * self.X / 2 + self.assertArrayType(value) self.assertAllClose(value, expected) def test_signal_setting_None(self): @@ -293,36 +371,51 @@ def test_evaluate_analytic(self): """Test for checking that with known operators that the Model returns the analyticlly known values.""" - test_operator_list = Array([self.X, self.Y, self.Z]) + test_operator_list = [self.X, self.Y, self.Z] signals = SignalList([Signal(1, j / 3) for j in range(3)]) simple_model = GeneratorModel( - operators=test_operator_list, static_operator=None, signals=signals, rotating_frame=None + operators=test_operator_list, + static_operator=None, + signals=signals, + rotating_frame=None, + array_library=self.array_library(), ) res = simple_model.evaluate(2) - self.assertAllClose(res, Array([[-0.5 + 0j, 1.0 + 0.5j], [1.0 - 0.5j, 0.5 + 0j]])) + self.assertArrayType(res) + self.assertAllClose(res, np.array([[-0.5 + 0j, 1.0 + 0.5j], [1.0 - 0.5j, 0.5 + 0j]])) - simple_model.static_operator = np.eye(2) + simple_model = GeneratorModel( + operators=test_operator_list, + static_operator=self.asarray(np.eye(2)), + signals=signals, + rotating_frame=None, + array_library=self.array_library(), + ) res = simple_model.evaluate(2) - self.assertAllClose(res, Array([[0.5 + 0j, 1.0 + 0.5j], [1.0 - 0.5j, 1.5 + 0j]])) + self.assertArrayType(res) + self.assertAllClose(res, np.array([[0.5 + 0j, 1.0 + 0.5j], [1.0 - 0.5j, 1.5 + 0j]])) def test_evaluate_in_frame_basis_analytic(self): """Tests evaluation in rotating frame against analytic values, both in specified basis and in frame basis.""" - test_operator_list = Array([self.X, self.Y, self.Z]) + test_operator_list = [self.X, self.Y, self.Z] signals = SignalList([Signal(1, j / 3) for j in range(3)]) + fop = np.array([[0, 1j], [1j, 0]]) simple_model = GeneratorModel( - operators=test_operator_list, static_operator=None, signals=signals, rotating_frame=None + operators=test_operator_list, + static_operator=self.asarray(np.eye(2)), + signals=signals, + rotating_frame=fop, + array_library=self.array_library(), ) - simple_model.static_operator = np.eye(2) - fop = Array([[0, 1j], [1j, 0]]) - simple_model.rotating_frame = fop res = simple_model(2.0) expected = ( expm(np.array(-2 * fop)) - @ (Array([[0.5 + 0j, 1.0 + 0.5j], [1.0 - 0.5j, 1.5 + 0j]]) - fop) + @ (np.array([[0.5 + 0j, 1.0 + 0.5j], [1.0 - 0.5j, 1.5 + 0j]]) - fop) @ expm(np.array(2 * fop)) ) + self.assertArrayType(res) self.assertAllClose(res, expected) simple_model.in_frame_basis = True @@ -334,36 +427,43 @@ def test_evaluate_in_frame_basis_analytic(self): @ np.array([[1, 1], [1, -1]]) / np.sqrt(2) ) - + self.assertArrayType(res) self.assertAllClose(res, expected) def test_order_of_assigning_properties(self): - """Tests whether setting the frame, static_operator, and signals - in the constructor is the same as constructing without - and then assigning them in any order. + """Tests whether setting signals in the constructor is the same as constructing without + and then assigning them. """ - paulis = Array([self.X, self.Y, self.Z]) - extra = Array(np.eye(2)) - state = Array([0.2, 0.5]) + paulis = np.array([self.X, self.Y, self.Z]) + extra = np.eye(2) + state = np.array([0.2, 0.5]) t = 2 sarr = [Signal(1, j / 3) for j in range(3)] sigvals = SignalList(sarr)(t) - farr = Array(np.array([[3j, 2j], [2j, 0]])) - farr2 = Array(np.array([[1j, 2], [-2, 3j]])) - evals, evect = wrap(np.linalg.eig)(farr) + farr = np.array([[3j, 2j], [2j, 0]]) + evals, evect = np.linalg.eig(farr) diafarr = np.diag(evals) paulis_in_frame_basis = np.conjugate(np.transpose(evect)) @ paulis @ evect ## Run checks without rotating frame for now - gm1 = GeneratorModel(operators=paulis, signals=sarr, static_operator=extra) - gm2 = GeneratorModel(operators=paulis, static_operator=extra) + gm1 = GeneratorModel( + operators=paulis, + signals=sarr, + static_operator=extra, + array_library=self.array_library(), + ) + gm2 = GeneratorModel( + operators=paulis, static_operator=extra, array_library=self.array_library() + ) gm2.signals = sarr - gm3 = GeneratorModel(operators=paulis, static_operator=extra) + gm3 = GeneratorModel( + operators=paulis, static_operator=extra, array_library=self.array_library() + ) gm3.signals = SignalList(sarr) # All should be the same, because there are no frames involved @@ -376,7 +476,14 @@ def test_order_of_assigning_properties(self): t31 = gm3.evaluate(t) gm3.in_frame_basis = True t32 = gm3.evaluate(t) - t_analytical = Array([[0.5, 1.0 + 0.5j], [1.0 - 0.5j, 1.5]]) + t_analytical = np.array([[0.5, 1.0 + 0.5j], [1.0 - 0.5j, 1.5]]) + + self.assertArrayType(t11) + self.assertArrayType(t12) + self.assertArrayType(t21) + self.assertArrayType(t22) + self.assertArrayType(t31) + self.assertArrayType(t32) self.assertAllClose(t11, t12) self.assertAllClose(t11, t21) @@ -389,7 +496,10 @@ def test_order_of_assigning_properties(self): ts1 = gm1.evaluate_rhs(t, state) gm1.in_frame_basis = True ts2 = gm1.evaluate_rhs(t, state) - ts_analytical = Array([0.6 + 0.25j, 0.95 - 0.1j]) + ts_analytical = np.array([0.6 + 0.25j, 0.95 - 0.1j]) + + self.assertArrayType(ts1) + self.assertArrayType(ts2) self.assertAllClose(ts1, ts2) self.assertAllClose(ts1, ts_analytical) @@ -398,31 +508,31 @@ def test_order_of_assigning_properties(self): self.assertAllClose(gm1(t) @ state, ts1) ## Now, run checks with frame - # If passing a frame in the first place, operators must be in frame basis.abs - # Testing at the same time whether having static_operatorrift = None is an issue. + # If passing a frame in the first place, operators must be in frame basis. + # Testing at the same time whether having static_operator = None is an issue. gm1 = GeneratorModel( - operators=paulis, signals=sarr, rotating_frame=RotatingFrame(farr), in_frame_basis=True + operators=paulis, + signals=sarr, + rotating_frame=RotatingFrame(farr), + in_frame_basis=True, + array_library=self.array_library(), + ) + gm2 = GeneratorModel( + operators=paulis, + rotating_frame=farr, + in_frame_basis=True, + array_library=self.array_library(), ) - gm2 = GeneratorModel(operators=paulis, rotating_frame=farr, in_frame_basis=True) gm2.signals = SignalList(sarr) - gm3 = GeneratorModel(operators=paulis, rotating_frame=farr, in_frame_basis=True) + gm3 = GeneratorModel( + operators=paulis, + rotating_frame=farr, + in_frame_basis=True, + array_library=self.array_library(), + ) gm3.signals = sarr - # Does adding a frame after make a difference? - # If so, does it make a difference if we add signals or the frame first? - gm4 = GeneratorModel(operators=paulis, in_frame_basis=True) - gm4.signals = sarr - gm4.rotating_frame = farr - gm5 = GeneratorModel(operators=paulis, in_frame_basis=True) - gm5.rotating_frame = farr - gm5.signals = sarr - gm6 = GeneratorModel(operators=paulis, signals=sarr, in_frame_basis=True) - gm6.rotating_frame = farr - # If we go to one frame, then transform back, does this make a difference? - gm7 = GeneratorModel(operators=paulis, signals=sarr, in_frame_basis=True) - gm7.rotating_frame = farr2 - gm7.rotating_frame = farr - - t_in_frame_actual = Array( + + t_in_frame_actual = ( np.diag(np.exp(-t * evals)) @ (np.tensordot(sigvals, paulis_in_frame_basis, axes=1) - diafarr) @ np.diag(np.exp(t * evals)) @@ -430,30 +540,25 @@ def test_order_of_assigning_properties(self): tf1 = gm1.evaluate(t) tf2 = gm2.evaluate(t) tf3 = gm3.evaluate(t) - tf4 = gm4.evaluate(t) - tf5 = gm5.evaluate(t) - tf6 = gm6.evaluate(t) - tf7 = gm7.evaluate(t) + self.assertArrayType(tf1) + self.assertArrayType(tf2) + self.assertArrayType(tf3) self.assertAllClose(t_in_frame_actual, tf1) self.assertAllClose(tf1, tf2) self.assertAllClose(tf1, tf3) - self.assertAllClose(tf1, tf4) - self.assertAllClose(tf1, tf5) - self.assertAllClose(tf1, tf6) - self.assertAllClose(tf1, tf7) def test_evaluate_rhs_lmult_equivalent_analytic(self): """Tests whether evaluate(t) @ state == evaluate_rhs(t,state) for analytically known values.""" - paulis = Array([self.X, self.Y, self.Z]) - extra = Array(np.eye(2)) + paulis = np.array([self.X, self.Y, self.Z]) + extra = np.eye(2) t = 2 - farr = Array(np.array([[3j, 2j], [2j, 0]])) - evals, evect = wrap(np.linalg.eig)(farr) + farr = np.array([[3j, 2j], [2j, 0]]) + evals, evect = np.linalg.eig(farr) diafarr = np.diag(evals) paulis_in_frame_basis = np.conjugate(np.transpose(evect)) @ paulis @ evect @@ -462,13 +567,13 @@ def test_evaluate_rhs_lmult_equivalent_analytic(self): sarr = [Signal(1, j / 3) for j in range(3)] sigvals = SignalList(sarr)(t) - t_in_frame_actual = Array( + t_in_frame_actual = ( np.diag(np.exp(-t * evals)) @ (np.tensordot(sigvals, paulis_in_frame_basis, axes=1) + extra_in_basis - diafarr) @ np.diag(np.exp(t * evals)) ) - state = Array([0.3, 0.1]) + state = np.array([0.3, 0.1]) state_in_frame_basis = np.conjugate(np.transpose(evect)) @ state gm1 = GeneratorModel( @@ -477,14 +582,17 @@ def test_evaluate_rhs_lmult_equivalent_analytic(self): rotating_frame=farr, static_operator=extra, in_frame_basis=True, + array_library=self.array_library(), ) - self.assertAllClose(gm1(t), t_in_frame_actual) + res = gm1(t) + self.assertArrayType(res) + self.assertAllClose(res, t_in_frame_actual) self.assertAllClose( gm1(t, state_in_frame_basis), t_in_frame_actual @ state_in_frame_basis, ) - t_not_in_frame_actual = Array( + t_not_in_frame_actual = ( expm(np.array(-t * farr)) @ (np.tensordot(sigvals, paulis, axes=1) + extra - farr) @ expm(np.array(t * farr)) @@ -496,28 +604,14 @@ def test_evaluate_rhs_lmult_equivalent_analytic(self): rotating_frame=farr, static_operator=extra, in_frame_basis=False, + array_library=self.array_library(), ) gm1.in_frame_basis = False - self.assertAllClose(gm2(t), t_not_in_frame_actual) + res = gm2(t) + self.assertArrayType(res) + self.assertAllClose(res, t_not_in_frame_actual) self.assertAllClose(gm1(t, state), t_not_in_frame_actual @ state) - # now, remove the frame - gm1.rotating_frame = None - gm2.rotating_frame = None - - t_expected = np.tensordot(sigvals, paulis, axes=1) + extra - - state_in_frame_basis = state - - self.assertAllClose(gm1._get_operators(True), gm1._get_operators(False)) - self.assertAllClose(gm1._get_static_operator(True), gm1._get_static_operator(False)) - - gm1.in_frame_basis = True - self.assertAllClose(gm1(t), t_expected) - self.assertAllClose(gm1(t, state), t_expected @ state_in_frame_basis) - self.assertAllClose(gm2(t), t_expected) - self.assertAllClose(gm2(t, state), t_expected @ state_in_frame_basis) - def test_evaluate_rhs_vectorized_pseudorandom(self): """Test for whether evaluating a model at m different states, each with an n-length statevector, by passing @@ -534,7 +628,12 @@ def test_evaluate_rhs_vectorized_pseudorandom(self): operators = rand.uniform(-1, 1, (k, n, n)) - gm = GeneratorModel(operators=operators, static_operator=None, signals=sig_list) + gm = GeneratorModel( + operators=operators, + static_operator=None, + signals=sig_list, + array_library=self.array_library(), + ) self.assertTrue(gm.evaluate_rhs(t, normal_states).shape == (n,)) self.assertTrue(gm.evaluate_rhs(t, vectorized_states).shape == (n, m)) for i in range(m): @@ -543,10 +642,16 @@ def test_evaluate_rhs_vectorized_pseudorandom(self): gm.evaluate_rhs(t, vectorized_states[:, i]), ) + # add pseudo frame to test farr = rand.uniform(-1, 1, (n, n)) farr = farr - np.conjugate(np.transpose(farr)) - - gm.rotating_frame = farr + gm = gm = GeneratorModel( + operators=operators, + static_operator=None, + signals=sig_list, + rotating_frame=farr, + array_library=self.array_library(), + ) self.assertTrue(gm.evaluate_rhs(t, normal_states).shape == (n,)) self.assertTrue(gm.evaluate_rhs(t, vectorized_states).shape == (n, m)) @@ -571,287 +676,74 @@ def test_evaluate_rhs_vectorized_pseudorandom(self): def test_evaluate_static(self): """Test evaluation of a GeneratorModel with only a static component.""" - static_model = GeneratorModel(static_operator=self.X) + static_model = GeneratorModel(static_operator=self.X, array_library=self.array_library()) self.assertAllClose(self.X, static_model(1.0)) # now with frame frame_op = -1j * (self.Z + 1.232 * self.Y) - static_model.rotating_frame = frame_op + static_model = GeneratorModel( + static_operator=self.X, rotating_frame=frame_op, array_library=self.array_library() + ) t = 1.2312 - expected = expm(-frame_op.data * t) @ (self.X - frame_op) @ expm(frame_op.data * t) - self.assertAllClose(expected, static_model(t)) + expected = expm(-frame_op * t) @ (self.X - frame_op) @ expm(frame_op * t) + res = static_model(t) + self.assertArrayType(res) + self.assertAllClose(expected, res) def test_get_operators_when_None(self): """Test getting operators when None.""" - model = GeneratorModel(static_operator=np.array([[1.0, 0.0], [0.0, -1.0]])) + model = GeneratorModel( + static_operator=np.array([[1.0, 0.0], [0.0, -1.0]]), array_library=self.array_library() + ) self.assertTrue(model.operators is None) -class TestGeneratorModelJax(TestGeneratorModel, TestJaxBase): - """Jax version of TestGeneratorModel tests.""" - - def test_jitable_funcs(self): - """Tests whether all functions are jitable. - Checks if having a frame makes a difference, as well as - all jax-compatible evaluation_modes.""" - self.jit_wrap(self.basic_model.evaluate)(1.0) - self.jit_wrap(self.basic_model.evaluate_rhs)(1, Array(np.array([0.2, 0.4]))) - - self.basic_model.rotating_frame = Array(np.array([[3j, 2j], [2j, 0]])) - - self.jit_wrap(self.basic_model.evaluate)(1) - self.jit_wrap(self.basic_model.evaluate_rhs)(1, Array(np.array([0.2, 0.4]))) - - self.basic_model.rotating_frame = None - - def test_gradable_funcs(self): - """Tests whether all functions are gradable. - Checks if having a frame makes a difference, as well as - all jax-compatible evaluation_modes.""" - self.jit_grad_wrap(self.basic_model.evaluate)(1.0) - self.jit_grad_wrap(self.basic_model.evaluate_rhs)(1.0, Array(np.array([0.2, 0.4]))) - - self.basic_model.rotating_frame = Array(np.array([[3j, 2j], [2j, 0]])) - - self.jit_grad_wrap(self.basic_model.evaluate)(1.0) - self.jit_grad_wrap(self.basic_model.evaluate_rhs)(1.0, Array(np.array([0.2, 0.4]))) - - self.basic_model.rotating_frame = None - - -class TestGeneratorModelSparse(QiskitDynamicsTestCase): - """Sparse-mode specific tests.""" +@partial(test_array_backends, array_libraries=["jax", "jax_sparse"]) +class TestGeneratorModelJAXTransformations: + """Test GeneratorModel under JAX transformations.""" def setUp(self): + """Build simple model elements.""" self.X = Operator.from_label("X").data self.Y = Operator.from_label("Y").data self.Z = Operator.from_label("Z").data - def test_switch_modes_and_evaluate(self): - """Test construction of a model, switching modes, and evaluating.""" - - model = GeneratorModel(static_operator=self.Z, operators=[self.X], signals=[1.0]) - self.assertAllClose(model(1.0), self.Z + self.X) + # define basic model elements + self.w = 2.0 + self.r = 0.5 + self.operators = [-1j * 2 * np.pi * self.Z / 2, -1j * 2 * np.pi * self.r * self.X / 2] + self.signals = [self.w, Signal(1.0, self.w)] - model.evaluation_mode = "sparse" - val = model(1.0) - self.validate_generator_eval(val, self.Z + self.X) - - model.evaluation_mode = "dense" - self.assertAllClose(model(1.0), self.Z + self.X) - - def test_frame_change_sparse(self): - """Test setting a frame after instantiation in sparse mode and evaluating.""" - model = GeneratorModel( - static_operator=self.Z, operators=[self.X], signals=[1.0], evaluation_mode="sparse" + self.basic_model = GeneratorModel( + operators=self.operators, signals=self.signals, array_library=self.array_library() ) - # test non-diagonal frame - model.rotating_frame = self.Z - expected = expm(1j * self.Z) @ ((1 + 1j) * self.Z + self.X) @ expm(-1j * self.Z) - val = model(1.0) - self.assertAllClose(val, expected) - - # test diagonal frame - model.rotating_frame = np.array([1.0, -1.0]) - val = model(1.0) - self.validate_generator_eval(val, expected) - - def test_switching_to_sparse_with_frame(self): - """Test switching to sparse with a frame already set.""" - - model = GeneratorModel( - static_operator=self.Z, - operators=[self.X], - signals=[1.0], - rotating_frame=np.array([1.0, -1.0]), + self.basic_model_w_frame = GeneratorModel( + operators=self.operators, + signals=self.signals, + rotating_frame=np.array([[3j, 2j], [2j, 0]]), + array_library=self.array_library(), ) - model.evaluation_mode = "sparse" - - expected = expm(1j * self.Z) @ ((1 + 1j) * self.Z + self.X) @ expm(-1j * self.Z) - val = model(1.0) - self.validate_generator_eval(val, expected) - - def validate_generator_eval(self, op, expected): - """Validate that op is sparse and agrees with expected.""" - self.assertTrue(issparse(op)) - self.assertAllClose(to_array(op), to_array(expected)) - - -class TestGeneratorModelSparseJax(TestGeneratorModelSparse, TestJaxBase): - """JAX version of sparse model tests.""" - - def validate_generator_eval(self, op, expected): - """Validate that op is sparse and agrees with expected.""" - self.assertTrue(type(op).__name__ == "BCOO") - self.assertAllClose(to_array(op), to_array(expected)) - - def test_jit_grad(self): - """Test jitting and gradding.""" - - model = GeneratorModel( - static_operator=-1j * self.Z, - operators=[-1j * self.X], - rotating_frame=self.Z, - evaluation_mode="sparse", - ) - - y = np.array([0.0, 1.0]) - - def func(a): - model_copy = model.copy() - model_copy.signals = [Signal(a)] - return model_copy(0.232, y) - - jitted_func = self.jit_wrap(func) - self.assertAllClose(jitted_func(1.2), func(1.2)) - - grad_jit_func = self.jit_grad_wrap(func) - grad_jit_func(1.2) - - -class Testtransfer_operator_functions(QiskitDynamicsTestCase): - """Tests for transfer_static_operator_between_frames and transfer_operators_between_frames.""" - - def test_all_None(self): - """Test all arguments being None.""" - - static_operator = transfer_static_operator_between_frames(None, None, None) - operators = transfer_operators_between_frames(None, None, None) - - self.assertTrue(static_operator is None) - self.assertTrue(operators is None) - - def test_array_inputs_diagonal_frame(self): - """Test correct handling when operators are arrays for diagonal frames.""" - - static_operator = -1j * np.array([[1.0, 0.0], [0.0, -1.0]]) - operators = -1j * np.array([[[0.0, 1.0], [1.0, 0.0]], [[0.0, -1j], [1j, 0.0]]]) - new_frame = -1j * np.array([1.0, -1.0]) - - out_static = transfer_static_operator_between_frames(static_operator, new_frame=new_frame) - out_operators = transfer_operators_between_frames(operators, new_frame=new_frame) - - self.assertTrue(isinstance(out_static, (np.ndarray, Array))) - self.assertTrue(isinstance(out_operators, (np.ndarray, Array))) - - self.assertAllClose(out_operators, operators) - self.assertAllClose(out_static, np.zeros((2, 2))) - - def test_array_inputs_pseudo_random(self): - """Test correct handling when operators are pseudo random arrays.""" - - rng = np.random.default_rng(34233) - num_terms = 3 - dim = 5 - b = 1.0 # bound on size of random terms - static_operator = rng.uniform(low=-b, high=b, size=(dim, dim)) + 1j * rng.uniform( - low=-b, high=b, size=(dim, dim) - ) - operators = rng.uniform(low=-b, high=b, size=(num_terms, dim, dim)) + 1j * rng.uniform( - low=-b, high=b, size=(dim, dim) - ) - old_frame = rng.uniform(low=-b, high=b, size=(dim, dim)) + 1j * rng.uniform( - low=-b, high=b, size=(dim, dim) - ) - old_frame = RotatingFrame(old_frame - old_frame.conj().transpose()) - - new_frame = rng.uniform(low=-b, high=b, size=(dim, dim)) + 1j * rng.uniform( - low=-b, high=b, size=(dim, dim) - ) - new_frame = RotatingFrame(new_frame - new_frame.conj().transpose()) - - out_static = transfer_static_operator_between_frames( - Array(static_operator), new_frame=new_frame, old_frame=old_frame - ) - out_operators = transfer_operators_between_frames( - Array(operators), new_frame=new_frame, old_frame=old_frame - ) - - self.assertTrue(isinstance(out_static, (np.ndarray, Array))) - self.assertTrue(isinstance(out_operators, (np.ndarray, Array))) - - expected_static = new_frame.operator_into_frame_basis( - (old_frame.operator_out_of_frame_basis(static_operator) + old_frame.frame_operator) - - new_frame.frame_operator - ) - expected_operators = [ - new_frame.operator_into_frame_basis(old_frame.operator_out_of_frame_basis(op)) - for op in operators - ] - - self.assertAllClose(out_static, expected_static) - self.assertAllClose(out_operators, expected_operators) - - -class Testtransfer_operator_functionsJax(Testtransfer_operator_functions, TestJaxBase): - """JAX version of Testtransfer_operator_functions.""" - - -class Testtransfer_operator_functionsSparse(QiskitDynamicsTestCase): - """Tests for transfer_static_operator_between_frames and transfer_operators_between_frames - for sparse cases. - """ - - def test_csr_inputs_diagonal_frame(self): - """Test correct handling when operators are csr matrices with a diagonal frame.""" - - static_operator = csr_matrix(-1j * np.array([[1.0, 0.0], [0.0, -1.0]])) - operators = [ - -1j * csr_matrix([[0.0, 1.0], [1.0, 0.0]]), - -1j * csr_matrix([[0.0, -1j], [1j, 0.0]]), - ] - new_frame = -1j * np.array([1.0, -1.0]) - - out_static = transfer_static_operator_between_frames(static_operator, new_frame=new_frame) - out_operators = transfer_operators_between_frames(operators, new_frame=new_frame) - - self.assertTrue(isinstance(out_static, csr_matrix)) - self.assertTrue(isinstance(out_operators, list) and issparse(out_operators[0])) - - self.assertAllCloseSparse(out_operators, operators) - self.assertAllCloseSparse(out_static, csr_matrix(np.zeros((2, 2)))) - - def test_csr_inputs_non_diagonal_frame(self): - """Test correct handling when operators are csr matrices with non-diagonal frames.""" + def test_jitable_funcs(self): + """Tests whether all functions are jitable. Checks if having a frame makes a difference, as + well as all jax-compatible evaluation_modes. + """ + from jax import jit - static_operator = csr_matrix(-1j * np.array([[1.0, 0.0], [0.0, -1.0]])) - operators = [ - -1j * csr_matrix([[0.0, 1.0], [1.0, 0.0]]), - -1j * csr_matrix([[0.0, -1j], [1j, 0.0]]), - ] - old_frame = np.array( - [ - [ - 0.0, - 1.0, - ], - [1.0, 0.0], - ] - ) - new_frame = -1j * np.array([1.0, -1.0]) + jit(self.basic_model.evaluate)(1.0) + jit(self.basic_model.evaluate_rhs)(1.0, np.array([0.2, 0.4])) - _, U = np.linalg.eigh(old_frame) - Uadj = U.conj().transpose() - - out_static = transfer_static_operator_between_frames( - static_operator, new_frame=new_frame, old_frame=old_frame - ) - out_operators = transfer_operators_between_frames( - operators, new_frame=new_frame, old_frame=old_frame - ) + jit(self.basic_model_w_frame.evaluate)(1.0) + jit(self.basic_model_w_frame.evaluate_rhs)(1.0, np.array([0.2, 0.4])) - self.assertTrue(isinstance(out_static, (np.ndarray, Array))) - self.assertTrue( - isinstance(out_operators, list) and isinstance(out_operators[0], (np.ndarray, Array)) - ) - - expected_ops = [U @ (op @ Uadj) for op in operators] - expected_static = ( - U @ to_array(static_operator) @ Uadj + (-1j * old_frame) - np.diag(new_frame) - ) + def test_gradable_funcs(self): + """Tests whether all functions are gradable. Checks if having a frame makes a difference, as + well as all jax-compatible evaluation_modes. + """ + self.jit_grad(self.basic_model.evaluate)(1.0) + self.jit_grad(self.basic_model.evaluate_rhs)(1.0, np.array([0.2, 0.4])) - self.assertAllClose(out_operators, expected_ops) - self.assertAllClose(out_static, expected_static) + self.jit_grad(self.basic_model_w_frame.evaluate)(1.0) + self.jit_grad(self.basic_model_w_frame.evaluate_rhs)(1.0, np.array([0.2, 0.4])) diff --git a/test/dynamics/models/test_hamiltonian_model.py b/test/dynamics/models/test_hamiltonian_model.py index adb859a28..99516e392 100644 --- a/test/dynamics/models/test_hamiltonian_model.py +++ b/test/dynamics/models/test_hamiltonian_model.py @@ -9,112 +9,117 @@ # 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 qiskit_dynamics.models.HamiltonianModel""" +from functools import partial + import numpy as np from scipy.linalg import expm -from scipy.sparse import csr_matrix from qiskit import QiskitError from qiskit.quantum_info.operators import Operator +from qiskit_dynamics.arraylias.alias import ArrayLike from qiskit_dynamics.models import HamiltonianModel from qiskit_dynamics.models.hamiltonian_model import is_hermitian from qiskit_dynamics.signals import Signal, SignalList -from qiskit_dynamics.array import Array -from qiskit_dynamics.type_utils import to_BCOO, to_csr -from ..common import QiskitDynamicsTestCase, TestJaxBase +from ..common import test_array_backends -class TestHamiltonianModelValidation(QiskitDynamicsTestCase): - """Test validation handling of HamiltonianModel.""" +@partial(test_array_backends, array_libraries=["numpy", "jax", "jax_sparse", "scipy_sparse"]) +class Testis_hermitian: + """Test is_hermitian validation function.""" - def test_operators_array_not_hermitian(self): - """Test raising error if operators are not Hermitian.""" + def test_cases(self): + """Test 2d array case.""" + self.assertTrue(is_hermitian(self.asarray([[1.0, 0.0], [0.0, 1.0]]))) + self.assertFalse(is_hermitian(self.asarray([[0.0, 1.0], [0.0, 0.0]]))) + self.assertFalse(is_hermitian(self.asarray([[0.0, 1j], [0.0, 0.0]]))) + self.assertTrue(is_hermitian(self.asarray([[0.0, 1j], [-1j, 0.0]]))) - with self.assertRaisesRegex(QiskitError, "operators must be Hermitian."): - HamiltonianModel(operators=[np.array([[0.0, 1.0], [0.0, 0.0]])]) - def test_operators_csr_not_hermitian(self): +@partial(test_array_backends, array_libraries=["numpy", "jax", "jax_sparse", "scipy_sparse"]) +class TestHamiltonianModelValidation: + """Test validation handling of HamiltonianModel.""" + + def test_operators_array_not_hermitian(self): """Test raising error if operators are not Hermitian.""" with self.assertRaisesRegex(QiskitError, "operators must be Hermitian."): - HamiltonianModel(operators=self.to_sparse([[[0.0, 1.0], [0.0, 0.0]]])) + HamiltonianModel(operators=[self.asarray([[0.0, 1.0], [0.0, 0.0]])]) def test_static_operator_not_hermitian(self): """Test raising error if static_operator is not Hermitian.""" with self.assertRaisesRegex(QiskitError, "static_operator must be Hermitian."): HamiltonianModel( - operators=[np.array([[0.0, 1.0], [1.0, 0.0]])], - static_operator=np.array([[0.0, 1.0], [0.0, 0.0]]), + operators=[self.asarray([[0.0, 1.0], [1.0, 0.0]])], + static_operator=self.asarray([[0.0, 1.0], [0.0, 0.0]]), ) def test_validate_false(self): """Verify setting validate=False avoids error raising.""" + if self.array_library() == "scipy_sparse": + operators = [self.asarray([[0.0, 1.0], [0.0, 0.0]])] + else: + operators = self.asarray([[[0.0, 1.0], [0.0, 0.0]]]) + ham_model = HamiltonianModel( - operators=[np.array([[0.0, 1.0], [0.0, 0.0]])], signals=[1.0], validate=False + operators=operators, signals=[1.0], validate=False, array_library=self.array_library() ) - self.assertAllClose(ham_model(1.0), -1j * np.array([[0.0, 1.0], [0.0, 0.0]])) - - def to_sparse(self, ops): - """Conversion of an operator or list of operators to the correct - sparse format for this class.""" - return to_csr(ops) - + res = ham_model(1.0) + self.assertArrayType(res) + self.assertAllClose(res, -1j * np.array([[0.0, 1.0], [0.0, 0.0]])) -class TestHamiltonianModelValidationJax(TestHamiltonianModelValidation, TestJaxBase): - """Test HamiltonianModel validation with JAX backend.""" - def to_sparse(self, ops): - return to_BCOO(ops) - - -class TestHamiltonianModel(QiskitDynamicsTestCase): +@partial(test_array_backends, array_libraries=["numpy", "jax", "jax_sparse", "scipy_sparse"]) +class TestHamiltonianModel: """Tests for HamiltonianModel.""" 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) - - # define a basic hamiltonian - w = 2.0 - r = 0.5 - operators = [2 * np.pi * self.Z / 2, 2 * np.pi * r * self.X / 2] - signals = [w, Signal(1.0, w)] - - self.w = w - self.r = r - self.basic_hamiltonian = HamiltonianModel(operators=operators, signals=signals) + """Build simple model elements.""" + self.X = Operator.from_label("X").data + self.Y = Operator.from_label("Y").data + self.Z = Operator.from_label("Z").data + + # define basic model elements + self.w = 2.0 + self.r = 0.5 + self.operators = [2 * np.pi * self.Z / 2, 2 * np.pi * self.r * self.X / 2] + self.signals = [self.w, Signal(1.0, self.w)] + + self.basic_hamiltonian = HamiltonianModel( + operators=self.operators, signals=self.signals, array_library=self.array_library() + ) def _basic_frame_evaluate_test(self, frame_operator, t): - """Routine for testing setting of valid frame operators using - basic_hamiltonian. - - Adapted from the version of this test in - test_operator_models.py, but relative to the way HamiltonianModel - modifies frame handling. + """Routine for testing setting of valid frame operators using the + basic_model. Args: - frame_operator (Array): now assumed to be a Hermitian operator H, with the + frame_operator (ArrayLike): now assumed to be a Hermitian operator H, with the frame being entered being F=-1j * H t (float): time of frame transformation """ - self.basic_hamiltonian.rotating_frame = frame_operator + basic_hamiltonian = HamiltonianModel( + operators=self.operators, + signals=self.signals, + rotating_frame=frame_operator, + array_library=self.array_library(), + ) # convert to 2d array if isinstance(frame_operator, Operator): - frame_operator = Array(frame_operator.data) - if isinstance(frame_operator, Array) and frame_operator.ndim == 1: + frame_operator = frame_operator.data + if isinstance(frame_operator, ArrayLike) and frame_operator.ndim == 1: frame_operator = np.diag(frame_operator) - value = self.basic_hamiltonian(t) / -1j + value = basic_hamiltonian(t) / -1j twopi = 2 * np.pi @@ -130,7 +135,7 @@ def _basic_frame_evaluate_test(self, frame_operator, t): + d_coeff * twopi * U @ self.X @ U.conj().transpose() / 2 - frame_operator ) - + self.assertArrayType(value) self.assertAllClose(value, expected) def test_diag_frame_operator_basic_hamiltonian(self): @@ -138,8 +143,8 @@ def test_diag_frame_operator_basic_hamiltonian(self): set up basic hamiltonian. """ - self._basic_frame_evaluate_test(Array([1.0, -1.0]), 1.123) - self._basic_frame_evaluate_test(Array([1.0, -1.0]), np.pi) + self._basic_frame_evaluate_test(np.array([1.0, -1.0]), 1.123) + self._basic_frame_evaluate_test(np.array([1.0, -1.0]), np.pi) def test_non_diag_frame_operator_basic_hamiltonian(self): """Test setting a non-diagonal frame operator for the internally @@ -155,25 +160,30 @@ def test_evaluate_no_frame_basic_hamiltonian(self): value = self.basic_hamiltonian(t) / -1j twopi = 2 * np.pi d_coeff = self.r * np.cos(2 * np.pi * self.w * t) - expected = twopi * self.w * self.Z.data / 2 + twopi * d_coeff * self.X.data / 2 + expected = twopi * self.w * self.Z / 2 + twopi * d_coeff * self.X / 2 + self.assertArrayType(value) self.assertAllClose(value, expected) def test_evaluate_in_frame_basis_basic_hamiltonian(self): """Test generator evaluation in frame basis in the basic_hamiltonian.""" - frame_op = (self.X + 0.2 * self.Y + 0.1 * self.Z).data + frame_op = self.X + 0.2 * self.Y + 0.1 * self.Z - # enter the frame given by the -1j * X - self.basic_hamiltonian.rotating_frame = frame_op - self.basic_hamiltonian.in_frame_basis = True + basic_hamiltonian = HamiltonianModel( + operators=self.operators, + signals=self.signals, + rotating_frame=frame_op, + in_frame_basis=True, + array_library=self.array_library(), + ) # get the frame basis used in model. Note that the Frame object # orders the basis according to the ordering of eigh _, U = np.linalg.eigh(frame_op) t = 3.21412 - value = self.basic_hamiltonian(t) / -1j + value = basic_hamiltonian(t) / -1j # compose the frame basis transformation with the exponential # frame rotation (this will be multiplied on the right) @@ -183,11 +193,9 @@ def test_evaluate_in_frame_basis_basic_hamiltonian(self): twopi = 2 * np.pi d_coeff = self.r * np.cos(2 * np.pi * self.w * t) expected = ( - Uadj - @ (twopi * self.w * self.Z.data / 2 + twopi * d_coeff * self.X.data / 2 - frame_op) - @ U + Uadj @ (twopi * self.w * self.Z / 2 + twopi * d_coeff * self.X / 2 - frame_op) @ U ) - + self.assertArrayType(value) self.assertAllClose(value, expected) def test_evaluate_pseudorandom(self): @@ -202,19 +210,19 @@ def test_evaluate_pseudorandom(self): rand_op = rng.uniform(low=-b, high=b, size=(dim, dim)) + 1j * rng.uniform( low=-b, high=b, size=(dim, dim) ) - frame_op = Array(rand_op + rand_op.conj().transpose()) + frame_op = rand_op + rand_op.conj().transpose() # random hermitian operators randoperators = rng.uniform(low=-b, high=b, size=(num_terms, dim, dim)) + 1j * rng.uniform( low=-b, high=b, size=(num_terms, dim, dim) ) - randoperators = Array(randoperators + randoperators.conj().transpose([0, 2, 1])) + randoperators = randoperators + randoperators.conj().transpose([0, 2, 1]) rand_coeffs = rng.uniform(low=-b, high=b, size=(num_terms)) + 1j * rng.uniform( low=-b, high=b, size=(num_terms) ) - rand_carriers = Array(rng.uniform(low=-b, high=b, size=(num_terms))) - rand_phases = Array(rng.uniform(low=-b, high=b, size=(num_terms))) + rand_carriers = rng.uniform(low=-b, high=b, size=(num_terms)) + rand_phases = rng.uniform(low=-b, high=b, size=(num_terms)) self._test_evaluate(frame_op, randoperators, rand_coeffs, rand_carriers, rand_phases) @@ -227,19 +235,19 @@ def test_evaluate_pseudorandom(self): rand_op = rng.uniform(low=-b, high=b, size=(dim, dim)) + 1j * rng.uniform( low=-b, high=b, size=(dim, dim) ) - frame_op = Array(rand_op + rand_op.conj().transpose()) + frame_op = rand_op + rand_op.conj().transpose() # random hermitian operators randoperators = rng.uniform(low=-b, high=b, size=(num_terms, dim, dim)) + 1j * rng.uniform( low=-b, high=b, size=(num_terms, dim, dim) ) - randoperators = Array(randoperators + randoperators.conj().transpose([0, 2, 1])) + randoperators = randoperators + randoperators.conj().transpose([0, 2, 1]) rand_coeffs = rng.uniform(low=-b, high=b, size=(num_terms)) + 1j * rng.uniform( low=-b, high=b, size=(num_terms) ) - rand_carriers = Array(rng.uniform(low=-b, high=b, size=(num_terms))) - rand_phases = Array(rng.uniform(low=-b, high=b, size=(num_terms))) + rand_carriers = rng.uniform(low=-b, high=b, size=(num_terms)) + rand_phases = rng.uniform(low=-b, high=b, size=(num_terms)) self._test_evaluate(frame_op, randoperators, rand_coeffs, rand_carriers, rand_phases) @@ -257,7 +265,11 @@ def env(t): sig_list.append(Signal(get_env_func(), freq, phase)) sig_list = SignalList(sig_list) model = HamiltonianModel( - operators=operators, static_operator=None, signals=sig_list, rotating_frame=frame_op + operators=operators, + static_operator=None, + signals=sig_list, + rotating_frame=frame_op, + array_library=self.array_library(), ) value = model(1.0) / -1j @@ -271,104 +283,72 @@ def env(t): self.assertAllClose(model._signals(1), coeffs) self.assertAllClose(model.operators, operators) + self.assertArrayType(value) self.assertAllClose(value, expected) def test_evaluate_static(self): """Test evaluation of a GeneratorModel with only a static component.""" - static_model = HamiltonianModel(static_operator=self.X) - self.assertAllClose(-1j * self.X, static_model(1.0)) + static_model = HamiltonianModel(static_operator=self.X, array_library=self.array_library()) + value = static_model(1.0) + self.assertArrayType(value) + self.assertAllClose(-1j * self.X, value) # now with frame frame_op = -1j * (self.Z + 1.232 * self.Y) - static_model.rotating_frame = frame_op + static_model = HamiltonianModel( + static_operator=self.X, rotating_frame=frame_op, array_library=self.array_library() + ) t = 1.2312 - expected = expm(-frame_op.data * t) @ (-1j * self.X - frame_op) @ expm(frame_op.data * t) - self.assertAllClose(expected, static_model(t)) + value = static_model(t) + expected = expm(-frame_op * t) @ (-1j * self.X - frame_op) @ expm(frame_op * t) + self.assertArrayType(value) + self.assertAllClose(expected, value) + +@partial(test_array_backends, array_libraries=["jax", "jax_sparse"]) +class TestHamiltonianModelJAXTransformations: + """Tests for jax transformations.""" -class TestHamiltonianModelJax(TestHamiltonianModel, TestJaxBase): - """Jax version of TestHamiltonianModel tests.""" + def setUp(self): + """Build simple model elements.""" + self.X = Operator.from_label("X").data + self.Y = Operator.from_label("Y").data + self.Z = Operator.from_label("Z").data + + # define basic model elements + self.w = 2.0 + self.r = 0.5 + self.operators = [2 * np.pi * self.Z / 2, 2 * np.pi * self.r * self.X / 2] + self.signals = [self.w, Signal(1.0, self.w)] + + self.basic_hamiltonian = HamiltonianModel( + operators=self.operators, signals=self.signals, array_library=self.array_library() + ) + + self.basic_hamiltonian_w_frame = HamiltonianModel( + operators=self.operators, + signals=self.signals, + rotating_frame=np.array([[3j, 2j], [2j, 0]]), + array_library=self.array_library(), + ) def test_jitable_funcs(self): """Tests whether all functions are jitable. Checks if having a frame makes a difference, as well as all jax-compatible evaluation_modes.""" - self.jit_wrap(self.basic_hamiltonian.evaluate)(1) - self.jit_wrap(self.basic_hamiltonian.evaluate_rhs)(1, Array(np.array([0.2, 0.4]))) - - self.basic_hamiltonian.rotating_frame = Array(np.array([[3j, 2j], [2j, 0]])) + from jax import jit - self.jit_wrap(self.basic_hamiltonian.evaluate)(1) - self.jit_wrap(self.basic_hamiltonian.evaluate_rhs)(1, Array(np.array([0.2, 0.4]))) - - self.basic_hamiltonian.rotating_frame = None + jit(self.basic_hamiltonian.evaluate)(1) + jit(self.basic_hamiltonian.evaluate_rhs)(1, np.array([0.2, 0.4])) + jit(self.basic_hamiltonian_w_frame.evaluate)(1) + jit(self.basic_hamiltonian_w_frame.evaluate_rhs)(1, np.array([0.2, 0.4])) def test_gradable_funcs(self): """Tests whether all functions are gradable. Checks if having a frame makes a difference, as well as all jax-compatible evaluation_modes.""" - self.jit_grad_wrap(self.basic_hamiltonian.evaluate)(1.0) - self.jit_grad_wrap(self.basic_hamiltonian.evaluate_rhs)(1.0, Array(np.array([0.2, 0.4]))) - - self.basic_hamiltonian.rotating_frame = Array(np.array([[3j, 2j], [2j, 0]])) - - self.jit_grad_wrap(self.basic_hamiltonian.evaluate)(1.0) - self.jit_grad_wrap(self.basic_hamiltonian.evaluate_rhs)(1.0, Array(np.array([0.2, 0.4]))) - - self.basic_hamiltonian.rotating_frame = None - - -class Testis_hermitian(QiskitDynamicsTestCase): - """Test is_hermitian validation function.""" - - def test_2d_array(self): - """Test 2d array case.""" - self.assertTrue(is_hermitian(Array([[1.0, 0.0], [0.0, 1.0]]))) - self.assertFalse(is_hermitian(Array([[0.0, 1.0], [0.0, 0.0]]))) - self.assertFalse(is_hermitian(Array([[0.0, 1j], [0.0, 0.0]]))) - self.assertTrue(is_hermitian(Array([[0.0, 1j], [-1j, 0.0]]))) - - def test_3d_array(self): - """Test 3d array case.""" - self.assertTrue(is_hermitian(Array([[[1.0, 0.0], [0.0, 1.0]]]))) - self.assertFalse(is_hermitian(Array([[[0.0, 1.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 0.0]]]))) - self.assertFalse(is_hermitian(Array([[[0.0, 1j], [0.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]]]))) - self.assertTrue(is_hermitian(Array([[[0.0, 1j], [-1j, 0.0]], [[0.0, 1.0], [1.0, 0.0]]]))) - - def test_csr_matrix(self): - """Test csr_matrix case.""" - self.assertTrue(is_hermitian(csr_matrix([[1.0, 0.0], [0.0, 1.0]]))) - self.assertFalse(is_hermitian(csr_matrix([[0.0, 1.0], [0.0, 0.0]]))) - self.assertFalse(is_hermitian(csr_matrix([[0.0, 1j], [0.0, 0.0]]))) - self.assertTrue(is_hermitian(csr_matrix([[0.0, 1j], [-1j, 0.0]]))) - - def test_list_csr_matrix(self): - """Test list of csr_matrix case.""" - self.assertTrue(is_hermitian([csr_matrix([[1.0, 0.0], [0.0, 1.0]])])) - self.assertFalse( - is_hermitian( - [csr_matrix([[0.0, 1.0], [0.0, 0.0]]), csr_matrix([[0.0, 1.0], [1.0, 0.0]])] - ) - ) - self.assertFalse( - is_hermitian( - [csr_matrix([[0.0, 1j], [0.0, 0.0]]), csr_matrix([[1.0, 0.0], [0.0, 1.0]])] - ) - ) - self.assertTrue( - is_hermitian( - [csr_matrix([[0.0, 1j], [-1j, 0.0]]), csr_matrix([[0.0, 1.0], [1.0, 0.0]])] - ) - ) - - -class Testis_hermitianBCOO(QiskitDynamicsTestCase, TestJaxBase): - """Test is_hermitian for jax BCOO case.""" - - def test_2d_cases(self): - """Test BCOO 2d cases.""" - self.assertTrue(is_hermitian(to_BCOO([[1.0, 0.0], [0.0, 1.0]]))) - self.assertFalse(is_hermitian(to_BCOO([[0.0, 1.0], [0.0, 0.0]]))) - self.assertFalse(is_hermitian(to_BCOO([[0.0, 1j], [0.0, 0.0]]))) - self.assertTrue(is_hermitian(to_BCOO([[0.0, 1j], [-1j, 0.0]]))) + self.jit_grad(self.basic_hamiltonian.evaluate)(1.0) + self.jit_grad(self.basic_hamiltonian.evaluate_rhs)(1.0, np.array([0.2, 0.4])) + self.jit_grad(self.basic_hamiltonian_w_frame.evaluate)(1.0) + self.jit_grad(self.basic_hamiltonian_w_frame.evaluate_rhs)(1.0, np.array([0.2, 0.4])) diff --git a/test/dynamics/models/test_lindblad_model.py b/test/dynamics/models/test_lindblad_model.py index 443303356..10d972d94 100644 --- a/test/dynamics/models/test_lindblad_model.py +++ b/test/dynamics/models/test_lindblad_model.py @@ -9,11 +9,11 @@ # 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,redundant-keyword-arg +# pylint: disable=invalid-name,redundant-keyword-arg,no-member -"""Tests for qiskit_dynamics.models.lindblad_models.py. Most -of the actual calculation checking is handled at the level of a -models.operator_collection.DenseLindbladOperatorCollection test.""" +"""Tests for qiskit_dynamics.models.lindblad_models.py.""" + +from functools import partial import numpy as np @@ -23,9 +23,7 @@ from qiskit.quantum_info.operators import Operator from qiskit_dynamics.models import LindbladModel from qiskit_dynamics.signals import Signal, SignalList -from qiskit_dynamics.array import Array -from qiskit_dynamics.type_utils import to_array -from ..common import QiskitDynamicsTestCase, TestJaxBase +from ..common import QiskitDynamicsTestCase, test_array_backends class TestLindbladModelErrors(QiskitDynamicsTestCase): @@ -151,10 +149,13 @@ def test_validate_false(self): self.assertAllClose(lindblad_model(1.0, np.eye(2)), np.zeros(2)) -class TestLindbladModel(QiskitDynamicsTestCase): - """Tests for LindbladModel.""" +class TestLindbladModel: + """Tests for LindbladModel. This class is turned into a proper test class through inheritance + below. + """ def setUp(self): + """Build simple model.""" self.X = Operator.from_label("X").data self.Y = Operator.from_label("Y").data self.Z = Operator.from_label("Z").data @@ -168,33 +169,34 @@ def setUp(self): self.w = w self.r = r - static_dissipators = Array([[[0.0, 0.0], [1.0, 0.0]]]) + static_dissipators = [[[0.0, 0.0], [1.0, 0.0]]] self.basic_lindblad = LindbladModel( hamiltonian_operators=ham_operators, hamiltonian_signals=ham_signals, static_dissipators=static_dissipators, - evaluation_mode=self.evaluation_mode, + array_library=self.array_library(), + vectorized=self.vectorized, ) @property - def evaluation_mode(self): - """Evaluation mode to use for tests, useful for inheritance.""" - return "dense" + def vectorized(self): + """Whether or not to run tests with vectorized LindbladModel.""" + return False def test_basic_lindblad_lmult(self): """Test lmult method of Lindblad generator OperatorModel.""" - A = Array([[1.0, 2.0], [3.0, 4.0]]) + A = np.array([[1.0, 2.0], [3.0, 4.0]]) t = 1.123 ham = ( 2 * np.pi * self.w * self.Z / 2 + 2 * np.pi * self.r * np.cos(2 * np.pi * self.w * t) * self.X / 2 ) - sm = Array([[0.0, 0.0], [1.0, 0.0]]) + sm = np.array([[0.0, 0.0], [1.0, 0.0]]) expected = self._evaluate_lindblad_rhs(A, ham, [sm]) - if "vectorized" in self.evaluation_mode: + if self.vectorized: expected = expected.flatten(order="F") A = A.flatten(order="F") @@ -207,7 +209,8 @@ def test_evaluate_only_dissipators(self): model = LindbladModel( dissipator_operators=[self.X], dissipator_signals=[1.0], - evaluation_mode=self.evaluation_mode, + array_library=self.array_library(), + vectorized=self.vectorized, ) rho = np.array([[1.0, 0.0], [0.0, 0.0]], dtype=complex) @@ -215,7 +218,7 @@ def test_evaluate_only_dissipators(self): expected = self._evaluate_lindblad_rhs( rho, ham=np.zeros((2, 2), dtype=complex), dissipators=[self.X] ) - if "vectorized" in self.evaluation_mode: + if self.vectorized: expected = expected.flatten(order="F") rho = rho.flatten(order="F") @@ -225,14 +228,16 @@ def test_evaluate_only_static_dissipators(self): """Test evaluation with just dissipators.""" model = LindbladModel( - static_dissipators=[self.X, self.Y], evaluation_mode=self.evaluation_mode + static_dissipators=[self.X, self.Y], + array_library=self.array_library(), + vectorized=self.vectorized, ) rho = np.array([[1.0, 0.0], [0.0, 0.0]], dtype=complex) expected = self._evaluate_lindblad_rhs( rho, ham=np.zeros((2, 2), dtype=complex), dissipators=[self.X, self.Y] ) - if "vectorized" in self.evaluation_mode: + if self.vectorized: expected = expected.flatten(order="F") rho = rho.flatten(order="F") @@ -241,11 +246,15 @@ def test_evaluate_only_static_dissipators(self): def test_evaluate_only_static_hamiltonian(self): """Test evaluation with just static hamiltonian.""" - model = LindbladModel(static_hamiltonian=self.X, evaluation_mode=self.evaluation_mode) + model = LindbladModel( + static_hamiltonian=self.X, + array_library=self.array_library(), + vectorized=self.vectorized, + ) rho = np.array([[1.0, 0.0], [0.0, 0.0]], dtype=complex) expected = self._evaluate_lindblad_rhs(rho, ham=self.X) - if "vectorized" in self.evaluation_mode: + if self.vectorized: expected = expected.flatten(order="F") rho = rho.flatten(order="F") @@ -257,12 +266,13 @@ def test_evaluate_only_hamiltonian_operators(self): model = LindbladModel( hamiltonian_operators=[self.X], hamiltonian_signals=[1.0], - evaluation_mode=self.evaluation_mode, + array_library=self.array_library(), + vectorized=self.vectorized, ) rho = np.array([[1.0, 0.0], [0.0, 0.0]], dtype=complex) expected = self._evaluate_lindblad_rhs(rho, ham=self.X) - if "vectorized" in self.evaluation_mode: + if self.vectorized: expected = expected.flatten(order="F") rho = rho.flatten(order="F") @@ -283,14 +293,14 @@ def test_lindblad_pseudorandom(self): randoperators = rng.uniform(low=-b, high=b, size=(num_ham, dim, dim)) + 1j * rng.uniform( low=-b, high=b, size=(num_ham, dim, dim) ) - rand_ham_ops = Array(randoperators + randoperators.conj().transpose([0, 2, 1])) + rand_ham_ops = randoperators + randoperators.conj().transpose([0, 2, 1]) # generate random hamiltonian coefficients rand_ham_coeffs = rng.uniform(low=-b, high=b, size=(num_ham)) + 1j * rng.uniform( low=-b, high=b, size=(num_ham) ) - rand_ham_carriers = Array(rng.uniform(low=-b, high=b, size=(num_ham))) - rand_ham_phases = Array(rng.uniform(low=-b, high=b, size=(num_ham))) + rand_ham_carriers = rng.uniform(low=-b, high=b, size=(num_ham)) + rand_ham_phases = rng.uniform(low=-b, high=b, size=(num_ham)) ham_sigs = [] for coeff, freq, phase in zip(rand_ham_coeffs, rand_ham_carriers, rand_ham_phases): @@ -299,23 +309,21 @@ def test_lindblad_pseudorandom(self): ham_sigs = SignalList(ham_sigs) # generate random static dissipators - rand_static_diss = Array( - rng.uniform(low=-b, high=b, size=(num_diss, dim, dim)) - + 1j * rng.uniform(low=-b, high=b, size=(num_diss, dim, dim)) - ) + rand_static_diss = rng.uniform( + low=-b, high=b, size=(num_diss, dim, dim) + ) + 1j * rng.uniform(low=-b, high=b, size=(num_diss, dim, dim)) # generate random dissipators - rand_diss = Array( - rng.uniform(low=-b, high=b, size=(num_diss, dim, dim)) - + 1j * rng.uniform(low=-b, high=b, size=(num_diss, dim, dim)) + rand_diss = rng.uniform(low=-b, high=b, size=(num_diss, dim, dim)) + 1j * rng.uniform( + low=-b, high=b, size=(num_diss, dim, dim) ) # random dissipator coefficients rand_diss_coeffs = rng.uniform(low=-b, high=b, size=(num_diss)) + 1j * rng.uniform( low=-b, high=b, size=(num_diss) ) - rand_diss_carriers = Array(rng.uniform(low=-b, high=b, size=(num_diss))) - rand_diss_phases = Array(rng.uniform(low=-b, high=b, size=(num_diss))) + rand_diss_carriers = rng.uniform(low=-b, high=b, size=(num_diss)) + rand_diss_phases = rng.uniform(low=-b, high=b, size=(num_diss)) diss_sigs = [] for coeff, freq, phase in zip(rand_diss_coeffs, rand_diss_carriers, rand_diss_phases): @@ -327,7 +335,7 @@ def test_lindblad_pseudorandom(self): rand_op = rng.uniform(low=-b, high=b, size=(dim, dim)) + 1j * rng.uniform( low=-b, high=b, size=(dim, dim) ) - frame_op = Array(rand_op - rand_op.conj().transpose()) + frame_op = rand_op - rand_op.conj().transpose() evect = np.linalg.eigh(1j * frame_op)[1] into_frame_basis = lambda x: evect.T.conj() @ x @ evect @@ -338,9 +346,10 @@ def test_lindblad_pseudorandom(self): static_dissipators=rand_static_diss, dissipator_operators=rand_diss, dissipator_signals=diss_sigs, - evaluation_mode=self.evaluation_mode, + rotating_frame=frame_op, + array_library=self.array_library(), + vectorized=self.vectorized, ) - lindblad_model.rotating_frame = frame_op t = rng.uniform(low=-b, high=b) # test storage of operators in class @@ -360,23 +369,22 @@ def test_lindblad_pseudorandom(self): lindblad_model.in_frame_basis = True self.assertAllClose( into_frame_basis(rand_diss), - to_array(lindblad_model.dissipator_operators), + lindblad_model.dissipator_operators, ) self.assertAllClose( into_frame_basis(rand_ham_ops), - to_array(lindblad_model.hamiltonian_operators), + lindblad_model.hamiltonian_operators, ) self.assertAllClose( into_frame_basis(-1j * frame_op), - to_array(lindblad_model.static_hamiltonian), + lindblad_model.static_hamiltonian, ) lindblad_model.in_frame_basis = False self.assertAllClose(-1j * frame_op, lindblad_model.static_hamiltonian) # evaluation tests - A = Array( - rng.uniform(low=-b, high=b, size=(dim, dim)) - + 1j * rng.uniform(low=-b, high=b, size=(dim, dim)) + A = rng.uniform(low=-b, high=b, size=(dim, dim)) + 1j * rng.uniform( + low=-b, high=b, size=(dim, dim) ) expected = self._evaluate_lindblad_rhs( @@ -401,7 +409,7 @@ def test_lindblad_pseudorandom(self): expected_in_frame_basis ) - if "vectorized" in self.evaluation_mode: + if self.vectorized: expected = expected.flatten(order="F") expected_in_frame_basis = expected_in_frame_basis.flatten(order="F") A = A.flatten(order="F") @@ -423,25 +431,26 @@ def test_dissipator_consistency(self): b = 1.0 # bound on size of random terms # generate random dissipators - rand_diss = Array( - rng.uniform(low=-b, high=b, size=(num_diss, dim, dim)) - + 1j * rng.uniform(low=-b, high=b, size=(num_diss, dim, dim)) + rand_diss = rng.uniform(low=-b, high=b, size=(num_diss, dim, dim)) + 1j * rng.uniform( + low=-b, high=b, size=(num_diss, dim, dim) ) static_model = LindbladModel( - static_dissipators=rand_diss, evaluation_mode=self.evaluation_mode + static_dissipators=rand_diss, + array_library=self.array_library(), + vectorized=self.vectorized, ) non_static_model = LindbladModel( dissipator_operators=rand_diss, dissipator_signals=[1.0] * num_diss, - evaluation_mode=self.evaluation_mode, + array_library=self.array_library(), + vectorized=self.vectorized, ) - rand_input = Array( - rng.uniform(low=-b, high=b, size=(dim, dim)) - + 1j * rng.uniform(low=-b, high=b, size=(dim, dim)) + rand_input = rng.uniform(low=-b, high=b, size=(dim, dim)) + 1j * rng.uniform( + low=-b, high=b, size=(dim, dim) ) - if "vectorized" in self.evaluation_mode: + if self.vectorized: rand_input = rand_input.flatten(order="F") self.assertAllClose(static_model(0.0, rand_input), non_static_model(0.0, rand_input)) @@ -464,8 +473,7 @@ def _evaluate_lindblad_rhs( Note: here we force everything into numpy arrays as these parts of the test are just for confirmation """ - # if a frame operator is given, transform the model pieces into - # the frame + # if a frame operator is given, transform the model pieces into the frame if frame_op is not None: frame_op = np.array(frame_op) U = expm(-t * frame_op) @@ -517,7 +525,8 @@ def test_get_operators_when_None(self): model = LindbladModel( static_hamiltonian=np.array([[1.0, 0.0], [0.0, -1.0]]), - evaluation_mode=self.evaluation_mode, + array_library=self.array_library(), + vectorized=self.vectorized, ) self.assertTrue(model.hamiltonian_operators is None) self.assertTrue(model.static_dissipators is None) @@ -525,161 +534,106 @@ def test_get_operators_when_None(self): model = LindbladModel( hamiltonian_operators=[np.array([[1.0, 0.0], [0.0, -1.0]])], - evaluation_mode=self.evaluation_mode, + array_library=self.array_library(), + vectorized=self.vectorized, ) self.assertTrue(model.static_hamiltonian is None) -class TestLindbladModelJax(TestLindbladModel, TestJaxBase): - """Jax version of TestLindbladModel tests.""" - - def test_jitable_funcs(self): - """Tests whether all functions are jitable. - Checks if having a frame makes a difference, as well as - all jax-compatible evaluation_modes.""" - rho = Array(np.array([[0.2, 0.4], [0.6, 0.8]])) - if "vectorized" in self.evaluation_mode: - rho = rho.flatten(order="F") - - self.jit_wrap(self.basic_lindblad.evaluate_rhs)(1.0, rho) - - self.basic_lindblad.rotating_frame = Array(np.array([[3j, 2j], [2j, 0]])) - self.jit_wrap(self.basic_lindblad.evaluate_rhs)(1.0, rho) - - def test_gradable_funcs(self): - """Tests whether all functions are gradable. - Checks if having a frame makes a difference, as well as - all jax-compatible evaluation_modes.""" - - rho = Array(np.array([[0.2, 0.4], [0.6, 0.8]])) - if "vectorized" in self.evaluation_mode: - rho = rho.flatten(order="F") - - self.jit_grad_wrap(self.basic_lindblad.evaluate_rhs)(1.0, rho) - - self.basic_lindblad.rotating_frame = Array(np.array([[3j, 2j], [2j, 0]])) - self.jit_grad_wrap(self.basic_lindblad.evaluate_rhs)(1.0, rho) - - -class TestLindbladModelSparse(TestLindbladModel): - """Sparse-mode tests.""" +@partial(test_array_backends, array_libraries=["numpy", "jax", "jax_sparse", "scipy_sparse"]) +class TestLindbladModelVectorized(TestLindbladModel): + """Vectorized version of tests.""" @property - def evaluation_mode(self): - """Evaluation mode to use for tests.""" - return "sparse" + def vectorized(self): + """Whether or not to use vectorized model.""" + return True - def test_switch_modes_and_evaluate(self): - """Test construction of a model, switching modes, and evaluating.""" - model = LindbladModel( - static_hamiltonian=self.Z, - hamiltonian_operators=[self.X], - hamiltonian_signals=[1.0], - evaluation_mode=self.evaluation_mode, - ) - rho = np.array([[1.0, 0.0], [0.0, 0.0]]) - ham = self.Z + self.X - expected = -1j * (ham @ rho - rho @ ham) - if "vectorized" in self.evaluation_mode: - expected = expected.flatten(order="F") - rho = rho.flatten(order="F") +# Tests for non-vectorized lindblad models +test_array_backends( + TestLindbladModel, array_libraries=["numpy", "jax", "jax_sparse", "scipy_sparse"] +) - self.assertAllClose(model(1.0, rho), expected) - if "vectorized" in self.evaluation_mode: - model.evaluation_mode = "dense_vectorized" - else: - model.evaluation_mode = "dense" +class TestLindbladModelJAXTransformations: + """JAX transformation tests for TestLindbladModel tests. This class is turned into a test class + below. + """ - self.assertAllClose(model(1.0, rho), expected) - - def test_frame_change_sparse(self): - """Test setting a frame after instantiation in sparse mode and evaluating.""" - model = LindbladModel( - static_hamiltonian=self.Z, - hamiltonian_operators=[self.X], - hamiltonian_signals=[1.0], - evaluation_mode=self.evaluation_mode, - ) + def setUp(self): + """Build simple model.""" + self.X = Operator.from_label("X").data + self.Y = Operator.from_label("Y").data + self.Z = Operator.from_label("Z").data - # test non-diagonal frame - model.rotating_frame = self.Z - rho = np.array([[1.0, 0.0], [0.0, 0.0]]) - ham = expm(1j * self.Z) @ self.X @ expm(-1j * self.Z) - expected = -1j * (ham @ rho - rho @ ham) - if "vectorized" in self.evaluation_mode: - expected = expected.flatten(order="F") - rho = rho.flatten(order="F") - self.assertAllClose(expected, model(1.0, rho)) + # define a basic hamiltonian + w = 2.0 + r = 0.5 + ham_operators = [2 * np.pi * self.Z / 2, 2 * np.pi * r * self.X / 2] + ham_signals = [w, Signal(1.0, w)] - # test diagonal frame - model.rotating_frame = np.array([1.0, -1.0]) - self.assertAllClose(expected, model(1.0, rho)) + self.w = w + self.r = r - def test_switching_to_sparse_with_frame(self): - """Test switching to sparse with a frame already set.""" + static_dissipators = [[[0.0, 0.0], [1.0, 0.0]]] - model = LindbladModel( - static_hamiltonian=self.Z, - hamiltonian_operators=[self.X], - hamiltonian_signals=[1.0], - rotating_frame=np.array([1.0, -1.0]), + self.basic_lindblad = LindbladModel( + hamiltonian_operators=ham_operators, + hamiltonian_signals=ham_signals, + static_dissipators=static_dissipators, + array_library=self.array_library(), + vectorized=self.vectorized, ) - model.evaluation_mode = self.evaluation_mode - - rho = np.array([[1.0, 0.0], [0.0, 0.0]]) - ham = expm(1j * self.Z) @ self.X @ expm(-1j * self.Z) - expected = -1j * (ham @ rho - rho @ ham) - if "vectorized" in self.evaluation_mode: - expected = expected.flatten(order="F") - rho = rho.flatten(order="F") - - self.assertAllClose(expected, model(1.0, rho)) - - -class TestLindbladModelJAXSparse(TestLindbladModelSparse, TestLindbladModelJax): - """JAX sparse-mode tests.""" - - -class TestLindbladModelDenseVectorized(TestLindbladModel): - """Test class for dense vectorized mode operation of LindbladModel.""" + self.rf_lindblad = LindbladModel( + hamiltonian_operators=ham_operators, + hamiltonian_signals=ham_signals, + static_dissipators=static_dissipators, + rotating_frame=np.array([[3j, 2j], [2j, 0]]), + array_library=self.array_library(), + vectorized=self.vectorized, + ) @property - def evaluation_mode(self): - return "dense_vectorized" + def vectorized(self): + """Whether or not to do vectorized tests.""" + return False + def test_jitable_funcs(self): + """Tests whether all functions are jitable. + Checks if having a frame makes a difference, as well as + all jax-compatible evaluation_modes.""" + rho = np.array([[0.2, 0.4], [0.6, 0.8]]) + if self.vectorized: + rho = rho.flatten(order="F") -class TestLindbladModelSparseVectorized(TestLindbladModelSparse): - """Test class for sparse vectorized mode operation of LindbladModel.""" + from jax import jit - @property - def evaluation_mode(self): - return "sparse_vectorized" + jit(self.basic_lindblad.evaluate_rhs)(1.0, rho) + jit(self.rf_lindblad.evaluate_rhs)(1.0, rho) + def test_gradable_funcs(self): + """Tests whether all functions are gradable. + Checks if having a frame makes a difference, as well as + all jax-compatible evaluation_modes.""" -class TestLindbladModelDenseVectorizedJax(TestLindbladModelJax): - """Test class for jax dense vectorized mode operation of LindbladModel.""" + rho = np.array([[0.2, 0.4], [0.6, 0.8]]) + if self.vectorized: + rho = rho.flatten(order="F") - @property - def evaluation_mode(self): - return "dense_vectorized" + self.jit_grad(self.basic_lindblad.evaluate_rhs)(1.0, rho) + self.jit_grad(self.rf_lindblad.evaluate_rhs)(1.0, rho) -class TestLindbladModelSparseVectorizedJax(TestLindbladModelJax): - """Test class for jax sparse vectorized mode operation of LindbladModel.""" +@partial(test_array_backends, array_libraries=["jax", "jax_sparse"]) +class TestLindbladModelJAXTransformationsVectorized(TestLindbladModelJAXTransformations): + """Test JAX transformations for vectorized lindblad model.""" @property - def evaluation_mode(self): - return "sparse_vectorized" - - -def get_const_func(const): - """Helper function for defining a constant function.""" + def vectorized(self): + return True - # pylint: disable=unused-argument - def env(t): - return const - return env +# test JAX transformations for non-vectorized +test_array_backends(TestLindbladModelJAXTransformations, array_libraries=["jax", "jax_sparse"])