Proposal: going to sparse matrix in tomography tools

[wave46]

Dear Cherab team.

The output matrices in, for example, admt_utils.py are usually very big and sparse and need lots of memory to store.
How do you think, if it would be profitable for both memory usage and computation efficiency to go to sparse matrix notation using scipy.sparse?

Thanks,
Ivan Kudashev

[jacklovell]

I did consider whether it was worthwhile making these sparse matrices when first implementing those utilities (which were initially copied from demo functionality scripts where optimisation was very much not a consideration). In the end I went with standard array types because it was a) what had already been implemented, b) simpler to debug and c) the inversion routines these matrices were used with (e.g scipy.optimize.nnls and cherab.tools.inversions.invert_regularised_sart) required non-sparse matrices anyway, so the memory savings would have been lost when they actually came to be used.

For the inversion grids of the JET and MAST-U inversions I was using these for the matrices weren't unreasonably large. I can see that if you want to do a high resolution ITER inversion for example then this may become an issue, but at that point O(N**3) inversion routines probably aren't suitable either so the requirement for non-sparse matrices is dropped.

Do the classes in scipy.sparse support the same operations as standard 2D arrays transparently? Would there be any modification needed to the utilities in admt_utils.py other than changing the initialisation of these arrays from np.zeros to a suitable sparse matrix initialiser? If that's the only change we could add a new boolean function argument sparse which would return sparse matrices when True and normal matrices when False, and it would default to False to maintain backwards compatibility. If more significant modifications are required I'd suggest adding new functions which geneate the sparse matrices: users would then call these new functions when they want sparse matrices and the original functions otherwise.

[wave46]

Thank you, @jacklovell for your answers.

I agree with you, that finally for using inversion tools one needs a dense matrix. Some of the classes in scipy.sparse (e.g. csr_matrix, 'lil_matrix') handle the same notation as normal matrices. However, it seems to me, that just changing initializing without changing the way how the matrices are defined will be even slower, because than each time the sparsity of the matrix is changed. Otherwise, one should first carefully store all of the indices and corresponding non-zero values as arrays and after that define a matrix.

So, summing the fact that one still needs dense matrices for the inversions and it that the initialization of the operators in admt_utils.py should be changed quite a lot, I think that dealing with sparse matrices can be left to user or, indeed, new functions should be made.

[jacklovell]

Been revisiting this during the development of the ADMT demo using a larger grid for the Generomak device, and I now agree that sparse support is worthwhile for the tomography routines. For inversions the runtime can be significantly reduced with sparse support: the Generomak demo uses 48 bolometer channels on a 2157 voxel grid and a sparse-aware NNLS inversion reduces the computation from 50 cpu-minutes (3 minutes on a 16-core machine) to 20 cpu-seconds (20 s wallclock as the algorithm is serial). So a very worthwhile gain.

A regularised NNLS inversion is pretty straightforward to implement, with the same API as the standard regularised NNLS inversion:

def invert_sparse_regularised_nnls(def invert_sparse_regularised_nnls(w_matrix, b_vector, alpha=0.01, tikhonov_matrix=None):
    import numpy as np
    from scipy.optimize import lsq_linear
    from scipy.sparse import csr_array, lil_array, eye_array
    m, n = w_matrix.shape
    if tikhonov_matrix is None:
        tikhonov_matrix = eye_array(n)
    c = lil_array((m+n, n))
    c[:m, :] = w_matrix
    c[m:, :] = alpha * tikhonov_matrix
    c = c.tocsr()  # 1 row per channel, each channel views small number of voxels
    d = np.zeros(m+n)
    d[:m] = b_vector
    # Normalize to avoid numerical issues.
    vmax = d.max()
    c /= vmax
    d /= vmax
    res = lsq_linear(c, d, bounds=(0, np.inf))
    return res.x, np.linalg.norm(res.fun) * vmax

This does require Scipy >= 0.17.0 for lsq_linear and Scipy > 1.8.0 for csr_array, lil_array and eye_array. Backwards compatibility with Scipy between 0.17 and 1.8 while avoiding upcoming deprecations in 1.18.0 can be achieved using:

try:
   from scipy.sparse import csr_array as csr, lil_array as lil, eye_array as eye
except ImportError:
    from scipy.sparse ipmort csr_matrix as csr, lil_matrix as lil, eye

And using csr, lil and eye instead of the explicit _array forms.

There are small numerical differences between the dense and sparse inversion results due to the different algorithms used, so in the interests of backwards compatibility I would propose having different functions for dense and sparse matrices.

The standard invert_regularised_lstsq can be converted to sparse using essentially the same function, just without setting the bounds argument to lsq_linear. The SART inversion routines implemented with Cython would need to be rewritten to support sparse matrices. This is a bit more involved.

I'm revisiting admt_utils.py as part of the Generomak demo in #427, and may end up rewriting some of it in a way that is transparent to matrix storage scheme so could make efficient use of sparse matrices. TBD.