patrick-kidger / lineax

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

Useful features #20

Closed thibmonsel closed 1 year ago

thibmonsel commented 1 year ago

It could probably useful to have some diagonal and triangular block matrices operators lineax.DiagonalBlockMatricesLinearOperator and lineax.TriangularlBlockMatricesLinearOperator ?

patrick-kidger commented 1 year ago

Our thinking on this front is to introduce a lineax.BlockLinearOperator, which you might call as e.g.

op1 = lineax.MatrixLinearOperator(...)
op2 = lineax.MatrixLinearOperator(...)
op = lineax.BlockLinearOperator([[op1, 0], [0, op2]])

which in the above case is diagonal. It uses nested linear operators to represent the nonzero parts and uses literal 0s to represent the zero parts of the overall operator.

Right now this is on the to-do list, see #4.