From a4d39eb7bd9c921acc1f51d88fdcbe51e6b61a11 Mon Sep 17 00:00:00 2001
From: Tobias Hangleiter <tobias.hangleiter@rwth-aachen.de>
Date: Tue, 2 Mar 2021 17:26:45 +0100
Subject: [PATCH] Factor out oper parsing, add qopt.DenseOperator detection

Used in both Basis and PulseSequence constructors
---
 filter_functions/basis.py          | 12 +-------
 filter_functions/pulse_sequence.py | 24 ++-------------
 filter_functions/util.py           | 49 ++++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 33 deletions(-)

diff --git a/filter_functions/basis.py b/filter_functions/basis.py
index 7f37877..08f2da9 100644
--- a/filter_functions/basis.py
+++ b/filter_functions/basis.py
@@ -173,21 +173,11 @@ def __new__(cls, basis_array: Sequence, traceless: Optional[bool] = None,
             except AttributeError:
                 pass
 
-            basis = np.empty((len(basis_array), *basis_array[0].shape), dtype=complex)
+            basis = util.parse_operators(basis_array, 'basis_array')
             if basis.shape[0] > np.product(basis.shape[1:]):
                 raise ValueError('Given overcomplete set of basis matrices. '
                                  'Not linearly independent.')
 
-            for i, elem in enumerate(basis_array):
-                if isinstance(elem, ndarray):   # numpy array
-                    basis[i] = elem
-                elif hasattr(elem, 'full'):     # qutip.Qobj
-                    basis[i] = elem.full()
-                elif hasattr(elem, 'todense'):  # sparse array
-                    basis[i] = elem.todense()
-                else:
-                    raise TypeError('At least one element invalid type!')
-
         basis = basis.view(cls)
         basis.btype = btype or 'Custom'
         basis.d = basis.shape[-1]
diff --git a/filter_functions/pulse_sequence.py b/filter_functions/pulse_sequence.py
index 2c4696f..e48d065 100644
--- a/filter_functions/pulse_sequence.py
+++ b/filter_functions/pulse_sequence.py
@@ -1134,28 +1134,8 @@ def _parse_Hamiltonian(H: Hamiltonian, n_dt: int, H_str: str) -> Tuple[Sequence[
         coeffs = args[0]
         identifiers = list(args[1])
 
-    # Parse opers and convert to squeezed ndarray if possible
-    parsed_opers = []
-    for oper in opers:
-        if isinstance(oper, ndarray):
-            parsed_opers.append(oper.squeeze())
-        elif hasattr(oper, 'full'):
-            # qutip.Qobj
-            parsed_opers.append(oper.full())
-        elif hasattr(oper, 'todense'):
-            # sparse object
-            parsed_opers.append(oper.todense())
-        else:
-            raise TypeError(f'Expected operators in {H_str} to be NumPy arrays or QuTiP Qobjs!')
-
-    # Check correct dimensions for the operators
-    if set(oper.ndim for oper in parsed_opers) != {2}:
-        raise ValueError(f'Expected all operators in {H_str} to be two-dimensional!')
-
-    if len(set(*set(oper.shape for oper in parsed_opers))) != 1:
-        raise ValueError(f'Expected operators in {H_str} to be square!')
-
-    parsed_opers = np.asarray(parsed_opers)
+    # Parse opers and convert to ndarray
+    parsed_opers = util.parse_operators(opers, H_str)
 
     if not all(hasattr(coeff, '__len__') for coeff in coeffs):
         raise TypeError(f'Expected coefficients in {H_str} to be a sequence')
diff --git a/filter_functions/util.py b/filter_functions/util.py
index 3f0bb26..3dab060 100644
--- a/filter_functions/util.py
+++ b/filter_functions/util.py
@@ -183,6 +183,55 @@ def wrapper(*args, **kwargs):
 parse_which_FF_parameter = parse_optional_parameters({'which': ('fidelity', 'generalized')})
 
 
+def parse_operators(opers: Sequence[Operator], err_loc: str) -> List[ndarray]:
+    """Parse a sequence of operators and convert to ndarray.
+
+    Parameters
+    ----------
+    opers: Sequence[Operator]
+        Sequence of operators.
+    err_loc: str
+        Some cosmetics for the exceptions to be raised.
+
+    Raises
+    ------
+    TypeError
+        If any operator is not a valid type.
+    ValueError
+        If not all operators are 2d and square.
+
+    Returns
+    -------
+    parse_opers: ndarray, shape (len(opers), *opers[0].shape)
+        The parsed ndarray.
+
+    """
+    parsed_opers = []
+    for oper in opers:
+        if isinstance(oper, ndarray):
+            parsed_opers.append(oper.squeeze())
+        elif hasattr(oper, 'full'):
+            # qutip.Qobj
+            parsed_opers.append(oper.full())
+        elif hasattr(oper, 'todense'):
+            # sparse object
+            parsed_opers.append(oper.todense())
+        elif hasattr(oper, 'data') and hasattr(oper, 'dexp'):
+            # qopt DenseMatrix
+            parsed_opers.append(oper.data)
+        else:
+            raise TypeError(f'Expected operators in {err_loc} to be NumPy arrays or QuTiP Qobjs!')
+
+    # Check correct dimensions of the operators
+    if set(oper.ndim for oper in parsed_opers) != {2}:
+        raise ValueError(f'Expected all operators in {err_loc} to be two-dimensional!')
+
+    if len(set(*set(oper.shape for oper in parsed_opers))) != 1:
+        raise ValueError(f'Expected operators in {err_loc} to be square!')
+
+    return np.asarray(parsed_opers)
+
+
 def _tensor_product_shape(shape_A: Sequence[int], shape_B: Sequence[int], rank: int):
     """Get shape of the tensor product between A and B of rank rank"""
     broadcast_shape = ()