Update model classes to use arraylias (#294)

Co-authored-by: Kento Ueda <[email protected]>

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,143 +76,103 @@ 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,
             operators=operators,
             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.")