HEOM Performance fixes

Author: pmenczel

qutip/core/cy/qobjevo.pyi

@@ -39,7 +39,7 @@ class QobjEvo:
     def arguments(self, new_args: dict[str, Any]) -> None: ...
     @overload
     def arguments(self, **new_args) -> None: ...
-    def compress(self) -> QobjEvo: ...
+    def compress(self, *, _skip_coeff=False, _skip_qobj=False) -> QobjEvo: ...
     def tidyup(self, atol: Number) -> QobjEvo: ...
     def copy(self) -> QobjEvo: ...
     def conj(self) -> QobjEvo: ...

qutip/core/cy/qobjevo.pyx

@@ -850,7 +850,7 @@ cdef class QobjEvo:
                 cleaned_elements.append(_EvoElement(qobj, coeff))
         return cleaned_elements
 
-    def compress(self):
+    def compress(self, *, _skip_coeff=False, _skip_qobj=False):
         """
         Look for redundance in the :obj:`.QobjEvo` components:
 
@@ -878,8 +878,10 @@ cdef class QobjEvo:
             else:
                 func_elements.append(element)
 
-        coeff_elements = self._compress_merge_coeff(coeff_elements)
-        coeff_elements = self._compress_merge_qobj(coeff_elements)
+        if not _skip_coeff:
+            coeff_elements = self._compress_merge_coeff(coeff_elements)
+        if not _skip_qobj:
+            coeff_elements = self._compress_merge_qobj(coeff_elements)
 
         cleaned_elements = coeff_elements + func_elements
 

qutip/core/direct_sum.py

@@ -24,6 +24,8 @@
 
 QobjLike = Union[Number, Qobj, QobjEvo]
 
+ARRAY_ONE = np.full((1,), 1, dtype=_data.base.idxint_dtype)
+
 
 def _is_like_ket(qobj: Any) -> bool:
     return (
@@ -376,7 +378,7 @@ def _do_direct_sum(
             _skip_check=True
         )
     if qobjevos:
-        result.compress()
+        result.compress(_skip_coeff=True)
     return result
 
 
@@ -419,7 +421,7 @@ def direct_component(sum_qobj, *index):
     if isinstance(sum_qobj, QobjEvo):
         result = sum_qobj.linear_map(lambda x: direct_component(x, *index),
                                      _skip_check=True)
-        result.compress()
+        result.compress(_skip_coeff=True)
         return result(0) if result.isconstant else result
 
     to_index, from_index = _check_component_index(sum_qobj._dims, index)
@@ -508,7 +510,7 @@ def set_direct_component(sum_qobj, component, *index):
         zeroed = sum_qobj.linear_map(
             lambda x: set_direct_component(x, None, *index),
             _skip_check=True)
-        zeroed.compress()
+        zeroed.compress(_skip_coeff=True)
         zeroed = zeroed(0) if zeroed.isconstant else zeroed
     else:
         zeroed = set_direct_component(sum_qobj, None, *index)
@@ -540,11 +542,21 @@ def _blow_up_qobj(x: Qobj, *, sum_dimensions, row, col, dtype):
     else:
         data_col = 0
 
+    xh, xw = x.shape
+    sh, sw = sum_dimensions.shape
+
     return Qobj(
-        _data.block_overwrite(
-            _data.zeros[dtype](*sum_dimensions.shape),
-            x.data, data_row, data_col, dtype=dtype
-        ), dims=sum_dimensions, copy=False
+        _data.block_build(
+            ARRAY_ONE,
+            ARRAY_ONE,
+            np.full((1,), x.data, dtype=_data.Data),
+            np.array([data_row, xh, sh - xh - data_row],
+                     dtype=_data.base.idxint_dtype),
+            np.array([data_col, xw, sw - xw - data_col],
+                     dtype=_data.base.idxint_dtype),
+            dtype=dtype
+        ),
+        dims=sum_dimensions, copy=False
     )
 
 

qutip/solver/heom/bofin_solvers.py

@@ -675,7 +675,7 @@ def __init__(self, H, bath, max_depth, *, odd_parity=False, options=None):
             else H  # already a liouvillian
         )
         if isinstance(self.L_sys, QobjEvo):
-            self.L_sys.compress()
+            self.L_sys.compress(_skip_coeff=True)
 
         # Dimensions that L_sys acts on
         self._sys_dims_ = self.L_sys._dims[0].oper
@@ -822,7 +822,7 @@ def _grad_prev(self, he_n, k):
         else:
             op = self._grad_prev_bosonic(he_n, k)
         if isinstance(op, QobjEvo):
-            op.compress()
+            op.compress(_skip_coeff=True)
         return op
 
     def _grad_prev_bosonic(self, he_n, k):
@@ -892,7 +892,7 @@ def _grad_next(self, he_n, k):
         else:
             op = self._grad_next_bosonic(he_n, k)
         if isinstance(op, QobjEvo):
-            op.compress()
+            op.compress(_skip_coeff=True)
         return op
 
     def _grad_next_bosonic(self, he_n, k):
@@ -928,7 +928,7 @@ def _calculate_rhs(self):
 
         ops = {}
         for he_n in self.ados.labels:
-            ops[(idx(he_n), idx(he_n))] = self._grad_n(he_n) + self.L_sys
+            ops[(idx(he_n), idx(he_n))] = self._grad_n(he_n)
             for k in range(len(self.ados.dims)):
                 next_he = self.ados.next(he_n, k)
                 if next_he is not None:
@@ -937,7 +937,17 @@ def _calculate_rhs(self):
                 if prev_he is not None:
                     ops[(idx(he_n), idx(prev_he))] = self._grad_prev(he_n, k)
 
-        return direct_sum_sparse(ops, self._rhs_dims)
+        if isinstance(self.L_sys, QobjEvo):
+            liouvillian_part = self.L_sys.linear_map(self._liouvillian_part)
+        else:
+            liouvillian_part = self._liouvillian_part(self.L_sys)
+
+        return liouvillian_part + direct_sum_sparse(ops, self._rhs_dims)
+
+    def _liouvillian_part(self, L):
+        return direct_sum_sparse(
+            {(idx, idx): L for idx in range(self._n_ados)}, self._rhs_dims
+        )
 
     def steady_state(
         self,