Feature: stacked

[mrava87]

This PR provides a first implementation of StackedDistributedArray as discussed in #76.

It contains the following new features:

• new array object StackedDistributedArray
• new base class MPIStackedLinearOperator that allows operating on multiple DistributedOperator performing usual operators (_matvec, _rmatvec, addition, chaining, ...)
• MPIStackedBlockDiag and MPIStackedVStack
• modified cg and cgls able to interact with both DistributedArray and StackedDistributedArray

[mrava87]

@rohanbabbar04 I finally managed to get some time to look into one of the two things which we didn't manage to implement, the StackedDistributedArray. This basically follows from the need to stack multiple distributed arrays to be able to solve an augmented system of equations without having to resort to the explicit normal equations.

Would you be interested to take a look and see what you think?

[rohanbabbar04]

My bad, I didn't get much time to implement it...

Thanks, it looks good. I will review it during the weekend

[rohanbabbar04]

Thanks @mrava87, It looks great....

We can perhaps add StackedBlockDiagonal and StackLinearOperator in this PR itself....

If you haven't worked in it I will work on it and push changes by the end of this week...

[mrava87]

@rohanbabbar04 that would be great! I was running out of time but indeed having those two was my goal before merging this PR. Please go ahead and I can review it once you push these new bits :D

[rohanbabbar04]

Ok so I saw the requirements of what we need to/could do...
I have already added the StackedBlockDiag...

Some things we will include in this PR(this PR could lead to a new release as it is a major change)

• Make the StackedLinearOperator class
• Extend the StackedBlockDiag and StackedVStack to StackedLinearOperator(also add some tests for each)
• Will be needed to update the cg and cgls solver to include the StackedLinearOperator.(should'nt be that difficult as it is a list of dist arrays)

[mrava87]

@rohanbabbar04 great! I will try to look into your changes during the week. If in the meantime you are able to work on the StackedLinearOperator this will be great.

A few comments to your points (which I also lifted up in the list of task in the first message)

Ok so I saw the requirements of what we need to/could do... I have already added the StackedBlockDiag...

Some things we will include in this PR(this PR could lead to a new release as it is a major change)

• Make the StackedLinearOperator class
• Extend the StackedBlockDiag and StackedVStack to StackedLinearOperator(also add some tests for each)
  Indeed, I have added some tests for the StackedDistributedArray but not yet for the operators

• Will be needed to update the cg and cgls solver to include the StackedLinearOperator.(should'nt be that difficult as it is a list of dist arrays)

This is already done. What I tried to do is to modify the solvers such that they are agnostic to whether one is working with DistributedArray or StackedDistributedArray... it required making a few changes, for example implementing in-place operations (e.g., iadd) to replace some of the lines that were working direction on local_array, eg

 x[:] += a * self.c.local_array  -> x += a * self.c

However in a few places I have been forced to do

self.c[:] = self.r.local_array + b * self.c.local_array  -> self.c = self.r + b * self.c

which I think may add a little extra cost as a brand new self.c will be made at each iteration, whilst before this was not the case.... any smarter idea is welcome :)

[rohanbabbar04]

@rohanbabbar04 great! I will try to look into your changes during the week. If in the meantime you are able to work on the StackedLinearOperator this will be great.

A few comments to your points (which I also lifted up in the list of task in the first message)

Ok so I saw the requirements of what we need to/could do... I have already added the StackedBlockDiag...
Some things we will include in this PR(this PR could lead to a new release as it is a major change)

• Make the StackedLinearOperator class
• Extend the StackedBlockDiag and StackedVStack to StackedLinearOperator(also add some tests for each)
  Indeed, I have added some tests for the StackedDistributedArray but not yet for the operators

• Will be needed to update the cg and cgls solver to include the StackedLinearOperator.(should'nt be that difficult as it is a list of dist arrays)
  This is already done. What I tried to do is to modify the solvers such that they are agnostic to whether one is working with DistributedArray or StackedDistributedArray... it required making a few changes, for example implementing in-place operations (e.g., iadd) to replace some of the lines that were working direction on local_array, eg

 x[:] += a * self.c.local_array  -> x += a * self.c

However in a few places I have been forced to do

self.c[:] = self.r.local_array + b * self.c.local_array  -> self.c = self.r + b * self.c

which I think may add a little extra cost as a brand new self.c will be made at each iteration, whilst before this was not the case.... any smarter idea is welcome :)

Sure, I will get the StackLinearOperator ready...
Let me check if I can think any thing better for self.c

[rohanbabbar04]

I have added some changes:

• Added MPIStackedLinearOperator
• Added and tested AdjointLinearOperator, TransposeLinearOperator, ScaledLinearOperator and ConjLinearOperator
• Prefixed the operators with MPI
• MPIStackedBlockDiag and MPIStackedVStack now extend the MPIStackedLinearOperator.
• Need to add and test the SumLinearOperator, ProductLinearOperator and PowerLinearOperator(will do it in 2-3 days and then @mrava87 you can review it.)

[mrava87]

@rohanbabbar04 let me know if/when I should go ahead with the review :)

