search

patrick-kidger / lineax

Linear solvers in JAX and Equinox. https://docs.kidger.site/lineax
Apache License 2.0
365 stars 24 forks source link

New solver based on Woodbury matrix identity #97

Open aidancrilly opened 5 months ago

aidancrilly commented 5 months ago

In issue #3 , a solver based on Woodbury matrix identity was given as a new solver feature for lineax.

Here I have made a first attempt at implementing it. It works slightly differently than the other base LinearOperators as it takes a LinearOperator as an argument (for A), with U, C and V as JAX arrays. This allows for the correct specialised solvers to be used for inverse(A) operations. E.g. A can be a DiagonalLinearOperator/TridiagonalLinearOperator/etc. - see test_Woodbury in test_operator.py

Things I am not certain of:

  1. Is passing solvers and states around in the SolverState of Woodbury the best approach. It would seem to me to minimise re-initialising solvers for multiple RHS. However perhaps there is a neater solution.
  2. I have not included a tag for Woodbury, it seems difficult to identify if matrices have this structure so it seems to me it should be left to the user. However, I am not sure if I have handled the operation-type linear operators (Add,Mul,Div,etc.) correctly in this instance.
  3. I had to add a pyright ignore in _solve.py under _linear_solve_transpose for allow_struct, not sure why. 
    ft.partial(eqxi.materialise_zeros, allow_struct=True), # pyright: ignore

    But it made me notice that materialise_zeros is from equinox 0.11.4 but pyproject.toml requires >=0.11.3. Does this require correcting?