GeoStat-Framework / pentapy

A Python toolbox for pentadiagonal linear systems
https://pentapy.readthedocs.io/
MIT License
14 stars 4 forks source link

Consider using np.asarray instead of np.array in solve #16

Closed derb12 closed 3 years ago

derb12 commented 3 years ago

Hi, thanks for the great library. I have a suggestion that could slightly improve performance when using the solvers with numpy array inputs.

In the solve function, you use np.array for the input matrix and right-hand-side array, which creates new arrays of the inputs even if they are already numpy arrays. I know NumPy's and Scipy's solvers do not modify the input arrays by default so they don't require copying, and it doesn't seem like your two solvers modify them either unless you call shift_banded before the solver. Therefore, I would suggest you use np.asarray whenever you do not require calling shift_banded (similar to how you use np.asanyarray within the tools functions), so that a new array is created only if the input is not already an array with the correct dtype.

I modified your 03_perform_simple.py example program and made it into a gist that shows the difference of the solver times for the PTRANS-I and PTRANS-II solvers when using np.asarray (https://gist.github.com/derb12/3e2670a509e86051ee9dfe5aa76a0a95). The modified solve function is solve_new. I also included just the time required to call np.array and np.asarray on the inputs. An image of the gist's output is below. For small N (< 10000), you can reduce time by ~5-10%, and for large N (> ~10000) the time reduction is ~30%.

image

If you're interested in making this change, feel free to copy and/or modify any of the code in the gist since it's basically just copies of your code anyway, or I can submit a pull request. You could simplify the code by always using np.asarray and simply calling shift_banded with copy=True, but I wanted to modify as little as possible for the gist.

MuellerSeb commented 3 years ago

Hey @derb12. Thanks for using Pentapy! That suggestion looks really good. Would you please open a PR with these changes? That would be great! Sebastian