Implementation of LSMR for iterative least squares. #86
Conversation
TODO: - Figure out early exiting - Add tests - Add type annotations - Add docstrings - Figure out complex dtypes - Clean up large number of state variables - Return actual status message on failure
There was a problem hiding this comment.
This looks really cool! Thank you for contributing it, I can see this must have been a lot of effort. :D I've gone through and left various comments, which should all be very small. In addition to those, can we add LSMR to the well-posed square tests?
On the exit conditions: how many does LSMR potentially have?
On complex dtypes: indeed, but I think we now have things almost entirely working. Should probably get that fixed on our end in the next few couple of weeks, and then we're just waiting on a new JAX release so that the upstream fixes there are generally available. (I think we should probably let all of that happen before merging this PR, so that we can add LSMR to the complex tests as well.)
|
The original version has 7 possible exit conditions:
I think we could get rid of 5-7 provided that maxiter is always finite. 2 and 3 could probably be combined into a standard "successful" exit assuming the user doesn't really care whether its the "true" solution or least squares. The main new one would be the condition number one. We could maybe shoehorn that into |
|
I know it's bad form to comment just when you want to 👍🏻 , but having LSMR in The problem here is one where I have a large (complex) linear system and I never want to construct the design matrix, so iterative methods are all I can really use. QR or SVD had huge memory overheads and were very slow. The Jax implementation for GMRES (in So this is a big 👍🏻 to this feature. |
Tests seem to be passing locally. I added
LSMRin the places it seemed to make sense (mainly places where I sawSVDwhich solves similar problems. Let me know if there are other tests that should be added.A few things that could use some input:
Resolves #85