Feature: in-place solvers

[mrava87]

This PR introduces two new functionalities to all of the solvers in the optimization module, namely:

• a method called memory_usage that can be used before running the setup and step methods to get an approximate estimate of the memory usage of the solver.
• a new optional parameter called preallocate that modifies the behavior of the solver compared to the current one, by preallocating all of the arrays in the setup method and performing all operations on arrays in-place (note that this does not apply to the Op.matvec/Op.rmatvec operations as those are currently not implemented in-place)

* approximate refers to the fact that the memory usage is calculated taking into account only the main variables of the solvers (ND-arrays of the size of either the model or the data) and not additional lists or arrays that store losses or other metrics.

[mrava87]

@cako I understand if you are busy and have no time for this, but if you can and are interested I'd appreciate any feedback (see https://github.com/PyLops/pylops_inplacesolvers for comparison with current solvers and in particular https://github.com/PyLops/pylops_inplacesolvers/blob/main/notebooks/Visualize.ipynb for some visualization of timings)

[cako]

Had a cursory look, and everything looked good. I did have a couple questions regarding intended use. Seems like preallocating is slower? And it doesnt change the memory consumption? If so, is there any reason to doing it? Regarding the memory usage computation, I think it's very useful.

[mrava87]

Had a cursory look, and everything looked good. I did have a couple questions regarding intended use. Seems like preallocating is slower? And it doesnt change the memory consumption? If so, is there any reason to doing it? Regarding the memory usage computation, I think it's very useful.

Thanks @cako!

Mmh not sure where you see that preallocating is slower. I re-run all experiments with N=3_000_000 (data and model size) and I get the below speed-ups:

Regarding memory, yes I do not expect the peak memory to change with pre-allocation as anyways the same variables are either created in the step method over and over again (current behavior) or upfront and filled (new behavior)...

In general, after some thinking, I decided to modify a bit the inner working of all the solvers when preallocate=False and revert back to the original codes (like they were still used with JAX arrays); this is mostly to avoid any backward compatibility problem as I noticed that numpy (and even more cupy) are quite strict in terms of data types when doing operations in place (eg +=) - basically an error will be raised if np.complex64 is added in place to a np.float32... whilst I think this is actually good and helps users to make sure their operators and solvers don't do unwanted internal casting (and that would be another incentive for using the new preallocate option), I do not want codes already written to break 😉

[mrava87]

@cako just checking you saw my reply?

[cako]

Yup! Sorry forgot to reply. But yes, your answer makes sense.

[mrava87]

Yup! Sorry forgot to reply. But yes, your answer makes sense.

Great, I am going to do one last look and polishing and then I will merge 😄