Adding custom linear solvers

[aadityacs]

Great library!

I think it would be nice if the library had a host of other linear solvers that can be useful for solving larger problems. For instance, I have a set of solvers like multi-grid solvers, intel's pypardiso integrated with JAX here.

I was wondering what would be the best way to go about integrating these solvers into the framework?

[tianjuxue]

Hello Aaditya, thanks for mentioning this and your linear solver examples. Indeed, the entire performance of this library JAX-FEM depends on scalable linear solvers. I just updated the code a little bit so that at least users have a chance to define their own favoriate solver, e.g., pardiso solver as in your example.

Please see the source code file solver.py as well as the user file example.py.

Let me know your thoughts! The next move for the library is really to go for massive parallel solvers (with multiple CPUs or even multiple GPUs), but I am a bit not sure what'd be the best practice at this moment.

[aadityacs]

I think we should be able to leverage already existing solvers on GPU(s) to speed up the computation.

For instance see this notebook where I am simply calling cupy code from jax to see a speed-up from 33 secs (on CPU) to 2 secs (on GPU)