From 3f19f258c216a4f5557d5c346ba707b2477ae026 Mon Sep 17 00:00:00 2001 From: Neill Lambert Date: Thu, 29 Aug 2024 17:24:36 +0900 Subject: [PATCH] refactor to just include TD on input terms Reverted back to using data objects of Ops, and made a new list to store only time-dependent ops for input ADOS. Solves slow RHS construction problem (partially) --- qutip/solver/heom/bofin_solvers.py | 178 +++++++++++++++-------------- 1 file changed, 93 insertions(+), 85 deletions(-) diff --git a/qutip/solver/heom/bofin_solvers.py b/qutip/solver/heom/bofin_solvers.py index f7e73620fa..745b354842 100644 --- a/qutip/solver/heom/bofin_solvers.py +++ b/qutip/solver/heom/bofin_solvers.py @@ -637,28 +637,31 @@ def __init__(self, H, bath, max_depth, *, options=None): _time_start = time() # pre-calculate identity matrix required by _grad_n - self._sId = Qobj(_data.identity(self._sup_shape, dtype="csr"),dims = self.L_sys.dims, superrep='super') + self._sId = _data.identity(self._sup_shape, dtype="csr") # pre-calculate superoperators required by _grad_prev and _grad_next: - Qs = [exp.Q for exp in self.ados.exponents] - self._spreQ = [spre(op) for op in Qs] - self._spostQ = [spost(op) for op in Qs] + Qs = [exp.Q.to("csr") for exp in self.ados.exponents] + self._spreQ = [spre(op).data for op in Qs] + self._spostQ = [spost(op).data for op in Qs] + self._s_pre_minus_post_Q = [ - (self._spreQ[k]-self._spostQ[k]) + _data.sub(self._spreQ[k], self._spostQ[k]) for k in range(self._n_exponents) ] self._s_pre_plus_post_Q = [ - (self._spreQ[k] + self._spostQ[k]) + _data.add(self._spreQ[k], self._spostQ[k]) for k in range(self._n_exponents) ] - self._spreQdag = [spre(op.dag()) for op in Qs] - self._spostQdag = [spost(op.dag()) for op in Qs] + self._spreQdag = [spre(op.dag()).data for op in Qs] + self._spostQdag = [spost(op.dag()).data for op in Qs] + + self._s_pre_minus_post_Qdag = [ - (self._spreQdag[k] - self._spostQdag[k]) + _data.sub(self._spreQdag[k], self._spostQdag[k])) for k in range(self._n_exponents) ] self._s_pre_plus_post_Qdag = [ - (self._spreQdag[k] - self._spostQdag[k]) + _data.add(self._spreQdag[k], self._spostQdag[k]) for k in range(self._n_exponents) ] @@ -716,7 +719,7 @@ def _grad_n(self, he_n): """ Get the gradient for the hierarchy ADO at level n. """ vk = self.ados.vk vk_sum = sum(he_n[i] * vk[i] for i in range(len(vk))) - op = self._sId * (-vk_sum) + op = _data.mul(self._sId, -vk_sum) #self._sId * (-vk_sum) return op def _grad_prev(self, he_n, k): @@ -728,35 +731,40 @@ def _grad_prev(self, he_n, k): def _grad_prev_bosonic(self, he_n, k): if self.ados.exponents[k].type == BathExponent.types.R: - op = ( - self._s_pre_minus_post_Q[k] * - -1j * he_n[k] * self.ados.ck[k] + op = _data.mul( + self._s_pre_minus_post_Q[k], + -1j * he_n[k] * self.ados.ck[k], ) + elif self.ados.exponents[k].type == BathExponent.types.I: - op = ( - self._s_pre_plus_post_Q[k] * - -1j * he_n[k] * 1j * self.ados.ck[k] + op = _data.mul( + self._s_pre_plus_post_Q[k], + -1j * he_n[k] * 1j * self.ados.ck[k], ) + elif self.ados.exponents[k].type == BathExponent.types.RI: - term1 = ( - self._s_pre_minus_post_Q[k] * - he_n[k] * -1j * self.ados.ck[k] + term1 = _data.mul( + self._s_pre_minus_post_Q[k], + he_n[k] * -1j * self.ados.ck[k], ) - term2 = ( - self._s_pre_plus_post_Q[k] * - he_n[k] * self.ados.ck2[k] + + term2 = _data.mul( + self._s_pre_plus_post_Q[k], + he_n[k] * self.ados.ck2[k], ) - op = term1+term2 + + op = _data.add(term1, term2) elif self.ados.exponents[k].type == BathExponent.types.Input: - op = ( - QobjEvo([self._s_pre_minus_post_Q[k] * - he_n[k], self.ados.ck[k]]) - ) + op = _data.mul( + self._s_pre_minus_post_Q[k], + he_n[k], + ) #omit ck here, it is included in ops_td + elif self.ados.exponents[k].type == BathExponent.types.Output: - op =( - self._s_pre_minus_post_Q[k] * - -1j * he_n[k] * self.ados.ck[k] + op = _data.mul( + self._s_pre_minus_post_Q[k], + he_n[k] * self.ados.ck[k], ) else: @@ -782,20 +790,20 @@ def _grad_prev_fermionic(self, he_n, k): sigma_bar_k = k + self.ados.sigma_bar_k_offset[k] if self.ados.exponents[k].type == BathExponent.types["+"]: - op = ( - (self._spreQdag[k]+1j * sign2 * ck[k]) * - ( - self._spostQdag[k] - - -1j * sign2 * sign1 * np.conj(ck[sigma_bar_k]) - ) + op = _data.sub( + _data.mul(self._spreQdag[k], -1j * sign2 * ck[k]), + _data.mul( + self._spostQdag[k], + -1j * sign2 * sign1 * np.conj(ck[sigma_bar_k]), + ), ) elif self.ados.exponents[k].type == BathExponent.types["-"]: op = _data.sub( - (self._spreQ[k] * -1j * sign2 * ck[k]) - - ( - self._spostQ[k] * - -1j * sign2 * sign1 * np.conj(ck[sigma_bar_k]) - ) + _data.mul(self._spreQ[k], -1j * sign2 * ck[k]), + _data.mul( + self._spostQ[k], + -1j * sign2 * sign1 * np.conj(ck[sigma_bar_k]), + ), ) else: raise ValueError( @@ -814,7 +822,7 @@ def _grad_next(self, he_n, k): def _grad_next_bosonic(self, he_n, k): if (self.ados.exponents[k].type != BathExponent.types.Input and self.ados.exponents[k].type != BathExponent.types.Output): - op = (self._s_pre_minus_post_Q[k]* -1j) + op = _data.mul(self._s_pre_minus_post_Q[k], -1j) #op = (self._s_pre_minus_post_Q[k]* -1j) return op else: return None @@ -832,12 +840,12 @@ def _grad_next_fermionic(self, he_n, k): if self.ados.exponents[k].type == BathExponent.types["+"]: if sign1 == -1: - op = (self._s_pre_minus_post_Q[k] * -1j * sign2) + op = _data.mul(self._s_pre_minus_post_Q[k], -1j * sign2) else: - op = (self._s_pre_plus_post_Q[k] * -1j * sign2) + op = _data.mul(self._s_pre_plus_post_Q[k], -1j * sign2) elif self.ados.exponents[k].type == BathExponent.types["-"]: if sign1 == -1: - op = (self._s_pre_minus_post_Qdag[k] * -1j * sign2) + op = _data.mul(self._s_pre_minus_post_Qdag[k], -1j * sign2) else: op = (self._s_pre_plus_post_Qdag[k] * -1j * sign2) else: @@ -868,13 +876,16 @@ def _rhs(self): prev_he = self.ados.prev(he_n, k) if prev_he is not None: op = self._grad_prev(he_n, k) - ops.add_op(he_n, prev_he, op) + if self.ados.exponents[k].type == BathExponent.types.Input: + ops.add_op(he_n, prev_he, op, self.ados.ck[k]) + else: + ops.add_op(he_n, prev_he, op) - return ops.gather() + return ops.gather(), ops.gather_td() def _calculate_rhs(self): """ Make the full RHS required by the solver. """ - rhs_mat = self._rhs() + rhs_mat, rhs_mat_td = self._rhs() rhs_dims = [ [self._sup_shape * self._n_ados], [self._sup_shape * self._n_ados] ] @@ -883,8 +894,8 @@ def _calculate_rhs(self): if self.L_sys.isconstant: # For the constant case, we just add the Liouvillian to the # diagonal blocks of the RHS matrix. - rhs_mat += Qobj(_data.kron(h_identity, self.L_sys(0).to("csr").data), dims=rhs_dims) - rhs = QobjEvo(rhs_mat) + rhs_mat += _data.kron(h_identity, self.L_sys(0).to("csr").data) + rhs = QobjEvo(Qobj(rhs_mat, dims=rhs_dims))+ rhs_mat_td else: # In the time dependent case, we construct the parameters # for the ODE gradient function under the assumption that @@ -895,7 +906,7 @@ def _calculate_rhs(self): # This assumption holds because only _grad_n dependents on # the system Liouvillian (and not _grad_prev or _grad_next) and # the bath coupling operators are not time-dependent. - rhs = rhs_mat#QobjEvo(Qobj(rhs_mat, dims=rhs_dims)) + rhs = QobjEvo(Qobj(rhs_mat, dims=rhs_dims)) + rhs_mat_td def _kron(x): return Qobj( @@ -909,7 +920,7 @@ def _kron(x): # check on the RHS creation. The base solver class will still # convert the RHS to the type required by the ODE integrator if # needed. - #assert isinstance(rhs_mat, _csr.CSR) + assert isinstance(rhs_mat, _csr.CSR) assert isinstance(rhs, QobjEvo) assert rhs.dims == rhs_dims @@ -1315,17 +1326,23 @@ def __init__(self, f_idx, block, nhe, rhs_dims): self._n_blocks = nhe self._f_idx = f_idx self._ops = [] + self._ops_td = [] self._rhs_dims = rhs_dims - def add_op(self, row_he, col_he, op): + def add_op(self, row_he, col_he, op, ck_td_factor = None): """ Add an block operator to the list. """ - self._ops.append( - (self._f_idx(row_he), self._f_idx(col_he), op) + if ck_td_factor == None: + self._ops.append( + (self._f_idx(row_he), self._f_idx(col_he), op) + ) + else: + self._ops_td.append( + (self._f_idx(row_he), self._f_idx(col_he), op, ck_td_factor) ) - def gather2(self): + def gather(self): """ Create the HEOM liouvillian from a sorted list of smaller sparse matrices. @@ -1354,33 +1371,31 @@ def gather2(self): self._n_blocks, self._block_size, ) - def gather(self): - """ Create the HEOM liouvillian from a sorted list of smaller sparse - matrices. + def gather_td(self): + """ Create the time-dependent parts of the HEOM liouvillian from a + sorted list of smaller sparse matrices and TD factors. .. note:: The list of operators contains tuples of the form - ``(row_idx, col_idx, op)``. The row_idx and col_idx give the - *block* row and column for each op. An operator with + ``(row_idx, col_idx, op, ck_td_factor)``. The row_idx and col_idx + give the *block* row and column for each op. An operator with block indices ``(N, M)`` is placed at position ``[N * block: (N + 1) * block, M * block: (M + 1) * block]`` - in the output matrix. + in the output matrix, with a time-dependent factor given + by the function ck_td_factor. + Returns ------- rhs : :obj:`Data` A combined matrix of shape ``(block * nhe, block * ne)``. """ - self._ops.sort() - ops = np.array(self._ops, dtype=[ - ("row", _data.base.idxint_dtype), - ("col", _data.base.idxint_dtype), - ("op", _data.CSR), - ]) - - + # TODO: checks on ck_td_factor + self._ops_td.sort() + + def _sum(op, loc): def _kron(x): @@ -1391,20 +1406,13 @@ def _kron(x): return op.linear_map(_kron) - def _sum2(op, row, col): - def _insert(x): - return Qobj(_csr._from_csr_blocks( - np.array([row],dtype= _data.base.idxint_dtype), np.array([col],dtype= _data.base.idxint_dtype), np.array([x.data],dtype=_data.CSR), - self._n_blocks, self._block_size, - ),dims=self._rhs_dims) - - return op.linear_map(_insert) - return np.sum([_sum(QobjEvo(op),_data.one_element_csr((self._n_blocks, self._n_blocks), (row, col))) for row, col, op in self._ops]) - #return np.sum([_sum2(QobjEvo(op),row,col) for row, col, op in ops]) - #return QobjEvo(ops["op"][0]).batch_linear_map(ops, self._rhs_dims, self._n_blocks) - - #def _sum_vectorized(ops, n_blocks, rhs_dims): + + return np.sum( + [_sum(QobjEvo([Qobj(op), tf_func]), + _data.one_element_csr((self._n_blocks, self._n_blocks), + (row, col))) + for row, col, op, tf_func in self._ops_td] + ) - #rhs = QobjEvo(Qobj(rhs_mat, dims=rhs_dims))