Speed up the HEOM RHS construct

[hodgestar]

Description

Adds a _from_csr_blocks constructor to speed up the construction of the HEOM RHS by a factor of ~4x.

What allows from_csr_blocks to be fast is that we can order the blocks in a way that makes it possible to write the rows of the output one after the other without a lot of work, and we never have to add a new non-zero element into a row or make a copy of the big output matrix.

The implementation in Python was slow just because there is a lot of Python looping and arithmetic when the number of operators being combined becomes large.

_from_csr_blocks is a private constructor for now. If use cases emerge we can consider making something similar part of the public interface later.

Related issues or PRs

• None

[hodgestar]

@nwlambert Would you mind giving this a try with some bigger realistic HEOM examples you have lying around? Eric's recent improvements to CSR mul and imul also help HEOM RHS construction since there is a lot of c * op happening. The 3.5x improvement from this branch is on top of that.

@Ericgig Would you mind giving your thoughts on the Cython interface both from a "is this good from a technical point of view" and from a "do we want this point of view". I guess we could also name it _from_csr_blocks if we don't want to expose it.

[Ericgig]

It doesn't look robust enough to be user facing as it is, but as a private function, if you get a 3.5x, I'd say we want it.

Technically, just give a type to i and it's all running in c code.

[coveralls]

Coverage: 75.328% (-0.06%) from 75.391% when pulling 58f82aa on hodgestar:feature/faster-heom-rhs-construction into 17e7958 on qutip:master.

[hodgestar]

It doesn't look robust enough to be user facing as it is, but as a private function, if you get a 3.5x, I'd say we want it.

I'll rename it to _from_csr_blocks and add some more checks on the ordering and shape of the ops if those aren't too expensive and add some tests. We can expose it publicly if that's ever useful.

Technically, just give a type to i and it's all running in c code.

Thanks! Adding the type to i did make it a little faster.

[Ericgig]

Seems fine.
This fixes the used data type to CSR and this should be documented.
I think you want csr.zeros instead of csr.empty.

[hodgestar]

This fixes the used data type to CSR and this should be documented.

It is documented in the docstring for block_ops? Or did you want more documentation somewhere?

[hodgestar]

@Ericgig Thank you for the review. Would you mind having a last look? I made quite a few small changes.

[Ericgig]

Better with zeros.

I would document in HEOMSolver say there that the Hamiltonian and bath operator will be coerced to csr.