Sparse matrix support as a linear operator

[kyoungseoun-chung]

Hi,

I am currently working on implementing your optimization module (conjugate gradient method) to solve a linear system.

The system that I am trying to solve is, simple 2D Poisson equation. Therefore, left-hand side matrix is a sparse matrix and my goal is to use pylops_distributed to reduce calculation time. (below is code snippets that I am testing)

    import numpy as np
    from pylops_distributed.optimization.cg import cg
    from pylops_distributed.basicoperators import MatrixMult
    import pylops_distributed
    from pyamg.gallery import poison

    Nx = 1024
    Ny = 1024
    Ntot = Nx * Ny
    A = poison((Nx, Ny), format='csr', dtype=np.float64)
    
    n_workers = 4
    client, _ = pylops_distributed.utils.backend.dask(n_workers=n_workers,
                                                      threads_per_worker=2)

    workers = int(np.sqrt(n_workers))

    A_da = da.from_array(A, chunks=(Nx//workers, Ny//workers))
    A_op = MatrixMult(A_da)

    b_da = da.zeros(Ntot, chunks=(Ntot//n_workers))

    x = cg(A_op, b_da)

    client.close()

The problem or question I have is that I cannot have any benefits of using distributed computation.

I am comparing this to scipy.linalg and pyamg libraries and pylops_distributed gives me about 100times slower calculation time.

Also, I observed the smaller the workers, the faster the calculation. Seems like there is a communication overhead.

I just wonder whether this is a known issue or not. (or whether pylops_distrubted is suitable to solve the problem I am tacking or not)

Thank you in advance!

Best regards,
Kyoungseoun Chung

[kyoungseoun-chung]

By the way, I don't know how to add a label for this issue. If this violates the contribution guideline, please let me know how to fix it. I guess this should belong to the question label...

[mrava87]

Hi @kyoungseoun-chung,
very sorry for noticing your issue before. Hopefully you made some progress already during this time :)

I think this is a classical case where distributed computing/dask/pylops-distributed (in the order of level of abstraction) is not really needed and so you may not get the benefit you wished to have. Let me try to explain a bit more what I mean.

Distributed computing, which is what dask does (and what pelops-distributed offers in the context of inverse problems), is beneficial when you have a very large problem that cannot sit in memory of a single computer, or even if this is not the case, when you expect to have very large computations interleaved by some smaller times for communication. On the other hand if your problem easily fits in memory and computations are fairly small, the overhead will kill any benefit of the distribution. I worry this is the case for your problem.

The first thing I noticed is that if your A is NxNy x NxNy and you choose chunks=(Nx//workers, Ny//workers) you are creating a massive amount of chunks, way smaller than what dask would do as default. Also when you say the smaller the worker the faster, I worry this is because you take way too small chunks and if Nworkers is small you have less chunks and so less overhead...

I also must admit I never used Dask with sparse arrays so I have little experience there but if I do:

import numpy as np
from pylops_distributed.optimization.cg import cg
from pylops_distributed.basicoperators import MatrixMult
import pylops_distributed
from pyamg.gallery import poisson

Nx = 1024
Ny = 1024
Ntot = Nx * Ny
A = poisson((Nx, Ny), format='csr', dtype=np.float64)

y = A * np.ones(A.shape[1])

I get quite instantaneous result. If I do add:

A_da = da.from_array(A)
x_da = da.ones(A.shape[1], chunks=4096)
y_da = A_da @ x_da
y_da.compute()

this takes a huge amount of time and gives me an error.

Even worse, if I do something much simpler:

from scipy.sparse import csc_matrix

row = np.array([0, 2, 2, 0, 1, 2])
col = np.array([0, 0, 1, 2, 2, 2])
data = np.array([1., 2, 3, 4, 5, 6])
A = csc_matrix((data, (row, col)), shape=(100, 100))
x = np.ones(100)
y = A @ x
A_da =  da.from_array(A)
x_da= da.from_array(x)
y_da = A_da @ x_da
y_da.compute()

I get the same error: ValueError: matmul: Input operand 0 does not have enough dimensions (has 0, gufunc core with signature (n?,k),(k,m?)->(n?,m?) requires 1) I worry something is wrong with sparse arrays and Dask, or we are doing something not right there. I would suggest you try to figure this out first and perhaps ask the Dask folks as if you can't get a good timing for the simple matrix-vector multiplication, everything else later will just be slow as you see...