Neaten up construction of block_rows, block_cols and block_ops arrays.

hodgestar · Mar 21, 2023 · c815c52 · c815c52
1 parent 0a41dab
commit c815c52
Showing 1 changed file with 6 additions and 13 deletions.
diff --git a/qutip/solver/heom/bofin_solvers.py b/qutip/solver/heom/bofin_solvers.py
@@ -1299,19 +1299,12 @@ def gather(self):
                 A combined matrix of shape ``(block * nhe, block * ne)``.
         """
         self._ops.sort()
-        block_rows = np.array(
-            [op[0] for op in self._ops],
-            dtype=_data.base.idxint_dtype,
-        )
-        block_cols = np.array(
-            [op[1] for op in self._ops],
-            dtype=_data.base.idxint_dtype,
-        )
-        block_ops = np.array(
-            [op[2] for op in self._ops],
-            dtype=object,
-        )
+        ops = np.array(self._ops, dtype=[
+            ("row", _data.base.idxint_dtype),
+            ("col", _data.base.idxint_dtype),
+            ("op", object),
+        ])
         return _csr._from_csr_blocks(
-            block_rows, block_cols, block_ops,
+            ops["row"], ops["col"], ops["op"],
             self._n_blocks, self._block_size,
         )