[rohanbabbar04]

@mrava87, most the things discussed are done...
But I encountered one issue while implementing the SumStackedLinearOperator and ProductStackedLinearOperator

import pylops_mpi

stacked_blockdiag = pylops_mpi.MPIStackedBlockDiag(ops=[....])
stacked_vstack = pylops_mpit.MPIStackedVStack(ops=[...])
s1= stacked_blockdiag + stacked_vstack
s1.H @ x

def _rmatvec(self, x: Union[DistributedArray, StackedDistributedArray]) -> Union[DistributedArray, StackedDistributedArray]:
        arr1 = self.args[0].rmatvec(x)
        arr2 = self.args[1].rmatvec(x)
        return arr1 + arr2

while applying in adjoint mode, arr1 will be a StackedDistributedArray whereas arr2 will be a DistributedArray, which we cannot sum. Any suggestions on this one?

After this we will need to modify the cg and cgls a little bit...

[mrava87]

Hi Rohan,
I am a bit unsure myself if this is possible/useful. I am tempted to say that for SumStackedLinearOperator also the forward is problematic as the input of MPIStackedBlockDiag is a StackedDistributedArray whilst the one of MPIStackedVStack is a DistributedArray. So, how would you implement this since we always assume that the input must be the same in the case one sums two operators (each operator is applied to the input and the outputs are summed).

For ProductStackedLinearOperator I think the scenario where you have MPIStackedBlockDiag *MPIStackedVStack should work because the first operator will take a DistributedArray and transform it in a StackedDistributedArray which the second operator can handle (same for adjoint); but the other order, MPIStackedVStack * MPIStackedBlockDiag, won't work and should be prevented with some checks, as here you would need a DistributedArray to be input of the second operator but the first by definition will always return a StackedDistributedArray... same you can't really have MPIStackedVStack *MPIStackedVStack

So, in short, I think we need:

• SumStackedLinearOperator: check that what you sum is of the same type, MPIStackedBlockDiag with MPIStackedBlockDiag or MPIStackedVStack with MPIStackedVStack
• ProductStackedLinearOperator: check that you have MPIStackedBlockDiag *MPIStackedVStack or MPIStackedBlockDiag * MPIStackedBlockDiag (all other cases are unfeasible

Makes sense?

[rohanbabbar04]

Make sense...
I'll just get it ready as well

[rohanbabbar04]

@mrava87 , I have made all the changes
We will need to add support for cg and cgls for MPIStackedLinearOperators...(especially handling StackedDistributedArray)

You can review all the changes we have done till then...

[rohanbabbar04]

@mrava87 , You can review this PR,

We can then merge it and open a new PR for handling StackedDistributedArray in cg and cgls. What do you think?

[mrava87]

@rohanbabbar04 sure let me review your changes.

Note that I already modified cg and cgls to work interchangeably with both DistributedArray and StackedDistributedArray (my bad for not adding a point in the list at the top of the PR). In fact, a tutorial already uses this. However I think we should look a bit more into it as a couple of changes that I made require re-initializing arrays whilst before they were working in place with .localarray… there may be a way to do it whilst still keeping interchangeability but I couldn’t find one yet, so let’s leave this for the new PR :)

[rohanbabbar04]

Sounds Great!!
I have noticed the changes, we will need to make changes in the docstrings of the cg and cgls. Would also like to add some tests for the StackedLinearOperators.(for cg and cgls solvers)...

We can do this in a new PR :)

[mrava87]

@rohanbabbar04 I did modify the docstrings of the solvers and did a final clean up of all the new codes... I merge now :)

Note that I created two new issues for the solvers, so we don't forget to add tests and look into any better way to avoid re-initializing arrays if not strictly needed :)