Add MPIBlockDiag and MPILinearOperator

[rohanbabbar04]

• Added MPIBlockDiag with _matvec and _rmatvec which now uses DistributedArray
• Added MPILinearOperator class
• Added basic classes(_ScaledLinearOperator, _AdjointLinearOperator, etc) and made them input and output DistributedArray.
• Added tests to verify above classes.
• Upgrade scipy to 1.8 and python to 3.8
• Updated docs for sphinx

[mrava87]

@rohanbabbar04, thanks for this PR. So far I struggle a bit to understand the rationale behind this operator. I had a few comments in the closed PR that I put here again, hopefully your answers (and example!) will help clarify my doubts.

• [Taken from closed PR] Looking at your example, I am not sure I understand how you envision/designed MPIBlockDiag… my idea is that if you have a block diagonal operator with N operators (N being the number of ranks), each rank will build one pylops operator and pass it to the MPIBlockDiag, so that once matvec/rmatvec is called each rank will compute their portion… so if you were to compare your new implementation with an equivalent pylops one, the pylops operator will have to take all of those operators and give the same result…. For simplicity take this example: you have 3 ranks and 3 5x5 matrices which you wrap into 3 MatrixMult operators. The first is filled with 1s, the second with 2s, the third with 3s.. now for pylops-MPI each rank makes one MatrixMult operator and passes it to MPIBlockDiag… the x is also distributed and can be just full of 1s for simplicity. Unless I don’t understand your example it doesn’t seem like it is built like this, rather each rank seems to make all of the ops…
• Based on the above, it seems like your x and y are not DistributedArray, instead just numpy arrays. Why? Again, I would think we should have distributed arrays passed to MPI operators, so that internally it is easy to get the portion needed
• In the PR you closed you had created a LinearOperator class and subclassed MPIBlockDiag. Was there any reason/benefit from doing that, or right now you do not see any use wrt MPIBlockDiag?

[rohanbabbar04]

@mrava87 Done the changes, if the code and tests are good to go, we can merge this
I will make a new PR for the example(Post Stack one)

[mrava87]

Good job! Still quite a few comments but I think we are getting closer to having a good MPIBlockDiag operator.

For now I think we can stick to this implementation. Later we can consider maybe adding also what you had originally, allowing users to choose (depending on the scenario) if they want to create&pass operators per rank or all of them from all ranks - the second may be more flexible for complex scenarios but I suspect the former will be the most used one as in most cases operators will be the same for all ranks, just filled with different 'content'. (or just doing the same think to different parts of the model and data vector)

[rohanbabbar04]

Hi @mrava87,

• I have added MPILinearOperator class and made MPIBlockDiag extend that.
• Added important classes(_ScaledLinearOperator, _AdjointLinearOperator, etc) and made them input and output DistributedArray.
• Since we are adding the above classes we need to test them. For the linearop tests(for Scaled, Prod and Sum) I have used MPIBlockDiag, as our linearop.

[mrava87]

This looks great!

I made just a few comments on the LinearOperator. Most should be easy to handle, but there is one or two where I cannot understand some of your code choices that would be good to discuss/agree upon before moving ahead.

After that, for me this is ready to be merged :)

[mrava87]

self.base_comm.allreduce(self.shape[1]) why?

You ask the user to pass shape and from the doc you say Shape of the MPI Linear Operator. which makes me think is the global shape. Also in MPIBlockDiag the shape is computed as shape = (base_comm.allreduce(self.nops), base_comm.allreduce(self.mops)) which means it is the global shape.. so why this?

[mrava87]

Same as above

[rohanbabbar04]

• I have made MPILinearOperator class a standalone class(not extending pylops.LinearOperator), have add some docstrings comment for different methods as a reason.
• Currently dot only supports matvec operations, not matmat.
• Thinking about asmpilinearoperator, currently reverted from self.base_comm.allreduce(self.shape[1]) to self.shape[1], but will have to handled for different operators as a direct self.Op._matvec(x.local_array) will possibly give an error...
• For the MPILinearOperator, there are mainly 3 parameters we primarily use which was in pylops.LinearOperator, which are Op, shape and dtype..

[mrava87]

Great, I think this is ready to go. We will see if things like matmat are needed in future but for now what we have in MPILinearOperator should be enough to be able to successfully solve our initial use-cases :)